feat(llm): return chat completion token ids via nvext

docs/components/frontend/nvext.md

@@ -33,9 +33,9 @@ Include `nvext` as a top-level field alongside standard OpenAI-compatible fields
 | `use_raw_prompt` | `bool` | `None` | Preprocessor | Bypasses the prompt template and passes the prompt directly to the tokenizer. |
 | `annotations` | `string[]` | `None` | Preprocessor | Triggers out-of-band information in the SSE stream via the `event:` field. |
 | `backend_instance_id` | `u64` | `None` | Router | Routes the request to a specific backend instance. |
-| `token_data` | `u32[]` | `None` | Preprocessor | Pre-tokenized prompt tokens. When provided with `backend_instance_id`, tokenization is skipped. |
+| `token_data` | `u32[]` | `None` | Preprocessor | Pre-tokenized prompt tokens. When provided, tokenization is skipped. `backend_instance_id` remains an independent routing hint. |
 | `max_thinking_tokens` | `u32` | `None` | Backend | Maximum thinking tokens allowed (passed through to backends). |
-| `extra_fields` | `string[]` | `None` | Response builder | Fields to include in the response `nvext`. Supported: `"worker_id"`, `"timing"`, `"routed_experts"`, `"engine_data"`, `"stop_reason"`. |
+| `extra_fields` | `string[]` | `None` | Response builder | Fields to include in the response `nvext`. Supported: `"worker_id"`, `"timing"`, `"routed_experts"`, `"engine_data"`, `"stop_reason"`, `"completion_token_ids"`. |
 | `prefill_worker_id` | `u64` | `None` | Router | Routes the request to a specific prefill worker (disaggregated serving). |
 | `decode_worker_id` | `u64` | `None` | Router | Routes the request to a specific decode worker (disaggregated serving). |
 | `agent_context` | object | `None` | Preprocessor | Passive session and trajectory identity for agent traces. See [Agent Context](#agent-context) below and [Agent Tracing](../../agents/agent-tracing.md). |
@@ -208,6 +208,7 @@ When the client requests response metadata via `extra_fields`, the response incl
 | `routed_experts` | `extra_fields: ["routed_experts"]` | Routed expert capture payload returned by SGLang-backed requests. |
 | `engine_data` | `extra_fields: ["engine_data"]` | Opaque backend-provided engine metadata. |
 | `stop_reason` | `extra_fields: ["stop_reason"]` | Backend-specific matched stop condition, returned under `nvext` because it is not part of the OpenAI completions schema. Dynamo currently serves this as a response-level field for single-choice requests; supporting `n > 1` will require an indexed per-choice shape. |
+| `completion_token_ids` | `extra_fields: ["completion_token_ids"]` | Generated token IDs accumulated across the chat-completions response and emitted on the final chunk. |
 | `token_ids` | Automatic (GAIE Stage 1) | Tokenized prompt for reuse in Stage 2 query-only mode. |
 
 ### Example response `nvext`

lib/llm/src/protocols/openai/chat_completions/delta.rs

@@ -10,7 +10,7 @@ use crate::{
         openai::{
             convert_backend_top_logprobs,
             delta_common::{self, DeltaGeneratorOptions},
-            nvext::NvExtProvider,
+            nvext::{NvExtProvider, NvExtResponse},
             token_to_utf8_bytes,
         },
     },
@@ -59,6 +59,9 @@ pub struct DeltaGenerator {
     options: DeltaGeneratorOptions,
     /// Request tracker for per-request metrics (shared with PreprocessedRequest).
     tracker: Arc<RequestTracker>,
+    /// Accumulated output token IDs across chunks, emitted on the final chunk
+    /// when `nvext.extra_fields` includes `completion_token_ids`.
+    accumulated_completion_token_ids: Vec<TokenIdType>,
 }
 
 impl DeltaGenerator {
@@ -75,6 +78,7 @@ impl DeltaGenerator {
             msg_counter: 0,
             options,
             tracker,
+            accumulated_completion_token_ids: Vec::new(),
         }
     }
 
@@ -257,6 +261,11 @@ impl crate::protocols::openai::DeltaGeneratorExt<NvCreateChatCompletionStreamRes
 
         self.usage.completion_tokens += token_length;
 
+        if self.options.response_fields.completion_token_ids && !delta.token_ids.is_empty() {
+            self.accumulated_completion_token_ids
+                .extend_from_slice(&delta.token_ids);
+        }
+
         // If backend provides completion_usage, use it to update usage stats
         // This is critical for prompt embeddings where prompt_tokens comes from
         // the embedding sequence length computed by the worker
@@ -311,27 +320,58 @@ impl crate::protocols::openai::DeltaGeneratorExt<NvCreateChatCompletionStreamRes
         // `NvExtResponseFieldSelection` (see `nvext.rs`). Both chat and
         // completions delta generators go through the same helper so the gating
         // rules stay in one place.
-        if let Some(nvext_response) = self.options.response_fields.build_response_nvext(
-            Some(&self.tracker),
-            delta.disaggregated_params.as_ref(),
-            finish_reason.is_some(),
-            delta.engine_data,
-            stop_reason,
-        ) && let Ok(nvext_json) = serde_json::to_value(&nvext_response)
+        if let Some(mut nvext_response) = self
+            .options
+            .response_fields
+            .build_response_nvext(
+                Some(&self.tracker),
+                delta.disaggregated_params.as_ref(),
+                finish_reason.is_some(),
+                delta.engine_data,
+                stop_reason,
+            )
+            .or_else(|| {
+                if self.options.response_fields.completion_token_ids && finish_reason.is_some() {
+                    Some(NvExtResponse {
+                        worker_id: None,
+                        timing: None,
+                        token_ids: None,
+                        routed_experts: None,
+                        engine_data: None,
+                        stop_reason: None,
+                        completion_token_ids: None,
+                    })
+                } else {
+                    None
+                }
+            })
         {
-            stream_response.nvext = Some(nvext_json);
-            if let Some(ref info) = nvext_response.worker_id {
-                tracing::debug!(
-                    "Injected worker_id into chat completion nvext: prefill={:?}, decode={:?}",
-                    info.prefill_worker_id,
-                    info.decode_worker_id
-                );
+            if self.options.response_fields.completion_token_ids && finish_reason.is_some() {
+                nvext_response.completion_token_ids =
+                    Some(self.accumulated_completion_token_ids.clone());
             }
-            if let Some(ref tokens) = nvext_response.token_ids {
-                tracing::debug!(
-                    "Injected token_ids into chat completion nvext: {} tokens",
-                    tokens.len()
-                );
+
+            if let Ok(nvext_json) = serde_json::to_value(&nvext_response) {
+                stream_response.nvext = Some(nvext_json);
+                if let Some(ref info) = nvext_response.worker_id {
+                    tracing::debug!(
+                        "Injected worker_id into chat completion nvext: prefill={:?}, decode={:?}",
+                        info.prefill_worker_id,
+                        info.decode_worker_id
+                    );
+                }
+                if let Some(ref tokens) = nvext_response.token_ids {
+                    tracing::debug!(
+                        "Injected token_ids into chat completion nvext: {} tokens",
+                        tokens.len()
+                    );
+                }
+                if let Some(ref tokens) = nvext_response.completion_token_ids {
+                    tracing::debug!(
+                        "Injected completion_token_ids into chat completion nvext: {} tokens",
+                        tokens.len()
+                    );
+                }
             }
         }
 
@@ -615,6 +655,48 @@ mod tests {
         assert!(nvext_json.get("routed_experts").is_none());
     }
 
+    #[test]
+    fn test_completion_token_ids_extra_field_emits_accumulated_ids_on_final_chunk() {
+        let request =
+            create_test_request_with_extra_fields(vec!["completion_token_ids".to_string()]);
+        let mut generator = request.response_generator("req-completion-ids".to_string());
+
+        let mut first_output = final_backend_output();
+        first_output.token_ids = vec![7];
+        first_output.tokens = vec![Some("A".to_string())];
+        first_output.text = Some("A".to_string());
+        first_output.finish_reason = None;
+        first_output.disaggregated_params = None;
+
+        let first_response = generator
+            .choice_from_postprocessor(first_output)
+            .expect("first choice generation");
+        assert!(
+            first_response.nvext.is_none(),
+            "completion_token_ids should be emitted only on the final chunk"
+        );
+
+        let mut final_output = final_backend_output();
+        final_output.token_ids = vec![8, 9];
+        final_output.tokens = vec![Some("B".to_string()), Some("C".to_string())];
+        final_output.text = Some("BC".to_string());
+        final_output.disaggregated_params = None;
+
+        let final_response = generator
+            .choice_from_postprocessor(final_output)
+            .expect("final choice generation");
+
+        let nvext_json = final_response
+            .nvext
+            .expect("nvext present for completion_token_ids request");
+        assert_eq!(
+            nvext_json.get("completion_token_ids"),
+            Some(&serde_json::json!([7, 8, 9]))
+        );
+        assert!(nvext_json.get("token_ids").is_none());
+        assert!(nvext_json.get("routed_experts").is_none());
+    }
+
     #[test]
     fn test_routed_experts_extra_field_emits_routed_experts() {
         use crate::protocols::openai::nvext::NvExt;

lib/llm/src/protocols/openai/nvext.rs

@@ -129,6 +129,11 @@ pub struct NvExtResponse {
     /// If `n > 1` is supported here, this needs an indexed/per-choice shape.
     #[serde(skip_serializing_if = "Option::is_none")]
     pub stop_reason: Option<serde_json::Value>,
+
+    /// Output token IDs generated by the engine.
+    /// Populated when the client requests `extra_fields: ["completion_token_ids"]`.
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub completion_token_ids: Option<Vec<u32>>,
 }
 
 pub(crate) fn merge_response_nvext(
@@ -165,6 +170,7 @@ pub struct NvExtResponseFieldSelection {
     pub routed_experts: bool,
     pub engine_data: bool,
     pub stop_reason: bool,
+    pub completion_token_ids: bool,
 }
 
 impl NvExtResponseFieldSelection {
@@ -182,6 +188,7 @@ impl NvExtResponseFieldSelection {
                     "routed_experts" => selection.routed_experts = true,
                     "engine_data" => selection.engine_data = true,
                     "stop_reason" => selection.stop_reason = true,
+                    "completion_token_ids" => selection.completion_token_ids = true,
                     _ => {}
                 }
             }
@@ -215,6 +222,8 @@ impl NvExtResponseFieldSelection {
     /// - `timing` requires the selection flag, `finish_reason_present == true`, **and** a tracker.
     /// - `engine_data` requires the selection flag **and** a non-`None` `engine_data_from_backend`.
     /// - `stop_reason` requires the selection flag **and** a non-`None` `stop_reason_from_backend`.
+    /// - `completion_token_ids` is accumulated by the chat-completions delta generator
+    ///   and attached to the final chunk after this helper returns.
     pub fn build_response_nvext(
         &self,
         tracker: Option<&std::sync::Arc<crate::protocols::common::timing::RequestTracker>>,
@@ -280,6 +289,7 @@ impl NvExtResponseFieldSelection {
             routed_experts,
             engine_data,
             stop_reason,
+            completion_token_ids: None,
         })
     }
 }
@@ -331,7 +341,7 @@ pub struct NvExt {
     /// Extra fields to be included in the response's nvext
     /// This is a list of field names that should be populated in the response
     /// Supported fields include "worker_id", "timing", "routed_experts", "engine_data",
-    /// "stop_reason", which map to fields in NvExtResponse.
+    /// "stop_reason", and "completion_token_ids", which map to fields in NvExtResponse.
     #[serde(default, skip_serializing_if = "Option::is_none")]
     #[builder(default, setter(strip_option))]
     pub extra_fields: Option<Vec<String>>,
@@ -764,6 +774,22 @@ mod tests {
         );
     }
 
+    #[test]
+    fn test_nvext_response_field_selection_completion_token_ids_only() {
+        let nvext = NvExt::builder()
+            .extra_fields(vec!["completion_token_ids".to_string()])
+            .build()
+            .unwrap();
+
+        assert_eq!(
+            NvExtResponseFieldSelection::from_nvext(Some(&nvext)),
+            NvExtResponseFieldSelection {
+                completion_token_ids: true,
+                ..Default::default()
+            }
+        );
+    }
+
     // Helpers for build_response_nvext tests -----------------------------
 
     fn sel_all_false() -> NvExtResponseFieldSelection {
@@ -966,6 +992,7 @@ mod tests {
             routed_experts: true,
             engine_data: false,
             stop_reason: false,
+            completion_token_ids: false,
         };
         let tracker = tracker_with_prefill_worker();
         let params = disagg_params_full();
@@ -1003,6 +1030,7 @@ mod tests {
                 routed_experts: true,
                 engine_data: false,
                 stop_reason: false,
+                completion_token_ids: false,
             }
         );
     }