[None][feat] Enable agent serving evaluation via trace-replay on Scaffolding - #14397
Merged
WeiHaocheng merged 52 commits intoJul 14, 2026
Merged
Conversation
KleinBlueC
requested review from
ZhanruiSunCh,
hchings,
laikhtewari,
lancelly,
nv-guomingz and
yuanjingx87
May 21, 2026 08:14
Contributor
📝 WalkthroughWalkthroughAdds Coder/Iter/Open-Deep/ToT agents and runners, multiple MCP servers (tavily, google, fetch_webpage, coder), execution tracing and replay with Pareto analyses/plots, executor decode-deferral and iteration-state logging, KV-retention defaults, benchmark wiring, config/docs updates, and minor build/ignore tweaks. ChangesEnd-to-end agents, MCP, trace/replay, executor updates
Sequence Diagram(s)(skipped) Estimated code review effort🎯 5 (Critical) | ⏱️ ~120 minutes Possibly related PRs
Suggested reviewers
✨ Finishing Touches🧪 Generate unit tests (beta)
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Note
Due to the large number of review comments, Critical severity comments were prioritized as inline comments.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
examples/scaffolding/benchmarks/coder_benchmark.py (1)
193-204:⚠️ Potential issue | 🟠 Major | ⚡ Quick winShutdown path drops
generation_worker, causing resource-leak risk.
create_coder_resources()returnsgeneration_worker, butasync_coder_benchmark()discards it and cleanup only shuts down MCP + LLM. Ensuregeneration_worker.shutdown()is always called.Suggested fix
-async def cleanup_coder_resources(llm, mcp_worker): +async def cleanup_coder_resources(llm, mcp_worker, generation_worker): """Cleanup Coder benchmark resources.""" await mcp_worker.async_shutdown() await shutdown_llm(llm) + generation_worker.shutdown() ... - llm, mcp_worker, _ = await create_coder_resources(args) + llm, mcp_worker, generation_worker = await create_coder_resources(args) ... - await cleanup_coder_resources(llm, mcp_worker) + await cleanup_coder_resources(llm, mcp_worker, generation_worker)🤖 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 `@examples/scaffolding/benchmarks/coder_benchmark.py` around lines 193 - 204, The shutdown path currently discards the returned generation_worker from create_coder_resources and never shuts it down; change the unpacking to capture it (e.g., llm, mcp_worker, generation_worker = await create_coder_resources(args)) and ensure generation_worker.shutdown() is invoked in the finally block (either by calling await generation_worker.shutdown() there or by extending cleanup_coder_resources to accept and shut down generation_worker and calling that). Update references to create_coder_resources, the local generation_worker variable, and cleanup_coder_resources/run_coder_benchmark_core as needed so the worker is always awaited and cleaned up.
🟡 Minor comments (17)
examples/scaffolding/contrib/Coder/run_coder.py-199-200 (1)
199-200:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winUse explicit runtime validation instead of
assertfor output checks.
assertis not reliable for runtime error handling and can be optimized out. Raise a concrete exception when outputs are missing.Suggested fix
- assert result.outputs[0].text is not None - print("\nFinal output:\n" + result.outputs[0].text) + if not result.outputs or result.outputs[0].text is None: + raise RuntimeError("Coder run returned no text output.") + print("\nFinal output:\n" + result.outputs[0].text)As per coding guidelines: “Use built-in Python exception types; use exceptions for error handling, not return values; … raise
ValueErrorinstead of assertions.”🤖 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 `@examples/scaffolding/contrib/Coder/run_coder.py` around lines 199 - 200, Replace the runtime assertion with explicit validation: check that result.outputs exists, has at least one item, and that result.outputs[0].text is not None (or non-empty) and if the check fails raise a ValueError with a clear message; update the block around result.outputs[0].text and the print call to perform this validation before printing so missing outputs produce a concrete exception instead of relying on assert.examples/scaffolding/contrib/tree_of_thought_research/config.yaml-1-3 (1)
1-3:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winUpdate stale runner labeling in this ToT config.
Comments still refer to IterResearch and
run_iter_research.py, which is misleading fortree_of_thought_research.Also applies to: 40-40
🤖 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 `@examples/scaffolding/contrib/tree_of_thought_research/config.yaml` around lines 1 - 3, Update the stale comments in the YAML header that reference "IterResearch" and "run_iter_research.py": replace the project label with "Tree of Thought Research" (or the canonical project name "tree_of_thought_research") and update the runner command reference to the correct runner script (e.g., "run_tree_of_thought_research.py" or the actual runner used by this example). Ensure both occurrences (the header comment and the runner comment) are updated so the config's comments accurately describe how to pass the file to the MCP servers and the correct runner.examples/scaffolding/contrib/open_deep_research/README.md-11-12 (1)
11-12:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winFix config path resolution in MCP startup snippets.
After
cdinto each MCP directory,--config examples/scaffolding/...is no longer a valid relative path. Use a directory-relative config path or run from repo root.Also applies to: 15-16, 19-20
🤖 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 `@examples/scaffolding/contrib/open_deep_research/README.md` around lines 11 - 12, The startup snippets use an invalid repo-root relative config path in the command "uv run tavily_search.py --config examples/scaffolding/contrib/open_deep_research/config.yaml"; after you cd into the MCP directory update the --config value to a directory-relative path (e.g. --config ../../contrib/open_deep_research/config.yaml) or change the instructions to run from the repo root, and apply the same replacement to the other two occurrences noted (the snippets at lines 15-16 and 19-20).examples/scaffolding/contrib/iter_research/README.md-6-7 (1)
6-7:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winFix the config path in MCP startup examples.
These commands
cdinto each tool directory, but still pass a repo-root-style--configpath. That path won’t resolve from the new working directory. Use a relative path from the tool directory (or run commands from repo root withoutcd).Also applies to: 10-11, 14-15
🤖 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 `@examples/scaffolding/contrib/iter_research/README.md` around lines 6 - 7, The example command cd's into examples/scaffolding/mcp/tavily_search but still uses a repo-root style --config path that will not resolve; update the --config argument in the uv run tavily_search.py invocation to a path relative to the tool directory (e.g., ../../contrib/iter_research/config.yaml) or remove the cd and run from the repo root, and make the same relative-path fix for the other affected examples (lines 10-11 and 14-15) so the --config flag points to a valid file from the current working directory.examples/scaffolding/mcp/coder/README.md-7-17 (1)
7-17:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winTool list is missing
python_interpreter.The exposed tools section is currently incomplete and can mislead clients wiring against the documented interface.
Suggested doc fix
- `update_plan` +- `python_interpreter` - `think` - `complete_task`🤖 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 `@examples/scaffolding/mcp/coder/README.md` around lines 7 - 17, The README's tool list is missing the python_interpreter tool; update the exposed tools bullet list in the examples/scaffolding/mcp/coder/README.md to include `python_interpreter` alongside `read_file`, `list_dir`, `grep_files`, `exec`, `shell`, `update_plan`, `think`, and `complete_task` so documentation matches the actual interface implemented (see the `python_interpreter` symbol in the Coder tools). Ensure the tool name is spelled exactly `python_interpreter` and placed consistently with the other bullets.examples/scaffolding/trace_replay/analysis/branch_summary.py-5-10 (1)
5-10:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winReplace en dashes with ASCII hyphens in the docstring.
This docstring contains ambiguous Unicode dash characters flagged by Ruff (
RUF002), which can fail lint.🤖 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 `@examples/scaffolding/trace_replay/analysis/branch_summary.py` around lines 5 - 10, The docstring in branch_summary.py uses Unicode en dashes (–) in the three bullet descriptions for by_branch_root, by_branch_depth, and by_system_prompt; replace those en dashes with ASCII hyphen-minus characters (-) in the docstring bullets so the lines read "* ``by_branch_root`` - ...", "* ``by_branch_depth`` - ...", and "* ``by_system_prompt`` - ...", then save the file.examples/scaffolding/trace_replay/analysis/__init__.py-35-60 (1)
35-60:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winSort
__all__to satisfy Ruff (RUF022).
__all__is currently unsorted and will trip lint in strict setups.🤖 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 `@examples/scaffolding/trace_replay/analysis/__init__.py` around lines 35 - 60, The exported names list __all__ is unsorted and triggers Ruff RUF022; fix it by alphabetically sorting the entries in the __all__ list (keep the exact string names and capitalization intact), e.g., reorder the elements like "ANNOTATED_TRACE_SUFFIX", "ANNOTATION_FIELDS", "ConversationSegments", ... so the list is lexicographically sorted; update the __all__ assignment in this module (search for the __all__ = [...] block) to the sorted sequence.examples/scaffolding/trace_replay/analysis/blocks.py-27-27 (1)
27-27:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winSort
_TokenNode.__slots__to satisfy Ruff (RUF023).The
__slots__tuple is not naturally sorted, which can trigger lint failures.🤖 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 `@examples/scaffolding/trace_replay/analysis/blocks.py` at line 27, The __slots__ tuple on _TokenNode is unsorted and triggers Ruff RUF023; update the _TokenNode.__slots__ definition to list its slot names in sorted order (e.g., alphabetically) so the tuple is deterministic and lint-compliant, ensuring you modify the __slots__ attribute in the _TokenNode class rather than other symbols.examples/scaffolding/trace_replay/README.md-30-56 (1)
30-56:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winAdd a language tag to the fenced code block.
The code fence at this section is missing a language specifier, which triggers markdown lint (
MD040).Suggested fix
-``` +```text examples/scaffolding/trace_replay/ ├── README.md -- this file ... -``` +```🤖 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 `@examples/scaffolding/trace_replay/README.md` around lines 30 - 56, The fenced code block in the README.md (the tree listing block) is missing a language tag which triggers MD040; update the opening fence from ``` to an explicit language tag such as ```text (or ```bash) in the README.md where the examples/scaffolding/trace_replay/ tree is shown so the block is properly recognized by markdown linters.examples/scaffolding/trace_replay/analysis/io.py-1-11 (1)
1-11:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winMissing NVIDIA copyright header.
All TensorRT-LLM source files must contain an NVIDIA copyright header. As per coding guidelines for
**/*.pyfiles.Proposed fix: add copyright header
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + """Filesystem helpers: locate trace files, load/write JSON."""🤖 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 `@examples/scaffolding/trace_replay/analysis/io.py` around lines 1 - 11, Add the required NVIDIA copyright header at the very top of this Python module (before the module docstring and imports). Ensure the header matches the project's standard Python file header used across other .py files and leave the existing module docstring, imports, and constants (OUTPUT_SUFFIX, ANNOTATED_TRACE_SUFFIX) unchanged; just prepend the canonical header text as the first lines of the file.examples/scaffolding/trace_replay/analysis/streams.py-1-18 (1)
1-18:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winMissing NVIDIA copyright header.
All TensorRT-LLM source files must contain an NVIDIA copyright header. As per coding guidelines for
**/*.pyfiles.Proposed fix: add copyright header
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + r"""Synthetic token streams + system-prompt registry.🤖 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 `@examples/scaffolding/trace_replay/analysis/streams.py` around lines 1 - 18, This file is missing the required NVIDIA copyright header; add the standard NVIDIA copyright/license header as the very first lines of examples/scaffolding/trace_replay/analysis/streams.py (before the module docstring) so it applies to the whole module, ensuring the header format matches other TensorRT-LLM Python files and includes the correct year/owner and license identifier; no code changes to SystemPromptRegistry or ConversationSegments are needed—only prepend the canonical NVIDIA header block to the top of the file.examples/scaffolding/trace_replay/analysis/compute_cache_hit_trace.py-1-19 (1)
1-19:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winMissing NVIDIA copyright header.
All TensorRT-LLM source files must contain an NVIDIA copyright header with the year of latest modification, using the Apache 2.0 license text. As per coding guidelines for
**/*.pyfiles.Proposed fix: add copyright header
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + r"""CLI: compute the ideal prefix KV-cache hit upper bound for scaffolding traces.🤖 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 `@examples/scaffolding/trace_replay/analysis/compute_cache_hit_trace.py` around lines 1 - 19, Add the required NVIDIA copyright header with the correct year of latest modification and the Apache 2.0 license boilerplate at the very top of the module file compute_cache_hit_trace.py (before the module docstring), matching the project's standard Python header format used across other `*.py` files.examples/scaffolding/trace_replay/pareto/trace_replay_pareto_aggregate.py-43-43 (1)
43-43:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winUnused import:
defaultdict.The
defaultdictimport fromcollectionsis never used in this file.Proposed fix
-from collections import defaultdict from datetime import datetime, timezone🤖 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 `@examples/scaffolding/trace_replay/pareto/trace_replay_pareto_aggregate.py` at line 43, Remove the unused import "defaultdict" from the top-level import statement in this file; locate the line that reads "from collections import defaultdict" and delete it so only necessary imports remain (no other code changes required).examples/scaffolding/trace_replay/pareto/trace_replay_client.py-264-306 (1)
264-306:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winFile handle
fpis not closed on exception before thetryblock.If an exception occurs between opening the file at line 268 and entering the try block at line 278, the file handle will leak. The
finallyblock at line 303 only runs if we reach line 278.Proposed fix: use context manager or move open inside try
output_path.parent.mkdir(parents=True, exist_ok=True) - fp = open(output_path, "w", encoding="utf-8") def _write_events(events): if not events: return for ev in events: fp.write(_json.dumps(ev) + "\n") fp.flush() total = 0 + fp = open(output_path, "w", encoding="utf-8") try: while not stop_event.is_set():🤖 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 `@examples/scaffolding/trace_replay/pareto/trace_replay_client.py` around lines 264 - 306, The file handle fp is opened before the try so it can leak if an exception occurs before entering the try; move the open(output_path, "w", encoding="utf-8") into the try (or wrap it with a context manager) so fp is created inside the try/finally scope and guaranteed to be closed in the existing finally, and keep the helper _write_events using that fp (or accept fp as a captured variable) so the existing finally block that calls await asyncio.to_thread(fp.close) always runs.tensorrt_llm/scaffolding/contrib/Coder/prompts.py-167-245 (1)
167-245:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winNormalize ambiguous unicode punctuation in prompt literals.
These lines contain smart quotes/dashes (
’,–,‑) flagged by RuffRUF001. Replacing them with ASCII equivalents avoids lint noise or CI failures.Also applies to: 274-276
🤖 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 `@tensorrt_llm/scaffolding/contrib/Coder/prompts.py` around lines 167 - 245, There are smart punctuation characters (’ – ‑) inside the multi-line prompt string literals in `tensorrt_llm/scaffolding/contrib/Coder/prompts.py` (the large assistant-guidelines/prompt templates and the later prompt strings around the subsequent examples); replace all smart quotes and dashes with ASCII equivalents (`'`, `-`, `--` or `"` as appropriate) in those prompt/template variables so the literals contain only standard ASCII punctuation and silence RUF001 lint warnings. Locate the problematic prompt/template string variables in this module (the main assistant guideline/multi-line prompt block and the later prompt templates referenced again near the end of the file) and perform a global replace of `’`→`'`, `–`→`-`, and `‑`→`-` inside those string literals only.tensorrt_llm/scaffolding/contrib/open_deep_research/README.md-52-64 (1)
52-64:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winThe MCP command examples use config paths that won’t resolve from the shown directories.
Because each snippet starts with
cd examples/scaffolding/mcp/...,--config examples/scaffolding/...is not cwd-correct.📝 Suggested fix
cd examples/scaffolding/mcp/tavily_search && uv run tavily_search.py \ - --config examples/scaffolding/contrib/open_deep_research/config.yaml + --config ../../contrib/open_deep_research/config.yaml @@ cd examples/scaffolding/mcp/google_scholar && uv run google_scholar.py \ - --config examples/scaffolding/contrib/open_deep_research/config.yaml + --config ../../contrib/open_deep_research/config.yaml @@ cd examples/scaffolding/mcp/fetch_webpage && uv run fetch_webpage.py \ - --config examples/scaffolding/contrib/open_deep_research/config.yaml + --config ../../contrib/open_deep_research/config.yaml🤖 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 `@tensorrt_llm/scaffolding/contrib/open_deep_research/README.md` around lines 52 - 64, The example MCP run commands in README.md use config paths that are incorrect for the shown cwd; update each command (the ones that run tavily_search.py, google_scholar.py and fetch_webpage.py) so the --config path is relative to the current directory (use ../contrib/open_deep_research/config.yaml instead of examples/scaffolding/contrib/open_deep_research/config.yaml) so the config resolves when you first run cd examples/scaffolding/mcp/<...> and then uv run <script>.tensorrt_llm/scaffolding/contrib/tree_of_thought_research/README.md-25-27 (1)
25-27:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winFix the example config path for ToT runner.
The command points to
iter_research/config.yaml; this README is fortree_of_thought_research.Proposed fix
python examples/scaffolding/contrib/tree_of_thought_research/run_tot_research.py \ - --config examples/scaffolding/contrib/iter_research/config.yaml \ + --config examples/scaffolding/contrib/tree_of_thought_research/config.yaml \ --prompt "Use current sources to compare TensorRT-LLM and vLLM for serving LLMs."🤖 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 `@tensorrt_llm/scaffolding/contrib/tree_of_thought_research/README.md` around lines 25 - 27, Update the example command in the README so the --config path points to the tree_of_thought_research config instead of iter_research; change the referenced config from "examples/scaffolding/contrib/iter_research/config.yaml" to the corresponding "examples/scaffolding/contrib/tree_of_thought_research/config.yaml" used by the run_tot_research.py example so the --config passed to run_tot_research.py matches this README's topic.
🧹 Nitpick comments (6)
examples/scaffolding/contrib/open_deep_research/run_open_deep_research.py (1)
30-67: ⚡ Quick winAdd explicit return types for function definitions.
parse_argumentsandmainshould use explicit return annotations (argparse.Namespace/None) per project typing requirements.As per coding guidelines, "Static type checking via mypy is opt-in by submodule and highly recommended; always annotate Python functions with return types (use
Noneif no return); make return type explicit".Also applies to: 107-107
🤖 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 `@examples/scaffolding/contrib/open_deep_research/run_open_deep_research.py` around lines 30 - 67, Add explicit return type annotations to the function definitions: annotate parse_arguments() to return argparse.Namespace and annotate main() to return None; update the def lines for parse_arguments and main accordingly (no other logic changes) so the module complies with the project's typing guideline requiring explicit function return types.examples/scaffolding/contrib/iter_research/run_iter_research.py (1)
23-46: ⚡ Quick winAdd explicit return type annotations to script functions.
parse_argumentsandmainshould declare return types to match the repository typing rule.As per coding guidelines, "Static type checking via mypy is opt-in by submodule and highly recommended; always annotate Python functions with return types (use
Noneif no return); make return type explicit".Also applies to: 86-87
🤖 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 `@examples/scaffolding/contrib/iter_research/run_iter_research.py` around lines 23 - 46, Add explicit return type annotations: update parse_arguments to declare -> argparse.Namespace (or import Namespace and use -> Namespace) and update main to declare -> None; adjust any type imports (e.g., from argparse import Namespace or import typing) if necessary to satisfy type checkers, and ensure the annotated signatures are used where parse_arguments() and main() are defined (functions named parse_arguments and main).tensorrt_llm/scaffolding/contrib/open_deep_research/supervisor.py (1)
143-147: ⚡ Quick winDemote dict-handling concern in
open_deep_research/supervisor.py: the scaffold treatstool_call.function.argumentsas a JSON string.
MCPCallTaskreceivestool_call.function.argumentsas-is and the MCP worker executesjson.loads(task.args); ifargumentswere a Pythondict, it would fail in that path as well, so the TypeError scenario is unlikely in this codebase. Consider aligning supervisor with the existing_parse_tool_argumentshelper only for consistency/future-proofing.🤖 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 `@tensorrt_llm/scaffolding/contrib/open_deep_research/supervisor.py` around lines 143 - 147, The supervisor currently blindly json.loads tool_call.function.arguments which duplicates parsing logic and may mis-handle dicts; update the code in open_deep_research/supervisor.py to use the shared parsing helper (_parse_tool_arguments) or a small wrapper that accepts either a JSON string or a dict and returns a dict, so MCPCallTask still receives the normalized arguments; locate the json.loads usage around tool_call.function.arguments and replace it with a call to _parse_tool_arguments (or equivalent) to ensure consistent parsing with MCPCallTask.tensorrt_llm/scaffolding/replay.py (1)
520-537: 💤 Low valueAdd
strict=Truetozip()calls in drop verification.The
conv_ids,prefixes, andkeepslists are built in parallel from the same_conversation_token_idsiteration (lines 482-491), so they should always match.♻️ Suggested fix
- for conv_id, prefix, keep in zip(conv_ids, prefixes, keeps): + for conv_id, prefix, keep in zip(conv_ids, prefixes, keeps, strict=True):Apply the same fix at line 566.
🤖 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 `@tensorrt_llm/scaffolding/replay.py` around lines 520 - 537, The parallel iteration over conv_ids, prefixes, and keeps in the drop verification uses zip(...) which can silently truncate mismatched iterables; update the zip calls that iterate these three lists (the one that creates probe GenerationTask with input_tokens=prefix and records via self._drop_path_stats.record, and the similar one later) to use zip(..., strict=True) so any length mismatch raises immediately; locate uses of zip(conv_ids, prefixes, keeps) around the probe creation and the later drop-verification block and add strict=True to each.tensorrt_llm/scaffolding/contrib/tree_of_thought_research/tot_research_controllers.py (1)
193-195: 💤 Low valueAdd
strict=Truetozip()for length validation.The
candidatesandeval_taskslists are built in lockstep (lines 183-186), so they should always match. Addingstrict=Truewill catch any future bug that breaks this invariant.♻️ Suggested fix
- for branch, eval_task in zip(candidates, eval_tasks): + for branch, eval_task in zip(candidates, eval_tasks, strict=True):🤖 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 `@tensorrt_llm/scaffolding/contrib/tree_of_thought_research/tot_research_controllers.py` around lines 193 - 195, The loop that pairs candidates and eval_tasks (for branch, eval_task in zip(candidates, eval_tasks)) silently allows length mismatches; update that zip call to zip(candidates, eval_tasks, strict=True) to enforce the invariant and surface any future mismatch bugs, keeping the body that sets branch.evaluation = eval_task.last_assistant_content() and branch.score = self._parse_score(branch.evaluation) unchanged.tensorrt_llm/scaffolding/controller.py (1)
44-47: 💤 Low valueConsider adding
strict=Truetozip().The three lists are caller-provided and should always match in length. Adding
strict=Truewould catch mismatches early.♻️ Suggested fix
- for controller, tasks, kwargs in zip(controllers, tasks_list, - kwargs_list): + for controller, tasks, kwargs in zip(controllers, tasks_list, + kwargs_list, strict=True):🤖 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 `@tensorrt_llm/scaffolding/controller.py` around lines 44 - 47, The loop over controllers, tasks_list, and kwargs_list uses zip(controllers, tasks_list, kwargs_list) which can silently truncate mismatched lengths; update that zip call to zip(controllers, tasks_list, kwargs_list, strict=True) to force a ValueError when lengths differ so mismatches are caught early; this change affects the loop that builds self.sub_gens by calling controller.process(tasks, **kwargs) and will surface caller errors immediately.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 1d50885d-31d4-4dfc-8375-1055cdce59c5
⛔ Files ignored due to path filters (2)
examples/scaffolding/contrib/mcp/weather/uv.lockis excluded by!**/*.lockexamples/scaffolding/contrib/open_deep_research/TavilyMCP/uv.lockis excluded by!**/*.lock
📒 Files selected for processing (124)
.extra-llm-api-config.yml.gitignore.pre-commit-config.yamlcpp/tensorrt_llm/kernels/cutlass_kernels/CMakeLists.txtexamples/scaffolding/benchmarks/__main__.pyexamples/scaffolding/benchmarks/coder_benchmark.pyexamples/scaffolding/contrib/Coder/README.mdexamples/scaffolding/contrib/Coder/run_coder.pyexamples/scaffolding/contrib/Coder/run_swebench.pyexamples/scaffolding/contrib/iter_research/README.mdexamples/scaffolding/contrib/iter_research/config.yamlexamples/scaffolding/contrib/iter_research/run_iter_research.pyexamples/scaffolding/contrib/mcp/README.mdexamples/scaffolding/contrib/mcp/websearch/.envexamples/scaffolding/contrib/mcp/websearch/README.mdexamples/scaffolding/contrib/mcp/websearch/main.pyexamples/scaffolding/contrib/mcp/websearch/pyproject.tomlexamples/scaffolding/contrib/mcp/websearch/websearch.pyexamples/scaffolding/contrib/open_deep_research/README.mdexamples/scaffolding/contrib/open_deep_research/TavilyMCP/README.mdexamples/scaffolding/contrib/open_deep_research/TavilyMCP/pyproject.tomlexamples/scaffolding/contrib/open_deep_research/TavilyMCP/travily.pyexamples/scaffolding/contrib/open_deep_research/config.yamlexamples/scaffolding/contrib/open_deep_research/run_deep_research.pyexamples/scaffolding/contrib/open_deep_research/run_open_deep_research.pyexamples/scaffolding/contrib/tree_of_thought_research/config.yamlexamples/scaffolding/contrib/tree_of_thought_research/run_tot_research.pyexamples/scaffolding/mcp/README.mdexamples/scaffolding/mcp/coder/README.mdexamples/scaffolding/mcp/coder/coder_mcp.pyexamples/scaffolding/mcp/coder/pyproject.tomlexamples/scaffolding/mcp/e2b/.env.exampleexamples/scaffolding/mcp/e2b/README.mdexamples/scaffolding/mcp/e2b/e2bserver.pyexamples/scaffolding/mcp/e2b/main.pyexamples/scaffolding/mcp/e2b/pyproject.tomlexamples/scaffolding/mcp/fetch_webpage/__init__.pyexamples/scaffolding/mcp/fetch_webpage/fetch_webpage.pyexamples/scaffolding/mcp/fetch_webpage/pyproject.tomlexamples/scaffolding/mcp/fetch_webpage/visit_controller.pyexamples/scaffolding/mcp/google_scholar/google_scholar.pyexamples/scaffolding/mcp/google_scholar/pyproject.tomlexamples/scaffolding/mcp/google_search/google_search.pyexamples/scaffolding/mcp/google_search/pyproject.tomlexamples/scaffolding/mcp/mcptest.pyexamples/scaffolding/mcp/tavily_search/__init__.pyexamples/scaffolding/mcp/tavily_search/pyproject.tomlexamples/scaffolding/mcp/tavily_search/tavily_controller.pyexamples/scaffolding/mcp/tavily_search/tavily_search.pyexamples/scaffolding/mcp/weather/pyproject.tomlexamples/scaffolding/mcp/weather/weather.pyexamples/scaffolding/mcp/wordllama/README.mdexamples/scaffolding/mcp/wordllama/main.pyexamples/scaffolding/mcp/wordllama/pyproject.tomlexamples/scaffolding/mcp/wordllama/wordllama_serve.pyexamples/scaffolding/trace_replay/README.mdexamples/scaffolding/trace_replay/analysis/README.mdexamples/scaffolding/trace_replay/analysis/__init__.pyexamples/scaffolding/trace_replay/analysis/aggregation.pyexamples/scaffolding/trace_replay/analysis/annotate.pyexamples/scaffolding/trace_replay/analysis/blocks.pyexamples/scaffolding/trace_replay/analysis/branch_summary.pyexamples/scaffolding/trace_replay/analysis/cache_hit.pyexamples/scaffolding/trace_replay/analysis/compute_cache_hit_trace.pyexamples/scaffolding/trace_replay/analysis/io.pyexamples/scaffolding/trace_replay/analysis/streams.pyexamples/scaffolding/trace_replay/metrics.pyexamples/scaffolding/trace_replay/pareto/README.mdexamples/scaffolding/trace_replay/pareto/_common.pyexamples/scaffolding/trace_replay/pareto/trace_replay_client.pyexamples/scaffolding/trace_replay/pareto/trace_replay_pareto_aggregate.pyexamples/scaffolding/trace_replay/plots/plot_trace_replay_agent_pareto.pyexamples/scaffolding/trace_replay/plots/plot_trace_replay_job_pareto.pyexamples/scaffolding/trace_replay/plots/plot_trace_replay_per_call_hit_curves.pyexamples/scaffolding/trace_replay/plots/plot_trace_replay_session_hit_pareto.pyexamples/scaffolding/trace_replay/plots/plot_trace_replay_session_hit_vs_time.pyexamples/scaffolding/trace_replay/plots/plot_trace_replay_token_pareto.pyexamples/scaffolding/trace_replay/run_trace_replay.pyexamples/scaffolding/trace_replay/trace_example/matplotlib__matplotlib-23412/matplotlib__matplotlib-23412.full.trace.jsonexamples/scaffolding/trace_replay/trace_example/matplotlib__matplotlib-23412/matplotlib__matplotlib-23412.trace.jsonpyproject.tomltensorrt_llm/_torch/pyexecutor/py_executor.pytensorrt_llm/_torch/pyexecutor/scheduler/adp_router.pytensorrt_llm/_torch/pyexecutor/scheduler/scheduler.pytensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.pytensorrt_llm/llmapi/llm.pytensorrt_llm/llmapi/llm_args.pytensorrt_llm/scaffolding/__init__.pytensorrt_llm/scaffolding/benchmark.pytensorrt_llm/scaffolding/contrib/Coder/__init__.pytensorrt_llm/scaffolding/contrib/Coder/coder.pytensorrt_llm/scaffolding/contrib/Coder/prompts.pytensorrt_llm/scaffolding/contrib/Coder/swebench.pytensorrt_llm/scaffolding/contrib/Coder/tools.pytensorrt_llm/scaffolding/contrib/iter_research/README.mdtensorrt_llm/scaffolding/contrib/iter_research/__init__.pytensorrt_llm/scaffolding/contrib/iter_research/agent.pytensorrt_llm/scaffolding/contrib/iter_research/prompts.pytensorrt_llm/scaffolding/contrib/iter_research/utils.pytensorrt_llm/scaffolding/contrib/open_deep_research/README.mdtensorrt_llm/scaffolding/contrib/open_deep_research/prompts.pytensorrt_llm/scaffolding/contrib/open_deep_research/researcher.pytensorrt_llm/scaffolding/contrib/open_deep_research/supervisor.pytensorrt_llm/scaffolding/contrib/open_deep_research/tools.pytensorrt_llm/scaffolding/contrib/tree_of_thought_research/README.mdtensorrt_llm/scaffolding/contrib/tree_of_thought_research/__init__.pytensorrt_llm/scaffolding/contrib/tree_of_thought_research/prompts.pytensorrt_llm/scaffolding/contrib/tree_of_thought_research/tools.pytensorrt_llm/scaffolding/contrib/tree_of_thought_research/tot_research_controllers.pytensorrt_llm/scaffolding/controller.pytensorrt_llm/scaffolding/execution_scope.pytensorrt_llm/scaffolding/execution_trace.pytensorrt_llm/scaffolding/replay.pytensorrt_llm/scaffolding/result.pytensorrt_llm/scaffolding/scaffolding_llm.pytensorrt_llm/scaffolding/task.pytensorrt_llm/scaffolding/task_collection.pytensorrt_llm/scaffolding/utils.pytensorrt_llm/scaffolding/worker.pytensorrt_llm/serve/openai_protocol.pytensorrt_llm/serve/openai_server.pytensorrt_llm/serve/postprocess_handlers.pytensorrt_llm/serve/resource_governor.pytests/test_tokenize_endpoint.py
💤 Files with no reviewable changes (10)
- examples/scaffolding/contrib/open_deep_research/TavilyMCP/pyproject.toml
- examples/scaffolding/contrib/mcp/README.md
- examples/scaffolding/contrib/mcp/websearch/.env
- examples/scaffolding/contrib/open_deep_research/TavilyMCP/README.md
- examples/scaffolding/contrib/mcp/websearch/pyproject.toml
- .gitignore
- examples/scaffolding/contrib/open_deep_research/run_deep_research.py
- examples/scaffolding/contrib/mcp/websearch/main.py
- examples/scaffolding/contrib/open_deep_research/TavilyMCP/travily.py
- examples/scaffolding/contrib/mcp/websearch/websearch.py
WeiHaocheng
reviewed
May 22, 2026
KleinBlueC
force-pushed
the
feat/scaffolding-trace-replay
branch
from
May 25, 2026 12:41
5000d85 to
8ac0979
Compare
Signed-off-by: Ryan Sun <rysun@nvidia.com> (cherry picked from commit f3869a8)
Signed-off-by: Yi Sun <yisun0618@gmail.com> (cherry picked from commit fc3bedb)
Signed-off-by: Yi Sun <yisun0618@gmail.com> (cherry picked from commit beef3b4)
Signed-off-by: Yi Sun <yisun0618@gmail.com> (cherry picked from commit 4d4197f)
(cherry picked from commit f696ee1) Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com>
(cherry picked from commit 1f3d60b) Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com>
Signed-off-by: Yi Sun <yisun0618@gmail.com> (cherry picked from commit 07445f8)
Signed-off-by: Yi Sun <yisun0618@gmail.com> (cherry picked from commit 1f5410d)
…fy scaffolding tracer (cherry picked from commit a1e77ac) Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com>
(cherry picked from commit f1b9e39) Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com>
Signed-off-by: Yi Sun <yisun0618@gmail.com> (cherry picked from commit c0c020b)
Signed-off-by: Yi Sun <yisun0618@gmail.com> (cherry picked from commit a950614)
Signed-off-by: kleinc <kleinc@nvidia.com> (cherry picked from commit 5d1fc08)
Contributor
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #58580 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #58580 [ run ] completed with state
|
The server registers a RequestValidationError handler that renders
validation failures as HTTP 400 with a {"error": ...} envelope (not the
FastAPI default 422), matching existing client/test contracts. The
tokenize end-to-end test wrongly asserted 422; fix it to expect 400.
Signed-off-by: KleinBlueC <102355066+KleinBlueC@users.noreply.github.com>
Signed-off-by: KleinBlueC <102355066+KleinBlueC@users.noreply.github.com>
Contributor
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #58746 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #58746 [ run ] completed with state
|
ChatWithMCPController now reads tool_call.id when creating MCPCallTask, but the test's dummy ToolCall only carried a function attribute. The resulting AttributeError was swallowed by the ScaffoldingLlm request handler, so test_scaffolding_with_chat_mcp_controller failed on the pre-seeded empty ScaffoldingOutput instead of surfacing the error. Signed-off-by: KleinBlueC <102355066+KleinBlueC@users.noreply.github.com>
Contributor
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #58799 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #58799 [ run ] completed with state |
Contributor
Author
|
/bot run --disable-fail-fast |
1 similar comment
Collaborator
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #58832 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #58832 [ run ] completed with state
|
Contributor
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #58853 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #58853 [ run ] completed with state
|
Contributor
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #58958 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #58958 [ run ] completed with state
|
Contributor
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #58976 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #58976 [ run ] completed with state |
WeiHaocheng
enabled auto-merge (squash)
July 13, 2026 14:35
WeiHaocheng
approved these changes
Jul 13, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
7 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
PR #14397 Change Summary
New Features
contrib/: Coder (sandboxed code-gen + SWE-bench), IterResearch (iterative research with apiary interpreter), and TreeOfThoughtResearchtrtllm-serveendpoints to support the above:POST /v1/tokenize(tokenize-only) andPOST /_control/kv_cache/truncate_tokens(token-level KV cache truncation)Description
I. TRT-LLM Internal Changes
Tiny, server-side only. All exist to give the scaffolding trace-replay client the hooks it needs.
serve/openai_protocol.py(+81)TokenizeRequest/TokenizeResponse/KVCacheTruncateTokensRequestPydantic modelsserve/openai_server.py(+54)POST /v1/tokenizeendpoint (tokenize-only, no executor work); defensivegetattron newKvCacheMetricsfields to avoid HTTP 500 on older bindingsserve/postprocess_handlers.py(+19/-5)/v1/completionschunks now usecmpl-{rsp.id}so clients can join with/perf_metrics, matching the chat-completion pathserve/resource_governor.py(+62)POST /_control/kv_cache/truncate_tokens(token-level KV cache truncation, bypasses chat template) + namespace alias for the legacytruncateroutetests/test_tokenize_endpoint.py(+263)/v1/tokenizeII. Scaffolding Changes (main body)
A. Trace Replay Framework (new)
Capture real agent request trajectories, then replay them offline for benchmarking, KV-cache hit analysis, and parallel-strategy comparison.
scaffolding/execution_trace.py,replay.py,execution_scope.py,utils.pyexamples/scaffolding/trace_replay/—run_trace_replay.py,metrics.py, ananalysis/package (aggregation, cache-hit, branch summary, etc.), plus a sample SWE-bench traceB. New Agents / Controllers (new, all under
contrib/)Each agent comes with its own example runner and
config.yaml.C. MCP Server Reorganization
Promotes
examples/scaffolding/contrib/mcp/→examples/scaffolding/mcp/, drops the oldwebsearch/, and adds new servers:coder/,fetch_webpage/,google_scholar/,google_search/,tavily_search/,wordllama/.D. Open Deep Research Refactor
Updates to the existing
open_deep_researchmodule — primarily reworking message format and system-message caching to play nicely with the replay engine. Example entry point renamed torun_open_deep_research.pyand gains aconfig.yaml.E. Framework Plumbing
Light edits to existing scaffolding modules (
scaffolding_llm.py,controller.py,task.py,worker.py,result.py,benchmark.py,task_collection.py) to wire in trace-capture hooks and replay-engine integration points.F. Cleanup
The final commits revert local experiment residue (scheduler logging tweaks, pareto-experiment server hacks) and delete redundant scripts / configs.
III. Reviewer Take-aways
tensorrt_llm/serve/, all additivecontrib/agents + MCP reorg. Existing users are unaffectedtests/test_tokenize_endpoint.py. The replay framework and new agents do not ship unit tests in this PRTest Coverage
PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.