From bbe28ecab32fbaad033e57ab80bb6f9da164358c Mon Sep 17 00:00:00 2001 From: Patrick Buckley Date: Thu, 26 Mar 2026 18:10:23 -0700 Subject: [PATCH] fix: cancel LLM judge daemon when user approves/denies tools (#187) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When the LLM judge is enabled and the user makes rapid approval decisions, judge daemon threads pile up competing for the inference server, causing timeouts. Pass a threading.Event from session to the judge — set on approval — so the daemon abandons remaining work within ~1s, including shutting down the executor to kill in-flight API calls. --- tests/test_judge.py | 45 ++++++++++++------- turnstone/core/judge.py | 91 ++++++++++++++++++++++++++------------- turnstone/core/session.py | 33 ++++++++++---- 3 files changed, 115 insertions(+), 54 deletions(-) diff --git a/tests/test_judge.py b/tests/test_judge.py index ad6ed36c..b2677f6d 100644 --- a/tests/test_judge.py +++ b/tests/test_judge.py @@ -4,6 +4,7 @@ from __future__ import annotations import json import time +from concurrent.futures import ThreadPoolExecutor from pathlib import Path from typing import Any from unittest.mock import MagicMock @@ -179,10 +180,13 @@ class TestErrorHandling: provider = _make_mock_provider(side_effect=RuntimeError("API error")) judge = _make_judge(provider) - result = judge._evaluate_single( - _make_item(), - [{"role": "user", "content": "test"}], - ) + with ThreadPoolExecutor(max_workers=1) as pool: + result = judge._evaluate_single( + _make_item(), + [{"role": "user", "content": "test"}], + cancel_event=None, + executor=pool, + ) assert result is None def test_provider_error_heuristic_still_returned(self): @@ -211,10 +215,13 @@ class TestErrorHandling: result_mock.content = "" judge = _make_judge(provider) - result = judge._evaluate_single( - _make_item(), - [{"role": "user", "content": "test"}], - ) + with ThreadPoolExecutor(max_workers=1) as pool: + result = judge._evaluate_single( + _make_item(), + [{"role": "user", "content": "test"}], + cancel_event=None, + executor=pool, + ) assert result is None @@ -254,10 +261,13 @@ class TestMultiTurnToolUse: provider.create_completion.side_effect = [turn1, turn2] judge = _make_judge(provider) - verdict = judge._evaluate_single( - _make_item(), - [{"role": "user", "content": "test"}], - ) + with ThreadPoolExecutor(max_workers=1) as pool: + verdict = judge._evaluate_single( + _make_item(), + [{"role": "user", "content": "test"}], + cancel_event=None, + executor=pool, + ) assert verdict is not None assert verdict.tier == "llm" assert provider.create_completion.call_count == 2 @@ -299,10 +309,13 @@ class TestMultiTurnToolUse: ] judge = _make_judge(provider) - judge._evaluate_single( - _make_item(), - [{"role": "user", "content": "test"}], - ) + with ThreadPoolExecutor(max_workers=1) as pool: + judge._evaluate_single( + _make_item(), + [{"role": "user", "content": "test"}], + cancel_event=None, + executor=pool, + ) # Should have called create_completion exactly _JUDGE_MAX_TURNS times assert provider.create_completion.call_count == 5 diff --git a/turnstone/core/judge.py b/turnstone/core/judge.py index 342036f6..bf9a83d2 100644 --- a/turnstone/core/judge.py +++ b/turnstone/core/judge.py @@ -853,6 +853,10 @@ If you used read_file to check a target, cite what you found.""" # --------------------------------------------------------------------------- +class _ExecutorPoisonedError(Exception): + """Raised when a timeout leaves the executor's worker thread stuck.""" + + class IntentJudge: """Session-scoped LLM judge for intent validation. @@ -912,18 +916,12 @@ class IntentJudge: self._model = session_model self._judge_context_window = context_window - # Executor for timeout-guarded API calls (1 thread — judge is serial) - self._executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="judge-api") - - def shutdown(self) -> None: - """Release the executor thread pool.""" - self._executor.shutdown(wait=False, cancel_futures=True) - def evaluate( self, items: list[dict[str, Any]], messages: list[dict[str, Any]], callback: Callable[[IntentVerdict], None], + cancel_event: threading.Event | None = None, ) -> list[IntentVerdict]: """Evaluate tool calls. Returns heuristic verdicts immediately. @@ -936,6 +934,10 @@ class IntentJudge: ``func_args``, ``approval_label``, ``call_id``). messages: Conversation history (OpenAI message format). callback: Called with each LLM verdict (or timeout/error fallback). + cancel_event: When set, the daemon judge thread abandons + remaining work. Callers should set this after the user + has already made an approval decision so the judge does + not keep consuming inference resources. Returns: List of heuristic verdicts (one per item), available immediately. @@ -958,7 +960,7 @@ class IntentJudge: # Spawn daemon thread for LLM judge thread = threading.Thread( target=self._run_judge, - args=(items, messages, heuristic_verdicts, callback), + args=(items, messages, heuristic_verdicts, callback, cancel_event), daemon=True, name="intent-judge", ) @@ -972,25 +974,44 @@ class IntentJudge: messages: list[dict[str, Any]], heuristic_verdicts: list[IntentVerdict], callback: Callable[[IntentVerdict], None], + cancel_event: threading.Event | None = None, ) -> None: """Daemon thread: run LLM judge for each item and invoke callback.""" - for item, h_verdict in zip(items, heuristic_verdicts, strict=True): - try: - llm_verdict = self._evaluate_single(item, messages) - # Arbitrate: only callback when LLM upgrades the heuristic - if llm_verdict and llm_verdict.confidence > h_verdict.confidence: - callback(llm_verdict) - # else: heuristic already delivered, no duplicate callback - except Exception: - log.exception( - "Judge evaluation failed for %s", - item.get("func_name", "?"), - ) + # Evaluation-scoped executor — avoids sharing mutable state with + # other daemon threads from concurrent evaluate() calls. + executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="judge-api") + try: + for idx, (item, h_verdict) in enumerate(zip(items, heuristic_verdicts, strict=True)): + if cancel_event and cancel_event.is_set(): + log.debug("judge.cancelled", remaining=len(items) - idx) + return + try: + llm_verdict = self._evaluate_single(item, messages, cancel_event, executor) + if cancel_event and cancel_event.is_set(): + return + # Arbitrate: only callback when LLM upgrades the heuristic + if llm_verdict and llm_verdict.confidence > h_verdict.confidence: + callback(llm_verdict) + # else: heuristic already delivered, no duplicate callback + except _ExecutorPoisonedError: + # Timeout left the worker stuck — replace the executor + # so subsequent items don't queue behind it. + executor.shutdown(wait=False, cancel_futures=True) + executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="judge-api") + except Exception: + log.exception( + "Judge evaluation failed for %s", + item.get("func_name", "?"), + ) + finally: + executor.shutdown(wait=False, cancel_futures=True) def _evaluate_single( self, item: dict[str, Any], messages: list[dict[str, Any]], + cancel_event: threading.Event | None, + executor: ThreadPoolExecutor, ) -> IntentVerdict | None: """Run LLM judge for a single tool call. Returns verdict or None.""" start = time.monotonic() @@ -1020,6 +1041,9 @@ class IntentJudge: result = None # will hold the last CompletionResult for turn in range(_JUDGE_MAX_TURNS): + if cancel_event and cancel_event.is_set(): + return None + turn_start = time.monotonic() is_last_turn = turn == _JUDGE_MAX_TURNS - 1 @@ -1043,7 +1067,7 @@ class IntentJudge: # 10 minutes — far too long for an advisory judge on local models. per_call_timeout = max(timeout_budget, 5.0) # at least 5s try: - future = self._executor.submit( + future = executor.submit( self._provider.create_completion, client=self._client, model=self._model, @@ -1053,17 +1077,24 @@ class IntentJudge: temperature=0.0, reasoning_effort="medium", ) - result = future.result(timeout=per_call_timeout) + # Poll in 1s increments so we notice cancellation promptly + # instead of blocking for the full per_call_timeout. + deadline = time.monotonic() + per_call_timeout + while True: + remaining = deadline - time.monotonic() + if cancel_event and cancel_event.is_set(): + future.cancel() + return None + if remaining <= 0: + raise TimeoutError + try: + result = future.result(timeout=min(remaining, 1.0)) + break + except TimeoutError: + pass # loop back to check remaining/cancel except TimeoutError: log.warning("Judge LLM call timed out on turn %d (%.0fs)", turn, per_call_timeout) - # Abandon the lingering API call and replace the executor so - # subsequent items in the batch don't queue behind it. - self._executor.shutdown(wait=False, cancel_futures=True) - self._executor = ThreadPoolExecutor( - max_workers=1, - thread_name_prefix="judge-api", - ) - return None + raise _ExecutorPoisonedError from None except Exception: log.exception("Judge LLM call failed on turn %d", turn) return None diff --git a/turnstone/core/session.py b/turnstone/core/session.py index 284132c1..7fa53108 100644 --- a/turnstone/core/session.py +++ b/turnstone/core/session.py @@ -319,6 +319,7 @@ class ChatSession: # Intent validation judge (lazy-initialized) self._judge_config: JudgeConfig | None = judge_config self._judge: IntentJudge | None = None + self._judge_cancel_event: threading.Event | None = None # MCP tool integration: merge external tools with built-in self._mcp_client = mcp_client self._mcp_refresh_cb: Any = None # Callable | None (avoid import) @@ -639,8 +640,8 @@ class ChatSession: def close(self) -> None: """Release resources (listener registrations, etc.).""" - if self._judge is not None: - self._judge.shutdown() + if self._judge_cancel_event is not None: + self._judge_cancel_event.set() if self._mcp_client and self._mcp_refresh_cb: self._mcp_client.remove_listener(self._mcp_refresh_cb) self._mcp_refresh_cb = None @@ -2074,20 +2075,24 @@ class ChatSession: def _evaluate_intent( self, items: list[dict[str, Any]], - ) -> None: + ) -> threading.Event | None: """Run intent validation on pending approval items. Attaches heuristic verdicts to items immediately. Spawns the async LLM judge that delivers final verdicts via UI callback. + + Returns a cancel event that, when set, tells the daemon judge + thread to abandon remaining work. Callers should set this + after the user has made an approval decision. """ judge = self._ensure_judge() if not judge: - return + return None # Only evaluate items that need approval and aren't errors pending = [it for it in items if it.get("needs_approval") and not it.get("error")] if not pending: - return + return None # Build func_args from tool-specific item keys so the heuristic # engine can pattern-match on argument content. @@ -2125,16 +2130,20 @@ class ChatSession: except Exception: log.debug("judge.verdict_delivery_failed", exc_info=True) + cancel_event = threading.Event() heuristic_verdicts = judge.evaluate( pending, list(self.messages), # snapshot — daemon thread must not see mutations callback=_on_verdict, + cancel_event=cancel_event, ) # Attach heuristic verdicts to items for the approval UI for item, verdict in zip(pending, heuristic_verdicts, strict=True): item["_heuristic_verdict"] = verdict.to_dict() + return cancel_event + def _evaluate_output(self, call_id: str, output: str, func_name: str) -> str: """Run the output guard on tool result text. @@ -2184,12 +2193,20 @@ class ChatSession: # Phase 1: prepare all tool calls items = [self._prepare_tool(tc) for tc in tool_calls] - # Intent validation (advisory, non-blocking) - self._evaluate_intent(items) + # Intent validation (advisory, non-blocking). + # Cancel any prior judge thread before spawning a new one. + if self._judge_cancel_event is not None: + self._judge_cancel_event.set() + judge_cancel = self._evaluate_intent(items) + self._judge_cancel_event = judge_cancel # track for close() # Phase 2: approve via UI self._emit_state("attention") - approved, user_feedback = self.ui.approve_tools(items) + try: + approved, user_feedback = self.ui.approve_tools(items) + finally: + if judge_cancel: + judge_cancel.set() # user decided (or disconnected) — stop judge self._emit_state("running") if not approved: # Mark all pending items as denied