Skip to content

fix(judge): protect unbounded judge transcript for non-litellm agents#837

Merged
rogeriochaves merged 7 commits into
mainfrom
fix/836-judge-transcript-size-protection
Jul 22, 2026
Merged

fix(judge): protect unbounded judge transcript for non-litellm agents#837
rogeriochaves merged 7 commits into
mainfrom
fix/836-judge-transcript-size-protection

Conversation

@Aryansharma28

Copy link
Copy Markdown
Contributor

Fixes #836

Problem

JudgeAgent's trace-size protection (structure-only digest + expand_trace/grep_trace discovery loop) is gated entirely on is_large_trace, which is computed only from spans collected by JudgeSpanCollector — i.e. spans produced by autotrack_litellm_calls. Any AgentAdapter that talks to its own backend directly (REST/SSE/gRPC/etc.) and never routes calls through litellm produces tool-call messages that never become spans.

For such agents, spans stays empty/small forever, so is_large_trace never trips — while the actual transcript block (built from input.messages via JudgeUtils.build_transcript_from_messages) has no size cap of any kind beyond stripping base64 media. Large tool-call content (SQL result sets, long documents, etc.) flows into the judge prompt completely unbounded.

The reporter hit this with a proprietary agent whose 27 tool calls (with large SQL result payloads) never became spans; the judge (gpt-4o-mini) silently dropped/ignored large portions of the unbounded transcript ("lost in the middle") and returned a confident but wrong faithfulness: 0.0 verdict claiming "no tool calls were made" — a false negative with no error signal at all.

Fix

  • is_large_trace is now is_large_span_trace OR is_large_transcript, where is_large_transcript is computed from the actual transcript's own estimated token count — independent of spans.
  • When the transcript itself is large, it's rendered as a structure-only skeleton (one line per message: index, role, tool-call name(s), estimated size — no content), the same pattern already used for large span digests.
  • Two new discovery tools, expand_transcript(indices) and grep_transcript(pattern), mirror expand_trace/grep_trace but operate on message position instead of span ID, so the judge can still reach full content on demand instead of it being force-fed (or silently dropped) on every call.
  • JudgeUtils.build_transcript_from_messages is refactored to expose render_transcript_lines, a per-message renderer shared by the full transcript builder and the new discovery tools, so expand_transcript shows byte-identical content to what the full transcript would have shown.
  • Small transcripts (the common case) are completely unaffected — same full inline rendering as before.

Testing

  • New tests/test_judge_transcript_size_protection.py reproduces the exact bug: empty spans (simulating a non-litellm agent) + a large transcript (27 tool calls with large payloads, mirroring the reporter's repro). Verified this test fails against unmodified main (4/5 cases) and passes after the fix.
  • Full existing judge test suite (test_judge_agent*.py, test_judge_discovery_*.py, test_judge_utilities.py, test_judge_force_verdict_hardening.py) passes unchanged — the span-based path is untouched.
  • Full pytest tests/ run: 1139 passed (the only failures are pre-existing e2e tests requiring live API keys, unrelated to this change).
  • pyright clean on all touched files.

Test plan

  • Reproduced the bug with a failing test against main
  • Fix makes the new test pass
  • Existing judge/discovery test suite passes unchanged
  • Full python test suite passes (excluding pre-existing env-gated e2e failures)
  • pyright clean

🤖 Generated with Claude Code

…gents

JudgeAgent's trace-size protection (structure-only digest + expand_trace/
grep_trace discovery) was gated entirely on JudgeSpanCollector spans, which
only ever contain spans from autotrack_litellm_calls. An AgentAdapter that
talks to its own backend directly (REST/SSE/etc., not via litellm) produces
tool-call messages that never become spans, so is_large_trace never tripped
for such agents while the raw `transcript` block built from input.messages
flowed into the judge prompt with no cap of any kind. A large enough
transcript could get "lost in the middle" for a small/cheap judge model,
producing a confident but wrong verdict instead of a visible failure.

Gate the transcript on its own estimated token count, independent of the
span digest. When it's large, render a structure-only skeleton (one line
per message: index, role, tool-call name, size) and offer the judge
expand_transcript/grep_transcript tools -- mirroring expand_trace/grep_trace
but keyed on message position instead of span ID -- so the full content is
still reachable on demand instead of force-fed on every call.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
@coderabbitai

