wuyang/zk-data-agent
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add session scoped runs and cancellation

Browse Source
This commit is contained in:
武阳
2026-05-07 16:26:20 +08:00
parent a4a0677ce4
commit 9d7d496443
6 changed files with 542 additions and 76 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/agent_runtime.py
+24
View File
@@ -1060,6 +1060,30 @@ class LocalCodingAgent:
tool_result = update.result
if tool_result is None:
raise RuntimeError(f'Tool executor returned no final result for {tool_call.name}')
if (
self.tool_context.cancel_event is not None
and self.tool_context.cancel_event.is_set()
):
result = AgentRunResult(
final_output='已取消当前请求。',
turns=turn_index,
tool_calls=tool_calls,
transcript=session.transcript(),
events=tuple(stream_events),
usage=total_usage,
total_cost_usd=total_cost_usd,
stop_reason='cancelled',
file_history=tuple(file_history),
session_id=session_id,
scratchpad_directory=(
str(scratchpad_directory)
if scratchpad_directory is not None
else None
),
)
result = self._persist_session(session, result)
self.last_run_result = result
return result
if self.plugin_runtime is not None:
self.plugin_runtime.record_tool_result(
tool_call.name,
src/agent_tool_core.py
+6
View File
@@ -40,6 +40,8 @@ class ToolExecutionContext:
scratchpad_directory: Path | None = None
python_env_dir: Path | None = None
extra_env: dict[str, str] = field(default_factory=dict)
cancel_event: Any | None = None
process_registry: Any | None = None
tool_registry: dict[str, 'AgentTool'] | None = None
search_runtime: 'SearchRuntime | None' = None
account_runtime: 'AccountRuntime | None' = None
@@ -127,6 +129,8 @@ def build_tool_context(
scratchpad_directory: Path | None = None,
python_env_dir: Path | None = None,
extra_env: dict[str, str] | None = None,
cancel_event: Any | None = None,
process_registry: Any | None = None,
tool_registry: dict[str, AgentTool] | None = None,
search_runtime: 'SearchRuntime | None' = None,
account_runtime: 'AccountRuntime | None' = None,
@@ -156,6 +160,8 @@ def build_tool_context(
else None
),
extra_env=dict(extra_env or {}),
cancel_event=cancel_event,
process_registry=process_registry,
tool_registry=tool_registry,
search_runtime=search_runtime,
account_runtime=account_runtime,
src/agent_tools.py
+86 -11
View File
@@ -1825,6 +1825,7 @@ def _grep_search(arguments: dict[str, Any], context: ToolExecutionContext) -> st
def _run_python_exec(arguments: dict[str, Any], context: ToolExecutionContext) -> str:
_ensure_process_execution_allowed(context, 'Python execution')
_raise_if_cancelled(context)
code = _optional_string(arguments, 'code')
script_path = _optional_string(arguments, 'script_path')
if bool(code) == bool(script_path):
@@ -1856,19 +1857,40 @@ def _run_python_exec(arguments: dict[str, Any], context: ToolExecutionContext) -
command = [interpreter, str(script), *raw_args]
mode = 'script'
process: subprocess.Popen[str] | None = None
try:
completed = subprocess.run(
process = subprocess.Popen(
command,
cwd=context.root,
input=stdin if stdin else None,
capture_output=True,
stdin=subprocess.PIPE if stdin else None,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
text=True,
timeout=timeout_seconds,
env=_build_subprocess_env(context),
)
except subprocess.TimeoutExpired as exc:
stdout = exc.stdout if isinstance(exc.stdout, str) else ''
stderr = exc.stderr if isinstance(exc.stderr, str) else ''
_register_child_process(context, process)
_register_child_process(context, process)
if stdin and process.stdin is not None:
process.stdin.write(stdin)
process.stdin.close()
process.stdin = None
deadline = time.monotonic() + timeout_seconds
while True:
_raise_if_cancelled(context, process)
remaining = deadline - time.monotonic()
if remaining <= 0:
raise subprocess.TimeoutExpired(command, timeout_seconds)
try:
stdout, stderr = process.communicate(timeout=min(remaining, 0.2))
break
except subprocess.TimeoutExpired:
continue
except subprocess.TimeoutExpired:
if process is not None:
_terminate_process(process)
stdout, stderr = process.communicate()
else:
stdout, stderr = '', ''
payload = [
'timed_out=true',
f'timeout_seconds={timeout_seconds:g}',
@@ -1894,11 +1916,17 @@ def _run_python_exec(arguments: dict[str, Any], context: ToolExecutionContext) -
'output_preview': _snapshot_text('\n'.join(payload).strip()),
},
)
except ToolExecutionError:
raise
finally:
if process is not None:
_unregister_child_process(context, process)
stdout = completed.stdout or ''
stderr = completed.stderr or ''
stdout = stdout or ''
stderr = stderr or ''
returncode = process.returncode if process is not None else 1
payload = [
f'exit_code={completed.returncode}',
f'exit_code={returncode}',
f'interpreter={interpreter}',
f'scratchpad={context.scratchpad_directory or ""}',
f'python_env={context.python_env_dir or ""}',
@@ -1915,7 +1943,7 @@ def _run_python_exec(arguments: dict[str, Any], context: ToolExecutionContext) -
'interpreter': interpreter,
'scratchpad_directory': str(context.scratchpad_directory) if context.scratchpad_directory else '',
'python_env_dir': str(context.python_env_dir) if context.python_env_dir else '',
'exit_code': completed.returncode,
'exit_code': returncode,
'stdout_preview': _snapshot_text(stdout),
'stderr_preview': _snapshot_text(stderr),
'output_preview': _snapshot_text('\n'.join(payload).strip()),
@@ -1923,6 +1951,48 @@ def _run_python_exec(arguments: dict[str, Any], context: ToolExecutionContext) -
)
def _raise_if_cancelled(
context: ToolExecutionContext,
process: subprocess.Popen[Any] | None = None,
) -> None:
cancel_event = context.cancel_event
if cancel_event is not None and cancel_event.is_set():
if process is not None:
_terminate_process(process)
raise ToolExecutionError('Run cancelled by user')
def _register_child_process(
context: ToolExecutionContext,
process: subprocess.Popen[Any],
) -> None:
registry = context.process_registry
if registry is not None:
registry.add(process)
def _unregister_child_process(
context: ToolExecutionContext,
process: subprocess.Popen[Any],
) -> None:
registry = context.process_registry
if registry is not None:
registry.discard(process)
def _terminate_process(process: subprocess.Popen[Any]) -> None:
if process.poll() is not None:
return
try:
process.terminate()
process.wait(timeout=1)
except (OSError, subprocess.TimeoutExpired):
try:
process.kill()
except OSError:
pass
def _run_python_package(arguments: dict[str, Any], context: ToolExecutionContext) -> str:
_ensure_process_execution_allowed(context, 'Python package management')
action = _require_string(arguments, 'action')
@@ -3364,6 +3434,10 @@ def _stream_bash(
try:
while selector.get_map():
if context.cancel_event is not None and context.cancel_event.is_set():
timeout_error = f'Command cancelled by user: {command}'
_terminate_process(process)
break
remaining = deadline - time.monotonic()
if remaining <= 0:
timeout_error = (
@@ -3398,6 +3472,7 @@ def _stream_bash(
stream=stream_name,
)
finally:
_unregister_child_process(context, process)
try:
selector.close()
except OSError: