fix(llm): preserve nvext through tool-call jail

[jthomson04]

Overview

Preserve response-level nvext metadata when the legacy streaming tool-call jail buffers and rewrites OpenAI chat chunks.

Summary

The jail unwraps NvCreateChatCompletionStreamResponse before parsing. It previously rebuilt emitted chunks with nvext: None, which removed requested metadata such as engine_data, completion token IDs, timing, worker IDs, and routed experts even when the backend returned it correctly.

The faulty path is present in Dynamo 1.3, Dynamo 1.4.1, and current main.

Details

• Buffer llm_metrics, nvext, and one clean response envelope only when an input contains metadata. Reuse the legacy jail's terminal-error latch as the single failure state.
• Merge completion_token_ids by ordered append. For every other top-level field, the latest supplied value for that field wins. Replace engine_data as one complete opaque value.
• Drain llm_metrics onto any emitted data chunk because the HTTP layer observes metrics before it removes payload-only usage chunks.
• Drain nvext onto a non-payload-usage output with a choice. If valid metadata remains at successful EOF, emit it once in a clean choice-less chunk before client usage and [DONE].
• Keep non-streaming behavior as one JSON response. DeltaAggregator absorbs the internal EOF metadata chunk into the final top-level nvext.
• Discard pending metadata after an annotated transport error. A failed response is not a valid training response.
• Select one tool-processing route after guided decoding and reuse it for request admission and postprocessing. Before backend dispatch, reject only legacy jailed requests with n > 1 when nvext.extra_fields requests engine_data, routed_experts, or stop_reason. Request-level metadata, Muse and Qwen unified routes, parser v2, and pass-through remain supported with n > 1.
• Share the synthetic-chunk scrub helper between legacy recovery, EOF clones, and parser v2. Synthetic chunks clear annotations, errors, usage, llm_metrics, and nvext.
• Document that parsed streams preserve raw completion token IDs once and in order, but parser rewriting prevents per-delta token-to-text alignment. Dynamo-aware streaming clients must inspect top-level nvext on every chunk.

This is independent of #13588 and #13640. Those PRs produce or expose training metadata; this fix prevents the legacy jail from deleting it.

Downstream context:

• feat(vllm-model): consume native Dynamo token data NVIDIA-NeMo/Gym#1784
• refactor(dynamo): move native token response handling to Gym NVIDIA-NeMo/RL#3763

After merge, please backport this fix to the maintained release/1.4.0 and release/1.3.0 branches. NeMo-RL currently uses ai-dynamo[vllm]==1.3.0.post1.

Where should the reviewer start?

OpenAIPreprocessor::apply_tool_calling_jail in lib/llm/src/preprocessor.rs, then jail_flushes_terminal_nvext_before_client_usage, jail_discards_pending_metadata_after_transport_error, and the two merge tests in lib/llm/tests/test_streaming_tool_parsers.rs.

Related Issues

• Confirmed — no related issue

Validation

Validated at 2875787bd81017d92eafd9beea8209766f0d87ba on canonical main 9a48f622e0169376a7f8cc0e697677d467e788c7.

• cargo test -p dynamo-llm --test test_streaming_tool_parsers — 45 passed.
• cargo test -p dynamo-llm --lib protocols::openai::chat_completions::tool_parser_v2::tests — 20 passed.
• cargo test -p dynamo-llm --lib protocols::openai::chat_completions::unified_parser::tests — 61 passed.
• cargo test -p dynamo-llm --test postprocessor_parsing_stream — 92 passed.
• The module-qualified operator regression passed and proved invalid legacy-jail metadata requests fail with InvalidArgument before the unreachable backend is called.
• Focused direct-SSE and DeltaAggregator metadata tests passed.
• cargo fmt --all -- --check — passed.
• cargo clippy --no-deps -p dynamo-llm --all-targets -- -D warnings — passed.
• Scoped pre-commit on all seven changed Rust and documentation files — passed.

---

Summary by CodeRabbit

• Bug Fixes
  • Improved streaming tool-call handling to preserve NVIDIA extension metadata across response chunks.
  • Prevented duplicated metrics, annotations, errors, usage details, and token IDs during recovery and reasoning transitions.
  • Ensured completion token IDs and engine metadata remain accurate in streamed responses.
• Tests
  • Expanded streaming coverage for metadata merging and duplicate-prevention behavior.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 66455368-bd27-44f2-90f9-97b133b1fdd5

