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

Support runtime guidance tool interrupts

Browse Source
This commit is contained in:
wuyang6
2026-06-23 15:20:59 +08:00
parent 712bccf814
commit 45792c8fd5
10 changed files with 955 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/agent_runtime.py
+404 -6
View File
@@ -6,6 +6,7 @@ from datetime import datetime, timezone
import json
from pathlib import Path
import threading
import time
from typing import Any, Callable
from uuid import uuid4
@@ -269,6 +270,45 @@ def _all_skill_tool_call_ids(transcript: tuple[dict[str, object], ...]) -> list[
return results
class _CombinedCancelEvent:
"""Expose a single cancel flag backed by run cancel plus per-tool interrupt."""
def __init__(
self,
run_cancel_event: Any | None,
tool_interrupt_event: threading.Event,
) -> None:
self._run_cancel_event = run_cancel_event
self._tool_interrupt_event = tool_interrupt_event
def is_set(self) -> bool:
run_cancelled = (
self._run_cancel_event is not None
and bool(self._run_cancel_event.is_set())
)
return run_cancelled or self._tool_interrupt_event.is_set()
def set(self) -> None:
self._tool_interrupt_event.set()
def wait(self, timeout: float | None = None) -> bool:
if self.is_set():
return True
if self._run_cancel_event is None:
return self._tool_interrupt_event.wait(timeout)
if timeout is None:
while not self.is_set():
self._tool_interrupt_event.wait(0.1)
return True
deadline = time.monotonic() + max(0.0, timeout)
while not self.is_set():
remaining = deadline - time.monotonic()
if remaining <= 0:
return self.is_set()
self._tool_interrupt_event.wait(min(0.1, remaining))
return True
@dataclass
class LocalCodingAgent:
model_config: ModelConfig
@@ -739,6 +779,7 @@ class LocalCodingAgent:
session,
stream_events,
turn_index=turn_index,
stage='before_model',
)
self._microcompact_session_if_needed(
session,
@@ -1036,6 +1077,32 @@ class LocalCodingAgent:
self.last_run_result = result
return result
if turn.tool_calls:
guidance_items = self._consume_runtime_guidance(
stream_events,
turn_index=turn_index,
stage='before_tools',
)
if guidance_items:
self._append_skipped_tool_results_for_guidance(
session,
stream_events,
turn=turn,
turn_index=turn_index,
guidance_items=guidance_items,
)
self._append_runtime_guidance_items(
session,
stream_events,
guidance_items,
turn_index=turn_index,
stage='before_tools',
)
assistant_response_segments.clear()
continuation_count = 0
last_content = turn.content
continue
if not turn.tool_calls:
assistant_response_segments.append(turn.content)
if self._is_empty_truncated_response(turn):
@@ -1092,6 +1159,16 @@ class LocalCodingAgent:
)
last_content = ''.join(assistant_response_segments)
continue
if self._inject_runtime_guidance(
session,
stream_events,
turn_index=turn_index,
stage='before_finish',
):
assistant_response_segments.clear()
continuation_count = 0
last_content = turn.content
continue
final_output = ''.join(assistant_response_segments)
if self._reached_auto_continuation_limit(
turn,
@@ -1130,6 +1207,7 @@ class LocalCodingAgent:
self.last_run_result = result
return result
replan_after_tool_guidance = False
for tool_call in turn.tool_calls:
assistant_response_segments.clear()
continuation_count = 0
@@ -1176,6 +1254,8 @@ class LocalCodingAgent:
self.last_run_result = result
return result
tool_result = None
tool_guidance_items: tuple[dict[str, object], ...] = ()
tool_guidance_interrupted = False
tool_message_index = session.start_tool(
name=tool_call.name,
tool_call_id=tool_call.id,
@@ -1279,11 +1359,19 @@ class LocalCodingAgent:
duplicate_skill_counts=duplicate_skill_counts,
)
elif tool_result is None:
tool_interrupt_event = threading.Event()
tool_context = replace(
self.tool_context,
cancel_event=_CombinedCancelEvent(
self.tool_context.cancel_event,
tool_interrupt_event,
),
)
for update in execute_tool_streaming(
self.tool_registry,
tool_call.name,
tool_call.arguments,
self.tool_context,
tool_context,
):
if update.kind == 'delta':
session.append_tool_delta(
@@ -1301,8 +1389,81 @@ class LocalCodingAgent:
'delta': update.content,
}
)
if not tool_guidance_items:
pending_guidance = self._consume_runtime_guidance(
stream_events,
turn_index=turn_index,
stage='during_tool',
)
if pending_guidance:
tool_guidance_items = pending_guidance
if self._runtime_guidance_should_interrupt_tool(
tool_call=tool_call,
guidance_items=pending_guidance,
):
tool_guidance_interrupted = True
tool_interrupt_event.set()
self._interrupt_running_tool_for_guidance(
stream_events,
turn_index=turn_index,
tool_call=tool_call,
guidance_items=pending_guidance,
)
else:
stream_events.append(
{
'type': 'runtime_guidance_deferred_after_tool',
'turn_index': turn_index,
'tool_name': tool_call.name,
'tool_call_id': tool_call.id,
'input_ids': self._runtime_guidance_item_ids(
pending_guidance
),
'message': (
'用户引导在工具执行中到达，运行时判断无需中断，'
'将等待当前工具完成后重规划'
),
}
)
continue
tool_result = update.result
if tool_result is not None and not tool_guidance_items:
tool_guidance_items = self._consume_runtime_guidance(
stream_events,
turn_index=turn_index,
stage='after_tool',
)
if tool_guidance_items:
stream_events.append(
{
'type': 'runtime_guidance_deferred_after_tool',
'turn_index': turn_index,
'tool_name': tool_call.name,
'tool_call_id': tool_call.id,
'input_ids': self._runtime_guidance_item_ids(
tool_guidance_items
),
'message': (
'用户引导在工具执行期间到达，但当前工具未提供中间输出；'
'已在工具完成后准备重规划'
),
}
)
if tool_result is not None and tool_guidance_items:
merged_metadata = dict(tool_result.metadata)
merged_metadata['runtime_guidance_input_ids'] = (
self._runtime_guidance_item_ids(tool_guidance_items)
)
if tool_guidance_interrupted:
merged_metadata['runtime_guidance_interrupted_tool'] = True
else:
merged_metadata['runtime_guidance_deferred_after_tool'] = True
tool_result = ToolExecutionResult(
name=tool_result.name,
ok=tool_result.ok,
content=tool_result.content,
metadata=merged_metadata,
)
if tool_result is None:
raise RuntimeError(f'Tool executor returned no final result for {tool_call.name}')
if (
@@ -1474,6 +1635,23 @@ class LocalCodingAgent:
)
if history_entry is not None:
file_history.append(history_entry)
if tool_guidance_items:
self._append_runtime_guidance_items(
session,
stream_events,
tool_guidance_items,
turn_index=turn_index,
stage=(
'during_tool_interrupted'
if tool_guidance_interrupted
else 'after_tool'
),
)
assistant_response_segments.clear()
continuation_count = 0
last_content = turn.content
replan_after_tool_guidance = True
break
if tool_result.metadata.get('stop_agent') is True:
stop_reason = str(
tool_result.metadata.get('stop_reason') or 'tool_requested_stop'
@@ -1561,6 +1739,9 @@ class LocalCodingAgent:
self.last_run_result = result
return result
if replan_after_tool_guidance:
continue
final_output = self._build_max_turns_output(
max_turns=self.runtime_config.max_turns,
last_content=last_content,
@@ -4286,10 +4467,34 @@ class LocalCodingAgent:
stream_events: list[dict[str, object]],
*,
turn_index: int,
) -> None:
stage: str,
) -> bool:
items = self._consume_runtime_guidance(
stream_events,
turn_index=turn_index,
stage=stage,
)
if not items:
return False
self._append_runtime_guidance_items(
session,
stream_events,
items,
turn_index=turn_index,
stage=stage,
)
return True
def _consume_runtime_guidance(
self,
stream_events: list[dict[str, object]],
*,
turn_index: int,
stage: str,
) -> tuple[dict[str, object], ...]:
provider = self.runtime_guidance_provider
if provider is None:
return
return ()
try:
items = tuple(
item
@@ -4302,12 +4507,22 @@ class LocalCodingAgent:
{
'type': 'runtime_guidance_error',
'turn_index': turn_index,
'stage': stage,
'error': str(exc)[:500],
}
)
return
if not items:
return
return ()
return items
def _append_runtime_guidance_items(
self,
session: AgentSessionState,
stream_events: list[dict[str, object]],
items: tuple[dict[str, object], ...],
*,
turn_index: int,
stage: str,
) -> None:
lines = [
'<runtime-guidance>',
(
@@ -4315,6 +4530,26 @@ class LocalCodingAgent:
'吸收这些要求；如果它们与当前目标冲突，先说明冲突并选择最合理的继续方式。'
),
]
if stage == 'before_tools':
lines.append(
'这批引导在工具执行前到达，运行时已经跳过上一批尚未执行的工具调用。'
'请重新判断：继续原计划、调整工具参数、改走新计划，或说明冲突。'
)
elif stage == 'before_finish':
lines.append(
'这批引导在最终回复前到达。请判断它是补充最终输出、修正当前结论，'
'还是应作为后续任务处理。'
)
elif stage == 'during_tool_interrupted':
lines.append(
'这批引导在工具执行中到达，运行时已经中断当前工具并停止继续执行同批旧工具计划。'
'请结合已有工具输出和这批引导重新规划；如果需要重跑工具，请使用新的参数。'
)
elif stage == 'after_tool':
lines.append(
'这批引导在工具执行期间或刚完成后到达。当前工具结果已保留，'
'运行时停止继续执行同批旧工具计划，请判断是否基于该结果继续、补充或重跑。'
)
item_ids: list[int] = []
for index, item in enumerate(items, start=1):
item_id = item.get('id')
@@ -4328,6 +4563,7 @@ class LocalCodingAgent:
metadata={
'kind': 'runtime_guidance',
'turn_index': turn_index,
'stage': stage,
'queue_item_ids': item_ids,
},
message_id=f'runtime_guidance_{turn_index}_{uuid4().hex[:8]}',
@@ -4337,12 +4573,174 @@ class LocalCodingAgent:
{
'type': 'runtime_guidance_injected',
'turn_index': turn_index,
'stage': stage,
'input_ids': item_ids,
'count': len(items),
'message': f'已将 {len(items)} 条用户引导注入当前任务',
}
)
@staticmethod
def _runtime_guidance_item_ids(
items: tuple[dict[str, object], ...],
) -> list[int]:
return [
int(item.get('id'))
for item in items
if isinstance(item.get('id'), int)
and not isinstance(item.get('id'), bool)
]
@staticmethod
def _runtime_guidance_should_interrupt_tool(
*,
tool_call: ToolCall,
guidance_items: tuple[dict[str, object], ...],
) -> bool:
combined = '\n'.join(
str(item.get('content') or '').strip().lower()
for item in guidance_items
)
if not combined:
return False
defer_markers = (
'完成后',
'结束后',
'最后',
'顺便',
'同时整理',
'补充输出',
'finish',
'after',
'then',
'also',
)
interrupt_markers = (
'停',
'中断',
'取消',
'暂停',
'别跑',
'别继续',
'不要继续',
'先别',
'错了',
'不对',
'改成',
'换成',
'重新',
'重来',
'不用',
'不是',
'别用',
'stop',
'cancel',
'interrupt',
'abort',
'wrong',
'instead',
'change to',
'switch to',
'restart',
'redo',
)
if any(marker in combined for marker in interrupt_markers):
return True
if any(marker in combined for marker in defer_markers):
return False
return tool_call.name in {
'bash',
'python_exec',
'jupyter_exec',
'remote_bash',
'remote_python_exec',
'Agent',
'delegate_agent',
} and len(combined) >= 120
def _interrupt_running_tool_for_guidance(
self,
stream_events: list[dict[str, object]],
*,
turn_index: int,
tool_call: ToolCall,
guidance_items: tuple[dict[str, object], ...],
) -> None:
process_registry = self.tool_context.process_registry
killed_processes = False
if process_registry is not None:
try:
process_registry.kill_all()
killed_processes = True
except Exception as exc: # noqa: BLE001
stream_events.append(
{
'type': 'runtime_guidance_error',
'turn_index': turn_index,
'stage': 'during_tool',
'tool_name': tool_call.name,
'tool_call_id': tool_call.id,
'error': str(exc)[:500],
}
)
stream_events.append(
{
'type': 'runtime_guidance_interrupt_tool',
'turn_index': turn_index,
'tool_name': tool_call.name,
'tool_call_id': tool_call.id,
'input_ids': self._runtime_guidance_item_ids(guidance_items),
'killed_processes': killed_processes,
'message': '用户引导要求调整当前任务，已中断当前工具并准备重规划',
}
)
def _append_skipped_tool_results_for_guidance(
self,
session: AgentSessionState,
stream_events: list[dict[str, object]],
*,
turn: AssistantTurn,
turn_index: int,
guidance_items: tuple[dict[str, object], ...],
) -> None:
input_ids = [
int(item.get('id'))
for item in guidance_items
if isinstance(item.get('id'), int)
and not isinstance(item.get('id'), bool)
]
for tool_call in turn.tool_calls:
tool_result = ToolExecutionResult(
name=tool_call.name,
ok=False,
content=(
'工具尚未执行时收到了用户运行中引导。运行时跳过了这次旧工具调用，'
'请结合后续 runtime guidance 重新规划。'
),
metadata={
'action': 'runtime_guidance_replan',
'runtime_guidance_skipped': True,
'runtime_guidance_input_ids': input_ids,
},
)
session.append_tool(
tool_call.name,
tool_call.id,
serialize_tool_result(tool_result),
display=False,
)
stream_events.append(
{
'type': 'runtime_guidance_replan_before_tools',
'turn_index': turn_index,
'input_ids': input_ids,
'tool_call_count': len(turn.tool_calls),
'tool_names': [tool_call.name for tool_call in turn.tool_calls],
'message': '用户引导在工具执行前到达，已跳过旧工具计划并要求模型重规划',
}
)
def render_system_prompt(self) -> str:
prompt_context = self.build_prompt_context()
parts = self.build_system_prompt_parts(prompt_context)