Code execute cancellation (#299)

* Hook cancelation token into code execution

* Add unit test for code cancellation

* actually save the merge

python/src/agnext/components/code_executor/_base.py

@@ -6,6 +6,8 @@ from __future__ import annotations
 from dataclasses import dataclass
 from typing import List, Protocol, runtime_checkable
 
+from agnext.core import CancellationToken
+
 
 @dataclass
 class CodeBlock:
@@ -27,7 +29,9 @@ class CodeResult:
 class CodeExecutor(Protocol):
     """Executes code blocks and returns the result."""
 
-    async def execute_code_blocks(self, code_blocks: List[CodeBlock]) -> CodeResult:
+    async def execute_code_blocks(
+        self, code_blocks: List[CodeBlock], cancellation_token: CancellationToken
+    ) -> CodeResult:
         """Execute code blocks and return the result.
 
         This method should be implemented by the code executor.

python/src/agnext/components/code_executor/_impl/local_commandline_code_executor.py

@@ -8,7 +8,7 @@ import warnings
 from hashlib import md5
 from pathlib import Path
 from string import Template
-from typing import Any, Callable, ClassVar, List, Sequence, Union, Optional
+from typing import Any, Callable, ClassVar, List, Sequence, Union
 
 from typing_extensions import ParamSpec
 
@@ -22,7 +22,6 @@ from .._func_with_reqs import (
 )
 from .command_line_code_result import CommandLineCodeResult
 from .utils import PYTHON_VARIANTS, get_file_name_from_content, lang_to_cmd, silence_pip  # type: ignore
-from ....core import CancellationToken
 
 __all__ = ("LocalCommandLineCodeExecutor",)
 
@@ -146,7 +145,7 @@ $functions"""
         """(Experimental) The working directory for the code execution."""
         return self._work_dir
 
-    async def _setup_functions(self, cancellation_token: Optional[CancellationToken]) -> None:
+    async def _setup_functions(self, cancellation_token: CancellationToken) -> None:
         func_file_content = build_python_functions_file(self._functions)
         func_file = self._work_dir / f"{self._functions_module}.py"
         func_file.write_text(func_file_content)
@@ -170,8 +169,7 @@ $functions"""
                     stderr=asyncio.subprocess.PIPE,
                 )
             )
-            if cancellation_token:
-                cancellation_token.link_future(task)
+            cancellation_token.link_future(task)
             try:
                 proc = await task
                 stdout, stderr = await asyncio.wait_for(proc.communicate(), self._timeout)
@@ -194,13 +192,13 @@ $functions"""
         self._setup_functions_complete = True
 
     async def execute_code_blocks(
-        self, code_blocks: List[CodeBlock], cancellation_token: Optional[CancellationToken] = None
+        self, code_blocks: List[CodeBlock], cancellation_token: CancellationToken
     ) -> CommandLineCodeResult:
         """(Experimental) Execute the code blocks and return the result.
 
         Args:
             code_blocks (List[CodeBlock]): The code blocks to execute.
-            cancellation_token (CancellationToken|None): an optional token to cancel the operation
+            cancellation_token (CancellationToken): a token to cancel the operation
 
         Returns:
             CommandLineCodeResult: The result of the code execution."""
@@ -211,7 +209,7 @@ $functions"""
         return await self._execute_code_dont_check_setup(code_blocks, cancellation_token)
 
     async def _execute_code_dont_check_setup(
-        self, code_blocks: List[CodeBlock], cancellation_token: Optional[CancellationToken]
+        self, code_blocks: List[CodeBlock], cancellation_token: CancellationToken
     ) -> CommandLineCodeResult:
         logs_all: str = ""
         file_names: List[Path] = []
@@ -262,8 +260,7 @@ $functions"""
                     stderr=asyncio.subprocess.PIPE,
                 )
             )
-            if cancellation_token:
-                cancellation_token.link_future(task)
+            cancellation_token.link_future(task)
             try:
                 proc = await task
                 stdout, stderr = await asyncio.wait_for(proc.communicate(), self._timeout)

python/src/agnext/components/tools/_code_execution.py

@@ -25,6 +25,8 @@ class PythonCodeExecutionTool(BaseTool[CodeExecutionInput, CodeExecutionResult])
 
     async def run(self, args: CodeExecutionInput, cancellation_token: CancellationToken) -> CodeExecutionResult:
         code_blocks = [CodeBlock(code=args.code, language="python")]
-        result = await self._executor.execute_code_blocks(code_blocks=code_blocks)
+        result = await self._executor.execute_code_blocks(
+            code_blocks=code_blocks, cancellation_token=cancellation_token
+        )
 
         return CodeExecutionResult(success=result.exit_code == 0, output=result.output)