Skip to main content

rig_core/providers/openai/completion/
streaming.rs

1use http::Request;
2use serde::{Deserialize, Serialize};
3use serde_json::json;
4use tracing::{Level, enabled, info_span};
5
6use crate::completion::{CompletionError, CompletionRequest, GetTokenUsage};
7use crate::http_client::HttpClientExt;
8use crate::json_utils::{self, merge};
9use crate::providers::internal::openai_chat_completions_compatible::{
10    self, CompatibleChoiceData, CompatibleChunk, CompatibleFinishReason, CompatibleStreamProfile,
11    CompatibleToolCallChunk,
12};
13use crate::providers::openai::completion::{
14    CompletionModelOptions, GenericCompletionModel, OpenAICompatibleProvider, Usage,
15};
16use crate::streaming;
17
18// ================================================================
19// OpenAI Completion Streaming API
20// ================================================================
21#[derive(Default, Deserialize, Debug)]
22pub(crate) struct StreamingFunction {
23    pub(crate) name: Option<String>,
24    #[serde(
25        default,
26        deserialize_with = "crate::json_utils::deserialize_json_string_or_value"
27    )]
28    pub(crate) arguments: Option<String>,
29}
30
31#[derive(Deserialize, Debug)]
32pub(crate) struct StreamingToolCall {
33    // Optional in several compatible dialects (e.g. Mistral); missing means
34    // a single in-flight tool call.
35    #[serde(default)]
36    pub(crate) index: usize,
37    pub(crate) id: Option<String>,
38    #[serde(default, deserialize_with = "json_utils::null_or_default")]
39    pub(crate) function: StreamingFunction,
40}
41
42impl From<&StreamingToolCall> for CompatibleToolCallChunk {
43    fn from(value: &StreamingToolCall) -> Self {
44        Self {
45            index: value.index,
46            id: value.id.clone(),
47            name: value.function.name.clone(),
48            arguments: value.function.arguments.clone(),
49        }
50    }
51}
52
53fn deserialize_delta_content<'de, D>(deserializer: D) -> Result<Option<String>, D::Error>
54where
55    D: serde::Deserializer<'de>,
56{
57    // Some compatible providers (e.g. Mistral's reasoning models) stream
58    // delta content as an array of content parts rather than a string.
59    let value = Option::<serde_json::Value>::deserialize(deserializer)?;
60    Ok(value.and_then(|value| match value {
61        serde_json::Value::String(text) => Some(text),
62        serde_json::Value::Array(parts) => {
63            let text = crate::providers::openai::completion::joined_text_parts(&parts);
64            (!text.is_empty()).then_some(text)
65        }
66        _ => None,
67    }))
68}
69
70#[derive(Deserialize, Debug)]
71struct StreamingDelta {
72    #[serde(default, deserialize_with = "deserialize_delta_content")]
73    content: Option<String>,
74    #[serde(default)]
75    reasoning_content: Option<String>,
76    // Not part of the official OpenAI API; some compatible providers (e.g.
77    // Groq) send the same payload under `reasoning`. A separate field rather
78    // than a serde alias so a delta carrying BOTH keys is not a
79    // duplicate-field error that drops the whole chunk.
80    #[serde(default)]
81    reasoning: Option<String>,
82    #[serde(default, deserialize_with = "json_utils::null_or_vec")]
83    tool_calls: Vec<StreamingToolCall>,
84    #[serde(default, deserialize_with = "json_utils::null_or_vec")]
85    reasoning_details: Vec<serde_json::Value>,
86}
87
88#[derive(Deserialize, Debug, PartialEq)]
89#[serde(rename_all = "snake_case")]
90pub enum FinishReason {
91    ToolCalls,
92    Stop,
93    ContentFilter,
94    Length,
95    #[serde(untagged)]
96    Other(String), // This will handle the deprecated function_call
97}
98
99#[derive(Deserialize, Debug)]
100struct StreamingChoice {
101    delta: StreamingDelta,
102    finish_reason: Option<FinishReason>,
103}
104
105#[derive(Deserialize, Debug)]
106struct StreamingCompletionChunk<U = Usage> {
107    id: Option<String>,
108    model: Option<String>,
109    choices: Vec<StreamingChoice>,
110    usage: Option<U>,
111}
112
113/// Final streaming response. `U` is the provider's streaming usage payload
114/// ([`Usage`] for OpenAI itself; providers with richer usage accounting, e.g.
115/// Mistral and DeepSeek, substitute their own via
116/// [`OpenAICompatibleProvider::StreamingUsage`].
117#[derive(Clone, Serialize, Deserialize)]
118pub struct StreamingCompletionResponse<U = Usage> {
119    pub usage: U,
120}
121
122impl<U> GetTokenUsage for StreamingCompletionResponse<U>
123where
124    U: GetTokenUsage,
125{
126    fn token_usage(&self) -> crate::completion::Usage {
127        self.usage.token_usage()
128    }
129}
130
131impl<Ext, H> GenericCompletionModel<Ext, H>
132where
133    crate::client::Client<Ext, H>: HttpClientExt + Clone + 'static,
134    Ext: crate::client::Provider
135        + OpenAICompatibleProvider
136        + Clone
137        + crate::wasm_compat::WasmCompatSend
138        + 'static,
139{
140    pub(crate) async fn stream(
141        &self,
142        completion_request: CompletionRequest,
143    ) -> Result<
144        streaming::StreamingCompletionResponse<StreamingCompletionResponse<Ext::StreamingUsage>>,
145        CompletionError,
146    > {
147        let preamble = completion_request.preamble.clone();
148        let options = CompletionModelOptions {
149            strict_tools: self.strict_tools,
150            tool_result_array_content: self.tool_result_array_content,
151            prompt_caching: self.prompt_caching,
152        };
153        let mut request = self.client.ext().build_completion_request(
154            self.model.clone(),
155            completion_request,
156            options,
157        )?;
158        self.client.ext().prepare_request(&mut request)?;
159
160        // Deliberately the configured model, not the per-request override:
161        // Azure's deployment URL is pinned to the model handle.
162        let path = self.client.ext().completion_path(&self.model);
163        let resolved_model = request.model.clone();
164        let mut request_as_json = serde_json::to_value(request)?;
165
166        // `merge` is shallow, so include_usage is inserted into any
167        // caller-supplied stream_options rather than merged over it: the
168        // caller's keys survive and the usage chunk is still requested.
169        if Ext::STREAM_INCLUDE_USAGE {
170            match request_as_json.get_mut("stream_options") {
171                Some(serde_json::Value::Object(options)) => {
172                    options
173                        .entry("include_usage")
174                        .or_insert(serde_json::Value::Bool(true));
175                }
176                Some(_) => {}
177                None => {
178                    request_as_json = merge(
179                        request_as_json,
180                        json!({"stream_options": {"include_usage": true}}),
181                    );
182                }
183            }
184        }
185        request_as_json = merge(request_as_json, json!({"stream": true}));
186        self.client
187            .ext()
188            .finalize_request_body_with_options(&mut request_as_json, options)?;
189
190        if enabled!(Level::TRACE) {
191            tracing::trace!(
192                target: "rig::completions",
193                "OpenAI Chat Completions streaming completion request: {}",
194                serde_json::to_string_pretty(&request_as_json)?
195            );
196        }
197
198        let req_body = serde_json::to_vec(&request_as_json)?;
199
200        let req = self
201            .client
202            .post(&path)?
203            .body(req_body)
204            .map_err(|e| CompletionError::HttpError(e.into()))?;
205
206        let span = if tracing::Span::current().is_disabled() {
207            info_span!(
208                target: "rig::completions",
209                "chat",
210                gen_ai.operation.name = "chat",
211                gen_ai.provider.name = Ext::PROVIDER_NAME,
212                gen_ai.request.model = resolved_model,
213                gen_ai.system_instructions = preamble,
214                gen_ai.response.id = tracing::field::Empty,
215                gen_ai.response.model = tracing::field::Empty,
216                gen_ai.usage.output_tokens = tracing::field::Empty,
217                gen_ai.usage.input_tokens = tracing::field::Empty,
218                gen_ai.usage.cache_read.input_tokens = tracing::field::Empty,
219                gen_ai.input.messages = tracing::field::Empty,
220                gen_ai.output.messages = tracing::field::Empty,
221            )
222        } else {
223            tracing::Span::current()
224        };
225
226        let client = self.client.clone();
227
228        tracing::Instrument::instrument(
229            openai_chat_completions_compatible::send_compatible_streaming_request(
230                client,
231                req,
232                OpenAICompatibleProfile::<Ext, Ext::StreamingUsage> {
233                    provider: self.client.ext().clone(),
234                    emits_complete_single_chunk_tool_calls:
235                        Ext::EMITS_COMPLETE_SINGLE_CHUNK_TOOL_CALLS,
236                    usage: std::marker::PhantomData,
237                },
238            ),
239            span,
240        )
241        .await
242    }
243}
244
245#[derive(Clone, Copy, Default)]
246struct OpenAICompatibleProfile<Ext = crate::providers::openai::OpenAICompletionsExt, U = Usage> {
247    provider: Ext,
248    emits_complete_single_chunk_tool_calls: bool,
249    usage: std::marker::PhantomData<U>,
250}
251
252impl<Ext, U> CompatibleStreamProfile for OpenAICompatibleProfile<Ext, U>
253where
254    Ext: OpenAICompatibleProvider + Clone + crate::wasm_compat::WasmCompatSend,
255    U: Clone
256        + Default
257        + GetTokenUsage
258        + serde::de::DeserializeOwned
259        + crate::wasm_compat::WasmCompatSend
260        + Unpin
261        + 'static,
262{
263    type Usage = U;
264    type Detail = serde_json::Value;
265    type FinalResponse = StreamingCompletionResponse<U>;
266
267    fn normalize_chunk(
268        &self,
269        data: &str,
270    ) -> Result<Option<CompatibleChunk<Self::Usage, Self::Detail>>, CompletionError> {
271        let data = match serde_json::from_str::<StreamingCompletionChunk<U>>(data) {
272            Ok(data) => data,
273            Err(error) => {
274                tracing::error!(?error, message = data, "Failed to parse SSE message");
275                return Ok(None);
276            }
277        };
278
279        Ok(Some(
280            openai_chat_completions_compatible::normalize_first_choice_chunk(
281                data.id,
282                data.model,
283                data.usage,
284                &data.choices,
285                |choice| CompatibleChoiceData {
286                    // `function_call` is the deprecated pre-tools finish reason
287                    // some compatible providers still emit for tool calls.
288                    finish_reason: match &choice.finish_reason {
289                        Some(FinishReason::ToolCalls) => CompatibleFinishReason::ToolCalls,
290                        Some(FinishReason::Other(other)) if other == "function_call" => {
291                            CompatibleFinishReason::ToolCalls
292                        }
293                        _ => CompatibleFinishReason::Other,
294                    },
295                    text: choice.delta.content.clone(),
296                    reasoning: choice
297                        .delta
298                        .reasoning_content
299                        .clone()
300                        .or_else(|| choice.delta.reasoning.clone()),
301                    tool_calls: openai_chat_completions_compatible::tool_call_chunks(
302                        &choice.delta.tool_calls,
303                    ),
304                    details: choice.delta.reasoning_details.clone(),
305                },
306            ),
307        ))
308    }
309
310    fn build_final_response(&self, usage: Self::Usage) -> Self::FinalResponse {
311        StreamingCompletionResponse { usage }
312    }
313
314    fn decorate_tool_call(
315        &self,
316        detail: &Self::Detail,
317        tool_calls: &mut std::collections::HashMap<usize, crate::streaming::RawStreamingToolCall>,
318    ) {
319        self.provider
320            .decorate_streaming_tool_call(detail, tool_calls);
321    }
322
323    fn uses_distinct_tool_call_eviction(&self) -> bool {
324        true
325    }
326
327    fn emits_complete_single_chunk_tool_calls(&self) -> bool {
328        self.emits_complete_single_chunk_tool_calls
329    }
330}
331
332pub async fn send_compatible_streaming_request<T>(
333    http_client: T,
334    req: Request<Vec<u8>>,
335) -> Result<streaming::StreamingCompletionResponse<StreamingCompletionResponse>, CompletionError>
336where
337    T: HttpClientExt + Clone + 'static,
338{
339    openai_chat_completions_compatible::send_compatible_streaming_request(
340        http_client,
341        req,
342        OpenAICompatibleProfile::<crate::providers::openai::OpenAICompletionsExt, Usage>::default(),
343    )
344    .await
345}
346
347#[cfg(test)]
348mod tests {
349    use super::*;
350    use crate::providers::internal::openai_chat_completions_compatible::test_support::{
351        assert_zero_arg_tool_call_is_emitted, sse_bytes_from_data_lines,
352    };
353
354    #[test]
355    fn test_streaming_function_deserialization() {
356        let json = r#"{"name": "get_weather", "arguments": "{\"location\":\"Paris\"}"}"#;
357        let function: StreamingFunction = serde_json::from_str(json).unwrap();
358        assert_eq!(function.name, Some("get_weather".to_string()));
359        assert_eq!(
360            function.arguments.as_ref().unwrap(),
361            r#"{"location":"Paris"}"#
362        );
363    }
364
365    #[test]
366    fn test_streaming_function_object_arguments() {
367        // Some OpenAI-compatible gateways send `arguments` as a JSON object
368        // instead of the spec-mandated JSON-encoded string. Accept it by
369        // re-serializing to the string form rather than dropping the chunk.
370        let json = r#"{"name": "list_dir", "arguments": {}}"#;
371        let function: StreamingFunction = serde_json::from_str(json).unwrap();
372        assert_eq!(function.name, Some("list_dir".to_string()));
373        assert_eq!(function.arguments.as_ref().unwrap(), "{}");
374
375        let json = r#"{"name": "get_weather", "arguments": {"city": "London"}}"#;
376        let function: StreamingFunction = serde_json::from_str(json).unwrap();
377        assert_eq!(function.arguments.as_ref().unwrap(), r#"{"city":"London"}"#);
378    }
379
380    #[test]
381    fn test_streaming_function_null_arguments() {
382        let json = r#"{"name": "list_dir", "arguments": null}"#;
383        let function: StreamingFunction = serde_json::from_str(json).unwrap();
384        assert!(function.arguments.is_none());
385
386        let json = r#"{"name": "list_dir"}"#;
387        let function: StreamingFunction = serde_json::from_str(json).unwrap();
388        assert!(function.arguments.is_none());
389    }
390
391    #[test]
392    fn test_streaming_tool_call_deserialization() {
393        let json = r#"{
394            "index": 0,
395            "id": "call_abc123",
396            "function": {
397                "name": "get_weather",
398                "arguments": "{\"city\":\"London\"}"
399            }
400        }"#;
401        let tool_call: StreamingToolCall = serde_json::from_str(json).unwrap();
402        assert_eq!(tool_call.index, 0);
403        assert_eq!(tool_call.id, Some("call_abc123".to_string()));
404        assert_eq!(tool_call.function.name, Some("get_weather".to_string()));
405    }
406
407    #[test]
408    fn test_streaming_tool_call_partial_deserialization() {
409        // Partial tool calls have no name and partial arguments
410        let json = r#"{
411            "index": 0,
412            "id": null,
413            "function": {
414                "name": null,
415                "arguments": "Paris"
416            }
417        }"#;
418        let tool_call: StreamingToolCall = serde_json::from_str(json).unwrap();
419        assert_eq!(tool_call.index, 0);
420        assert!(tool_call.id.is_none());
421        assert!(tool_call.function.name.is_none());
422        assert_eq!(tool_call.function.arguments.as_ref().unwrap(), "Paris");
423    }
424
425    #[test]
426    fn test_streaming_tool_call_missing_function_deserialization() {
427        let json = r#"{
428            "index": 0,
429            "id": "call_abc123"
430        }"#;
431        let tool_call: StreamingToolCall = serde_json::from_str(json).unwrap();
432        assert_eq!(tool_call.index, 0);
433        assert_eq!(tool_call.id, Some("call_abc123".to_string()));
434        assert!(tool_call.function.name.is_none());
435        assert!(tool_call.function.arguments.is_none());
436    }
437
438    #[test]
439    fn test_streaming_tool_call_null_function_deserialization() {
440        let json = r#"{
441            "index": 0,
442            "id": "call_abc123",
443            "function": null
444        }"#;
445        let tool_call: StreamingToolCall = serde_json::from_str(json).unwrap();
446        assert_eq!(tool_call.index, 0);
447        assert_eq!(tool_call.id, Some("call_abc123".to_string()));
448        assert!(tool_call.function.name.is_none());
449        assert!(tool_call.function.arguments.is_none());
450    }
451
452    #[test]
453    fn test_streaming_delta_with_tool_calls() {
454        let json = r#"{
455            "content": null,
456            "tool_calls": [{
457                "index": 0,
458                "id": "call_xyz",
459                "function": {
460                    "name": "search",
461                    "arguments": ""
462                }
463            }]
464        }"#;
465        let delta: StreamingDelta = serde_json::from_str(json).unwrap();
466        assert!(delta.content.is_none());
467        assert_eq!(delta.tool_calls.len(), 1);
468        assert_eq!(delta.tool_calls[0].id, Some("call_xyz".to_string()));
469    }
470
471    #[test]
472    fn test_streaming_delta_with_null_tool_calls() {
473        let json = r#"{
474            "content": "Hello",
475            "tool_calls": null
476        }"#;
477        let delta: StreamingDelta = serde_json::from_str(json).unwrap();
478        assert_eq!(delta.content, Some("Hello".to_string()));
479        assert!(delta.tool_calls.is_empty());
480    }
481
482    #[test]
483    fn test_streaming_chunk_deserialization() {
484        let json = r#"{
485            "choices": [{
486                "delta": {
487                    "content": "Hello",
488                    "tool_calls": []
489                }
490            }],
491            "usage": {
492                "prompt_tokens": 10,
493                "completion_tokens": 5,
494                "total_tokens": 15
495            }
496        }"#;
497        let chunk: StreamingCompletionChunk = serde_json::from_str(json).unwrap();
498        assert_eq!(chunk.choices.len(), 1);
499        assert_eq!(chunk.choices[0].delta.content, Some("Hello".to_string()));
500        assert!(chunk.usage.is_some());
501    }
502
503    #[test]
504    fn test_streaming_chunk_with_multiple_tool_call_deltas() {
505        // Simulates multiple partial tool call chunks arriving
506        let json_start = r#"{
507            "choices": [{
508                "delta": {
509                    "content": null,
510                    "tool_calls": [{
511                        "index": 0,
512                        "id": "call_123",
513                        "function": {
514                            "name": "get_weather",
515                            "arguments": ""
516                        }
517                    }]
518                }
519            }],
520            "usage": null
521        }"#;
522
523        let json_chunk1 = r#"{
524            "choices": [{
525                "delta": {
526                    "content": null,
527                    "tool_calls": [{
528                        "index": 0,
529                        "id": null,
530                        "function": {
531                            "name": null,
532                            "arguments": "{\"loc"
533                        }
534                    }]
535                }
536            }],
537            "usage": null
538        }"#;
539
540        let json_chunk2 = r#"{
541            "choices": [{
542                "delta": {
543                    "content": null,
544                    "tool_calls": [{
545                        "index": 0,
546                        "id": null,
547                        "function": {
548                            "name": null,
549                            "arguments": "ation\":\"NYC\"}"
550                        }
551                    }]
552                }
553            }],
554            "usage": null
555        }"#;
556
557        // Verify each chunk deserializes correctly
558        let start_chunk: StreamingCompletionChunk = serde_json::from_str(json_start).unwrap();
559        assert_eq!(start_chunk.choices[0].delta.tool_calls.len(), 1);
560        assert_eq!(
561            start_chunk.choices[0].delta.tool_calls[0]
562                .function
563                .name
564                .as_ref()
565                .unwrap(),
566            "get_weather"
567        );
568
569        let chunk1: StreamingCompletionChunk = serde_json::from_str(json_chunk1).unwrap();
570        assert_eq!(chunk1.choices[0].delta.tool_calls.len(), 1);
571        assert_eq!(
572            chunk1.choices[0].delta.tool_calls[0]
573                .function
574                .arguments
575                .as_ref()
576                .unwrap(),
577            "{\"loc"
578        );
579
580        let chunk2: StreamingCompletionChunk = serde_json::from_str(json_chunk2).unwrap();
581        assert_eq!(chunk2.choices[0].delta.tool_calls.len(), 1);
582        assert_eq!(
583            chunk2.choices[0].delta.tool_calls[0]
584                .function
585                .arguments
586                .as_ref()
587                .unwrap(),
588            "ation\":\"NYC\"}"
589        );
590    }
591
592    #[tokio::test]
593    async fn test_streaming_usage_only_chunk_is_not_ignored() {
594        use crate::test_utils::MockStreamingClient;
595        use futures::StreamExt;
596
597        // Some providers emit a final "usage-only" chunk where `choices` is empty.
598        let client = MockStreamingClient {
599            sse_bytes: sse_bytes_from_data_lines([
600                "{\"choices\":[{\"delta\":{\"content\":\"Hello\",\"tool_calls\":[]}}],\"usage\":null}",
601                "{\"choices\":[],\"usage\":{\"prompt_tokens\":10,\"completion_tokens\":5,\"total_tokens\":15}}",
602                "[DONE]",
603            ]),
604        };
605
606        let req = http::Request::builder()
607            .method("POST")
608            .uri("http://localhost/v1/chat/completions")
609            .body(Vec::new())
610            .unwrap();
611
612        let mut stream = send_compatible_streaming_request(client, req)
613            .await
614            .unwrap();
615
616        let mut final_usage = None;
617        while let Some(chunk) = stream.next().await {
618            if let streaming::StreamedAssistantContent::Final(res) = chunk.unwrap() {
619                final_usage = Some(res.usage);
620                break;
621            }
622        }
623
624        let usage = final_usage.expect("expected a final response with usage");
625        assert_eq!(usage.prompt_tokens, 10);
626        assert_eq!(usage.total_tokens, 15);
627    }
628
629    #[tokio::test]
630    async fn test_streaming_reasoning_content_and_text_chunks_are_incremental() {
631        use crate::test_utils::MockStreamingClient;
632        use futures::StreamExt;
633
634        let client = MockStreamingClient {
635            sse_bytes: sse_bytes_from_data_lines([
636                "{\"id\":\"cmpl-1\",\"model\":\"Qwen/Qwen3-4B\",\"choices\":[{\"delta\":{\"reasoning_content\":\"think \",\"tool_calls\":[]},\"finish_reason\":null}],\"usage\":null}",
637                "{\"id\":\"cmpl-1\",\"model\":\"Qwen/Qwen3-4B\",\"choices\":[{\"delta\":{\"reasoning_content\":\"more\",\"tool_calls\":[]},\"finish_reason\":null}],\"usage\":null}",
638                "{\"id\":\"cmpl-1\",\"model\":\"Qwen/Qwen3-4B\",\"choices\":[{\"delta\":{\"content\":\"hel\",\"tool_calls\":[]},\"finish_reason\":null}],\"usage\":null}",
639                "{\"id\":\"cmpl-1\",\"model\":\"Qwen/Qwen3-4B\",\"choices\":[{\"delta\":{\"content\":\"lo\",\"tool_calls\":[]},\"finish_reason\":\"stop\"}],\"usage\":null}",
640                "{\"choices\":[],\"usage\":{\"prompt_tokens\":4,\"completion_tokens\":6,\"total_tokens\":10}}",
641                "[DONE]",
642            ]),
643        };
644
645        let req = http::Request::builder()
646            .method("POST")
647            .uri("http://localhost/v1/chat/completions")
648            .body(Vec::new())
649            .unwrap();
650
651        let mut stream = send_compatible_streaming_request(client, req)
652            .await
653            .unwrap();
654
655        let mut reasoning_chunks = Vec::new();
656        let mut text_chunks = Vec::new();
657        let mut final_usage = None;
658
659        while let Some(chunk) = stream.next().await {
660            match chunk.unwrap() {
661                streaming::StreamedAssistantContent::ReasoningDelta { reasoning, .. } => {
662                    reasoning_chunks.push(reasoning)
663                }
664                streaming::StreamedAssistantContent::Text(text) => text_chunks.push(text.text),
665                streaming::StreamedAssistantContent::Final(response) => {
666                    final_usage = Some(response.usage)
667                }
668                _ => {}
669            }
670        }
671
672        assert_eq!(
673            reasoning_chunks,
674            vec!["think ".to_string(), "more".to_string()]
675        );
676        assert_eq!(text_chunks, vec!["hel".to_string(), "lo".to_string()]);
677
678        let usage = final_usage.expect("expected final usage");
679        assert_eq!(usage.prompt_tokens, 4);
680        assert_eq!(usage.total_tokens, 10);
681        let token_usage = usage.token_usage();
682        assert_eq!(token_usage.output_tokens, 6);
683    }
684
685    #[tokio::test]
686    async fn test_streaming_cached_input_tokens_populated() {
687        use crate::test_utils::MockStreamingClient;
688        use futures::StreamExt;
689
690        // Usage chunk includes prompt_tokens_details with cached_tokens.
691        let client = MockStreamingClient {
692            sse_bytes: sse_bytes_from_data_lines([
693                "{\"choices\":[{\"delta\":{\"content\":\"Hi\",\"tool_calls\":[]}}],\"usage\":null}",
694                "{\"choices\":[],\"usage\":{\"prompt_tokens\":100,\"completion_tokens\":10,\"total_tokens\":110,\"prompt_tokens_details\":{\"cached_tokens\":80}}}",
695                "[DONE]",
696            ]),
697        };
698
699        let req = http::Request::builder()
700            .method("POST")
701            .uri("http://localhost/v1/chat/completions")
702            .body(Vec::new())
703            .unwrap();
704
705        let mut stream = send_compatible_streaming_request(client, req)
706            .await
707            .unwrap();
708
709        let mut final_response = None;
710        while let Some(chunk) = stream.next().await {
711            if let streaming::StreamedAssistantContent::Final(res) = chunk.unwrap() {
712                final_response = Some(res);
713                break;
714            }
715        }
716
717        let res = final_response.expect("expected a final response");
718
719        // Verify provider-level usage has the cached_tokens
720        assert_eq!(
721            res.usage
722                .prompt_tokens_details
723                .as_ref()
724                .unwrap()
725                .cached_tokens,
726            80
727        );
728
729        // Verify core Usage also has cached_input_tokens via GetTokenUsage
730        let core_usage = res.token_usage();
731        assert_eq!(core_usage.cached_input_tokens, 80);
732        assert_eq!(core_usage.input_tokens, 100);
733        assert_eq!(core_usage.total_tokens, 110);
734    }
735
736    /// Reproduces the bug where a proxy/gateway sends multiple parallel tool
737    /// calls all sharing `index: 0` but with distinct `id` values.  Without
738    /// the fix, rig merges both calls into one corrupted entry.
739    #[tokio::test]
740    async fn test_duplicate_index_different_id_tool_calls() {
741        use crate::test_utils::MockStreamingClient;
742        use futures::StreamExt;
743
744        // Simulate a gateway that sends two tool calls both at index 0.
745        // First tool call: id="call_aaa", name="command", args={"cmd":"ls"}
746        // Second tool call: id="call_bbb", name="git", args={"action":"log"}
747        let client = MockStreamingClient {
748            sse_bytes: sse_bytes_from_data_lines([
749                "{\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":\"call_aaa\",\"function\":{\"name\":\"command\",\"arguments\":\"\"}}]},\"finish_reason\":null}],\"usage\":null}",
750                "{\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":null,\"function\":{\"name\":null,\"arguments\":\"{\\\"cmd\\\"\"}}]},\"finish_reason\":null}],\"usage\":null}",
751                "{\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":null,\"function\":{\"name\":null,\"arguments\":\":\\\"ls\\\"}\"}}]},\"finish_reason\":null}],\"usage\":null}",
752                "{\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":\"call_bbb\",\"function\":{\"name\":\"git\",\"arguments\":\"\"}}]},\"finish_reason\":null}],\"usage\":null}",
753                "{\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":null,\"function\":{\"name\":null,\"arguments\":\"{\\\"action\\\"\"}}]},\"finish_reason\":null}],\"usage\":null}",
754                "{\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":null,\"function\":{\"name\":null,\"arguments\":\":\\\"log\\\"}\"}}]},\"finish_reason\":null}],\"usage\":null}",
755                "{\"choices\":[{\"delta\":{\"tool_calls\":[]},\"finish_reason\":\"tool_calls\"}],\"usage\":null}",
756                "{\"choices\":[],\"usage\":{\"prompt_tokens\":20,\"completion_tokens\":10,\"total_tokens\":30}}",
757                "[DONE]",
758            ]),
759        };
760
761        let req = http::Request::builder()
762            .method("POST")
763            .uri("http://localhost/v1/chat/completions")
764            .body(Vec::new())
765            .unwrap();
766
767        let mut stream = send_compatible_streaming_request(client, req)
768            .await
769            .unwrap();
770
771        let mut collected_tool_calls = Vec::new();
772        while let Some(chunk) = stream.next().await {
773            if let streaming::StreamedAssistantContent::ToolCall {
774                tool_call,
775                internal_call_id: _,
776            } = chunk.unwrap()
777            {
778                collected_tool_calls.push(tool_call);
779            }
780        }
781
782        assert_eq!(
783            collected_tool_calls.len(),
784            2,
785            "expected 2 separate tool calls, got {collected_tool_calls:?}"
786        );
787
788        assert_eq!(collected_tool_calls[0].id, "call_aaa");
789        assert_eq!(collected_tool_calls[0].function.name, "command");
790        assert_eq!(
791            collected_tool_calls[0].function.arguments,
792            serde_json::json!({"cmd": "ls"})
793        );
794
795        assert_eq!(collected_tool_calls[1].id, "call_bbb");
796        assert_eq!(collected_tool_calls[1].function.name, "git");
797        assert_eq!(
798            collected_tool_calls[1].function.arguments,
799            serde_json::json!({"action": "log"})
800        );
801    }
802
803    #[tokio::test]
804    async fn test_tool_call_id_chunk_without_function_is_preserved() {
805        use crate::test_utils::MockStreamingClient;
806        use futures::StreamExt;
807
808        let client = MockStreamingClient {
809            sse_bytes: sse_bytes_from_data_lines([
810                "{\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":\"call_abc123\"}]},\"finish_reason\":null}],\"usage\":null}",
811                "{\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":null,\"function\":{\"name\":\"lookup\",\"arguments\":\"\"}}]},\"finish_reason\":null}],\"usage\":null}",
812                "{\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":null,\"function\":{\"name\":null,\"arguments\":\"{\\\"id\\\":1}\"}}]},\"finish_reason\":null}],\"usage\":null}",
813                "{\"choices\":[{\"delta\":{\"tool_calls\":[]},\"finish_reason\":\"tool_calls\"}],\"usage\":null}",
814                "[DONE]",
815            ]),
816        };
817
818        let req = http::Request::builder()
819            .method("POST")
820            .uri("http://localhost/v1/chat/completions")
821            .body(Vec::new())
822            .unwrap();
823
824        let mut stream = send_compatible_streaming_request(client, req)
825            .await
826            .unwrap();
827
828        let mut collected_tool_calls = Vec::new();
829        while let Some(chunk) = stream.next().await {
830            if let streaming::StreamedAssistantContent::ToolCall {
831                tool_call,
832                internal_call_id: _,
833            } = chunk.unwrap()
834            {
835                collected_tool_calls.push(tool_call);
836            }
837        }
838
839        assert_eq!(
840            collected_tool_calls.len(),
841            1,
842            "expected id-only chunk to be retained for later tool-call deltas"
843        );
844        assert_eq!(collected_tool_calls[0].id, "call_abc123");
845        assert_eq!(collected_tool_calls[0].function.name, "lookup");
846        assert_eq!(
847            collected_tool_calls[0].function.arguments,
848            serde_json::json!({"id": 1})
849        );
850    }
851
852    /// Reproduces the bug where a provider (e.g. GLM-4 via OpenAI-compatible
853    /// endpoint) sends a unique `id` on every SSE delta chunk for the same
854    /// logical tool call.  Without the fix, each chunk triggers an eviction,
855    /// yielding incomplete fragments as "completed" tool calls.
856    #[tokio::test]
857    async fn test_unique_id_per_chunk_single_tool_call() {
858        use crate::test_utils::MockStreamingClient;
859        use futures::StreamExt;
860
861        // Each chunk carries a different id but they all represent delta
862        // fragments of the SAME tool call at index 0.
863        let client = MockStreamingClient {
864            sse_bytes: sse_bytes_from_data_lines([
865                "{\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":\"chatcmpl-tool-aaa\",\"function\":{\"name\":\"web_search\",\"arguments\":\"null\"}}]},\"finish_reason\":null}],\"usage\":null}",
866                "{\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":\"chatcmpl-tool-bbb\",\"function\":{\"name\":\"\",\"arguments\":\"{\\\"query\\\": \\\"META\"}}]},\"finish_reason\":null}],\"usage\":null}",
867                "{\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":\"chatcmpl-tool-ccc\",\"function\":{\"name\":\"\",\"arguments\":\" Platforms news\\\"}\"}}]},\"finish_reason\":null}],\"usage\":null}",
868                "{\"choices\":[{\"delta\":{\"tool_calls\":[]},\"finish_reason\":\"tool_calls\"}],\"usage\":null}",
869                "{\"choices\":[],\"usage\":{\"prompt_tokens\":15,\"completion_tokens\":8,\"total_tokens\":23}}",
870                "[DONE]",
871            ]),
872        };
873
874        let req = http::Request::builder()
875            .method("POST")
876            .uri("http://localhost/v1/chat/completions")
877            .body(Vec::new())
878            .unwrap();
879
880        let mut stream = send_compatible_streaming_request(client, req)
881            .await
882            .unwrap();
883
884        let mut collected_tool_calls = Vec::new();
885        while let Some(chunk) = stream.next().await {
886            if let streaming::StreamedAssistantContent::ToolCall {
887                tool_call,
888                internal_call_id: _,
889            } = chunk.unwrap()
890            {
891                collected_tool_calls.push(tool_call);
892            }
893        }
894
895        assert_eq!(
896            collected_tool_calls.len(),
897            1,
898            "expected 1 tool call (all chunks are fragments of the same call), got {collected_tool_calls:?}"
899        );
900
901        assert_eq!(collected_tool_calls[0].function.name, "web_search");
902        // The arguments should be the fully accumulated string, not fragments
903        let args_str = match &collected_tool_calls[0].function.arguments {
904            serde_json::Value::String(s) => s.clone(),
905            v => v.to_string(),
906        };
907        assert!(
908            args_str.contains("META Platforms news"),
909            "expected accumulated arguments containing the full query, got: {args_str}"
910        );
911    }
912
913    #[tokio::test]
914    async fn test_zero_arg_tool_call_normalized_on_finish_reason() {
915        use crate::test_utils::MockStreamingClient;
916
917        let client = MockStreamingClient {
918            sse_bytes: sse_bytes_from_data_lines([
919                "{\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":\"call_123\",\"function\":{\"name\":\"ping\",\"arguments\":\"\"}}]},\"finish_reason\":null}],\"usage\":null}",
920                "{\"choices\":[{\"delta\":{\"tool_calls\":[]},\"finish_reason\":\"tool_calls\"}],\"usage\":null}",
921                "[DONE]",
922            ]),
923        };
924
925        let req = http::Request::builder()
926            .method("POST")
927            .uri("http://localhost/v1/chat/completions")
928            .body(Vec::new())
929            .unwrap();
930
931        let stream = send_compatible_streaming_request(client, req)
932            .await
933            .unwrap();
934
935        assert_zero_arg_tool_call_is_emitted(stream, "call_123", "ping", true).await;
936    }
937
938    #[tokio::test]
939    async fn test_zero_arg_tool_call_is_preserved_at_eof() {
940        use crate::test_utils::MockStreamingClient;
941
942        let client = MockStreamingClient {
943            sse_bytes: sse_bytes_from_data_lines([
944                "{\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":\"call_123\",\"function\":{\"name\":\"ping\",\"arguments\":\"\"}}]},\"finish_reason\":null}],\"usage\":null}",
945            ]),
946        };
947
948        let req = http::Request::builder()
949            .method("POST")
950            .uri("http://localhost/v1/chat/completions")
951            .body(Vec::new())
952            .unwrap();
953
954        let stream = send_compatible_streaming_request(client, req)
955            .await
956            .unwrap();
957
958        assert_zero_arg_tool_call_is_emitted(stream, "call_123", "ping", true).await;
959    }
960}