From 50df8b6833d49b1aab2df2ff525977f6aa6be6d0 Mon Sep 17 00:00:00 2001 From: r266-tech Date: Mon, 23 Mar 2026 01:00:46 +0800 Subject: [PATCH] fix: clean up temp directory in JupyterCodeExecutor.stop() JupyterCodeExecutor created a temp directory via tempfile.mkdtemp() when no output_dir was provided, but never cleaned it up in stop(). This was inconsistent with other executors (LocalCommandLineCodeExecutor, DockerCommandLineCodeExecutor, AzureContainerCodeExecutor) which all properly clean up their temp directories. Changes: - Use tempfile.TemporaryDirectory() instead of tempfile.mkdtemp() when output_dir is None (auto-cleanup on context exit) - Add cleanup of self._temp_dir in stop() method - Add tests verifying cleanup behavior for both cases Fixes #7217 --- .../jupyter/_jupyter_code_executor.py | 14 +++++++--- .../test_jupyter_code_executor.py | 26 +++++++++++++++++++ 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..50f2e786b981 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 @@ -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( 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..cb82cb66b23b 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,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()"