coderabbitai Bot commented Jul 22, 2026

Copy link
Copy Markdown

Review Change Stack

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

  • @coderabbitai resume to resume automatic reviews.
  • @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

  • ▶️ Resume reviews
  • 🔍 Trigger review

Walkthrough

Large transcripts now use bounded skeleton prompts with expand_transcript and grep_transcript discovery tools. Shared rendering supports indexed message access, while JudgeAgent integrates transcript discovery alongside trace discovery and adds regression and end-to-end coverage.

Changes

Transcript discovery

Layer / File(s) Summary
Transcript rendering and discovery utilities
python/scenario/_judge/judge_utils.py, python/scenario/_judge/transcript_tools.py, python/scenario/_judge/__init__.py
Shared message rendering supports transcript skeletons, selective expansion, case-insensitive grep, bounded results, and public utility exports.
JudgeAgent transcript protection and tool execution
python/scenario/judge_agent.py
Large transcripts use skeleton prompts and transcript discovery schemas; discovery execution can expand or grep the working message set.
Transcript protection regression and end-to-end coverage
python/tests/test_judge_transcript_size_protection.py, python/examples/test_judge_transcript_size_protection_e2e.py
Tests cover large non-litellm transcripts, bounded prompts, discovery behavior, malformed inputs, error paths, dispatch handling, and successful needle recovery.

Sequence Diagram(s)

sequenceDiagram
  participant AgentInput
  participant JudgeAgent
  participant transcript_tools
  participant JudgeModel
  AgentInput->>JudgeAgent: working_messages
  JudgeAgent->>transcript_tools: build_transcript_skeleton(working_messages)
  transcript_tools-->>JudgeAgent: skeleton transcript
  JudgeAgent->>JudgeModel: prompt with skeleton and discovery tools
  JudgeModel->>JudgeAgent: expand_transcript or grep_transcript call
  JudgeAgent->>transcript_tools: execute discovery tool
  transcript_tools-->>JudgeAgent: bounded transcript content
  JudgeAgent->>JudgeModel: discovery result
Loading

Poem

I’m a rabbit with transcripts tucked tight,
Skeletons keep the judge’s prompt light.
I grep through the rows,
Expand what it knows,
And hop through protected content tonight.

🚥 Pre-merge checks | ✅ 5
✅ Passed checks (5 passed)
Check name Status Explanation
Title check ✅ Passed The title clearly matches the main change: adding transcript protection for non-LiteLLM judges.
Description check ✅ Passed The description directly explains the transcript-size protection fix and added discovery tools.
Linked Issues check ✅ Passed The changes implement #836 by basing protection on transcript size and adding expand_transcript/grep_transcript.
Out of Scope Changes check ✅ Passed The refactor and tests stay focused on transcript protection and related discovery behavior.
Docstring Coverage ✅ Passed No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Commit unit tests in branch fix/836-judge-transcript-size-protection

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

Comment @coderabbitai help to get the list of available commands.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 5

🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@python/scenario/_judge/judge_utils.py`:
- Around line 180-242: Replace broad Any, Dict, List, and Sequence annotations
with the concrete message and tool-call types already used by the judge APIs,
using built-in dict[...] and list[...] generics. In
python/scenario/_judge/judge_utils.py:180-242, narrow _render_message_line,
render_transcript_lines, and related helpers; in
python/scenario/_judge/transcript_tools.py:20-60 and 77-217, type
_IndexedMessage.message and update helper annotations, marking budget constants
as Final[int]; in python/scenario/judge_agent.py:852-861 and 1020-1022, narrow
working_messages and tool_call to their concrete types.

In `@python/scenario/_judge/transcript_tools.py`:
- Around line 97-120: The skeleton-building logic around the indexed-message
loop must enforce a total token or character budget, including tool-call
metadata, so it cannot exceed the judge context. Cap displayed tool names and
metadata per message, stop adding entries once the budget is reached, and append
a concise indication of omitted contiguous index ranges so omitted messages
remain discoverable.
- Around line 146-162: Limit requested indices and render selected transcript
entries incrementally in the transcript indexing function, stopping when the
character budget is reached instead of joining all selected content before
truncation; preserve invalid-index reporting. In
python/scenario/_judge/transcript_tools.py lines 146-162, update the
selection/rendering flow accordingly. In python/scenario/judge_agent.py lines
815-826, add a matching maxItems constraint to the indices tool schema.

In `@python/tests/test_judge_transcript_size_protection.py`:
- Around line 31-34: Add explicit, parameterized type annotations throughout the
helper, fixture, and callback definitions covered by this change, including
message payloads, tool arguments, fixture yields, and the mock_completion
interface. Replace bare list/dict annotations with concrete generic collection
types and annotate missing parameters and return values, including
create_mock_collector. Leave pytest test methods unannotated to match repository
convention.
- Around line 102-111: Update the setup_config fixture to capture
ScenarioConfig.default_config before replacing it, then restore that saved value
during teardown instead of unconditionally assigning None. Preserve the existing
context_scenario cleanup and test configuration setup.
🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

  • Push a commit to this branch (recommended)
  • Create a new PR with the fixes

ℹ️ Review info
⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: 4ea0b0eb-211c-4c26-bcd7-fb177cc59beb

📥 Commits

Reviewing files that changed from the base of the PR and between e018d76 and a7494ed.

📒 Files selected for processing (5)
  • python/scenario/_judge/__init__.py
  • python/scenario/_judge/judge_utils.py
  • python/scenario/_judge/transcript_tools.py
  • python/scenario/judge_agent.py
  • python/tests/test_judge_transcript_size_protection.py

Comment thread python/scenario/_judge/judge_utils.py Outdated
Comment thread python/scenario/_judge/transcript_tools.py Outdated
Comment thread python/scenario/_judge/transcript_tools.py
Comment thread python/tests/test_judge_transcript_size_protection.py Outdated
Comment thread python/tests/test_judge_transcript_size_protection.py Outdated
Self-review caught a crash: expand_transcript's indices arg is reachable
from an LLM tool call, and sorted(set(indices) - set(valid_range)) raises
TypeError the moment a non-integer sneaks into that list (e.g. a lax
litellm-routed provider that doesn't enforce the "type": "integer" schema
as strictly as OpenAI, sending "2" instead of 2). Verified the crash
reproduces, then coerced/filtered indices defensively instead of trusting
the schema to hold at runtime -- matching the "never raises" contract this
file already applies to tool-call argument rendering.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🧹 Nitpick comments (1)
python/tests/test_judge_transcript_size_protection.py (1)

310-314: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Assert the concrete ignored indices.

The test only checks the generic marker, so it would pass even if the ignored values were omitted. Assert "Ignored non-integer indices: [None, {}]" to pin the diagnostic contract. Based on the provided expand_transcript contract, ignored values are part of the returned diagnostic.

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@python/tests/test_judge_transcript_size_protection.py` around lines 310 -
314, Update the assertions in the expand_transcript test to verify the complete
ignored-index diagnostic, including the concrete values “[None, {}]”, rather
than only checking the generic marker. Preserve the existing assistant output
assertion and validate the returned message exactly as specified by the
diagnostic contract.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@python/tests/test_judge_transcript_size_protection.py`:
- Around line 310-314: Update the assertions in the expand_transcript test to
verify the complete ignored-index diagnostic, including the concrete values
“[None, {}]”, rather than only checking the generic marker. Preserve the
existing assistant output assertion and validate the returned message exactly as
specified by the diagnostic contract.

ℹ️ Review info
⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: 7a0addb9-7405-491b-9edc-131d0568de57

📥 Commits

Reviewing files that changed from the base of the PR and between a7494ed and 45bd5bd.

📒 Files selected for processing (2)
  • python/scenario/_judge/transcript_tools.py
  • python/tests/test_judge_transcript_size_protection.py
🚧 Files skipped from review as they are similar to previous changes (1)
  • python/scenario/_judge/transcript_tools.py

Addresses CodeRabbit findings on #837:

- build_transcript_skeleton had no cap of its own. A transcript of
  thousands of tiny messages (streaming/high-frequency tool-call agent)
  produces a skeleton whose *total* size can exceed the same
  token_threshold that triggered structure-only rendering in the first
  place -- the exact failure mode this fix exists to close. Verified: 5000
  short messages produced a 32852-token skeleton against an 8192 default
  threshold. Now bounded to a head+tail window (start frames the scenario,
  end is what the judge is deciding on) with the elided middle still
  reachable via expand_transcript/grep_transcript.
- expand_transcript's index matching checked membership against a List
  (O(N*M) for N messages and M requested indices); switched to a set for
  O(1) lookups.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@python/tests/test_judge_transcript_size_protection.py`:
