ewilliams28/turnstone
1
0
Fork 0
mirror of https://github.com/turnstonelabs/turnstone.git synced 2026-08-12 23:12:23 -06:00
Code Issues Packages Projects Releases Wiki Activity

fix: cancel LLM judge daemon when user approves/denies tools (#187)

Browse Source 
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.
This commit is contained in:
Patrick Buckley
2026-03-26 18:10:23 -07:00
committed by GitHub
parent 93a9fd3c28
commit bbe28ecab3
3 changed files with 115 additions and 54 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/test_judge.py
+29 -16
View File
@@ -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
turnstone/core/judge.py
+61 -30
View File
@@ -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
turnstone/core/session.py
+25 -8
View File
@@ -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