📥 Commits

Reviewing files that changed from the base of the PR and between 10e1fc1 and 18335e0.

📒 Files selected for processing (2)

• lib/llm/src/preprocessor.rs
• lib/llm/tests/test_streaming_tool_parsers.rs

Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review.

---

Walkthrough

The tool-call jail now preserves and merges llm_metrics and nvext across buffered response chunks. Synthetic recovery and flush chunks use shared metadata scrubbing to prevent duplicated per-chunk metadata. Streaming tests verify merged token IDs and preserved engine data.

Changes

Response metadata preservation

Layer / File(s)	Summary
Buffered metadata accumulation and emission lib/llm/src/preprocessor.rs, lib/llm/tests/test_streaming_tool_parsers.rs	The tool-call jail accumulates and merges llm_metrics and nvext, restores them after processing, and emits merged token IDs and engine data. Tests verify one merged nvext value.
Synthetic chunk metadata scrubbing lib/llm/src/preprocessor.rs	Parser recovery and end-of-stream synthetic chunks use the shared scrubber to remove copied per-chunk metadata.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Merge Risk: ⚪ Minimal · up to 18335

This change preserves response metadata while buffered tool-call chunks are rewritten. The fix is localized and validated by the reported tests and checks; no actionable merge-blocking risk remains after normal review.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 75.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 4 functions across 2 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title clearly and concisely describes the main change: preserving nvext metadata through the LLM tool-call jail.
Description check	✅ Passed	The description is complete and follows the repository template. It includes an overview, implementation details, reviewer starting points, related-issue status, and validation results.

• Fix all pre-merge checks with AI

---

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

[devin-ai-integration]

✅ Devin Review: No Issues Found

Devin Review analyzed this PR and found no bugs or issues to report.

[jthomson04]

Automated review of the nvext-through-the-jail change. The direction is right, but I think the core mechanism — "stamp the accumulated nvext onto the next emitted data chunk" — has a case where it doesn't deliver the metadata it's meant to preserve, plus a merge-semantics issue and a missing end-of-stream flush. Details inline.

One finding falls outside the diff, so it can't be an inline comment:

---

lib/llm/src/protocols/openai/chat_completions.rs:300-302 — stream_choice_chunk_from_template is the one synthetic-chunk site this PR didn't consolidate, and it still leaks nvext.

The PR's doc comment now says the scrub helper is "kept as one function so the call sites cannot drift apart", but stream_choice_chunk_from_template clears inner.usage and llm_metrics on the cloned envelope and not nvext — the exact drift scrub_synthetic_chunk_metadata exists to prevent.

Repro: DYN_ENABLE_EXPERIMENTAL_PARSERS_V2=1, Qwen3-Coder family, stream_options.include_usage unset (so no empty-choices chunk ever refreshes template), engine drops the terminal finish_reason. tool_parser_v2::apply_stream's EOF backstop calls finish_unterminated_choices(..., template) where template is a clone of the last content chunk, nvext included. The trailing ToolCalls chunk therefore re-emits completion_token_ids: [16], DeltaAggregator append-merges it, and the aggregated response gets [..., 16, 16] — plus a duplicate engine_data/timing on the wire for streaming clients.

Note tool_parser_v2.rs:1065 and :1118 already assert llm_metrics.is_none() on these chunks; the nvext counterpart is missing.

---

🤖 Generated with Claude Code

[jthomson04]

The out-of-diff parser-v2 finding is also fixed in f3bf4f3. I moved the existing synthetic-chunk scrub helper into the shared chat-completion module with crate-only visibility. Both legacy recovery/EOF clones and parser-v2 stream_choice_chunk_from_template now clear annotations, errors, usage, llm_metrics, and nvext. The existing parser-v2 synthetic-terminal test now supplies nvext and proves that the synthetic response clears it; this test failed before the shared scrub fix.

[jthomson04]

/ok to test f3bf4f3

[github-actions]

🌿 Fern Docs Preview: https://nvidia-preview-01a044fb-c443-76ef-ab35-9e9b29a40c18.docs.buildwithfern.com/dynamo/dev

[jthomson04]

/ok to test 57acc75

[rmccorm4]

Left one small comment, and otherwise needs merge conflicts fixed

[jthomson04]

/ok to test 2875787

[rmccorm4]

one nit, otherwise LGTM: https://github.com/ai-dynamo/dynamo/pull/13730/changes#r3875689803