- Around line 327-331: Replace the cast(Any, messages) calls in the transcript
skeleton test cases with the project’s concrete parameterized chat-message type,
and annotate the messages fixtures accordingly. Pass the typed messages directly
to build_transcript_skeleton at all referenced cases so Pyright validates their
structure without Any.
🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

  • Push a commit to this branch (recommended)
  • Create a new PR with the fixes

ℹ️ Review info
⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: 619362a4-9e9e-4e5c-a9c5-e0b480a12db7

📥 Commits

Reviewing files that changed from the base of the PR and between 45bd5bd and 76c56de.

📒 Files selected for processing (2)
  • python/scenario/_judge/transcript_tools.py
  • python/tests/test_judge_transcript_size_protection.py
🚧 Files skipped from review as they are similar to previous changes (1)
  • python/scenario/_judge/transcript_tools.py

Comment thread python/tests/test_judge_transcript_size_protection.py Outdated
Adds direct-unit coverage for the guard clauses that weren't exercised by
the earlier repro tests: empty-transcript short-circuits on all three
functions, expand_transcript's empty/out-of-range index requests, grep_
transcript's no-match and >MAX_GREP_MATCHES branches, and the char-budget
truncation path. Also covers _execute_discovery_tool's JSON-decode-error
and unknown-tool-name branches through the real discovery loop, since
those sit directly next to the expand_transcript/grep_transcript dispatch
arms this fix added.

scenario/_judge/transcript_tools.py (the new module) is now at 100% line
coverage. judge_utils.py's and judge_agent.py's remaining coverage gaps
are pre-existing, in code this PR didn't touch (base64/audio media
truncation edge cases, voice modality resolution, deep verdict-parsing
branches) -- out of scope for #836.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

@Aryansharma28 Aryansharma28 left a comment

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Deep self-review of #837

Since I authored this fix, I ran it adversarially against itself: exercised every guard clause directly, tried to break expand_transcript/grep_transcript with malformed LLM tool-call input, checked the combined-large-span-and-large-transcript path, and triaged every CodeRabbit finding rather than rubber-stamping or dismissing them. Three real issues surfaced and are fixed in follow-up commits (45bd5bd, 76c56de). No open findings remain.

Verdict: no blockers. Ready to merge once CI's unrelated infra issue (see below) is resolved or waived.


[P1] expand_transcript crashed on non-integer indices from a lax LLM provider — FIXED (45bd5bd)

Why: expand_transcript's indices argument is reachable from an LLM tool call. The tool schema declares "items": {"type": "integer"}, but not every litellm-routed provider enforces function-calling argument types as strictly as OpenAI. I reproduced the crash directly:

expand_transcript(messages, indices=["1", 5])
# TypeError: '<' not supported between instances of 'str' and 'int'

sorted(set(indices) - set(valid_range)) raised the moment a non-integer entered the list — an unhandled exception that would have propagated up through _execute_discovery_tool_run_discovery_loopcall(), crashing the entire judge call instead of degrading gracefully.

