diff --git a/docs/fern/pages/use-cases/reinforcement-learning/implementation-guide.md b/docs/fern/pages/use-cases/reinforcement-learning/implementation-guide.md index 8ebd27206d6d..efc3b9bdec70 100644 --- a/docs/fern/pages/use-cases/reinforcement-learning/implementation-guide.md +++ b/docs/fern/pages/use-cases/reinforcement-learning/implementation-guide.md @@ -193,7 +193,7 @@ The OpenAI-compatible completion routes provide a cross-backend token-in/token-o | Feature | Request | Response | Notes | |---|---|---|---| | Token input | Set `prompt` to an integer array on `/v1/completions`, or set `nvext.token_data` on a chat or completion request. | Standard completion response | `nvext.token_data` bypasses frontend tokenization. | -| Completion token IDs | Add `"completion_token_ids"` to `nvext.extra_fields`. | `nvext.completion_token_ids` | Requires one prompt and one generated choice. Streaming responses contain token deltas; non-streaming responses contain the concatenated IDs. | +| Completion token IDs | Add `"completion_token_ids"` to `nvext.extra_fields`. | `nvext.completion_token_ids` | Requires one prompt and one generated choice. Streaming responses contain ordered token deltas; non-streaming responses contain the concatenated IDs. A tool parser can buffer and rewrite output, so token IDs on parsed streams are not aligned with each rewritten `delta.content`. | | Completion log probabilities | Set `logprobs` on `/v1/completions`, or set `logprobs: true` and `top_logprobs` on `/v1/chat/completions`. | Standard `choices[].logprobs` | The selected engine must support the requested log probability mode. | | Prompt log probabilities | Set top-level `prompt_logprobs` and add `"prompt_logprobs"` to `nvext.extra_fields`. | `nvext.prompt_logprobs` on the final response | The first prompt position is `null` because it has no preceding-token probability. | | Prefix-cache salt | Set `nvext.cache_salt`. | No response field | vLLM includes the opaque salt in prompt cache keys, separating reuse between requests with different salts. | @@ -201,6 +201,10 @@ The OpenAI-compatible completion routes provide a cross-backend token-in/token-o | Raw engine metadata | Add `"engine_data"` to `nvext.extra_fields`. | `nvext.engine_data` | Backend-specific and not a stable cross-backend schema. Prefer named fields when available. | | SGLang `meta_info` upload | Set `nvext.metadata_upload.url`. | Out-of-band object per choice | Requires an RL-enabled SGLang worker and fsspec support. | +On the legacy chat tool-parser path, a successful streaming response can include a choice-less `nvext` frame after the last generated choice and before the client usage frame and `[DONE]`. Dynamo-aware streaming clients must inspect top-level `nvext` on every frame. Non-streaming requests still return one JSON response because Dynamo aggregates this internal frame into the final top-level `nvext`. + +When Dynamo combines buffered metadata, it appends `completion_token_ids` in stream order. For each other top-level field, the latest supplied value for that field wins; a field that is not supplied again keeps its earlier value. Dynamo replaces `engine_data` as one complete value and does not merge its nested fields. The legacy tool parser requires `n: 1` when `extra_fields` requests `engine_data`, `routed_experts`, or `stop_reason`. This limit does not apply to parser v2 or to request-level metadata such as timing and worker IDs. + See [NVIDIA Request Extensions](../../developer-guide/additional-resources/nvidia-request-extensions-nvext.md) for the complete `nvext` reference. Backend RL flags also select engine-specific behavior: diff --git a/lib/llm/src/http/service/metrics.rs b/lib/llm/src/http/service/metrics.rs index c5f94d810386..4eaea51efdae 100644 --- a/lib/llm/src/http/service/metrics.rs +++ b/lib/llm/src/http/service/metrics.rs @@ -3181,6 +3181,41 @@ mod tests { "internal metrics leaked to client SSE: {wire}" ); + // A choice-less Dynamo metadata frame is client-visible. + let metadata: crate::protocols::openai::chat_completions::NvCreateChatCompletionStreamResponse = + serde_json::from_value(serde_json::json!({ + "id": "chatcmpl-x", "object": "chat.completion.chunk", "created": 1, + "model": "test-model", "choices": [], + "nvext": {"engine_data": {"prompt_token_ids": [1, 2]}} + })) + .unwrap(); + let metadata = Annotated { + id: None, + data: Some(metadata), + event: None, + comment: None, + error: None, + }; + let mut http_queue_guard = None; + let event = process_chat_response_using_event_converter_and_observe_metrics( + EventConverter::from(metadata), + &mut collector, + &mut http_queue_guard, + ReasoningField::default(), + ) + .expect("conversion ok") + .expect("nvext chunk should yield a client event"); + let sse = Sse::new(futures::stream::once(async move { + Ok::<_, std::convert::Infallible>(event) + })); + let body = sse.into_response().into_body(); + let bytes = axum::body::to_bytes(body, usize::MAX).await.unwrap(); + let wire = String::from_utf8_lossy(&bytes); + assert!( + wire.contains("engine_data") && wire.contains("prompt_token_ids"), + "nvext metadata did not reach client SSE: {wire}" + ); + // (2) Payload-only usage chunk (event = payload_usage, carries usage data). let usage: crate::protocols::openai::chat_completions::NvCreateChatCompletionStreamResponse = serde_json::from_value(serde_json::json!({ @@ -3198,10 +3233,11 @@ mod tests { }; let mut http_queue_guard = None; - let result = process_response_using_event_converter_and_observe_metrics( + let result = process_chat_response_using_event_converter_and_observe_metrics( EventConverter::from(payload_usage), &mut collector, &mut http_queue_guard, + ReasoningField::default(), ) .expect("conversion ok"); assert!( diff --git a/lib/llm/src/preprocessor.rs b/lib/llm/src/preprocessor.rs index 3898a4b7d379..0ecc63c6f1d7 100644 --- a/lib/llm/src/preprocessor.rs +++ b/lib/llm/src/preprocessor.rs @@ -79,11 +79,17 @@ use crate::protocols::{ TokenIdType, common::{ OutputOptionsProvider, SamplingOptionsProvider, StopConditionsProvider, - extensions::{AgentHints, NvExtProvider, request_cache_salt, routing_constraints_to_kv}, + extensions::{ + AgentHints, NvExtProvider, merge_response_nvext, request_cache_salt, + routing_constraints_to_kv, + }, }, openai::{ DeltaGeneratorExt, - chat_completions::{NvCreateChatCompletionRequest, NvCreateChatCompletionStreamResponse}, + chat_completions::{ + NvCreateChatCompletionRequest, NvCreateChatCompletionStreamResponse, + scrub_synthetic_chunk_metadata, + }, completions::{NvCreateCompletionRequest, NvCreateCompletionResponse}, embeddings::{NvCreateEmbeddingRequest, NvCreateEmbeddingResponse}, }, @@ -146,6 +152,44 @@ impl ImageDimFetchFailure { } } +fn validate_legacy_jail_nvext_choice_count( + n: u8, + extra_fields: Option<&[String]>, + is_legacy_jail: bool, +) -> Result<()> { + if n <= 1 || !is_legacy_jail { + return Ok(()); + } + + const CHOICE_SPECIFIC_FIELDS: [&str; 3] = ["engine_data", "routed_experts", "stop_reason"]; + if let Some(field) = extra_fields.and_then(|fields| { + fields + .iter() + .find(|field| CHOICE_SPECIFIC_FIELDS.contains(&field.as_str())) + }) { + return Err(invalid_argument_error(format!( + "legacy tool-call parsing requires n = 1 when nvext.extra_fields requests choice-specific field `{field}`" + ))); + } + + Ok(()) +} + +#[derive(Clone, Debug, PartialEq, Eq)] +enum ToolProcessingRoute { + MuseUnified(String), + QwenUnified(&'static str), + ParserV2(String), + LegacyJail(Option), + PassThrough, +} + +impl ToolProcessingRoute { + fn uses_legacy_jail(&self) -> bool { + matches!(self, Self::LegacyJail(_)) + } +} + fn tool_content_part_as_user( part: &ChatCompletionRequestToolMessageContentPart, ) -> Cow<'_, ChatCompletionRequestUserMessageContentPart> { @@ -390,40 +434,6 @@ struct ChoiceReasoningState { parser_finished: bool, } -/// Strip every per-chunk field from a response that is being reused as the -/// envelope for a synthetic end-of-stream chunk. -/// -/// Both end-of-stream flushes build their chunk by cloning the last chunk they -/// saw, because that is the only way to carry the buffered bytes out on a -/// correctly-shaped response. The clone is not a new generation: it produced no -/// tokens and re-channels bytes the parser was already holding. So everything -/// describing the *original* chunk's generation has to go, or it is reported -/// twice: -/// -/// - `event` / `comment` — the annotation channel carries per-chunk payloads -/// such as generated `token_ids`. -/// - `error` — an error annotation must not be replayed on a later chunk. -/// - `usage` / `llm_metrics` — `metrics.rs` sums `chunk_tokens` and samples an -/// ITL point per chunk carrying `llm_metrics`. -/// - `nvext` — per-chunk NVIDIA extensions. `merge_response_nvext` append-merges -/// `completion_token_ids`, so a repeat turns `[42]` into `[42, 42]`, and -/// fields it does not append (such as `prompt_logprobs`) are overwritten. -/// -/// Kept as one function so the two call sites cannot drift apart: they already -/// did, which is how `nvext` survived on both. -fn scrub_synthetic_chunk_metadata( - response: &mut Annotated, -) -> Option<()> { - response.event = None; - response.comment = None; - response.error = None; - let data = response.data.as_mut()?; - data.inner.usage = None; - data.llm_metrics = None; - data.nvext = None; - Some(()) -} - /// Estimates reasoning-token usage from the parser-classified Chat Completion stream. /// /// This is intentionally chunk-granular: if one decoded chunk contains both reasoning @@ -3561,6 +3571,79 @@ impl OpenAIPreprocessor { stream } + fn tool_processing_route( + &self, + request: &NvCreateChatCompletionRequest, + guided_tool_constraint: &crate::protocols::openai::GuidedToolConstraint, + ) -> anyhow::Result { + use crate::protocols::openai::chat_completions::{tool_parser_v2, unified_parser}; + + let uses_tool_call_structural_tag = guided_tool_constraint.uses_structural_tag(); + if let Some(family) = tool_parser_v2::unified_family( + self.tool_call_parser.as_deref(), + self.runtime_config.reasoning_parser.as_deref(), + ) && !uses_tool_call_structural_tag + && matches!( + request.inner.tool_choice.as_ref(), + None | Some(ChatCompletionToolChoiceOption::Auto) + | Some(ChatCompletionToolChoiceOption::None) + ) + { + return Ok(ToolProcessingRoute::MuseUnified(family)); + } + + if let Some(family) = unified_parser::selected_family( + self.tool_call_parser.as_deref(), + self.runtime_config.reasoning_parser.as_deref(), + ) { + return Ok(ToolProcessingRoute::QwenUnified(family)); + } + + let effective_tool_call_parser = self.tool_call_parser.clone().or_else(|| { + self.runtime_config + .reasoning_parser + .as_deref() + .filter(|parser| matches!(*parser, "kimi_k3" | "kimi-k3")) + .map(str::to_string) + }); + let parser_unwraps_all_kimi_k3_responses = effective_tool_call_parser + .as_deref() + .is_some_and(|parser| matches!(parser, "kimi_k3" | "kimi-k3")); + let tool_call_parsing_enabled = Self::tool_call_parsing_enabled(request); + let has_tools = request + .inner + .tools + .as_ref() + .is_some_and(|tools| !tools.is_empty()); + let should_jail = if tool_call_parsing_enabled || parser_unwraps_all_kimi_k3_responses { + Self::should_apply_tool_jail( + effective_tool_call_parser.as_ref(), + request.inner.tool_choice.as_ref(), + has_tools, + )? + } else { + false + }; + + if !should_jail { + return Ok(ToolProcessingRoute::PassThrough); + } + + if let Some(parser_name) = effective_tool_call_parser.as_deref() + && tool_parser_v2::enabled() + && tool_parser_v2::supports_family(parser_name) + && !uses_tool_call_structural_tag + && matches!( + request.inner.tool_choice.as_ref(), + None | Some(ChatCompletionToolChoiceOption::Auto) + ) + { + Ok(ToolProcessingRoute::ParserV2(parser_name.to_string())) + } else { + Ok(ToolProcessingRoute::LegacyJail(effective_tool_call_parser)) + } + } + pub fn postprocessor_parsing_stream( &self, stream: S, @@ -3579,11 +3662,13 @@ impl OpenAIPreprocessor { self.runtime_config.reasoning_parser.as_deref(), uses_tool_call_structural_tag, )?; + let tool_processing_route = self.tool_processing_route(request, &guided_tool_constraint)?; self.postprocessor_parsing_stream_with_constraint( stream, request, prompt_injected_reasoning, guided_tool_constraint, + tool_processing_route, ) } @@ -3593,13 +3678,14 @@ impl OpenAIPreprocessor { request: &NvCreateChatCompletionRequest, prompt_injected_reasoning: bool, guided_tool_constraint: crate::protocols::openai::GuidedToolConstraint, + tool_processing_route: ToolProcessingRoute, ) -> anyhow::Result< impl Stream> + Send + 'static, > where S: Stream> + Send + 'static, { - use crate::protocols::openai::chat_completions::unified_parser; + use crate::protocols::openai::chat_completions::{tool_parser_v2, unified_parser}; let uses_tool_call_structural_tag = guided_tool_constraint.uses_structural_tag(); let defer_reasoning_for_nonempty_content = Self::wants_reasoning_as_content_when_empty(request.chat_template_args.as_ref()); @@ -3634,16 +3720,7 @@ impl OpenAIPreprocessor { // JSON for named/required, native markup for auto/none/structural-tag), so // it does not need the same entry gate. // - if let Some(family) = tool_parser_v2::unified_family( - self.tool_call_parser.as_deref(), - self.runtime_config.reasoning_parser.as_deref(), - ) && !uses_tool_call_structural_tag - && matches!( - request.inner.tool_choice.as_ref(), - None | Some(ChatCompletionToolChoiceOption::Auto) - | Some(ChatCompletionToolChoiceOption::None) - ) - { + if let ToolProcessingRoute::MuseUnified(family) = &tool_processing_route { let tool_definitions = request.inner.tools.as_ref().map(|tools| { tools .iter() @@ -3654,9 +3731,13 @@ impl OpenAIPreprocessor { }) .collect() }); - let unified: Pin + Send>> = Box::pin( - tool_parser_v2::apply_unified_stream(stream, tool_definitions, family, true), - ); + let unified: Pin + Send>> = + Box::pin(tool_parser_v2::apply_unified_stream( + stream, + tool_definitions, + family.clone(), + true, + )); return Ok(Self::apply_unified_response_policies( unified, Self::tool_call_parsing_enabled(request), @@ -3664,10 +3745,7 @@ impl OpenAIPreprocessor { )); } - if let Some(family) = unified_parser::selected_family( - self.tool_call_parser.as_deref(), - self.runtime_config.reasoning_parser.as_deref(), - ) { + if let ToolProcessingRoute::QwenUnified(family) = &tool_processing_route { let tool_definitions = request.inner.tools.as_ref().map(|tools| { tools .iter() @@ -3798,41 +3876,7 @@ impl OpenAIPreprocessor { stream }; - // Check if tools are present and if we should apply jail - let has_tools = request - .inner - .tools - .as_ref() - .is_some_and(|tools| !tools.is_empty()); - - // K3's reasoning-only path still emits XTML response/message wrappers. - // vLLM strips those in its K3 reasoner when no tool parser is active; - // reuse the Rust K3 tool parser as the wrapper decoder so configuring - // only `--dyn-reasoning-parser kimi_k3` remains safe too. - let effective_tool_call_parser = self.tool_call_parser.clone().or_else(|| { - self.runtime_config - .reasoning_parser - .as_deref() - .filter(|parser| matches!(*parser, "kimi_k3" | "kimi-k3")) - .map(str::to_string) - }); - - // A parser describes model syntax, but request semantics decide whether - // tool calls are allowed. Keep K3's wrapper decoder active because that - // model wraps ordinary assistant content in XTML even without tools. - let parser_unwraps_all_kimi_k3_responses = effective_tool_call_parser - .as_deref() - .is_some_and(|parser| matches!(parser, "kimi_k3" | "kimi-k3")); let tool_call_parsing_enabled = Self::tool_call_parsing_enabled(request); - let should_jail = if tool_call_parsing_enabled || parser_unwraps_all_kimi_k3_responses { - Self::should_apply_tool_jail( - effective_tool_call_parser.as_ref(), - request.inner.tool_choice.as_ref(), - has_tools, - )? - } else { - false - }; // Convert OpenAI tools to parser ToolDefinition format before applying jail let tool_definitions = request.inner.tools.as_ref().map(|tools| { @@ -3846,57 +3890,30 @@ impl OpenAIPreprocessor { .collect() }); - // When DYN_ENABLE_EXPERIMENTAL_PARSERS_V2 is set, supported families - // (Qwen3-Coder, DeepSeek-V4) stream through the dynamo-parsers-v2 parser - // instead of the jail, which is never built for them on this path. - // tool_choice=required/named and structural-tag still use the jail's - // Immediate mode, since those rely on guided-decoded JSON rather than the - // native markup the v2 parser reads. See tool_parser_v2::apply_stream. - use crate::protocols::openai::chat_completions::tool_parser_v2; - - // Guided JSON does NOT go to the jail. We installed the grammar that produced - // this output, so the shape is already known: a named choice's payload is the - // argument object itself, and a required choice's is an array of - // `{name, parameters}`. Routing comes from the SHARED constraint predicate, not - // from `tool_choice`, because a Kimi K3 forced request is indistinguishable - // there yet installs no JSON schema at all. That predicate is read once at the - // top of this function, into `guided_tool_streaming`. - - let parser_name = effective_tool_call_parser.as_deref(); - let use_parsers_v2 = tool_parser_v2::enabled() - && parser_name.is_some_and(tool_parser_v2::supports_family) - && !uses_tool_call_structural_tag - && matches!( - request.inner.tool_choice.as_ref(), - None | Some(dynamo_protocols::types::ChatCompletionToolChoiceOption::Auto) - ); - - // Apply jail conditionally let transformed_stream: Pin + Send>> = - if should_jail && use_parsers_v2 { - Box::pin(tool_parser_v2::apply_stream( - stream, - tool_definitions, - parser_name - .expect("use_parsers_v2 implies a parser name") - .to_string(), - )) - } else if should_jail { - // A forced tool_choice installed a JSON grammar, so the jail may release - // calls as they arrive instead of buffering to the closing brace. The - // jail keeps its own native fallback, so a backend that ignores the - // grammar (MiniMax M2 emits XML under `required`) still parses normally. - // Same request-scoped decision the unified path above was given. - Box::pin(Self::apply_tool_calling_jail( - effective_tool_call_parser, - request.inner.tool_choice.clone(), - tool_definitions, - uses_tool_call_structural_tag, - guided_tool_streaming, - stream, - )) - } else { - Box::pin(stream) + match tool_processing_route { + ToolProcessingRoute::ParserV2(parser_name) => Box::pin( + tool_parser_v2::apply_stream(stream, tool_definitions, parser_name), + ), + ToolProcessingRoute::LegacyJail(effective_tool_call_parser) => { + // A forced tool_choice installed a JSON grammar, so the jail may release + // calls as they arrive instead of buffering to the closing brace. The + // jail keeps its own native fallback, so a backend that ignores the + // grammar (MiniMax M2 emits XML under `required`) still parses normally. + // Same request-scoped decision the unified path above was given. + Box::pin(Self::apply_tool_calling_jail( + effective_tool_call_parser, + request.inner.tool_choice.clone(), + tool_definitions, + uses_tool_call_structural_tag, + guided_tool_streaming, + stream, + )) + } + ToolProcessingRoute::PassThrough => Box::pin(stream), + ToolProcessingRoute::MuseUnified(_) | ToolProcessingRoute::QwenUnified(_) => { + unreachable!("unified routes return before legacy response processing") + } }; Ok(Self::apply_tool_call_response_policy( @@ -4420,9 +4437,9 @@ impl OpenAIPreprocessor { /// `Annotated`, runs the moved jail, and /// re-wraps the result. /// - /// `nvext` is not populated on the streaming tool-call path (only the unary - /// aggregator/anthropic paths set it), so the jail never needs to preserve - /// it and re-wrapped chunks carry `nvext: None`. + /// The parser can buffer and rewrite several input chunks before it emits an + /// output. Completion token IDs therefore describe the ordered buffered + /// group, not the rewritten text in one output delta. pub fn apply_tool_calling_jail( tool_call_parser: Option, tool_choice: Option, @@ -4454,12 +4471,18 @@ impl OpenAIPreprocessor { // and `observe_current_osl` takes the latest `output_tokens`. (The // annotation form on data-less usage chunks rides through untouched via // `event`/`comment`.) + // + // `nvext` uses the unary aggregator's merge rules: completion token IDs + // are appended, while the latest supplied value wins for every other + // top-level field. `engine_data` is replaced as one complete value. #[derive(Default)] - struct PendingMetrics { - template: Option, + struct PendingDynamoMetadata { + metrics_template: Option, chunk_tokens: usize, + nvext: Option, + response_template: Option, } - let pending = Arc::new(Mutex::new(PendingMetrics::default())); + let pending = Arc::new(Mutex::new(PendingDynamoMetadata::default())); let pending_in = Arc::clone(&pending); // Per-choice recovery state — allocated only for glm47 since only that @@ -4517,12 +4540,37 @@ impl OpenAIPreprocessor { std::future::ready(!is_error) }); - // dynamo `Annotated` -> jail `Annotated` (buffer llm_metrics) + // dynamo `Annotated` -> jail `Annotated` (buffer Dynamo metadata) let jail_input = stream.map(move |mut a| { - if let Some(metrics) = a.data.as_mut().and_then(|nv| nv.llm_metrics.take()) { - let mut p = pending_in.lock().expect("jail metrics buffer poisoned"); - p.chunk_tokens = p.chunk_tokens.saturating_add(metrics.chunk_tokens); - p.template = Some(metrics); + let has_metadata = a + .data + .as_ref() + .is_some_and(|nv| nv.llm_metrics.is_some() || nv.nvext.is_some()); + if has_metadata { + let mut p = pending_in + .lock() + .expect("jail Dynamo metadata buffer poisoned"); + if let Some(nv) = a.data.as_mut() { + if p.response_template.is_none() { + p.response_template = Some( + dynamo_protocols::types::CreateChatCompletionStreamResponse { + id: nv.inner.id.clone(), + object: nv.inner.object.clone(), + created: nv.inner.created, + model: nv.inner.model.clone(), + choices: Vec::new(), + usage: None, + service_tier: nv.inner.service_tier.clone(), + system_fingerprint: nv.inner.system_fingerprint.clone(), + }, + ); + } + if let Some(metrics) = nv.llm_metrics.take() { + p.chunk_tokens = p.chunk_tokens.saturating_add(metrics.chunk_tokens); + p.metrics_template = Some(metrics); + } + merge_response_nvext(&mut p.nvext, nv.nvext.take()); + } } // Buffer input content only for glm47 (truncation recovery). // Only retain from the last marker onward to bound @@ -4564,7 +4612,7 @@ impl OpenAIPreprocessor { } }); - // jail `Annotated` -> dynamo `Annotated` (re-attach llm_metrics) + // jail `Annotated` -> dynamo `Annotated` (re-attach Dynamo metadata) // The crate encodes the opt-in in WHICH entry point you call, so pick here and // box both arms to one type. `apply_tool_calling_jail` keeps the published // five-argument signature for everyone else. @@ -4593,196 +4641,249 @@ impl OpenAIPreprocessor { jail_input, )) }; - let terminal_error_out = Arc::clone(&terminal_error); - jailed - .flat_map(move |a| { - // Stamp the accumulated metrics onto the next emitted data chunk; - // data-less/synthesized chunks carry it forward (or `None`). - let llm_metrics = a.data.as_ref().and_then(|_| { - let mut p = pending.lock().expect("jail metrics buffer poisoned"); - let chunk_tokens = p.chunk_tokens; - p.chunk_tokens = 0; - p.template.take().map(|mut metrics| { - metrics.chunk_tokens = chunk_tokens; - metrics - }) + let pending_out = Arc::clone(&pending); + let pending_eof = Arc::clone(&pending); + let jailed_output = jailed.flat_map(move |a| { + // Metrics can ride on payload-only usage chunks because the HTTP + // layer observes them before removing the chunk. Client-visible + // nvext must wait for a non-payload-usage output with a choice. + let has_choices = a.data.as_ref().is_some_and(|data| !data.choices.is_empty()); + let is_payload_usage = a.event.as_deref() == Some(ANNOTATION_PAYLOAD_USAGE); + let (llm_metrics, nvext) = a.data.as_ref().map_or((None, None), |_| { + let mut p = pending_out + .lock() + .expect("jail Dynamo metadata buffer poisoned"); + let chunk_tokens = p.chunk_tokens; + p.chunk_tokens = 0; + let metrics = p.metrics_template.take().map(|mut metrics| { + metrics.chunk_tokens = chunk_tokens; + metrics }); - let mut nv_chunk = Annotated { - data: a.data.map(|inner| NvCreateChatCompletionStreamResponse { - inner, - nvext: None, - llm_metrics, - }), - id: a.id, - event: a.event, - comment: a.comment, - error: a.error.map(DynamoError::msg), + let nvext = if has_choices && !is_payload_usage { + p.nvext.take() + } else { + None }; + (metrics, nvext) + }); + let mut nv_chunk = Annotated { + data: a.data.map(|inner| NvCreateChatCompletionStreamResponse { + inner, + nvext, + llm_metrics, + }), + id: a.id, + event: a.event, + comment: a.comment, + error: a.error.map(DynamoError::msg), + }; - // glm47: on finish_reason=length, recover the last incomplete - // block. rfind skips complete blocks so earlier parsed - // calls are never duplicated. Recovered content WILL contain raw - // markup — callers that require "no tool tags in content" must - // filter on finish_reason=length. - // - // TODO: this recovery runs inside apply_tool_calling_jail, which - // the v2 path bypasses (use_parsers_v2 branch above). Adding - // "glm47" to V2_FAMILIES in tool_parser_v2.rs silently disables - // streaming recovery while aggregator.rs keeps running. At that - // point hoist this above the jail/v2 branch — it only needs - // buffered input text + finish_reason, both available there. - // Pass 1 (immutable): compute the recovery tail per choice and - // whether the jail already released it as content on this chunk. - // We collect into a Vec so we can release the immutable borrow on - // nv_chunk before mutating it in pass 2. - let recoveries: Vec = - if is_glm47 { - let mut cr = choice_recovery.lock().expect("choice recovery poisoned"); - nv_chunk - .data - .iter() - .flat_map(|data| data.inner.choices.iter()) - .filter_map(|choice| { - let state = cr.entry(choice.index).or_default(); - if !state.recovered - && let Some(ChatCompletionMessageContent::Text(t)) = - &choice.delta.content - { - state.emitted_text.push_str(t); - // Bound like input_text: retain only the suffix from - // the last marker onward — all the contains(&tail) - // check needs. - let mut keep_from = - match state.emitted_text.rfind(glm47_start.as_str()) { - Some(pos) => pos, - None => state - .emitted_text - .len() - .saturating_sub(glm47_start.len() - 1), - }; - while keep_from > 0 - && !state.emitted_text.is_char_boundary(keep_from) - { - keep_from -= 1; - } - state.emitted_text.drain(..keep_from); - } - if state.recovered - || !matches!( - choice.finish_reason, - Some(dynamo_protocols::types::FinishReason::Length) - ) - { - return None; - } - let tail = state.input_text.rfind(glm47_start.as_str()).and_then( - |pos| { - let t = &state.input_text[pos..]; - if !t.contains(glm47_end.as_str()) { - Some(t.to_string()) - } else { - None - } - }, - )?; - let tail_already_emitted = state.emitted_text.contains(&tail); - state.recovered = true; - Some(PendingRecovery { - choice_idx: choice.index, - tail, - tail_already_emitted, - }) - }) - .collect() - } else { - vec![] - }; - - // Pass 2 (mutable): when the jail already released the tail verbatim, - // suppress the finish chunk's content entirely. The recovery chunk - // carries just the marker-onwards tail, matching the non-streaming - // path (rfind result only, no post-call prose). - for pr in &recoveries { - if !pr.tail_already_emitted { - continue; - } - if let Some(ref mut data) = nv_chunk.data { - for rc in data - .inner - .choices - .iter_mut() - .filter(|c| c.index == pr.choice_idx) + // glm47: on finish_reason=length, recover the last incomplete + // block. rfind skips complete blocks so earlier parsed + // calls are never duplicated. Recovered content WILL contain raw + // markup — callers that require "no tool tags in content" must + // filter on finish_reason=length. + // + // TODO: this recovery runs inside apply_tool_calling_jail, which + // the v2 path bypasses (use_parsers_v2 branch above). Adding + // "glm47" to V2_FAMILIES in tool_parser_v2.rs silently disables + // streaming recovery while aggregator.rs keeps running. At that + // point hoist this above the jail/v2 branch — it only needs + // buffered input text + finish_reason, both available there. + // Pass 1 (immutable): compute the recovery tail per choice and + // whether the jail already released it as content on this chunk. + // We collect into a Vec so we can release the immutable borrow on + // nv_chunk before mutating it in pass 2. + let recoveries: Vec = if is_glm47 { + let mut cr = choice_recovery.lock().expect("choice recovery poisoned"); + nv_chunk + .data + .iter() + .flat_map(|data| data.inner.choices.iter()) + .filter_map(|choice| { + let state = cr.entry(choice.index).or_default(); + if !state.recovered + && let Some(ChatCompletionMessageContent::Text(t)) = + &choice.delta.content { - // The jail released the truncated block verbatim as content - // on this chunk, potentially preceded by post-call prose. - // glm47's parser drops post-call prose deliberately, so - // suppress the whole content here and let the recovery - // chunk carry just the marker-onwards tail — matching batch. - rc.delta.content = None; + state.emitted_text.push_str(t); + // Bound like input_text: retain only the suffix from + // the last marker onward — all the contains(&tail) + // check needs. + let mut keep_from = match state.emitted_text.rfind(glm47_start.as_str()) + { + Some(pos) => pos, + None => state + .emitted_text + .len() + .saturating_sub(glm47_start.len() - 1), + }; + while keep_from > 0 && !state.emitted_text.is_char_boundary(keep_from) { + keep_from -= 1; + } + state.emitted_text.drain(..keep_from); } + if state.recovered + || !matches!( + choice.finish_reason, + Some(dynamo_protocols::types::FinishReason::Length) + ) + { + return None; + } + let tail = + state + .input_text + .rfind(glm47_start.as_str()) + .and_then(|pos| { + let t = &state.input_text[pos..]; + if !t.contains(glm47_end.as_str()) { + Some(t.to_string()) + } else { + None + } + })?; + let tail_already_emitted = state.emitted_text.contains(&tail); + state.recovered = true; + Some(PendingRecovery { + choice_idx: choice.index, + tail, + tail_already_emitted, + }) + }) + .collect() + } else { + vec![] + }; + + // Pass 2 (mutable): when the jail already released the tail verbatim, + // suppress the finish chunk's content entirely. The recovery chunk + // carries just the marker-onwards tail, matching the non-streaming + // path (rfind result only, no post-call prose). + for pr in &recoveries { + if !pr.tail_already_emitted { + continue; + } + if let Some(ref mut data) = nv_chunk.data { + for rc in data + .inner + .choices + .iter_mut() + .filter(|c| c.index == pr.choice_idx) + { + // The jail released the truncated block verbatim as content + // on this chunk, potentially preceded by post-call prose. + // glm47's parser drops post-call prose deliberately, so + // suppress the whole content here and let the recovery + // chunk carry just the marker-onwards tail — matching batch. + rc.delta.content = None; } } + } - // Pass 3: emit a recovery chunk per affected choice carrying just - // the truncated tail (marker onwards, no post-call prose). - let recovery_chunks: Vec<_> = recoveries - .into_iter() - .filter_map(|pr| { - let PendingRecovery { - choice_idx, tail, .. - } = pr; - tracing::warn!( - choice_index = choice_idx, - recovered_bytes = tail.len(), - "glm47 streaming: partial emitted as content \ + // Pass 3: emit a recovery chunk per affected choice carrying just + // the truncated tail (marker onwards, no post-call prose). + let recovery_chunks: Vec<_> = recoveries + .into_iter() + .filter_map(|pr| { + let PendingRecovery { + choice_idx, tail, .. + } = pr; + tracing::warn!( + choice_index = choice_idx, + recovered_bytes = tail.len(), + "glm47 streaming: partial emitted as content \ on length finish" - ); - let mut rec = nv_chunk.clone(); - rec.id = None; - rec.event = None; - rec.comment = None; - rec.error = None; - let rd = rec.data.as_mut()?; - rd.inner.usage = None; - rd.llm_metrics = None; - rd.inner.choices.retain(|c| c.index == choice_idx); - for rc in &mut rd.inner.choices { - rc.delta.content = - Some(ChatCompletionMessageContent::Text(tail.clone())); - rc.delta.tool_calls = None; - rc.finish_reason = None; - rc.logprobs = None; - } - Some(rec) - }) - .collect(); + ); + let mut rec = nv_chunk.clone(); + rec.id = None; + scrub_synthetic_chunk_metadata(&mut rec); + let rd = rec.data.as_mut()?; + rd.inner.choices.retain(|c| c.index == choice_idx); + for rc in &mut rd.inner.choices { + rc.delta.content = Some(ChatCompletionMessageContent::Text(tail.clone())); + rc.delta.tool_calls = None; + rc.finish_reason = None; + rc.logprobs = None; + } + Some(rec) + }) + .collect(); - futures::stream::iter(recovery_chunks.into_iter().chain(std::iter::once(nv_chunk))) - }) - // See the `terminal_error` comment above: once the upstream error was - // latched, `take_while` drops everything the jail still emits from that - // point on (its finalize output has no way to know the request already - // failed), and `chain` substitutes the latched error as the stream's - // final and only item from there. When no error occurred, `terminal_error` - // is never populated, `take_while` never stops early, and `chain`'s - // `filter_map` drops the `None` it reads back — this stage is then a - // no-op passthrough. - .take_while(move |_| { - let stop = terminal_error_out + futures::stream::iter(recovery_chunks.into_iter().chain(std::iter::once(nv_chunk))) + }); + + // Once an upstream error is latched, drop any output the jail synthesized + // while it observed the shortened stream. The wrapper below then emits only + // the original error and discards all pending metadata and usage. + let terminal_error_out = Arc::clone(&terminal_error); + let jailed_output = jailed_output.take_while(move |_| { + let stop = terminal_error_out + .lock() + .expect("jail terminal error poisoned") + .is_some(); + std::future::ready(!stop) + }); + + let with_eof_metadata = async_stream::stream! { + tokio::pin!(jailed_output); + while let Some(response) = jailed_output.next().await { + yield response; + } + + let terminal_error = { + terminal_error .lock() .expect("jail terminal error poisoned") - .is_some(); - std::future::ready(!stop) - }) - .chain( - stream::once(async move { - terminal_error + .take() + }; + if let Some(error) = terminal_error { + { + let mut p = pending_eof .lock() - .expect("jail terminal error poisoned") - .take() - }) - .filter_map(std::future::ready), - ) + .expect("jail Dynamo metadata buffer poisoned"); + p.metrics_template = None; + p.chunk_tokens = 0; + p.nvext = None; + p.response_template = None; + } + yield error; + return; + } + + let eof_metadata = { + let mut p = pending_eof + .lock() + .expect("jail Dynamo metadata buffer poisoned"); + let chunk_tokens = p.chunk_tokens; + p.chunk_tokens = 0; + let llm_metrics = p.metrics_template.take().map(|mut metrics| { + metrics.chunk_tokens = chunk_tokens; + metrics + }); + let nvext = p.nvext.take(); + if llm_metrics.is_none() && nvext.is_none() { + None + } else { + p.response_template.take().map(|inner| Annotated { + data: Some(NvCreateChatCompletionStreamResponse { + inner, + nvext, + llm_metrics, + }), + id: None, + event: None, + comment: None, + error: None, + }) + } + }; + if let Some(response) = eof_metadata { + yield response; + } + }; + + Self::hold_usage_until_stream_end(with_eof_metadata) } /// Whether the selected tool-call or reasoning parser depends on the @@ -5421,7 +5522,7 @@ impl OpenAIPreprocessor { // See `scrub_synthetic_chunk_metadata`: this chunk produced no // tokens, so every per-chunk field from the envelope it was // cloned from has to be dropped rather than reported twice. - scrub_synthetic_chunk_metadata(&mut response)?; + scrub_synthetic_chunk_metadata(&mut response); let data = response.data.as_mut()?; // Rebuild the choice list from the flushed indices rather than // reusing the envelope's own choices: with `n > 1` the last @@ -5647,7 +5748,8 @@ impl OpenAIPreprocessor { /// been emitted. A truncated upstream stream can end without /// `finish_reason`, so reasoning recovery happens at EOF; forwarding usage /// immediately would put that recovered content after the chunk clients - /// treat as the stream trailer. + /// treat as the stream trailer. A transport error discards the pending + /// trailer because the response did not complete successfully. fn hold_usage_until_stream_end( stream: S, ) -> impl Stream> + Send @@ -5657,7 +5759,14 @@ impl OpenAIPreprocessor { async_stream::stream! { tokio::pin!(stream); let mut pending_usage = None; + let mut transport_failed = false; while let Some(response) = stream.next().await { + if response.error.is_some() { + transport_failed = true; + pending_usage = None; + yield response; + continue; + } let is_usage_only = response.data.as_ref().is_some_and(|data| { data.inner.choices.is_empty() && data.inner.usage.is_some() }); @@ -5669,7 +5778,7 @@ impl OpenAIPreprocessor { yield response; } } - if let Some(usage) = pending_usage { + if !transport_failed && let Some(usage) = pending_usage { yield usage; } } @@ -5801,7 +5910,7 @@ impl OpenAIPreprocessor { // is a clone of an already-counted chunk, so it goes through // the shared scrub rather than repeating a partial copy of // it here. - scrub_synthetic_chunk_metadata(&mut response)?; + scrub_synthetic_chunk_metadata(&mut response); let data = response.data.as_mut()?; let mut template = data.inner.choices.first()?.clone(); template.delta.role = None; @@ -5933,6 +6042,16 @@ impl &mut common_request, prompt_injected_reasoning, )?; + let tool_processing_route = + self.tool_processing_route(&request, &guided_tool_constraint)?; + validate_legacy_jail_nvext_choice_count( + request.inner.n.unwrap_or(1), + request + .nvext + .as_ref() + .and_then(|nvext| nvext.extra_fields.as_deref()), + tool_processing_route.uses_legacy_jail(), + )?; tracing::trace!(request = ?common_request, prompt_injected_reasoning, "Pre-processed request"); let trace_state = crate::request_trace::build_request_end_trace_state( @@ -5991,6 +6110,7 @@ impl &request, prompt_injected_reasoning, guided_tool_constraint, + tool_processing_route, )?; let transformed_stream = Self::normalize_chat_stream_roles(transformed_stream); @@ -6334,6 +6454,32 @@ mod tests { ); } + #[test] + fn legacy_jail_rejects_multiple_choices_for_choice_specific_nvext() { + let engine_data = vec!["engine_data".to_string()]; + let request_level = vec!["timing".to_string(), "worker_id".to_string()]; + + assert!(validate_legacy_jail_nvext_choice_count(2, Some(&engine_data), true).is_err()); + assert!(validate_legacy_jail_nvext_choice_count(1, Some(&engine_data), true).is_ok()); + assert!(validate_legacy_jail_nvext_choice_count(2, Some(&request_level), true).is_ok()); + + for route in [ + ToolProcessingRoute::MuseUnified("muse_glimmer".to_string()), + ToolProcessingRoute::QwenUnified("qwen3"), + ToolProcessingRoute::ParserV2("qwen3_coder".to_string()), + ToolProcessingRoute::PassThrough, + ] { + assert!( + validate_legacy_jail_nvext_choice_count( + 2, + Some(&engine_data), + route.uses_legacy_jail(), + ) + .is_ok() + ); + } + } + fn chat_stream_chunk( index: u32, role: Option, @@ -8373,17 +8519,17 @@ mod tests { ); } - struct UnreachableCompletionBackend; + struct UnreachableBackend; #[async_trait] impl AsyncEngine, ManyOut>, Error> - for UnreachableCompletionBackend + for UnreachableBackend { async fn generate( &self, _request: SingleIn, ) -> Result>, Error> { - panic!("over-budget completion must be rejected before backend dispatch") + panic!("request must be rejected before backend dispatch") } } @@ -8415,7 +8561,7 @@ mod tests { }; let next: Arc< dyn AsyncEngine, ManyOut>, Error>, - > = Arc::new(UnreachableCompletionBackend); + > = Arc::new(UnreachableBackend); let result = Operator::generate(preprocessor.as_ref(), PipelineContext::new(request), next).await; @@ -8428,6 +8574,48 @@ mod tests { assert_eq!(dynamo_err.error_type(), ErrorType::InvalidArgument); } + #[tokio::test] + async fn chat_operator_rejects_invalid_legacy_jail_request_before_dispatch() { + let mut mdc = ModelDeploymentCard::load_from_disk( + "tests/data/sample-models/mock-llama-3.1-8b-instruct", + None, + ) + .unwrap(); + mdc.runtime_config.tool_call_parser = Some("hermes".to_string()); + let preprocessor = OpenAIPreprocessor::new(mdc).unwrap(); + let request: NvCreateChatCompletionRequest = serde_json::from_value(serde_json::json!({ + "model": "test-model", + "messages": [{"role": "user", "content": "What is the weather?"}], + "n": 2, + "max_tokens": 4, + "tools": [{ + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the weather", + "parameters": {"type": "object", "properties": {}} + } + }], + "tool_choice": "auto", + "nvext": {"extra_fields": ["engine_data"]} + })) + .unwrap(); + let next: Arc< + dyn AsyncEngine, ManyOut>, Error>, + > = Arc::new(UnreachableBackend); + + let result = + Operator::generate(preprocessor.as_ref(), PipelineContext::new(request), next).await; + let Err(err) = result else { + panic!("invalid legacy-jail request should fail admission"); + }; + let dynamo_err = err + .downcast_ref::() + .expect("error should preserve the DynamoError type"); + assert_eq!(dynamo_err.error_type(), ErrorType::InvalidArgument); + assert!(dynamo_err.to_string().contains("legacy tool-call parsing")); + } + fn test_prompt_formatter(template: &str) -> Arc { let template: dynamo_renderer::ChatTemplate = serde_json::from_value(serde_json::json!({ "chat_template": template diff --git a/lib/llm/src/protocols/openai/chat_completions.rs b/lib/llm/src/protocols/openai/chat_completions.rs index ac0c51cdafc8..41b0f523c7fa 100644 --- a/lib/llm/src/protocols/openai/chat_completions.rs +++ b/lib/llm/src/protocols/openai/chat_completions.rs @@ -285,6 +285,21 @@ pub struct NvCreateChatCompletionStreamResponse { pub llm_metrics: Option, } +/// Synthetic chunks reuse a real response envelope but consume no backend data. +/// Clear copied transport and per-chunk fields so clients do not count them twice. +pub(crate) fn scrub_synthetic_chunk_metadata( + response: &mut Annotated, +) -> Option<()> { + response.event = None; + response.comment = None; + response.error = None; + let data = response.data.as_mut()?; + data.inner.usage = None; + data.llm_metrics = None; + data.nvext = None; + Some(()) +} + /// Build one synthetic stream choice from an existing response template. /// /// Both streaming tool-call paths use this constructor when an engine omits a @@ -299,8 +314,6 @@ pub(super) fn stream_choice_chunk_from_template( finish_reason: Option, ) -> Annotated { let mut response = template.clone(); - response.inner.usage = None; - response.llm_metrics = None; #[allow(deprecated)] let choice = ChatChoiceStream { index, @@ -316,13 +329,15 @@ pub(super) fn stream_choice_chunk_from_template( logprobs: None, }; response.inner.choices = vec![choice]; - Annotated { + let mut chunk = Annotated { data: Some(response), id: None, event: None, comment: None, error: None, - } + }; + scrub_synthetic_chunk_metadata(&mut chunk); + chunk } /// Implements `NvExtProvider` for `NvCreateChatCompletionRequest`, diff --git a/lib/llm/src/protocols/openai/chat_completions/aggregator.rs b/lib/llm/src/protocols/openai/chat_completions/aggregator.rs index 934bb9c16241..fba52d0bae75 100644 --- a/lib/llm/src/protocols/openai/chat_completions/aggregator.rs +++ b/lib/llm/src/protocols/openai/chat_completions/aggregator.rs @@ -1777,7 +1777,33 @@ mod tests { annotated_delta2.data.as_mut().expect("delta data").nvext = Some(serde_json::json!({ "stop_reason": 128001 })); - let stream = Box::pin(stream::iter(vec![annotated_delta1, annotated_delta2])); + let mut metadata = annotated_delta2.clone(); + let metadata_data = metadata.data.as_mut().expect("metadata data"); + metadata_data.inner.choices.clear(); + metadata_data.nvext = Some(serde_json::json!({ + "engine_data": { + "prompt_token_ids": [1, 2], + "completion_token_ids": [10, 11], + "completion_logprobs": [-0.1, -0.2], + } + })); + + let mut usage = metadata.clone(); + let usage_data = usage.data.as_mut().expect("usage data"); + usage_data.nvext = None; + usage_data.inner.usage = Some(dynamo_protocols::types::CompletionUsage { + prompt_tokens: 2, + completion_tokens: 2, + total_tokens: 4, + ..Default::default() + }); + + let stream = Box::pin(stream::iter(vec![ + annotated_delta1, + annotated_delta2, + metadata, + usage, + ])); let response = DeltaAggregator::apply(stream, ParsingOptions::default()) .await .expect("aggregate stream"); @@ -1785,10 +1811,19 @@ mod tests { assert_eq!( response.nvext, Some(serde_json::json!({ - "engine_data": { "trace_id": "abc" }, + "engine_data": { + "prompt_token_ids": [1, 2], + "completion_token_ids": [10, 11], + "completion_logprobs": [-0.1, -0.2], + }, "stop_reason": 128001, })) ); + assert_eq!(response.inner.choices.len(), 1); + assert_eq!( + response.inner.usage.expect("aggregated usage").total_tokens, + 4 + ); } #[allow(deprecated)] diff --git a/lib/llm/src/protocols/openai/chat_completions/tool_parser_v2.rs b/lib/llm/src/protocols/openai/chat_completions/tool_parser_v2.rs index 9c64f4734191..5ceed0754c3b 100644 --- a/lib/llm/src/protocols/openai/chat_completions/tool_parser_v2.rs +++ b/lib/llm/src/protocols/openai/chat_completions/tool_parser_v2.rs @@ -1121,7 +1121,8 @@ mod tests { )]); let mut finished = HashSet::new(); let mut tool_emitted = HashSet::from([3]); - let template = usage_chunk().data.expect("usage response data"); + let mut template = usage_chunk().data.expect("usage response data"); + template.nvext = Some(serde_json::json!({"completion_token_ids": [42]})); let responses = finish_unterminated_choices(&mut states, &mut finished, &mut tool_emitted, &template); @@ -1140,6 +1141,10 @@ mod tests { response.llm_metrics.is_none(), "terminal chunk must not repeat LLM metrics" ); + assert!( + response.nvext.is_none(), + "terminal chunk must not repeat nvext" + ); assert_eq!(response.inner.choices.len(), 1); assert_eq!(response.inner.choices[0].index, 3); assert_eq!( @@ -1270,7 +1275,10 @@ mod tests { .chunks(8) .map(|b| chunk(std::str::from_utf8(b).unwrap(), false)) .collect(); - chunks.push(usage_chunk()); + let mut usage = usage_chunk(); + usage.data.as_mut().expect("usage data").nvext = + Some(serde_json::json!({"completion_token_ids": [42]})); + chunks.push(usage); let out: Vec<_> = apply_unified_stream(stream::iter(chunks), None, "muse_glimmer".to_string(), true) @@ -1308,6 +1316,24 @@ mod tests { finish_position < usage_position, "synthesized finish chunk must precede usage" ); + let finish_data = out[finish_position] + .data + .as_ref() + .expect("synthesized finish data"); + assert!( + finish_data.nvext.is_none(), + "synthetic terminal chunk must not duplicate nvext" + ); + assert_eq!( + out.iter() + .filter(|response| response + .data + .as_ref() + .is_some_and(|data| data.nvext.is_some())) + .count(), + 1, + "only the real usage chunk may carry nvext" + ); } // Reasoning + content, no tools: the backstop must NOT invent a finish_reason diff --git a/lib/llm/tests/test_streaming_tool_parsers.rs b/lib/llm/tests/test_streaming_tool_parsers.rs index 0660f4f5a8f6..7955c8759f5b 100644 --- a/lib/llm/tests/test_streaming_tool_parsers.rs +++ b/lib/llm/tests/test_streaming_tool_parsers.rs @@ -27,8 +27,12 @@ across backends. */ use dynamo_llm::preprocessor::OpenAIPreprocessor; +use dynamo_llm::protocols::common::metrics::LLMMetricAnnotation; use dynamo_llm::protocols::openai::chat_completions::NvCreateChatCompletionStreamResponse; -use dynamo_protocols::types::{ChatChoiceStream, ChatCompletionMessageContent, FinishReason}; +use dynamo_protocols::types::{ + ChatChoiceStream, ChatCompletionMessageContent, ChatCompletionToolChoiceOption, + CompletionUsage, FinishReason, +}; use dynamo_runtime::protocols::annotated::Annotated; use futures::{Stream, StreamExt, stream}; use std::pin::Pin; @@ -1916,87 +1920,125 @@ mod tests { ); } - // The jail moved to dynamo-parsers and operates on the shared - // `Create` payload, so the boundary adapter (apply_tool_calling_jail) must - // buffer the dynamo-only typed `llm_metrics` and re-attach it. This asserts - // the buffered chunk_tokens sum and latest output_tokens survive the jail on - // a tool-call stream (they'd all be None without the buffer/re-attach). - #[tokio::test] - async fn jail_preserves_llm_metrics_across_buffered_tool_call() { - use dynamo_llm::protocols::common::metrics::LLMMetricAnnotation; - use dynamo_llm::protocols::openai::chat_completions::NvCreateChatCompletionStreamResponse; - use dynamo_protocols::types::{ - ChatChoiceStream, ChatCompletionMessageContent, ChatCompletionStreamResponseDelta, - CreateChatCompletionStreamResponse, Role, - }; - use dynamo_runtime::protocols::annotated::Annotated; - use futures::StreamExt; - - fn chunk( - text: &str, - chunk_tokens: usize, - output_tokens: usize, - ) -> Annotated { - #[allow(deprecated)] - let choice = ChatChoiceStream { - index: 0, - delta: ChatCompletionStreamResponseDelta { - role: Some(Role::Assistant), - content: Some(ChatCompletionMessageContent::Text(text.to_string())), - tool_calls: None, - function_call: None, - refusal: None, - reasoning_content: None, - }, - finish_reason: None, - logprobs: None, - }; - Annotated { - data: Some(NvCreateChatCompletionStreamResponse { - inner: CreateChatCompletionStreamResponse { - id: "id".to_string(), - object: "chat.completion.chunk".to_string(), - created: 0, - model: "m".to_string(), - choices: vec![choice], - usage: None, - service_tier: None, - system_fingerprint: None, - }, - nvext: None, - llm_metrics: Some(LLMMetricAnnotation { - input_tokens: 7, - output_tokens, - chunk_tokens, - cached_tokens: None, - prefill_worker_id: None, - prefill_dp_rank: None, - prefill_worker_type: None, - decode_worker_id: None, - decode_dp_rank: None, - decode_worker_type: None, - tokenize_latency: None, - detokenize_total_latency: None, - detokenize_count: None, - ..Default::default() - }), - }), - id: None, - event: None, - comment: None, - error: None, - } - } + fn metadata_chunk( + text: &str, + finish_reason: Option, + chunk_tokens: usize, + output_tokens: usize, + nvext: serde_json::Value, + ) -> Annotated { + let mut chunk = make_chunk(text, finish_reason); + let data = chunk.data.as_mut().expect("chunk data"); + data.nvext = Some(nvext); + data.llm_metrics = Some(LLMMetricAnnotation { + input_tokens: 7, + output_tokens, + chunk_tokens, + cached_tokens: None, + image_count: 0, + video_count: 0, + audio_count: 0, + image_tokens: None, + prefill_worker_id: None, + prefill_dp_rank: None, + prefill_worker_type: None, + decode_worker_id: None, + decode_dp_rank: None, + decode_worker_type: None, + tokenize_latency: None, + detokenize_total_latency: None, + detokenize_count: None, + }); + chunk + } + + fn metadata_usage_chunk() -> Annotated { + let mut chunk = make_chunk("", None); + let data = chunk.data.as_mut().expect("usage data"); + data.inner.choices.clear(); + data.inner.usage = Some(CompletionUsage { + prompt_tokens: 7, + completion_tokens: 7, + total_tokens: 14, + ..Default::default() + }); + chunk.event = Some(dynamo_llm::preprocessor::ANNOTATION_PAYLOAD_USAGE.to_string()); + chunk + } - // Hermes tool call split across two metric-bearing chunks -> the jail - // buffers both, then emits one tool-call chunk. + fn metadata_only_chunk( + finish_reason: Option, + chunk_tokens: usize, + output_tokens: usize, + nvext: serde_json::Value, + ) -> Annotated { + let mut chunk = metadata_chunk("", finish_reason, chunk_tokens, output_tokens, nvext); + let choice = &mut chunk.data.as_mut().expect("metadata data").inner.choices[0]; + choice.delta.content = None; + choice.delta.role = None; + chunk + } + + fn assert_buffered_metrics( + out: &[Annotated], + expected_chunk_tokens: usize, + expected_output_tokens: usize, + ) { + let total_chunk_tokens: usize = out + .iter() + .filter_map(|a| a.data.as_ref().and_then(|d| d.llm_metrics.as_ref())) + .map(|m| m.chunk_tokens) + .sum(); + assert_eq!(total_chunk_tokens, expected_chunk_tokens); + let max_osl = out + .iter() + .filter_map(|a| a.data.as_ref().and_then(|d| d.llm_metrics.as_ref())) + .map(|m| m.output_tokens) + .max(); + assert_eq!(max_osl, Some(expected_output_tokens)); + } + + fn metadata_outputs( + out: &[Annotated], + ) -> Vec<&NvCreateChatCompletionStreamResponse> { + out.iter() + .filter_map(|response| response.data.as_ref().filter(|data| data.nvext.is_some())) + .collect() + } + + // Immediate mode holds all generated choices until EOF. A usage chunk can + // therefore leave the jail before the parsed tool call, but it must not take + // the client-visible nvext that belongs to the generated choice. + #[tokio::test] + async fn jail_keeps_nvext_pending_across_usage_chunk() { + let engine_data = serde_json::json!({ + "prompt_token_ids": [1, 2], + "completion_token_ids": [10, 11, 12, 13, 14, 15, 16], + "completion_logprobs": [-0.1, -0.2, -0.3, -0.4, -0.5, -0.6, -0.7], + }); let chunks = vec![ - chunk("\n{\"name\": \"get_weather\", \"arg", 3, 3), - chunk("uments\": {\"location\": \"SF\"}}\n", 4, 7), + metadata_chunk( + "[{\"name\": \"get_weather\", \"parameters\": {\"loc", + None, + 3, + 3, + serde_json::json!({ "completion_token_ids": [10, 11, 12] }), + ), + metadata_chunk( + "ation\": \"SF\"}}", + None, + 4, + 7, + serde_json::json!({ + "completion_token_ids": [13, 14, 15, 16], + "engine_data": engine_data, + }), + ), + metadata_usage_chunk(), ]; let out: Vec<_> = OpenAIPreprocessor::apply_tool_calling_jail( Some("hermes".to_string()), - None, + Some(ChatCompletionToolChoiceOption::Required), None, false, false, @@ -2005,25 +2047,118 @@ mod tests { .collect() .await; - let total_chunk_tokens: usize = out + assert_buffered_metrics(&out, 7, 7); + let usage = out .iter() - .filter_map(|a| a.data.as_ref().and_then(|d| d.llm_metrics.as_ref())) - .map(|m| m.chunk_tokens) - .sum(); + .find_map(|a| { + a.data + .as_ref() + .filter(|d| d.inner.choices.is_empty() && d.inner.usage.is_some()) + }) + .expect("usage output"); + assert!(usage.nvext.is_none(), "usage output must not consume nvext"); + + let metadata = metadata_outputs(&out); + assert_eq!(metadata.len(), 1, "nvext must be emitted exactly once"); + assert!(!metadata[0].inner.choices.is_empty()); + let nvext = metadata[0].nvext.as_ref().expect("nvext"); assert_eq!( - total_chunk_tokens, 7, - "buffered chunk_tokens (3+4) must survive the jail; got {total_chunk_tokens}" + nvext["completion_token_ids"], + serde_json::json!([10, 11, 12, 13, 14, 15, 16]) ); - let max_osl = out - .iter() - .filter_map(|a| a.data.as_ref().and_then(|d| d.llm_metrics.as_ref())) - .map(|m| m.output_tokens) - .max(); + assert_eq!(nvext["engine_data"], engine_data); + } + + #[tokio::test] + async fn jail_flushes_terminal_nvext_before_client_usage() { + let final_engine_data = serde_json::json!({ + "prompt_token_ids": [1, 2], + "completion_token_ids": [30, 31], + "completion_logprobs": [-0.1, -0.2], + }); + let mut usage = metadata_usage_chunk(); + usage.event = None; + let early = metadata_only_chunk( + None, + 1, + 1, + serde_json::json!({ + "completion_token_ids": [30], + "engine_data": {"phase": "partial"}, + "worker_id": {"decode_worker_id": 9}, + }), + ); + let terminal = metadata_only_chunk( + Some(FinishReason::Stop), + 1, + 2, + serde_json::json!({ + "completion_token_ids": [31], + "engine_data": final_engine_data, + "timing": {"total_ms": 12.5}, + "prompt_logprobs": [null, {"token": 1}], + }), + ); + let chunks = vec![early, terminal, usage]; + let out: Vec<_> = OpenAIPreprocessor::apply_tool_calling_jail( + Some("hermes".to_string()), + Some(ChatCompletionToolChoiceOption::Required), + None, + false, + true, + Box::pin(futures::stream::iter(chunks)), + ) + .collect() + .await; + + assert_buffered_metrics(&out, 2, 2); + let metadata = metadata_outputs(&out); + assert_eq!(metadata.len(), 1, "nvext must be emitted exactly once"); + assert!(metadata[0].inner.choices.is_empty()); + let nvext = metadata[0].nvext.as_ref().expect("nvext"); + assert_eq!(nvext["completion_token_ids"], serde_json::json!([30, 31])); + assert_eq!(nvext["engine_data"], final_engine_data); + assert_eq!( + nvext["worker_id"], + serde_json::json!({"decode_worker_id": 9}) + ); + assert_eq!(nvext["timing"], serde_json::json!({"total_ms": 12.5})); assert_eq!( - max_osl, - Some(7), - "final cumulative output_tokens must survive the jail" + nvext["prompt_logprobs"], + serde_json::json!([null, {"token": 1}]) ); + + assert!(out.last().is_some_and(|response| { + response + .data + .as_ref() + .is_some_and(|data| data.inner.usage.is_some()) + })); + } + + #[tokio::test] + async fn jail_discards_pending_metadata_after_transport_error() { + let terminal = metadata_only_chunk( + Some(FinishReason::Stop), + 1, + 1, + serde_json::json!({"engine_data": {"request_id": "failed"}}), + ); + let chunks = vec![terminal, Annotated::from_error("transport failed")]; + + let out: Vec<_> = OpenAIPreprocessor::apply_tool_calling_jail( + Some("hermes".to_string()), + Some(ChatCompletionToolChoiceOption::Required), + None, + false, + false, + Box::pin(futures::stream::iter(chunks)), + ) + .collect() + .await; + + assert!(out.iter().any(|response| response.error.is_some())); + assert!(metadata_outputs(&out).is_empty()); } } @@ -2106,13 +2241,15 @@ fn content_texts(chunks: &[Annotated]) -> // The latch must fire on the first finish chunk and not re-fire on the terminal one. #[tokio::test] async fn test_glm47_streaming_truncated_data_less_terminal_chunk() { - let chunks = vec![ + let mut chunks = vec![ make_glm47_chunk( Some("get_weathercityBos"), None, ), make_glm47_chunk(None, Some(FinishReason::Length)), ]; + chunks[1].data.as_mut().unwrap().nvext = + Some(serde_json::json!({"completion_token_ids": [42]})); let out = run_glm47_jail(chunks).await; let texts = content_texts(&out); @@ -2136,6 +2273,19 @@ async fn test_glm47_streaming_truncated_data_less_terminal_chunk() { 1, "finish_reason=length must appear exactly once in output" ); + let nvext_chunks: Vec<_> = out + .iter() + .filter_map(|a| a.data.as_ref().and_then(|d| d.nvext.as_ref())) + .collect(); + assert_eq!( + nvext_chunks.len(), + 1, + "the synthetic recovery chunk must not repeat nvext" + ); + assert_eq!( + nvext_chunks[0]["completion_token_ids"], + serde_json::json!([42]) + ); } // Prose follows a complete tool call, then a truncated second call — all in one jailed