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()"