Fix: coerce/filter indices defensively (int() with try/except, malformed entries reported back to the LLM instead of raising) — matching the "never raises" contract this file already applies to tool-call argument rendering (_render_tool_arguments's docstring). Covered by test_non_integer_indices_do_not_raise / test_unparseable_indices_are_reported_not_raised.

[P2] build_transcript_skeleton had no cap on its own total size (CodeRabbit) — FIXED (76c56de)

Why: CodeRabbit correctly flagged that the skeleton renders one line per message with no total budget. A transcript of thousands of small messages (a streaming/high-frequency tool-call agent) has a tiny per-message cost but an unbounded total — undercutting the entire point of this fix, since the skeleton is the thing that's supposed to stay under token_threshold. I verified: 5000 short messages produced a 32852-token skeleton against an 8192 default threshold — 4x over.

Fix: bounded to TOOL_RESULT_TOKEN_BUDGET (~4096 tokens) with a head+tail window (start frames the scenario, end is what the judge is deciding on) and the elided middle still reachable via expand_transcript/grep_transcript. Verified the fix brings the same 5000-message case down to 4099 tokens while preserving [0] and [4999]. Covered by TestSkeletonItselfIsBounded.

[P3] expand_transcript index matching was O(N×M) (CodeRabbit, partial) — FIXED (76c56de)

Why: membership-checked entry.index in parsed_indices against a List[int] rather than a set. Cheap, real, no downside to fix.

What I did NOT apply from CodeRabbit's suggestion: the "incrementally render + stop early" rewrite and a maxItems schema cap. The total-transcript-string cost this would guard against is already unconditionally paid every judge call (JudgeUtils.build_transcript_from_messages runs regardless of size to decide is_large_transcript), so the marginal protection is small relative to the added complexity, and expand_trace's existing span_ids schema has no equivalent cap either — adding one only to the new tool would be an asymmetry without a matching real-world risk.

Triaged and intentionally not applied

  • CodeRabbit: tighten Any/Dict/List to concrete generics across judge_utils.py/transcript_tools.py/judge_agent.py. These match the pre-existing convention in this exact package — _execute_discovery_tool(self, tool_call: Any, spans: Sequence[Any]) used Any for both parameters before this PR touched it. Tightening only the new code would make it inconsistent with the file it extends.
  • CodeRabbit: "use jsonify instead of json.dumps" (ast-grep). False positive — this is a Flask-specific lint rule firing on a non-Flask codebase; there's no Flask response object anywhere near this code.
  • CodeRabbit: add concrete type annotations to test helpers. The existing sibling test file (test_judge_agent_progressive.py) already uses the same loose-annotation style I matched (def create_large_trace() -> list:, def mock_litellm_response(tool_name: str, args: dict): with no return type). Diverging only in the new test file would reduce consistency, not improve it.
  • CodeRabbit: restore ScenarioConfig.default_config to its prior value instead of None in the test fixture teardown. Valid observation, but it's copy-pasted verbatim from the same setup_config fixture already duplicated across every judge test file in this suite. Fixing it in isolation in one file doesn't actually improve test isolation (the other files still reset to None), and a suite-wide fix is out of scope for #836.

Verified directly (adversarial, not just read)

  • Reproduced the original bug against unmodified main: 4/5 new tests fail without the fix, pass with it.
  • Combined-large-trigger case: built a scenario where both span digest and transcript exceed threshold simultaneously — confirmed all four discovery tools (expand_trace, grep_trace, expand_transcript, grep_transcript) are offered together with no crash.
  • Empty-messages, malformed tool_calls, out-of-range/unparseable indices, no-grep-match, >20-grep-match, and truncation-at-budget paths all exercised directly.
  • scenario/_judge/transcript_tools.py (the new module) is at 100% line coverage. Remaining gaps in judge_utils.py/judge_agent.py are pre-existing, in code this PR didn't touch (base64/audio media truncation edge cases, voice modality resolution, deep verdict-parsing branches) — out of scope for #836.
  • pyright clean on every touched file.
  • Full pytest tests/ (1139 tests): passes; only pre-existing e2e failures requiring live API keys, unrelated to this change.

Things I couldn't verify from the diff alone

  • CI's Test (Examples) step is failing on this PR with AuthenticationError: Incorrect API key provided for every example test. This is unambiguous: it fails before any of this PR's code runs, on every single example regardless of content, and main's most recent successful run (2026-07-20) passed this exact step — this reads as the OPENAI_API_KEY repo secret being invalid/rotated as of today, not something introduced by this change. I don't have access to repo secrets to confirm or fix it. Flagging so a maintainer with secret access can check, rather than blocking this PR on it.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🧹 Nitpick comments (1)
python/tests/test_judge_transcript_size_protection.py (1)

420-500: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Extract shared helper to remove duplication between the two dispatch tests.

test_malformed_tool_call_arguments_do_not_crash_the_loop and test_hallucinated_tool_name_reports_unknown_tool repeat nearly identical MagicMock tool-call construction and an almost-identical mock_completion closure. A small helper (e.g. _mock_tool_call_response(tool_name, arguments, call_id) plus a shared "capture tool response content" driver) would remove this duplication.

♻️ Sketch of a shared helper
+def _mock_tool_call_response(tool_name: str, arguments: str, call_id: str) -> MagicMock:
+    response = MagicMock()
+    response.choices = [MagicMock()]
+    tool_call = MagicMock()
+    tool_call.id = call_id
+    tool_call.function.name = tool_name
+    tool_call.function.arguments = arguments
+    response.choices[0].message.tool_calls = [tool_call]
+    response.choices[0].message.content = None
+    response.choices[0].message.role = "assistant"
+    return response
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@python/tests/test_judge_transcript_size_protection.py` around lines 420 -
500, Extract the duplicated tool-call response setup and completion-capture
logic from test_malformed_tool_call_arguments_do_not_crash_the_loop and
test_hallucinated_tool_name_reports_unknown_tool into small shared test helpers,
such as a MagicMock response builder and a completion driver. Keep each test’s
distinct tool name, arguments, call ID, expected captured content, call count,
and result assertions unchanged.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@python/tests/test_judge_transcript_size_protection.py`:
- Around line 420-500: Extract the duplicated tool-call response setup and
completion-capture logic from
test_malformed_tool_call_arguments_do_not_crash_the_loop and
test_hallucinated_tool_name_reports_unknown_tool into small shared test helpers,
such as a MagicMock response builder and a completion driver. Keep each test’s
distinct tool name, arguments, call ID, expected captured content, call count,
and result assertions unchanged.

ℹ️ Review info
⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: 6f5e4403-a124-4428-a94e-c9d24de593d4

📥 Commits

Reviewing files that changed from the base of the PR and between 76c56de and 47b86a4.

📒 Files selected for processing (1)
  • python/tests/test_judge_transcript_size_protection.py

Adds examples/test_judge_transcript_size_protection_e2e.py, following this
repo's existing examples/ pattern for live-LLM tests. A NonLitellmSalesAgent
returns 27 large tool-call/tool-result messages in one turn -- never
touching litellm itself, so JudgeSpanCollector stays empty -- with a single
fact buried in the middle. A real JudgeAgent (gpt-4.1-mini) must find that
fact through the transcript's own discovery tools to pass.

Verified locally against a live OpenAI key: PASSED. The judge correctly
located and used the buried fact despite the large non-litellm transcript,
confirming the fix works end-to-end with a real model, not just in mocked
unit tests.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
@github-actions github-actions Bot added the low-risk-change PR qualifies as low-risk per policy and can be merged without manual review label Jul 22, 2026
github-actions[bot]
github-actions Bot previously approved these changes Jul 22, 2026

@github-actions github-actions Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Approved by automation: PR qualifies as low-risk-change under the documented policy.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@python/examples/test_judge_transcript_size_protection_e2e.py`:
- Around line 63-71: Update the final assistant message in the test setup so it
is generic and does not repeat NEEDLE_FACT or any equivalent identifying fact.
Keep NEEDLE_FACT present only in the buried tool result, while preserving the
existing success assertions and transcript-discovery flow.
🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

  • Push a commit to this branch (recommended)
  • Create a new PR with the fixes

ℹ️ Review info
⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: 9fee658d-882b-461b-a3e8-93ff5e806389

📥 Commits

Reviewing files that changed from the base of the PR and between 47b86a4 and f4ffd95.

📒 Files selected for processing (1)
  • python/examples/test_judge_transcript_size_protection_e2e.py

Comment thread python/examples/test_judge_transcript_size_protection_e2e.py
…ummary

CodeRabbit review (PR #837): the final assistant message repeated
NEEDLE_FACT verbatim, so the judge could pass by reading that summary
alone (e.g. a cheap expand_transcript on just the last message) without
ever searching the buried tool result -- defeating the point of the test.

Made the summary generic so NEEDLE_FACT exists only inside the buried
tool-call result. Re-verified live against gpt-4.1-mini: still PASSED,
now a genuine proof that the judge searches the transcript rather than
reading a conveniently-placed restatement.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
@github-actions github-actions Bot removed the low-risk-change PR qualifies as low-risk per policy and can be merged without manual review label Jul 22, 2026
…ure state

Addresses the remaining unresolved CodeRabbit findings on PR #837:

- scenario/_judge/transcript_tools.py: _IndexedMessage.message is now
  ChatCompletionMessageParam instead of Any; the three budget constants
  are Final[int].
- scenario/_judge/judge_utils.py: _render_message_line's msg param is now
  ChatCompletionMessageParam instead of Any.
- scenario/judge_agent.py: working_messages (the parameter this fix
  introduced) is now Sequence[ChatCompletionMessageParam] instead of
  Sequence[Any] in both _run_discovery_loop and _execute_discovery_tool.
- tests/test_judge_transcript_size_protection.py: helper functions
  (create_mock_collector, create_large_non_litellm_transcript,
  create_base_input, mock_litellm_response, setup_config) now have
  concrete parameter/return types instead of bare list/dict/no-annotation;
  cast(Any, messages) narrowed to cast(List[ChatCompletionMessageParam],
  messages) everywhere the underlying data is well-formed (the deliberately
  malformed `indices` casts in the malformed-input tests stay Any, since
  narrowing those would misrepresent what's being tested).
- setup_config fixture now restores ScenarioConfig.default_config to its
  prior value instead of unconditionally resetting to None.

Not applied: switching to lowercase builtin generics (dict[...]/list[...])
in judge_utils.py/transcript_tools.py -- every other file in this package
consistently imports List/Dict from typing (verified across
trace_tools.py, judge_agent.py, judge_span_digest_formatter.py); doing it
only in the new code would make this PR's files inconsistent with the
package it extends. Also not applied: rewriting judge_utils.py's
pre-existing base64/media-truncation functions' Any usage -- that's
unchanged logic outside this PR's diff, not something #836 touches.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
@github-actions

Copy link
Copy Markdown
Contributor

Automated low-risk assessment

This PR was evaluated against the repository's Low-Risk Pull Requests procedure and does not qualify as low risk.

This PR changes runtime behavior of the JudgeAgent (how transcript size is computed and rendered), adds new discovery tools (expand_transcript/grep_transcript) and wiring into the LLM tool invocation flow, and refactors transcript rendering utilities. These are not mere docs/tests/UI changes and affect core agent logic and model interactions, so they do not meet the "low risk" criteria.

This PR requires a manual review before merging.

@Aryansharma28

Copy link
Copy Markdown
Contributor Author

All 7 CodeRabbit review threads are now resolved:

  • "Cap the skeleton itself" — fixed in 76c56de
  • "Bound expansion before rendering" — partially fixed in 76c56de (O(1) set membership); the incremental-render rewrite wasn't applied (reasoning in my earlier review comment)
  • "Ensure test requires transcript discovery" — fixed in 903b626
  • "Tighten typings" (judge_utils.py/transcript_tools.py/judge_agent.py) — fixed in f3e4eb7 for the code this PR actually introduced; pre-existing untouched code and the builtin-generics style switch were deliberately left alone (would make new code inconsistent with the rest of the package)
  • "Add concrete types to test helpers" — fixed in f3e4eb7
  • "Restore previous default configuration" in the test fixture — fixed in f3e4eb7
  • "Replace cast(Any, messages) with a concrete type" — fixed in f3e4eb7 everywhere the underlying data is well-formed; the indices casts in the two deliberately-malformed-input tests stay Any since those tests exist specifically to pass data that violates the type contract

Full test suite (113 tests across the judge suite) + pyright: clean after every change. python-ci's only remaining red check is Test (Examples), unrelated to this PR (pre-existing broken OPENAI_API_KEY repo secret, flagged in my earlier review).

@rogeriochaves
rogeriochaves merged commit 6604334 into main Jul 22, 2026
16 of 18 checks passed
@rogeriochaves
rogeriochaves deleted the fix/836-judge-transcript-size-protection branch July 22, 2026 16:23
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

JudgeAgent trace-size protection only covers litellm spans — raw transcript block is unbounded for non-litellm agent adapters

2 participants