From 2163cbe24e2583d4c497cb9fb91fbbdc6a7412dd Mon Sep 17 00:00:00 2001
From: Varun Chawla <varun_6april@hotmail.com>
Date: Sat, 7 Feb 2026 11:21:41 -0800
Subject: [PATCH] fix: clean up temporary output directory in
 JupyterCodeExecutor.stop()

---
 .../jupyter/_jupyter_code_executor.py         | 15 ++++++++---
 .../test_jupyter_code_executor.py             | 25 +++++++++++++++++++
 2 files changed, 36 insertions(+), 4 deletions(-)

diff --git a/python/packages/autogen-ext/src/autogen_ext/code_executors/jupyter/_jupyter_code_executor.py b/python/packages/autogen-ext/src/autogen_ext/code_executors/jupyter/_jupyter_code_executor.py
index 2476b5a3349f..0eac0dbd2888 100644
--- a/python/packages/autogen-ext/src/autogen_ext/code_executors/jupyter/_jupyter_code_executor.py
+++ b/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:
diff --git a/python/packages/autogen-ext/tests/code_executors/test_jupyter_code_executor.py b/python/packages/autogen-ext/tests/code_executors/test_jupyter_code_executor.py
index b6789d0b5e41..22fcc64491af 100644
--- a/python/packages/autogen-ext/tests/code_executors/test_jupyter_code_executor.py
+++ b/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()"