fix: clean up temporary output 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
 
@@ -306,6 +308,11 @@ async def stop(self) -> None:
             self.kernel_context = None
 
         self._client = None
+
+        if self._temp_dir is not None:
+            self._temp_dir.cleanup()
+            self._temp_dir = None
+
         self._started = False
 
     def _to_config(self) -> JupyterCodeExecutorConfig:

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

@@ -225,3 +225,28 @@ 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 temporary output directory is cleaned up when stop() is called."""
+    executor = JupyterCodeExecutor()
+    temp_dir = executor.output_dir
+    assert temp_dir.exists()
+
+    await executor.start()
+    await executor.stop()
+
+    assert not temp_dir.exists(), "Temporary output directory should be cleaned up after stop()"
+
+
+@pytest.mark.asyncio
+async def test_user_provided_output_dir_not_cleaned_up(tmp_path: Path) -> None:
+    """Test that user-provided output directory is NOT cleaned up when stop() is called."""
+    executor = JupyterCodeExecutor(output_dir=tmp_path)
+    assert executor.output_dir == tmp_path
+
+    await executor.start()
+    await executor.stop()
+
+    assert tmp_path.exists(), "User-provided output directory should not be cleaned up after stop()"