diff --git a/tests/test_output_guard_judge.py b/tests/test_output_guard_judge.py index 43a4716d..4e05a491 100644 --- a/tests/test_output_guard_judge.py +++ b/tests/test_output_guard_judge.py @@ -23,6 +23,10 @@ def _make_provider( """Build a mock LLMProvider whose create_completion returns the given content.""" provider = MagicMock() provider.provider_name = "openai" + # The judge reads context_window at construction for its oversize guard. + caps = MagicMock() + caps.context_window = 200_000 + provider.get_capabilities = MagicMock(return_value=caps) def _create_completion(**_kwargs: Any) -> Any: if delay: @@ -243,6 +247,58 @@ class TestEvaluateFailurePaths: assert stragglers == [], f"non-daemon worker survived evaluate(): {stragglers}" +class TestOversizeGuard: + """A tool output that would overflow the judge model's context window must + not silently fall to heuristic-only via an opaque provider 400 — it is + detected up front and surfaced as a labelled llm_error the operator sees.""" + + def test_oversize_output_skips_llm_and_returns_labeled_error(self) -> None: + # ``content`` would parse to a clean verdict IF the provider were + # called — so a labelled oversize error proves the call was skipped. + judge = _make_judge(content='{"risk_level": "low", "flags": [], "reasoning": "x"}') + judge._judge_context_window = 50 # tiny window forces the guard to trip + v = judge.evaluate("Z" * 2000, func_name="web_fetch", call_id="c1") + assert not v.succeeded + assert "output_too_large_for_judge_window" in v.error + assert v.judge_model # model recorded so the audit row is attributable + + def test_output_within_window_is_judged_normally(self) -> None: + judge = _make_judge(content='{"risk_level": "low", "flags": [], "reasoning": "x"}') + v = judge.evaluate("a small, safe output", func_name="bash", call_id="c1") + assert v.succeeded + assert "too_large" not in v.error + + def test_guard_threshold_scales_with_resolved_window(self) -> None: + """The same output that overflows a tiny window passes a large one — + the guard is keyed to the judge model, not a fixed cap.""" + payload = "Z" * 4000 # assembled prompt overflows a 200-tok window, fits 200k + small = _make_judge(content='{"risk_level": "low", "flags": [], "reasoning": "x"}') + small._judge_context_window = 200 + big = _make_judge(content='{"risk_level": "low", "flags": [], "reasoning": "x"}') + big._judge_context_window = 200_000 + assert not small.evaluate(payload, call_id="c1").succeeded + assert big.evaluate(payload, call_id="c1").succeeded + + def test_missing_provider_capabilities_falls_back_to_default_window(self) -> None: + """If the provider can't report a context window, construction uses the + conservative default rather than crashing or an unbounded window.""" + from turnstone.core.output_guard_judge import _DEFAULT_JUDGE_CONTEXT_WINDOW + + provider = _make_provider(content='{"risk_level": "none", "flags": []}') + provider.get_capabilities = MagicMock(side_effect=RuntimeError("no caps")) + config = JudgeConfig(output_guard_llm=True) + client = MagicMock() + client.base_url = "http://test" + client.api_key = "k" + judge = OutputGuardJudge( + config=config, + session_provider=provider, + session_client=client, + session_model="test-model", + ) + assert judge._judge_context_window == _DEFAULT_JUDGE_CONTEXT_WINDOW + + class TestAliasResolution: def test_unknown_alias_falls_back_to_session_model(self) -> None: # Registry says alias does not exist; judge should fall back. @@ -390,14 +446,24 @@ class TestFenceEscape: assert "Heuristic stage flagged:" not in prompt assert "Heuristic annotations:" not in prompt - def test_user_prompt_truncates_long_tool_args(self) -> None: + def test_user_prompt_does_not_default_truncate_tool_args(self) -> None: + """tool_args lowers whole — no default cap. A pathologically large call + is caught by evaluate()'s window backstop, not by clipping a normal + argument into a misleading prefix.""" long_args = '{"query": "' + ("x" * 1000) + '"}' prompt = OutputGuardJudge._user_prompt( "the output", func_name="search", tool_args=long_args ) - assert "...(truncated)" in prompt - # Original full 1000+ chars must not appear. - assert long_args not in prompt + assert long_args in prompt + assert "chars omitted" not in prompt + + def test_user_prompt_never_truncates_the_output_under_review(self) -> None: + """The fenced output is the content being judged and must reach the + judge whole.""" + big_output = "Z" * 20_000 + prompt = OutputGuardJudge._user_prompt(big_output, func_name="web_fetch") + assert big_output in prompt + assert "chars omitted" not in prompt def test_user_prompt_skips_heuristic_section_when_clean(self) -> None: # risk='none' and empty flags → no "Heuristic stage flagged" line. diff --git a/turnstone/core/output_guard_judge.py b/turnstone/core/output_guard_judge.py index 62bb89ce..461c903e 100644 --- a/turnstone/core/output_guard_judge.py +++ b/turnstone/core/output_guard_judge.py @@ -58,6 +58,19 @@ if TYPE_CHECKING: log = get_logger(__name__) +# Prompt-size guard. A tool output large enough to overflow the judge model's +# context window would come back as an opaque provider 400 and fall silently to +# heuristic-only; we detect it up front instead (see ``evaluate``). +# ``_CHARS_PER_TOKEN`` mirrors ``judge._CHARS_PER_TOKEN`` — the same token +# estimate both judges budget with, duplicated to keep this module free of a +# runtime judge import; keep the two in sync if either is retuned. ``0.9`` +# leaves headroom for the 512-token response plus estimation error. The +# default window is a conservative fallback used only when the provider can't +# report capabilities — the resolved model's real window is preferred. +_CHARS_PER_TOKEN = 3.5 +_MAX_PROMPT_RATIO = 0.9 +_DEFAULT_JUDGE_CONTEXT_WINDOW = 32_768 + # --------------------------------------------------------------------------- # Verdict @@ -259,17 +272,30 @@ class OutputGuardJudge: # Alias resolution mirrors IntentJudge.__init__ at judge.py:917-960. # An empty / unset alias falls through to the session model silently; # a set-but-unknown alias logs a warning and also falls through. + # Judge model's context window drives the oversize-output guard in + # ``evaluate``. On the alias path it comes from the registry's + # ModelConfig — NOT ``provider.get_capabilities()``, which returns a + # static 200000 for every model absent from its table (i.e. every local + # / self-hosted judge), so the guard keyed off it would never trip for + # the small-window local judges it exists to protect. resolved = False if config.output_guard_model and model_registry is not None: try: if model_registry.has_alias(config.output_guard_model): - client, model_name, _ = model_registry.resolve(config.output_guard_model) + client, model_name, model_cfg = model_registry.resolve( + config.output_guard_model + ) self._provider = model_registry.get_provider(config.output_guard_model) self._client_factory_args = self._extract_client_config( client, self._provider.provider_name ) self._model = model_name self._judge_model_alias = config.output_guard_model + # getattr guard: a malformed ModelConfig must not abort + # resolution — degrade the window, keep the model. + self._judge_context_window = getattr( + model_cfg, "context_window", _DEFAULT_JUDGE_CONTEXT_WINDOW + ) resolved = True except Exception: log.debug( @@ -292,6 +318,15 @@ class OutputGuardJudge: ) self._model = session_model self._judge_model_alias = "" + # Session-model fallback: no ModelConfig in hand, so use the static + # capability table (correct for commercial models; a conservative + # default for anything it doesn't know). + try: + self._judge_context_window = self._provider.get_capabilities( + self._model + ).context_window + except Exception: + self._judge_context_window = _DEFAULT_JUDGE_CONTEXT_WINDOW # Lazy-init in _create_client(); reused across evaluate() calls. # Session swaps the entire OutputGuardJudge on credential / model @@ -405,6 +440,33 @@ class OutputGuardJudge: }, ] + # Oversize guard. The heuristic stage has already run and its verdict + # stands regardless; what's at stake here is only the opted-in LLM tier. + # A prompt that overflows the judge window returns an opaque provider + # 400, which would fall to heuristic-only with no trace that the LLM was + # even attempted. Detect it up front: skip the doomed call, log a + # warning, and return a LABELLED error verdict so the skip surfaces as a + # distinct ``llm_error`` audit row (reason = "output_too_large…") the + # operator can see, rather than a silent no-op. + prompt_chars = sum(len(str(m["content"])) for m in judge_messages) + est_tokens = int(prompt_chars / _CHARS_PER_TOKEN) + if est_tokens > self._judge_context_window * _MAX_PROMPT_RATIO: + log.warning( + "output_guard_judge.output_too_large", + call_id=call_id, + func_name=func_name, + output_chars=len(output), + est_prompt_tokens=est_tokens, + judge_context_window=self._judge_context_window, + ) + return self._error_verdict( + verdict_id, + call_id, + start, + f"output_too_large_for_judge_window: ~{est_tokens} tok " + f"> {self._judge_context_window} window", + ) + try: client = self._create_client() except Exception as e: @@ -513,9 +575,11 @@ class OutputGuardJudge: verdict + heuristic annotations) precede the fence. The system prompt classifies each field's trust level: framework-supplied fields are TRUSTED; ``tool_args`` is UNTRUSTED (caller-supplied, - may contain injection); fenced output is UNTRUSTED. Tool args - are truncated to 500 chars to bound prompt cost while preserving - shape. + may contain injection); fenced output is UNTRUSTED. Neither + ``tool_args`` nor the fenced output is truncated here — both lower + whole. A pathologically large call is caught by the window backstop in + ``evaluate`` (which skips the LLM tier honestly rather than feeding it a + silently-clipped prefix), never by a default cap on a normal argument. """ nonce = fence.mint_nonce() @@ -525,8 +589,7 @@ class OutputGuardJudge: if tool_description: lines.append(f"Description: {tool_description}") if tool_args: - truncated = tool_args if len(tool_args) <= 500 else tool_args[:500] + "...(truncated)" - lines.append(f"Called with: {truncated}") + lines.append(f"Called with: {tool_args}") if heuristic_risk != "none" or heuristic_flags: flags_str = ", ".join(heuristic_flags) if heuristic_flags else "(none)" lines.append(