fix: clean up temp directory in JupyterCodeExecutor.stop()

python/packages/autogen-ext/src/autogen_ext/code_executors/jupyter/_jupyter_code_executor.py

@@ -145,11 +145,13 @@ def __init__(
         if timeout < 1:
             raise ValueError("Timeout must be greater than or equal to 1.")
 
-        self._output_dir: Path = Path(tempfile.mkdtemp()) if output_dir is None else Path(output_dir)
-        self._output_dir.mkdir(exist_ok=True, parents=True)
-
         self._temp_dir: Optional[tempfile.TemporaryDirectory[str]] = None
-        self._temp_dir_path: Optional[Path] = None
+        if output_dir is None:
+            self._temp_dir = tempfile.TemporaryDirectory()
+            self._output_dir = Path(self._temp_dir.name)
+        else:
+            self._output_dir = Path(output_dir)
+        self._output_dir.mkdir(exist_ok=True, parents=True)
 
         self._started = False
 
@@ -308,6 +310,10 @@ async def stop(self) -> None:
         self._client = None
         self._started = False
 
+        if self._temp_dir is not None:
+            self._temp_dir.cleanup()
+            self._temp_dir = None
+
     def _to_config(self) -> JupyterCodeExecutorConfig:
         """Convert current instance to config object"""
         return JupyterCodeExecutorConfig(

python/packages/autogen-ext/tests/code_executors/test_jupyter_code_executor.py

@@ -225,3 +225,29 @@ async def test_runtime_error_not_started() -> None:
     code_blocks = [CodeBlock(code="print('hello world!')", language="python")]
     with pytest.raises(RuntimeError, match="Executor must be started before executing cells"):
         await executor.execute_code_blocks(code_blocks, CancellationToken())
+
+
+@pytest.mark.asyncio
+async def test_temp_dir_cleanup_on_stop() -> None:
+    """Test that temp directory is cleaned up when output_dir is not specified."""
+    executor = JupyterCodeExecutor()
+    temp_dir_path = executor.output_dir
+    assert temp_dir_path.exists()
+
+    await executor.start()
+    code_blocks = [CodeBlock(code="print('hello world!')", language="python")]
+    await executor.execute_code_blocks(code_blocks, CancellationToken())
+    await executor.stop()
+
+    assert not temp_dir_path.exists(), "Temp directory should be cleaned up after stop()"
+
+
+@pytest.mark.asyncio
+async def test_explicit_output_dir_not_cleaned_up(tmp_path: Path) -> None:
+    """Test that explicitly provided output_dir is NOT cleaned up after stop()."""
+    async with JupyterCodeExecutor(output_dir=tmp_path) as executor:
+        await executor.start()
+        code_blocks = [CodeBlock(code="print('hello world!')", language="python")]
+        await executor.execute_code_blocks(code_blocks, CancellationToken())
+
+    assert tmp_path.exists(), "Explicit output_dir should NOT be cleaned up after stop()"