diff --git a/config.py b/config.py index 2510127..ff91e1f 100644 --- a/config.py +++ b/config.py @@ -95,6 +95,7 @@ def default_config() -> dict: "auto_test": True, # run Test Engineer after every coding task }, "tool_approval": True, # False = auto-approve tools; True = ask for non-safe tools + "tool_timeout_seconds": 600, # per-call deadline for a tool call (0 = no deadline) "ui": { "default": "plain", # plain | tui "refresh_interval_ms": 1500, diff --git a/core.py b/core.py index b5cfeb5..92cca39 100644 --- a/core.py +++ b/core.py @@ -35,6 +35,59 @@ } +# ── Per-tool execution deadline & interruptibility ─────────────────────────── +# Every tool call runs on a worker thread that the conversation loop polls in +# short slices. Before this, a tool that never returned (a wedged subprocess, a +# socket read with no timeout) blocked the loop forever, and interrupt() — which +# only sets self._cancel — could not break it, because nothing was polled while +# the tool ran. A call that overruns `tool_timeout_seconds`, or the user's +# interrupt, is now abandoned instead of awaited (the worker is a daemon thread). +_SEQUENTIAL_INTERRUPT_POLL_SECONDS = 0.1 +_DEFAULT_TOOL_TIMEOUT_SECONDS = 600.0 +# Tools that supervise their own liveness and legitimately block for a whole +# long-running operation (subagent/delegation batches carry their own timeouts; +# browser_task has per-action Playwright timeouts): no generic per-call deadline, +# but they still run on a polled worker so an interrupt can abandon them. +_TOOL_DEADLINE_EXEMPT_TOOLS = frozenset({ + "delegate_task", "spawn_subagent", + "start_coding_session", "start_tracked_coding_task", + "browser_task", +}) + + +def _resolve_tool_timeout(agent) -> float | None: + """Per-call deadline in seconds for a non-exempt tool; ``None`` disables it. + + Reads ``tool_timeout_seconds`` from the agent config; ``0`` or a negative + value disables the deadline. A missing or invalid value falls back to + ``_DEFAULT_TOOL_TIMEOUT_SECONDS``. + """ + cfg = getattr(agent, "cfg", None) + raw = _DEFAULT_TOOL_TIMEOUT_SECONDS + if isinstance(cfg, dict) and "tool_timeout_seconds" in cfg: + raw = cfg["tool_timeout_seconds"] + try: + seconds = float(raw) + except (TypeError, ValueError): + seconds = _DEFAULT_TOOL_TIMEOUT_SECONDS + return None if seconds <= 0 else seconds + + +def _tool_timeout_message(name: str, timeout_s: float) -> str: + return ( + f"Error executing tool '{name}': timed out after {timeout_s:.1f}s and was abandoned. " + "Do not assume its side effects happened — it may still be running in the background. " + "Retry with a narrower call, or raise `tool_timeout_seconds` in config.yaml." + ) + + +def _tool_interrupt_message(name: str, elapsed: float) -> str: + return ( + f"Error: tool '{name}' was abandoned after {elapsed:.1f}s because the run was " + "interrupted by the user. Its side effects (if any) are unknown." + ) + + # Cached at startup — re-detect only on explicit request _SYSTEM_CAPS: dict = {} @@ -1296,79 +1349,90 @@ def _run_conversation_loop(self, user_input: str, image_path: str | None = None) except Exception: pass - from concurrent.futures import ThreadPoolExecutor - def run_one(call): - name, args = call["name"], call["arguments"] - t0 = time.time() - try: - res = self._execute_tool(name, args) - except Exception as e: - import traceback - res = f"Error executing tool: {e}\n{traceback.format_exc()}" - t_elapsed = time.time() - t0 - return res, t_elapsed - + # Each parallel call runs on its own daemon worker and + # reports through a queue. A wedged tool can then be + # abandoned instead of blocking the batch — an executor's + # ``with`` block joins its workers on exit, so one hung + # call used to freeze the whole round. This loop polls for + # results, interruption and the per-call deadline. if non_blocked_parallel_calls: - import concurrent.futures - with ThreadPoolExecutor(max_workers=len(non_blocked_parallel_calls)) as executor: - future_to_call = {} - for idx, call in non_blocked_parallel_calls: - future_to_call[executor.submit(run_one, call)] = (idx, call) + import queue as _queue - futures = list(future_to_call.keys()) - completed_results = {} + timeout_s = _resolve_tool_timeout(self) + results_q: _queue.Queue = _queue.Queue() + started_at: dict = {} - while futures: - if self._cancel.is_set(): - for f in futures: - f.cancel() - break - - done, futures = concurrent.futures.wait( - futures, timeout=0.1, return_when=concurrent.futures.FIRST_COMPLETED + def _worker(call, _q=results_q): + name, args = call["name"], call["arguments"] + t0 = time.time() + try: + res = self._execute_tool(name, args) + except Exception as e: + import traceback + res = f"Error executing tool: {e}\n{traceback.format_exc()}" + _q.put((call["id"], res, time.time() - t0)) + + pending: dict = {} + for idx, call in non_blocked_parallel_calls: + pending[call["id"]] = call + started_at[call["id"]] = time.time() + threading.Thread( + target=_worker, args=(call,), + name=f"koza-tool-{call['name']}", daemon=True, + ).start() + + completed_results = {} + while pending: + if self._cancel.is_set(): + break + try: + call_id, res, t_elapsed = results_q.get( + timeout=_SEQUENTIAL_INTERRUPT_POLL_SECONDS ) + except _queue.Empty: + call_id = None - for f in done: - idx, call = future_to_call[f] - name, args = call["name"], call["arguments"] + if call_id is not None and call_id in pending: + call = pending.pop(call_id) + name, args = call["name"], call["arguments"] + res = self._guardrail.after_call(name, args, res) + completed_results[call_id] = (res, t_elapsed) + if hasattr(self, "tool_progress_callback") and self.tool_progress_callback: try: - res, t_elapsed = f.result() - except concurrent.futures.CancelledError: - res = "Tool execution cancelled due to user interrupt." - t_elapsed = 0.0 - except Exception as e: - import traceback - res = f"Error executing tool: {e}\n{traceback.format_exc()}" - t_elapsed = 0.0 - - res = self._guardrail.after_call(name, args, res) - completed_results[call["id"]] = (res, t_elapsed) - - if hasattr(self, "tool_progress_callback") and self.tool_progress_callback: - try: - self.tool_progress_callback("tool.completed", name, str(res)[:300], args, duration=t_elapsed) - except Exception: - pass - - for idx, call in non_blocked_parallel_calls: - call_id = call["id"] - if call_id in completed_results: - res, t_elapsed = completed_results[call_id] - else: - res = "Tool execution cancelled or skipped." - t_elapsed = 0.0 + self.tool_progress_callback("tool.completed", name, str(res)[:300], args, duration=t_elapsed) + except Exception: + pass + + # Abandon any call that overran its deadline. + if timeout_s is not None and pending: + now = time.time() + for cid in list(pending): + if now - started_at[cid] >= timeout_s: + overrun = pending.pop(cid) + completed_results[cid] = ( + _tool_timeout_message(overrun["name"], timeout_s), + timeout_s, + ) + + for idx, call in non_blocked_parallel_calls: + call_id = call["id"] + if call_id in completed_results: + res, t_elapsed = completed_results[call_id] + else: + res = "Tool execution cancelled or skipped." + t_elapsed = 0.0 - name, args = call["name"], call["arguments"] - yield {"type": "tool_done", "name": name, "result": res, "elapsed": t_elapsed} - res_str = str(res) - self.messages.append({ - "role": "tool", - "tool_call_id": call_id, - "name": name, - "content": res_str, - }) - if "PERMANENT FAILURE" in res_str: - permanent_failure = True + name, args = call["name"], call["arguments"] + yield {"type": "tool_done", "name": name, "result": res, "elapsed": t_elapsed} + res_str = str(res) + self.messages.append({ + "role": "tool", + "tool_call_id": call_id, + "name": name, + "content": res_str, + }) + if "PERMANENT FAILURE" in res_str: + permanent_failure = True for idx, call in parallel_calls: if call["id"] in blocked_results: @@ -1431,13 +1495,9 @@ def run_one(call): self.tool_progress_callback("tool.started", name, preview, args) except Exception: pass - t0 = time.time() - try: - result = self._execute_tool(name, args) - except Exception as e: - import traceback - result = f"Error executing tool: {e}\n{traceback.format_exc()}" - t_elapsed = time.time() - t0 + # Per-call deadline + interrupt poll: a wedged tool is + # abandoned instead of blocking the whole loop forever. + result, t_elapsed, _outcome = self._run_tool_with_deadline(name, args) # Check loop guardrail after call result = self._guardrail.after_call(name, args, result) @@ -1457,6 +1517,16 @@ def run_one(call): }) if "PERMANENT FAILURE" in result_str: permanent_failure = True + if _outcome == "interrupted": + for remaining_call in calls[i + 1:]: + self.messages.append({ + "role": "tool", + "tool_call_id": remaining_call["id"], + "name": remaining_call["name"], + "content": "Process interrupted by user.", + }) + yield {"type": "interrupted"} + return if permanent_failure: # Tell the model once, then stop — do not retry @@ -1640,6 +1710,46 @@ def _auto_save_credentials(self, user_input: str) -> None: except Exception: pass + def _run_tool_with_deadline(self, name: str, args: dict) -> tuple[str, float, str]: + """Run one tool call on a daemon worker, polled for interrupt and deadline. + + Returns ``(result, elapsed, outcome)`` where ``outcome`` is ``"done"``, + ``"timeout"`` or ``"interrupted"``. A tool that never returns is abandoned + (the worker is a daemon thread and is never joined again) instead of + blocking the conversation loop forever. + """ + import time + + timeout_s = None + if name not in _TOOL_DEADLINE_EXEMPT_TOOLS: + timeout_s = _resolve_tool_timeout(self) + + box: dict = {} + started = time.time() + + def _work() -> None: + t0 = time.time() + try: + box["result"] = self._execute_tool(name, args) + except Exception as e: + import traceback + box["result"] = f"Error executing tool: {e}\n{traceback.format_exc()}" + box["elapsed"] = time.time() - t0 + + worker = threading.Thread(target=_work, name=f"koza-tool-{name}", daemon=True) + worker.start() + + while True: + worker.join(_SEQUENTIAL_INTERRUPT_POLL_SECONDS) + if not worker.is_alive(): + return str(box.get("result", "")), float(box.get("elapsed", 0.0)), "done" + + elapsed = time.time() - started + if self._cancel.is_set(): + return _tool_interrupt_message(name, elapsed), elapsed, "interrupted" + if timeout_s is not None and elapsed >= timeout_s: + return _tool_timeout_message(name, timeout_s), timeout_s, "timeout" + def _execute_tool(self, name: str, args: dict) -> str: from tools.middleware import MiddlewareChain from tools.registry import ALL_HANDLERS diff --git a/tests/test_tool_deadline.py b/tests/test_tool_deadline.py new file mode 100644 index 0000000..26b54b9 --- /dev/null +++ b/tests/test_tool_deadline.py @@ -0,0 +1,252 @@ +"""A tool call must run under a deadline and stay interruptible. + +Koza used to call the tool handler inline in the conversation loop: a tool that +never returned (a wedged subprocess, a socket read with no timeout) blocked the +whole loop forever, and ``interrupt()`` — which only sets ``self._cancel`` — +could not break it, because nothing was polled while the tool ran. Every call now +runs on a daemon worker that the loop polls, so an overrunning or interrupted +call is abandoned with a visible result instead of hanging the turn. +""" +import os +import sys +import threading +import time + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from core import ( + _DEFAULT_TOOL_TIMEOUT_SECONDS, + Agent, + ToolLoopGuardrail, + _resolve_tool_timeout, +) + +# ── unit: _run_tool_with_deadline ──────────────────────────────────────────── + + +def _agent(cfg=None, tool_fn=None): + agent = Agent.__new__(Agent) + agent._cancel = threading.Event() + agent.cfg = cfg or {} + agent._execute_tool = tool_fn or (lambda name, args: f"ran {name}") + return agent + + +def test_normal_call_returns_done_with_its_result(): + agent = _agent(tool_fn=lambda name, args: f"ran {name}") + + result, _elapsed, outcome = agent._run_tool_with_deadline("read_file", {"path": "a"}) + + assert outcome == "done" + assert result == "ran read_file" + + +def test_wedged_tool_is_abandoned_at_the_deadline(): + never = threading.Event() + agent = _agent(cfg={"tool_timeout_seconds": 0.3}, tool_fn=lambda name, args: never.wait()) + + t0 = time.monotonic() + result, elapsed, outcome = agent._run_tool_with_deadline("run_command", {"command": "sleep"}) + wall = time.monotonic() - t0 + + assert outcome == "timeout", "an overrunning call must be abandoned, not awaited" + assert wall < 5, "the call must not hang the caller past its deadline" + assert "timed out after 0.3s" in result + assert abs(elapsed - 0.3) < 0.1 + + +def test_interrupt_abandons_a_wedged_tool(): + never = threading.Event() + agent = _agent(cfg={"tool_timeout_seconds": 30}, tool_fn=lambda name, args: never.wait()) + + def _interrupt_soon(): + time.sleep(0.2) + agent._cancel.set() + + threading.Thread(target=_interrupt_soon, daemon=True).start() + + result, _elapsed, outcome = agent._run_tool_with_deadline("run_command", {"command": "sleep"}) + + assert outcome == "interrupted", "an interrupt must break a tool that never returns" + assert "interrupted by the user" in result + + +def test_exempt_tools_have_no_deadline(): + agent = _agent( + cfg={"tool_timeout_seconds": 0.2}, + tool_fn=lambda name, args: (time.sleep(0.5), "batch-done")[1], + ) + + result, _elapsed, outcome = agent._run_tool_with_deadline("delegate_task", {}) + + assert outcome == "done" + assert result == "batch-done" + + +def test_zero_config_disables_the_deadline_but_keeps_interruptibility(): + agent = _agent( + cfg={"tool_timeout_seconds": 0}, + tool_fn=lambda name, args: (time.sleep(0.35), "slow-ok")[1], + ) + + result, _elapsed, outcome = agent._run_tool_with_deadline("run_command", {}) + + assert outcome == "done" + assert result == "slow-ok" + + +def test_resolve_tool_timeout_reads_config_with_a_safe_fallback(): + assert _resolve_tool_timeout(_agent(cfg={})) == _DEFAULT_TOOL_TIMEOUT_SECONDS + assert _resolve_tool_timeout(_agent(cfg={"tool_timeout_seconds": 42})) == 42.0 + assert _resolve_tool_timeout(_agent(cfg={"tool_timeout_seconds": 0})) is None + assert _resolve_tool_timeout(_agent(cfg={"tool_timeout_seconds": -5})) is None + assert _resolve_tool_timeout(_agent(cfg={"tool_timeout_seconds": "bad"})) == _DEFAULT_TOOL_TIMEOUT_SECONDS + + +# ── end to end: the conversation loop keeps moving ─────────────────────────── + + +class _FakeProvider: + name = "fake" + supports_vision = False + supports_thinking = False + _model = "fake-model" + + def __init__(self, rounds): + self._rounds = list(rounds) + self.rounds_started = 0 + + def stream_chat(self, messages, tools=None, cancel_event=None): + self.rounds_started += 1 + if not self._rounds: + return + yield from self._rounds.pop(0) + + +def _tool_chunk(index, name, args_chunk, call_id=None): + return { + "__tool_chunk__": True, + "index": index, + "id": call_id or f"{name}-{index}", + "name": name, + "args_chunk": args_chunk, + } + + +def _finish(reason="tool_calls"): + return {"__finish_reason__": reason} + + +def _make_agent(rounds, tool_fn=None, cfg=None): + agent = Agent.__new__(Agent) + agent.provider = _FakeProvider(rounds) + agent.messages = [{"role": "system", "content": "test"}] + agent._cancel = threading.Event() + agent._busy = False + agent.cfg = cfg or {} + agent._router = None + agent._guardrail = ToolLoopGuardrail() + agent.permission_callback = None + agent.tool_progress_callback = None + agent._pre_fetch_links = lambda user_input: iter(()) + agent._refresh_memory_context = lambda *a, **k: None + agent._resolve_available_tools = lambda *a, **k: [] + agent._execute_tool = tool_fn or (lambda name, args: f"ran {name}") + return agent + + +def _events(agent, text="do something"): + return list(agent._run_conversation_loop(text)) + + +def _tool_results(agent): + return [m for m in agent.messages if m.get("role") == "tool"] + + +def _last_assistant(agent): + return [m for m in agent.messages if m.get("role") == "assistant"][-1] + + +def test_loop_does_not_hang_when_a_sequential_tool_wedges(): + never = threading.Event() + + def tool(name, args): + never.wait() # never returns + return "never" + + agent = _make_agent( + [ + [_tool_chunk(0, "run_command", '{"command": "sleep"}'), _finish("tool_calls")], + ["recovered", _finish("stop")], + ], + tool_fn=tool, + cfg={"tool_timeout_seconds": 0.3}, + ) + + t0 = time.monotonic() + _events(agent) + wall = time.monotonic() - t0 + + assert wall < 5, "the turn must not hang on a wedged tool" + tool_msgs = _tool_results(agent) + assert len(tool_msgs) == 1 + assert "timed out after" in tool_msgs[0]["content"] + assert agent.provider.rounds_started == 2, "the loop must continue after a timeout" + assert _last_assistant(agent)["content"] == "recovered" + + +def test_loop_stops_when_an_interrupt_abandons_a_tool(): + ref = {} + + def tool(name, args): + ref["agent"]._cancel.set() + threading.Event().wait() # never returns + return "never" + + agent = _make_agent( + [[_tool_chunk(0, "run_command", '{"command": "sleep"}'), _finish("tool_calls")]], + tool_fn=tool, + cfg={"tool_timeout_seconds": 30}, + ) + ref["agent"] = agent + + events = _events(agent) + + assert {"type": "interrupted"} in events + assert agent.provider.rounds_started == 1, "an interrupted turn must not start a new round" + tool_msgs = _tool_results(agent) + assert len(tool_msgs) == 1 + assert "interrupted by the user" in tool_msgs[0]["content"] + + +def test_parallel_batch_abandons_a_wedged_call_at_the_deadline(): + never = threading.Event() + + def tool(name, args): + if name == "web_search": + never.wait() # only the search wedges + return f"ok:{name}" + + agent = _make_agent( + [ + [ + _tool_chunk(0, "read_file", '{"path": "a.py"}'), + _tool_chunk(1, "web_search", '{"query": "q"}'), + _finish("tool_calls"), + ], + ["done", _finish("stop")], + ], + tool_fn=tool, + cfg={"tool_timeout_seconds": 0.4}, + ) + + t0 = time.monotonic() + _events(agent) + wall = time.monotonic() - t0 + + assert wall < 5, "one wedged parallel call must not freeze the batch" + by_name = {m["name"]: m["content"] for m in _tool_results(agent)} + assert by_name["read_file"] == "ok:read_file" + assert "timed out after" in by_name["web_search"] + assert agent.provider.rounds_started == 2 + assert _last_assistant(agent)["content"] == "done"