Skip to main content

rig_core/providers/anthropic/
streaming.rs

1use async_stream::stream;
2use futures::StreamExt;
3use serde::{Deserialize, Serialize};
4use serde_json::{Value, json};
5use tracing::{Level, enabled};
6use tracing_futures::Instrument;
7
8use super::completion::{
9    AnthropicCompatibleProvider, AnthropicCompletionRequest, AnthropicRequestParams, CacheTtl,
10    Content, GenericCompletionModel, Usage,
11};
12use crate::completion::{CompletionError, CompletionRequest, GetTokenUsage};
13use crate::http_client::sse::{Event, GenericEventSource};
14use crate::http_client::{self, HttpClientExt};
15use crate::message::ReasoningContent;
16use crate::streaming::{
17    self, RawStreamingChoice, RawStreamingToolCall, StreamingResult, ToolCallDeltaContent,
18};
19use crate::telemetry::{CompletionOperation, CompletionSpanBuilder, SpanCombinator};
20use crate::wasm_compat::{WasmCompatSend, WasmCompatSync};
21use std::collections::HashMap;
22
23/// Build the Anthropic streaming request body.
24///
25/// Derives it from the *same* typed [`AnthropicCompletionRequest`] the blocking
26/// path builds (in `completion.rs`), rather than re-assembling the body by hand.
27/// The previous hand-rolled `json!` body had drifted from the blocking one and
28/// silently dropped `output_schema` (structured-output config); reaching for the
29/// typed request fixes that and keeps the two in lockstep. The one intentional
30/// streaming-only difference — an explicit `tool_choice: auto` when tools are
31/// advertised but the caller left the choice unset — is re-applied below so the
32/// streaming request bytes stay stable.
33fn create_streaming_request_body(
34    request_model: String,
35    mut completion_request: CompletionRequest,
36    max_tokens: u64,
37    prompt_caching: bool,
38    automatic_caching: bool,
39    automatic_caching_ttl: Option<CacheTtl>,
40) -> Result<Value, CompletionError> {
41    // The typed request's `TryFrom` requires `max_tokens`; feed it the value the
42    // caller already resolved (the request's own value, else the model default).
43    completion_request.max_tokens = Some(max_tokens);
44
45    let request = AnthropicCompletionRequest::try_from(AnthropicRequestParams {
46        model: &request_model,
47        request: completion_request,
48        prompt_caching,
49        automatic_caching,
50        automatic_caching_ttl,
51    })?;
52
53    let mut body = serde_json::to_value(&request)?;
54    if let Some(map) = body.as_object_mut() {
55        // `AnthropicCompletionRequest` has no `stream` field (the blocking path
56        // omits it, defaulting to non-streaming); set it for the streaming endpoint.
57        map.insert("stream".to_string(), Value::Bool(true));
58
59        // Preserve the streaming path's long-standing `tool_choice` shape, which
60        // emitted `tool_choice` *iff* a non-empty tool set was advertised (Anthropic
61        // rejects `tool_choice` without `tools`). The blocking typed request instead
62        // serializes any caller-set `tool_choice` regardless of tools and omits it
63        // when unset, so reconcile here:
64        //   - tools present, choice unset -> add the explicit `auto` the streaming
65        //     wire has always carried (equivalent to Anthropic's default);
66        //   - tools absent -> drop a caller-set `tool_choice` that would otherwise
67        //     be sent without `tools` and rejected.
68        if map.contains_key("tools") {
69            map.entry("tool_choice")
70                .or_insert_with(|| json!({ "type": "auto" }));
71        } else {
72            map.remove("tool_choice");
73        }
74    }
75
76    Ok(body)
77}
78
79#[derive(Debug, Deserialize)]
80#[serde(tag = "type", rename_all = "snake_case")]
81pub enum StreamingEvent {
82    MessageStart {
83        message: MessageStart,
84    },
85    ContentBlockStart {
86        index: usize,
87        content_block: Content,
88    },
89    ContentBlockDelta {
90        index: usize,
91        delta: ContentDelta,
92    },
93    ContentBlockStop {
94        index: usize,
95    },
96    MessageDelta {
97        delta: MessageDelta,
98        usage: PartialUsage,
99    },
100    MessageStop,
101    Ping,
102    #[serde(other)]
103    Unknown,
104}
105
106#[derive(Debug, Deserialize)]
107pub struct MessageStart {
108    pub id: String,
109    pub role: String,
110    pub content: Vec<Content>,
111    pub model: String,
112    pub stop_reason: Option<String>,
113    pub stop_sequence: Option<String>,
114    pub usage: Usage,
115}
116
117#[derive(Debug, Deserialize)]
118#[serde(tag = "type", rename_all = "snake_case")]
119pub enum ContentDelta {
120    TextDelta {
121        text: String,
122    },
123    InputJsonDelta {
124        partial_json: String,
125    },
126    ThinkingDelta {
127        thinking: String,
128    },
129    SignatureDelta {
130        signature: String,
131    },
132    CitationsDelta {
133        citation: super::completion::Citation,
134    },
135    /// Forward-compatibility fallback. Any delta type Anthropic adds in the
136    /// future that this crate does not yet model deserializes here so the
137    /// surrounding [`StreamingEvent`] still parses.
138    #[serde(other)]
139    Unknown,
140}
141
142#[derive(Debug, Deserialize)]
143pub struct MessageDelta {
144    pub stop_reason: Option<String>,
145    pub stop_sequence: Option<String>,
146}
147
148#[derive(Debug, Deserialize, Clone, Serialize, Default)]
149pub struct PartialUsage {
150    pub output_tokens: usize,
151    #[serde(default)]
152    pub input_tokens: Option<usize>,
153    #[serde(default)]
154    pub cache_creation_input_tokens: Option<u64>,
155    #[serde(default)]
156    pub cache_read_input_tokens: Option<u64>,
157}
158
159impl GetTokenUsage for PartialUsage {
160    fn token_usage(&self) -> crate::completion::Usage {
161        let mut usage = crate::completion::Usage::new();
162
163        usage.input_tokens = self.input_tokens.unwrap_or_default() as u64;
164        usage.output_tokens = self.output_tokens as u64;
165        usage.cached_input_tokens = self.cache_read_input_tokens.unwrap_or(0);
166        usage.cache_creation_input_tokens = self.cache_creation_input_tokens.unwrap_or(0);
167        usage.total_tokens = usage.input_tokens
168            + usage.cached_input_tokens
169            + usage.cache_creation_input_tokens
170            + usage.output_tokens;
171        usage
172    }
173}
174
175#[derive(Default)]
176struct ToolCallState {
177    name: String,
178    id: String,
179    internal_call_id: String,
180    input_json: String,
181}
182
183struct ServerToolUseState {
184    name: String,
185    id: String,
186    initial_input: Value,
187    input_json: String,
188}
189
190#[derive(Default)]
191struct ThinkingState {
192    thinking: String,
193    signature: String,
194}
195
196#[derive(Clone, Debug, Deserialize, Serialize)]
197pub struct StreamingCompletionResponse {
198    pub usage: PartialUsage,
199}
200
201impl GetTokenUsage for StreamingCompletionResponse {
202    fn token_usage(&self) -> crate::completion::Usage {
203        let mut usage = crate::completion::Usage::new();
204        usage.input_tokens = self.usage.input_tokens.unwrap_or(0) as u64;
205        usage.output_tokens = self.usage.output_tokens as u64;
206        usage.cached_input_tokens = self.usage.cache_read_input_tokens.unwrap_or(0);
207        usage.cache_creation_input_tokens = self.usage.cache_creation_input_tokens.unwrap_or(0);
208        usage.total_tokens = usage.input_tokens
209            + usage.cached_input_tokens
210            + usage.cache_creation_input_tokens
211            + usage.output_tokens;
212
213        usage
214    }
215}
216
217impl<Ext, T> GenericCompletionModel<Ext, T>
218where
219    T: HttpClientExt + Clone + Default + 'static,
220    Ext: AnthropicCompatibleProvider + Clone + WasmCompatSend + WasmCompatSync + 'static,
221{
222    pub(crate) async fn stream(
223        &self,
224        completion_request: CompletionRequest,
225    ) -> Result<streaming::StreamingCompletionResponse<StreamingCompletionResponse>, CompletionError>
226    {
227        let request_model = completion_request
228            .model
229            .clone()
230            .unwrap_or_else(|| self.model.clone());
231        let span = CompletionSpanBuilder::new(
232            Ext::PROVIDER_NAME,
233            &request_model,
234            CompletionOperation::ChatStreaming,
235        )
236        .system_instructions(
237            completion_request.preamble.as_deref(),
238            completion_request.record_telemetry_content,
239        )
240        .build();
241        let max_tokens = if let Some(tokens) = completion_request.max_tokens {
242            tokens
243        } else if let Some(tokens) = self.default_max_tokens {
244            tokens
245        } else {
246            return Err(CompletionError::RequestError(
247                "`max_tokens` must be set for Anthropic".into(),
248            ));
249        };
250
251        let body = create_streaming_request_body(
252            request_model,
253            completion_request,
254            max_tokens,
255            self.prompt_caching,
256            self.automatic_caching,
257            self.automatic_caching_ttl.clone(),
258        )?;
259
260        if enabled!(Level::TRACE) {
261            tracing::trace!(
262                target: "rig::completions",
263                "Anthropic completion request: {}",
264                serde_json::to_string_pretty(&body)?
265            );
266        }
267
268        let body: Vec<u8> = serde_json::to_vec(&body)?;
269
270        let req = self
271            .client
272            .post("/v1/messages")?
273            .body(body)
274            .map_err(http_client::Error::Protocol)?;
275
276        let stream = GenericEventSource::new(self.client.clone(), req);
277
278        // Use our SSE decoder to directly handle Server-Sent Events format
279        let stream: StreamingResult<StreamingCompletionResponse> = Box::pin(stream! {
280            let mut current_tool_call: Option<ToolCallState> = None;
281            let mut server_tool_uses: HashMap<usize, ServerToolUseState> = HashMap::new();
282            let mut current_thinking: Option<ThinkingState> = None;
283            let mut sse_stream = Box::pin(stream);
284            let mut input_tokens = 0;
285            let mut final_usage = None;
286
287            let mut text_content = String::new();
288
289            while let Some(sse_result) = sse_stream.next().await {
290                match sse_result {
291                    Ok(Event::Open) => {}
292                    Ok(Event::Message(sse)) => {
293                        // Parse the SSE data as a StreamingEvent
294                        match serde_json::from_str::<StreamingEvent>(&sse.data) {
295                            Ok(event) => {
296                                match &event {
297                                    StreamingEvent::MessageStart { message } => {
298                                        input_tokens = message.usage.input_tokens;
299
300                                        let span = tracing::Span::current();
301                                        span.record("gen_ai.response.id", &message.id);
302                                        span.record("gen_ai.response.model", &message.model);
303                                    },
304                                    StreamingEvent::MessageDelta { delta, usage } => {
305                                        if delta.stop_reason.is_some() {
306                                            // cache_creation_input_tokens and cache_read_input_tokens
307                                            // are cumulative totals on message_delta.usage per the
308                                            // Anthropic streaming API spec — use them directly.
309                                            let usage = PartialUsage {
310                                                 output_tokens: usage.output_tokens,
311                                                 input_tokens: usize::try_from(input_tokens).ok(),
312                                                 cache_creation_input_tokens: usage.cache_creation_input_tokens,
313                                                 cache_read_input_tokens: usage.cache_read_input_tokens
314                                            };
315
316                                            let span = tracing::Span::current();
317                                            span.record_token_usage(&usage);
318                                            final_usage = Some(usage);
319                                            break;
320                                        }
321                                    }
322                                    _ => {}
323                                }
324
325                                if let Some(result) = handle_event(
326                                    &event,
327                                    &mut current_tool_call,
328                                    &mut server_tool_uses,
329                                    &mut current_thinking,
330                                ) {
331                                    if let Ok(RawStreamingChoice::Message(ref text)) = result {
332                                        text_content += text;
333                                    }
334                                    yield result;
335                                }
336                            },
337                            Err(e) => {
338                                if !sse.data.trim().is_empty() {
339                                    yield Err(CompletionError::ResponseError(
340                                        format!("Failed to parse JSON: {} (Data: {})", e, sse.data)
341                                    ));
342                                }
343                            }
344                        }
345                    },
346                    Err(e) => {
347                        yield Err(CompletionError::from_stream_transport(e));
348                        break;
349                    }
350                }
351            }
352
353            // Ensure event source is closed when stream ends
354            sse_stream.close();
355
356            yield Ok(RawStreamingChoice::FinalResponse(StreamingCompletionResponse {
357                usage: final_usage.unwrap_or_default()
358            }))
359        }.instrument(span));
360
361        Ok(streaming::StreamingCompletionResponse::stream(stream))
362    }
363}
364
365fn handle_event(
366    event: &StreamingEvent,
367    current_tool_call: &mut Option<ToolCallState>,
368    server_tool_uses: &mut HashMap<usize, ServerToolUseState>,
369    current_thinking: &mut Option<ThinkingState>,
370) -> Option<Result<RawStreamingChoice<StreamingCompletionResponse>, CompletionError>> {
371    match event {
372        StreamingEvent::ContentBlockDelta { index, delta } => match delta {
373            ContentDelta::TextDelta { text } => {
374                if current_tool_call.is_none() {
375                    return Some(Ok(RawStreamingChoice::Message(text.clone())));
376                }
377                None
378            }
379            ContentDelta::InputJsonDelta { partial_json } => {
380                if let Some(server_tool_use) = server_tool_uses.get_mut(index) {
381                    server_tool_use.input_json.push_str(partial_json);
382                    return None;
383                }
384
385                if let Some(tool_call) = current_tool_call {
386                    tool_call.input_json.push_str(partial_json);
387                    // Emit the delta so UI can show progress
388                    return Some(Ok(RawStreamingChoice::ToolCallDelta {
389                        id: tool_call.id.clone(),
390                        internal_call_id: tool_call.internal_call_id.clone(),
391                        content: ToolCallDeltaContent::Delta(partial_json.clone()),
392                    }));
393                }
394                None
395            }
396            ContentDelta::ThinkingDelta { thinking } => {
397                current_thinking
398                    .get_or_insert_with(ThinkingState::default)
399                    .thinking
400                    .push_str(thinking);
401
402                Some(Ok(RawStreamingChoice::ReasoningDelta {
403                    id: None,
404                    reasoning: thinking.clone(),
405                }))
406            }
407            ContentDelta::SignatureDelta { signature } => {
408                current_thinking
409                    .get_or_insert_with(ThinkingState::default)
410                    .signature
411                    .push_str(signature);
412
413                // Don't yield signature chunks, they will be included in the final Reasoning
414                None
415            }
416            ContentDelta::CitationsDelta { citation } => {
417                Some(Ok(RawStreamingChoice::TextAdditionalParams(json!({
418                    "citations": [citation]
419                }))))
420            }
421            ContentDelta::Unknown => None,
422        },
423        StreamingEvent::ContentBlockStart {
424            index,
425            content_block,
426        } => match content_block {
427            Content::Text { citations, .. } => {
428                let additional_params = (!citations.is_empty()).then(|| {
429                    json!({
430                        "citations": citations
431                    })
432                });
433                Some(Ok(RawStreamingChoice::TextStart { additional_params }))
434            }
435            Content::ServerToolUse { id, name, input } => {
436                server_tool_uses.insert(
437                    *index,
438                    ServerToolUseState {
439                        name: name.clone(),
440                        id: id.clone(),
441                        initial_input: input.clone(),
442                        input_json: String::new(),
443                    },
444                );
445                None
446            }
447            raw @ (Content::WebSearchToolResult { .. }
448            | Content::CodeExecutionToolResult { .. }) => Some(Ok(RawStreamingChoice::TextStart {
449                additional_params: Some(json!({
450                    super::completion::ANTHROPIC_RAW_CONTENT_KEY: raw
451                })),
452            })),
453            Content::ToolUse { id, name, .. } => {
454                let internal_call_id = crate::id::generate();
455                *current_tool_call = Some(ToolCallState {
456                    name: name.clone(),
457                    id: id.clone(),
458                    internal_call_id: internal_call_id.clone(),
459                    input_json: String::new(),
460                });
461                Some(Ok(RawStreamingChoice::ToolCallDelta {
462                    id: id.clone(),
463                    internal_call_id,
464                    content: ToolCallDeltaContent::Name(name.clone()),
465                }))
466            }
467            Content::Thinking { .. } => {
468                *current_thinking = Some(ThinkingState::default());
469                None
470            }
471            Content::RedactedThinking { data } => Some(Ok(RawStreamingChoice::Reasoning {
472                id: None,
473                content: ReasoningContent::Redacted { data: data.clone() },
474            })),
475            // Handle other content types - they don't need special handling
476            _ => None,
477        },
478        StreamingEvent::ContentBlockStop { index } => {
479            if let Some(thinking_state) = Option::take(current_thinking)
480                && !thinking_state.thinking.is_empty()
481            {
482                let signature = if thinking_state.signature.is_empty() {
483                    None
484                } else {
485                    Some(thinking_state.signature)
486                };
487
488                return Some(Ok(RawStreamingChoice::Reasoning {
489                    id: None,
490                    content: ReasoningContent::Text {
491                        text: thinking_state.thinking,
492                        signature,
493                    },
494                }));
495            }
496
497            if let Some(server_tool_use) = server_tool_uses.remove(index) {
498                let input = if server_tool_use.input_json.is_empty() {
499                    if server_tool_use.initial_input.is_null() {
500                        json!({})
501                    } else {
502                        server_tool_use.initial_input
503                    }
504                } else {
505                    match serde_json::from_str(&server_tool_use.input_json) {
506                        Ok(json_value) => json_value,
507                        Err(e) => return Some(Err(CompletionError::from(e))),
508                    }
509                };
510
511                return Some(Ok(RawStreamingChoice::TextStart {
512                    additional_params: Some(json!({
513                        super::completion::ANTHROPIC_RAW_CONTENT_KEY: Content::ServerToolUse {
514                            id: server_tool_use.id,
515                            name: server_tool_use.name,
516                            input,
517                        }
518                    })),
519                }));
520            }
521
522            if let Some(tool_call) = Option::take(current_tool_call) {
523                let json_str = if tool_call.input_json.is_empty() {
524                    "{}"
525                } else {
526                    &tool_call.input_json
527                };
528                match serde_json::from_str(json_str) {
529                    Ok(json_value) => {
530                        let raw_tool_call =
531                            RawStreamingToolCall::new(tool_call.id, tool_call.name, json_value)
532                                .with_internal_call_id(tool_call.internal_call_id);
533                        Some(Ok(RawStreamingChoice::ToolCall(raw_tool_call)))
534                    }
535                    Err(e) => Some(Err(CompletionError::from(e))),
536                }
537            } else {
538                None
539            }
540        }
541        // Ignore other event types or handle as needed
542        StreamingEvent::MessageStart { .. }
543        | StreamingEvent::MessageDelta { .. }
544        | StreamingEvent::MessageStop
545        | StreamingEvent::Ping
546        | StreamingEvent::Unknown => None,
547    }
548}
549
550#[cfg(test)]
551mod tests {
552    use super::super::completion::{
553        CLAUDE_OPUS_4_8, CacheControl, CacheTtl, Message, SystemContent,
554        apply_prompt_cache_control, build_tool_definitions, resolve_top_level_cache_control,
555    };
556    use super::*;
557    use crate::OneOrMany;
558    use crate::completion::Message as RigMessage;
559    use crate::completion::request::Document as RigDocument;
560    use async_stream::stream;
561    use futures::StreamExt;
562
563    #[cfg(not(all(target_arch = "wasm32", target_os = "unknown")))]
564    fn to_stream_result(
565        stream: impl futures::Stream<
566            Item = Result<RawStreamingChoice<StreamingCompletionResponse>, CompletionError>,
567        > + Send
568        + 'static,
569    ) -> crate::streaming::StreamingResult<StreamingCompletionResponse> {
570        Box::pin(stream)
571    }
572
573    #[cfg(all(target_arch = "wasm32", target_os = "unknown"))]
574    fn to_stream_result(
575        stream: impl futures::Stream<
576            Item = Result<RawStreamingChoice<StreamingCompletionResponse>, CompletionError>,
577        > + 'static,
578    ) -> crate::streaming::StreamingResult<StreamingCompletionResponse> {
579        Box::pin(stream)
580    }
581
582    #[test]
583    fn test_streaming_tool_build_marks_final_combined_tool() {
584        let mut additional_params = json!({
585            "tools": [{
586                "name": "provider_tool",
587                "description": "Provider tool",
588                "input_schema": {"type": "object"}
589            }]
590        });
591
592        let mut tools = build_tool_definitions(
593            vec![crate::completion::ToolDefinition {
594                name: "rig_tool".to_string(),
595                description: "Rig tool".to_string(),
596                parameters: json!({"type": "object", "properties": {}}),
597            }],
598            &mut additional_params,
599        )
600        .unwrap();
601        let mut system: Vec<SystemContent> = Vec::new();
602        let mut messages: Vec<Message> = Vec::new();
603        apply_prompt_cache_control(&mut system, &mut messages, &mut tools, true, None).unwrap();
604
605        assert_eq!(tools.len(), 2);
606        assert!(tools[0].get("cache_control").is_none());
607        assert_eq!(tools[1]["name"], "provider_tool");
608        assert_eq!(tools[1]["cache_control"]["type"], "ephemeral");
609    }
610
611    #[test]
612    fn streaming_request_keeps_documents_after_leading_system_messages() {
613        let request = CompletionRequest {
614            model: None,
615            preamble: None,
616            chat_history: OneOrMany::many(vec![
617                RigMessage::system("System prompt"),
618                RigMessage::assistant("Earlier assistant turn"),
619                RigMessage::system("Mid-conversation instruction"),
620                RigMessage::user("Prompt"),
621            ])
622            .unwrap(),
623            documents: vec![RigDocument {
624                id: "doc1".to_string(),
625                text: "Document text.".to_string(),
626                additional_props: Default::default(),
627            }],
628            tools: vec![],
629            temperature: None,
630            max_tokens: Some(64),
631            tool_choice: None,
632            additional_params: None,
633            output_schema: None,
634            record_telemetry_content: false,
635        };
636
637        let body = create_streaming_request_body(
638            CLAUDE_OPUS_4_8.to_string(),
639            request,
640            64,
641            false,
642            false,
643            None,
644        )
645        .expect("streaming request body should build");
646
647        assert_eq!(body["system"][0]["text"], "System prompt");
648        assert_eq!(body["system"][1]["text"], "Mid-conversation instruction");
649        let messages = body["messages"]
650            .as_array()
651            .expect("messages should be array");
652        assert_eq!(messages.len(), 3);
653        assert_eq!(messages[0]["role"], "user");
654        assert!(
655            messages[0].to_string().contains("<file id: doc1>"),
656            "document message should follow top-level system: {messages:?}"
657        );
658        assert_eq!(messages[1]["role"], "assistant");
659        assert_eq!(messages[2]["role"], "user");
660        assert_eq!(
661            messages
662                .iter()
663                .filter(|message| message.to_string().contains("<file id: doc1>"))
664                .count(),
665            1,
666            "document message should appear exactly once: {messages:?}"
667        );
668    }
669
670    #[test]
671    fn streaming_body_is_blocking_body_plus_stream_flag_and_carries_output_schema() {
672        let schema: schemars::Schema = serde_json::from_value(json!({
673            "title": "WeatherResponse",
674            "type": "object",
675            "properties": { "city": { "type": "string" } }
676        }))
677        .expect("schema should deserialize");
678
679        let request = CompletionRequest {
680            model: None,
681            preamble: Some("You are helpful".to_string()),
682            chat_history: OneOrMany::one(RigMessage::user("What's the weather?")),
683            documents: vec![],
684            tools: vec![],
685            temperature: Some(0.5),
686            max_tokens: Some(64),
687            tool_choice: None,
688            additional_params: None,
689            output_schema: Some(schema),
690            record_telemetry_content: false,
691        };
692
693        let streaming_body = create_streaming_request_body(
694            CLAUDE_OPUS_4_8.to_string(),
695            request.clone(),
696            64,
697            false,
698            false,
699            None,
700        )
701        .expect("streaming request body should build");
702
703        // The streaming endpoint flag is set.
704        assert_eq!(streaming_body["stream"], serde_json::Value::Bool(true));
705
706        // Regression: `output_schema` now reaches the streaming wire as
707        // `output_config` (the hand-rolled body dropped it entirely, so this
708        // assertion would have failed before the typed-request unification).
709        assert_eq!(
710            streaming_body["output_config"]["format"]["type"],
711            "json_schema"
712        );
713        assert!(
714            streaming_body["output_config"]["format"]["schema"].is_object(),
715            "streaming body must carry the structured-output schema: {streaming_body}"
716        );
717
718        // Unification invariant: the streaming body is exactly the blocking body
719        // (built via the same typed request) plus `stream: true`. Pins the two
720        // wire formats together so a future edit can't reintroduce drift.
721        let blocking = AnthropicCompletionRequest::try_from(AnthropicRequestParams {
722            model: CLAUDE_OPUS_4_8,
723            request,
724            prompt_caching: false,
725            automatic_caching: false,
726            automatic_caching_ttl: None,
727        })
728        .expect("blocking request body should build");
729        let mut expected = serde_json::to_value(&blocking).expect("serialize blocking body");
730        expected
731            .as_object_mut()
732            .expect("body is an object")
733            .insert("stream".to_string(), serde_json::Value::Bool(true));
734
735        assert_eq!(streaming_body, expected);
736    }
737
738    #[test]
739    fn streaming_body_keeps_explicit_tool_choice_auto_when_tools_present_but_unset() {
740        let request = CompletionRequest {
741            model: None,
742            preamble: None,
743            chat_history: OneOrMany::one(RigMessage::user("Add 2 and 3")),
744            documents: vec![],
745            tools: vec![crate::completion::ToolDefinition {
746                name: "add".to_string(),
747                description: "Add x and y".to_string(),
748                parameters: json!({
749                    "type": "object",
750                    "properties": { "x": { "type": "integer" } }
751                }),
752            }],
753            temperature: None,
754            max_tokens: Some(64),
755            tool_choice: None,
756            additional_params: None,
757            output_schema: None,
758            record_telemetry_content: false,
759        };
760
761        let body = create_streaming_request_body(
762            CLAUDE_OPUS_4_8.to_string(),
763            request,
764            64,
765            false,
766            false,
767            None,
768        )
769        .expect("streaming request body should build");
770
771        // Tools advertised + `tool_choice` unset must still carry the explicit
772        // `auto` the streaming wire format has always sent (parity with recorded
773        // fixtures), even though the blocking typed request omits it.
774        assert_eq!(body["tool_choice"], json!({ "type": "auto" }));
775        assert!(body["tools"].is_array());
776    }
777
778    #[test]
779    fn streaming_body_drops_tool_choice_when_no_tools_are_advertised() {
780        // The typed request serializes a caller-set `tool_choice` regardless of
781        // whether tools are present, but the streaming path has always emitted
782        // `tool_choice` *only* alongside a non-empty tool set (Anthropic rejects it
783        // otherwise). A `tool_choice` set with no tools must not reach the wire.
784        let request = CompletionRequest {
785            model: None,
786            preamble: None,
787            chat_history: OneOrMany::one(RigMessage::user("Hi")),
788            documents: vec![],
789            tools: vec![],
790            temperature: None,
791            max_tokens: Some(64),
792            tool_choice: Some(crate::message::ToolChoice::Auto),
793            additional_params: None,
794            output_schema: None,
795            record_telemetry_content: false,
796        };
797
798        let body = create_streaming_request_body(
799            CLAUDE_OPUS_4_8.to_string(),
800            request,
801            64,
802            false,
803            false,
804            None,
805        )
806        .expect("streaming request body should build");
807
808        assert!(
809            body.get("tool_choice").is_none(),
810            "tool_choice must be omitted when no tools are advertised: {body}"
811        );
812        assert!(body.get("tools").is_none());
813    }
814
815    #[test]
816    fn test_streaming_prompt_cache_control_uses_raw_top_level_ttl() {
817        let mut additional_params = json!({
818            "cache_control": {"type": "ephemeral", "ttl": "1h"}
819        });
820        let top_level_cache_control =
821            resolve_top_level_cache_control(false, None, &mut additional_params).unwrap();
822        let mut tools = build_tool_definitions(
823            vec![crate::completion::ToolDefinition {
824                name: "rig_tool".to_string(),
825                description: "Rig tool".to_string(),
826                parameters: json!({"type": "object", "properties": {}}),
827            }],
828            &mut additional_params,
829        )
830        .unwrap();
831        let mut system = vec![SystemContent::Text {
832            text: "System prompt".to_string(),
833            cache_control: None,
834        }];
835        let mut messages: Vec<Message> = Vec::new();
836
837        apply_prompt_cache_control(
838            &mut system,
839            &mut messages,
840            &mut tools,
841            true,
842            top_level_cache_control.as_ref(),
843        )
844        .unwrap();
845
846        assert_eq!(tools[0]["cache_control"]["type"], "ephemeral");
847        assert_eq!(tools[0]["cache_control"]["ttl"], "1h");
848        match &system[0] {
849            SystemContent::Text {
850                cache_control: Some(CacheControl::Ephemeral { ttl }),
851                ..
852            } => assert_eq!(ttl.as_ref(), Some(&CacheTtl::OneHour)),
853            other => panic!("expected system cache_control, got {other:?}"),
854        }
855        assert!(additional_params.get("cache_control").is_none());
856    }
857
858    fn handle_event(
859        event: &StreamingEvent,
860        current_tool_call: &mut Option<ToolCallState>,
861        current_thinking: &mut Option<ThinkingState>,
862    ) -> Option<Result<RawStreamingChoice<StreamingCompletionResponse>, CompletionError>> {
863        let mut server_tool_uses = HashMap::new();
864        super::handle_event(
865            event,
866            current_tool_call,
867            &mut server_tool_uses,
868            current_thinking,
869        )
870    }
871
872    #[test]
873    fn test_thinking_delta_deserialization() {
874        let json = r#"{"type": "thinking_delta", "thinking": "Let me think about this..."}"#;
875        let delta: ContentDelta = serde_json::from_str(json).unwrap();
876
877        match delta {
878            ContentDelta::ThinkingDelta { thinking } => {
879                assert_eq!(thinking, "Let me think about this...");
880            }
881            _ => panic!("Expected ThinkingDelta variant"),
882        }
883    }
884
885    #[test]
886    fn test_signature_delta_deserialization() {
887        let json = r#"{"type": "signature_delta", "signature": "abc123def456"}"#;
888        let delta: ContentDelta = serde_json::from_str(json).unwrap();
889
890        match delta {
891            ContentDelta::SignatureDelta { signature } => {
892                assert_eq!(signature, "abc123def456");
893            }
894            _ => panic!("Expected SignatureDelta variant"),
895        }
896    }
897
898    #[test]
899    fn test_thinking_delta_streaming_event_deserialization() {
900        let json = r#"{
901            "type": "content_block_delta",
902            "index": 0,
903            "delta": {
904                "type": "thinking_delta",
905                "thinking": "First, I need to understand the problem."
906            }
907        }"#;
908
909        let event: StreamingEvent = serde_json::from_str(json).unwrap();
910
911        match event {
912            StreamingEvent::ContentBlockDelta { index, delta } => {
913                assert_eq!(index, 0);
914                match delta {
915                    ContentDelta::ThinkingDelta { thinking } => {
916                        assert_eq!(thinking, "First, I need to understand the problem.");
917                    }
918                    _ => panic!("Expected ThinkingDelta"),
919                }
920            }
921            _ => panic!("Expected ContentBlockDelta event"),
922        }
923    }
924
925    #[test]
926    fn test_signature_delta_streaming_event_deserialization() {
927        let json = r#"{
928            "type": "content_block_delta",
929            "index": 0,
930            "delta": {
931                "type": "signature_delta",
932                "signature": "ErUBCkYICBgCIkCaGbqC85F4"
933            }
934        }"#;
935
936        let event: StreamingEvent = serde_json::from_str(json).unwrap();
937
938        match event {
939            StreamingEvent::ContentBlockDelta { index, delta } => {
940                assert_eq!(index, 0);
941                match delta {
942                    ContentDelta::SignatureDelta { signature } => {
943                        assert_eq!(signature, "ErUBCkYICBgCIkCaGbqC85F4");
944                    }
945                    _ => panic!("Expected SignatureDelta"),
946                }
947            }
948            _ => panic!("Expected ContentBlockDelta event"),
949        }
950    }
951
952    #[test]
953    fn test_handle_thinking_delta_event() {
954        let event = StreamingEvent::ContentBlockDelta {
955            index: 0,
956            delta: ContentDelta::ThinkingDelta {
957                thinking: "Analyzing the request...".to_string(),
958            },
959        };
960
961        let mut tool_call_state = None;
962        let mut thinking_state = None;
963        let result = handle_event(&event, &mut tool_call_state, &mut thinking_state);
964
965        assert!(result.is_some());
966        let choice = result.unwrap().unwrap();
967
968        match choice {
969            RawStreamingChoice::ReasoningDelta { id, reasoning, .. } => {
970                assert_eq!(id, None);
971                assert_eq!(reasoning, "Analyzing the request...");
972            }
973            _ => panic!("Expected ReasoningDelta choice"),
974        }
975
976        // Verify thinking state was updated
977        assert!(thinking_state.is_some());
978        assert_eq!(thinking_state.unwrap().thinking, "Analyzing the request...");
979    }
980
981    #[test]
982    fn test_handle_signature_delta_event() {
983        let event = StreamingEvent::ContentBlockDelta {
984            index: 0,
985            delta: ContentDelta::SignatureDelta {
986                signature: "test_signature".to_string(),
987            },
988        };
989
990        let mut tool_call_state = None;
991        let mut thinking_state = None;
992        let result = handle_event(&event, &mut tool_call_state, &mut thinking_state);
993
994        // SignatureDelta should not yield anything (returns None)
995        assert!(result.is_none());
996
997        // But signature should be captured in thinking state
998        assert!(thinking_state.is_some());
999        assert_eq!(thinking_state.unwrap().signature, "test_signature");
1000    }
1001
1002    #[test]
1003    fn test_handle_redacted_thinking_content_block_start_event() {
1004        let event = StreamingEvent::ContentBlockStart {
1005            index: 0,
1006            content_block: Content::RedactedThinking {
1007                data: "redacted_blob".to_string(),
1008            },
1009        };
1010        let mut tool_call_state = None;
1011        let mut thinking_state = None;
1012        let result = handle_event(&event, &mut tool_call_state, &mut thinking_state);
1013
1014        assert!(result.is_some());
1015        match result.unwrap().unwrap() {
1016            RawStreamingChoice::Reasoning {
1017                content: ReasoningContent::Redacted { data },
1018                ..
1019            } => {
1020                assert_eq!(data, "redacted_blob");
1021            }
1022            _ => panic!("Expected Redacted reasoning chunk"),
1023        }
1024    }
1025
1026    #[test]
1027    fn test_handle_text_delta_event() {
1028        let event = StreamingEvent::ContentBlockDelta {
1029            index: 0,
1030            delta: ContentDelta::TextDelta {
1031                text: "Hello, world!".to_string(),
1032            },
1033        };
1034
1035        let mut tool_call_state = None;
1036        let mut thinking_state = None;
1037        let result = handle_event(&event, &mut tool_call_state, &mut thinking_state);
1038
1039        assert!(result.is_some());
1040        let choice = result.unwrap().unwrap();
1041
1042        match choice {
1043            RawStreamingChoice::Message(text) => {
1044                assert_eq!(text, "Hello, world!");
1045            }
1046            _ => panic!("Expected Message choice"),
1047        }
1048    }
1049
1050    #[test]
1051    fn test_handle_text_block_start_event() {
1052        let event = StreamingEvent::ContentBlockStart {
1053            index: 0,
1054            content_block: Content::Text {
1055                text: String::new(),
1056                citations: Vec::new(),
1057                cache_control: None,
1058            },
1059        };
1060
1061        let mut tool_call_state = None;
1062        let mut thinking_state = None;
1063        let result = handle_event(&event, &mut tool_call_state, &mut thinking_state);
1064
1065        assert!(result.is_some());
1066        let choice = result.unwrap().unwrap();
1067        assert!(matches!(
1068            choice,
1069            RawStreamingChoice::TextStart {
1070                additional_params: None
1071            }
1072        ));
1073    }
1074
1075    #[test]
1076    fn test_thinking_delta_does_not_interfere_with_tool_calls() {
1077        // Thinking deltas should still be processed even if a tool call is in progress
1078        let event = StreamingEvent::ContentBlockDelta {
1079            index: 0,
1080            delta: ContentDelta::ThinkingDelta {
1081                thinking: "Thinking while tool is active...".to_string(),
1082            },
1083        };
1084
1085        let mut tool_call_state = Some(ToolCallState {
1086            name: "test_tool".to_string(),
1087            id: "tool_123".to_string(),
1088            internal_call_id: crate::id::generate(),
1089            input_json: String::new(),
1090        });
1091        let mut thinking_state = None;
1092
1093        let result = handle_event(&event, &mut tool_call_state, &mut thinking_state);
1094
1095        assert!(result.is_some());
1096        let choice = result.unwrap().unwrap();
1097
1098        match choice {
1099            RawStreamingChoice::ReasoningDelta { reasoning, .. } => {
1100                assert_eq!(reasoning, "Thinking while tool is active...");
1101            }
1102            _ => panic!("Expected ReasoningDelta choice"),
1103        }
1104
1105        // Tool call state should remain unchanged
1106        assert!(tool_call_state.is_some());
1107    }
1108
1109    #[test]
1110    fn test_handle_input_json_delta_event() {
1111        let event = StreamingEvent::ContentBlockDelta {
1112            index: 0,
1113            delta: ContentDelta::InputJsonDelta {
1114                partial_json: "{\"arg\":\"value".to_string(),
1115            },
1116        };
1117
1118        let mut tool_call_state = Some(ToolCallState {
1119            name: "test_tool".to_string(),
1120            id: "tool_123".to_string(),
1121            internal_call_id: crate::id::generate(),
1122            input_json: String::new(),
1123        });
1124        let mut thinking_state = None;
1125
1126        let result = handle_event(&event, &mut tool_call_state, &mut thinking_state);
1127
1128        // Should emit a ToolCallDelta
1129        assert!(result.is_some());
1130        let choice = result.unwrap().unwrap();
1131
1132        match choice {
1133            RawStreamingChoice::ToolCallDelta {
1134                id,
1135                internal_call_id: _,
1136                content,
1137            } => {
1138                assert_eq!(id, "tool_123");
1139                match content {
1140                    ToolCallDeltaContent::Delta(delta) => assert_eq!(delta, "{\"arg\":\"value"),
1141                    _ => panic!("Expected Delta content"),
1142                }
1143            }
1144            _ => panic!("Expected ToolCallDelta choice, got {:?}", choice),
1145        }
1146
1147        // Verify the input_json was accumulated
1148        assert!(tool_call_state.is_some());
1149        let state = tool_call_state.unwrap();
1150        assert_eq!(state.input_json, "{\"arg\":\"value");
1151    }
1152
1153    #[test]
1154    fn test_tool_call_accumulation_with_multiple_deltas() {
1155        let mut tool_call_state = Some(ToolCallState {
1156            name: "test_tool".to_string(),
1157            id: "tool_123".to_string(),
1158            internal_call_id: crate::id::generate(),
1159            input_json: String::new(),
1160        });
1161        let mut thinking_state = None;
1162
1163        // First delta
1164        let event1 = StreamingEvent::ContentBlockDelta {
1165            index: 0,
1166            delta: ContentDelta::InputJsonDelta {
1167                partial_json: "{\"location\":".to_string(),
1168            },
1169        };
1170        let result1 = handle_event(&event1, &mut tool_call_state, &mut thinking_state);
1171        assert!(result1.is_some());
1172
1173        // Second delta
1174        let event2 = StreamingEvent::ContentBlockDelta {
1175            index: 0,
1176            delta: ContentDelta::InputJsonDelta {
1177                partial_json: "\"Paris\",".to_string(),
1178            },
1179        };
1180        let result2 = handle_event(&event2, &mut tool_call_state, &mut thinking_state);
1181        assert!(result2.is_some());
1182
1183        // Third delta
1184        let event3 = StreamingEvent::ContentBlockDelta {
1185            index: 0,
1186            delta: ContentDelta::InputJsonDelta {
1187                partial_json: "\"temp\":\"20C\"}".to_string(),
1188            },
1189        };
1190        let result3 = handle_event(&event3, &mut tool_call_state, &mut thinking_state);
1191        assert!(result3.is_some());
1192
1193        // Verify accumulated JSON
1194        assert!(tool_call_state.is_some());
1195        let state = tool_call_state.as_ref().unwrap();
1196        assert_eq!(
1197            state.input_json,
1198            "{\"location\":\"Paris\",\"temp\":\"20C\"}"
1199        );
1200
1201        // Final ContentBlockStop should emit complete tool call
1202        let stop_event = StreamingEvent::ContentBlockStop { index: 0 };
1203        let final_result = handle_event(&stop_event, &mut tool_call_state, &mut thinking_state);
1204        assert!(final_result.is_some());
1205
1206        match final_result.unwrap().unwrap() {
1207            RawStreamingChoice::ToolCall(RawStreamingToolCall {
1208                id,
1209                name,
1210                arguments,
1211                ..
1212            }) => {
1213                assert_eq!(id, "tool_123");
1214                assert_eq!(name, "test_tool");
1215                assert_eq!(
1216                    arguments.get("location").unwrap().as_str().unwrap(),
1217                    "Paris"
1218                );
1219                assert_eq!(arguments.get("temp").unwrap().as_str().unwrap(), "20C");
1220            }
1221            other => panic!("Expected ToolCall, got {:?}", other),
1222        }
1223
1224        // Tool call state should be taken
1225        assert!(tool_call_state.is_none());
1226    }
1227
1228    #[test]
1229    fn test_citations_delta_streaming_event_deserialization() {
1230        let json = r#"{
1231            "type": "content_block_delta",
1232            "index": 0,
1233            "delta": {
1234                "type": "citations_delta",
1235                "citation": {
1236                    "type": "char_location",
1237                    "cited_text": "The grass is green.",
1238                    "document_index": 0,
1239                    "document_title": "Example",
1240                    "start_char_index": 0,
1241                    "end_char_index": 20
1242                }
1243            }
1244        }"#;
1245
1246        let event: StreamingEvent = serde_json::from_str(json).unwrap();
1247        let StreamingEvent::ContentBlockDelta { index, delta } = event else {
1248            panic!("expected ContentBlockDelta");
1249        };
1250        assert_eq!(index, 0);
1251        let ContentDelta::CitationsDelta { citation } = delta else {
1252            panic!("expected CitationsDelta");
1253        };
1254        let crate::providers::anthropic::completion::Citation::CharLocation {
1255            start_char_index,
1256            end_char_index,
1257            ..
1258        } = citation
1259        else {
1260            panic!("expected CharLocation");
1261        };
1262        assert_eq!(start_char_index, 0);
1263        assert_eq!(end_char_index, 20);
1264    }
1265
1266    #[test]
1267    fn test_search_result_citations_delta_streaming_event_deserialization() {
1268        let json = r#"{
1269            "type": "content_block_delta",
1270            "index": 0,
1271            "delta": {
1272                "type": "citations_delta",
1273                "citation": {
1274                    "type": "search_result_location",
1275                    "cited_text": "API requests require a key.",
1276                    "source": "https://docs.example.com/api-reference",
1277                    "title": "API Reference",
1278                    "search_result_index": 0,
1279                    "start_block_index": 0,
1280                    "end_block_index": 1
1281                }
1282            }
1283        }"#;
1284
1285        let event: StreamingEvent = serde_json::from_str(json).unwrap();
1286        let StreamingEvent::ContentBlockDelta { delta, .. } = event else {
1287            panic!("expected ContentBlockDelta");
1288        };
1289        let ContentDelta::CitationsDelta { citation } = delta else {
1290            panic!("expected CitationsDelta");
1291        };
1292        assert!(matches!(
1293            citation,
1294            crate::providers::anthropic::completion::Citation::SearchResultLocation {
1295                search_result_index: 0,
1296                start_block_index: 0,
1297                end_block_index: 1,
1298                ..
1299            }
1300        ));
1301    }
1302
1303    #[test]
1304    fn test_web_search_result_citations_delta_streaming_event_deserialization() {
1305        let json = r#"{
1306            "type": "content_block_delta",
1307            "index": 0,
1308            "delta": {
1309                "type": "citations_delta",
1310                "citation": {
1311                    "type": "web_search_result_location",
1312                    "cited_text": "Claude Shannon was a mathematician.",
1313                    "url": "https://example.com/shannon",
1314                    "title": "Claude Shannon",
1315                    "encrypted_index": "encrypted-reference"
1316                }
1317            }
1318        }"#;
1319
1320        let event: StreamingEvent = serde_json::from_str(json).unwrap();
1321        let StreamingEvent::ContentBlockDelta { delta, .. } = event else {
1322            panic!("expected ContentBlockDelta");
1323        };
1324        let ContentDelta::CitationsDelta { citation } = delta else {
1325            panic!("expected CitationsDelta");
1326        };
1327        assert!(matches!(
1328            citation,
1329            crate::providers::anthropic::completion::Citation::WebSearchResultLocation {
1330                ref url,
1331                ref encrypted_index,
1332                ..
1333            } if url == "https://example.com/shannon"
1334                && encrypted_index == "encrypted-reference"
1335        ));
1336    }
1337
1338    #[test]
1339    fn test_web_search_result_citations_delta_allows_null_title() {
1340        let json = r#"{
1341            "type": "content_block_delta",
1342            "index": 0,
1343            "delta": {
1344                "type": "citations_delta",
1345                "citation": {
1346                    "type": "web_search_result_location",
1347                    "cited_text": "Claude Shannon was a mathematician.",
1348                    "url": "https://example.com/shannon",
1349                    "title": null,
1350                    "encrypted_index": "encrypted-reference"
1351                }
1352            }
1353        }"#;
1354
1355        let event: StreamingEvent = serde_json::from_str(json).unwrap();
1356        let StreamingEvent::ContentBlockDelta { delta, .. } = event else {
1357            panic!("expected ContentBlockDelta");
1358        };
1359        let ContentDelta::CitationsDelta { citation } = delta else {
1360            panic!("expected CitationsDelta");
1361        };
1362        assert!(matches!(
1363            citation,
1364            crate::providers::anthropic::completion::Citation::WebSearchResultLocation {
1365                title: None,
1366                ..
1367            }
1368        ));
1369    }
1370
1371    #[test]
1372    fn test_text_content_block_start_allows_null_citations() {
1373        // The Anthropic Messages API emits an explicit `"citations": null` on the
1374        // first text `content_block_start` event. `#[serde(default)]` alone covers
1375        // a missing field but not an explicit null, so this must deserialize to an
1376        // empty citation list rather than failing the whole stream (see #1971).
1377        let json = r#"{
1378            "type": "content_block_start",
1379            "index": 0,
1380            "content_block": {
1381                "type": "text",
1382                "text": "",
1383                "citations": null
1384            }
1385        }"#;
1386
1387        let event: StreamingEvent = serde_json::from_str(json).unwrap();
1388        let StreamingEvent::ContentBlockStart { content_block, .. } = event else {
1389            panic!("expected ContentBlockStart");
1390        };
1391        let Content::Text {
1392            text, citations, ..
1393        } = content_block
1394        else {
1395            panic!("expected text content block");
1396        };
1397        assert_eq!(text, "");
1398        assert!(citations.is_empty());
1399    }
1400
1401    #[test]
1402    fn test_web_search_content_block_start_events_deserialize() {
1403        let server_tool_use = r#"{
1404            "type": "content_block_start",
1405            "index": 1,
1406            "content_block": {
1407                "type": "server_tool_use",
1408                "id": "srvtoolu_01",
1409                "name": "web_search",
1410                "input": {
1411                    "query": "claude shannon birth date"
1412                }
1413            }
1414        }"#;
1415        let event: StreamingEvent = serde_json::from_str(server_tool_use).unwrap();
1416        assert!(matches!(
1417            event,
1418            StreamingEvent::ContentBlockStart {
1419                content_block: Content::ServerToolUse {
1420                    ref id,
1421                    ref name,
1422                    ref input
1423                },
1424                ..
1425            } if id == "srvtoolu_01"
1426                && name == "web_search"
1427                && input["query"] == "claude shannon birth date"
1428        ));
1429
1430        let web_search_tool_result = r#"{
1431            "type": "content_block_start",
1432            "index": 2,
1433            "content_block": {
1434                "type": "web_search_tool_result",
1435                "tool_use_id": "srvtoolu_01",
1436                "content": [{
1437                    "type": "web_search_result",
1438                    "url": "https://example.com/shannon",
1439                    "title": "Claude Shannon",
1440                    "encrypted_content": "encrypted-content"
1441                }]
1442            }
1443        }"#;
1444        let event: StreamingEvent = serde_json::from_str(web_search_tool_result).unwrap();
1445        assert!(matches!(
1446            event,
1447            StreamingEvent::ContentBlockStart {
1448                content_block: Content::WebSearchToolResult {
1449                    ref tool_use_id,
1450                    ref content
1451                },
1452                ..
1453            } if tool_use_id == "srvtoolu_01"
1454                && content[0]["encrypted_content"] == "encrypted-content"
1455        ));
1456    }
1457
1458    #[test]
1459    fn test_code_execution_tool_result_block_is_preserved() {
1460        let event: StreamingEvent = serde_json::from_value(serde_json::json!({
1461            "type": "content_block_start",
1462            "index": 1,
1463            "content_block": {
1464                "type": "code_execution_tool_result",
1465                "tool_use_id": "srvtoolu_01",
1466                "content": {
1467                    "type": "code_execution_result",
1468                    "return_code": 0,
1469                    "stdout": "42\n",
1470                    "stderr": "",
1471                    "content": []
1472                }
1473            }
1474        }))
1475        .unwrap();
1476        let mut tool_call_state = None;
1477        let mut server_tool_uses = HashMap::new();
1478        let mut thinking_state = None;
1479
1480        let choice = super::handle_event(
1481            &event,
1482            &mut tool_call_state,
1483            &mut server_tool_uses,
1484            &mut thinking_state,
1485        )
1486        .expect("code_execution_tool_result block should produce raw metadata")
1487        .unwrap();
1488
1489        let RawStreamingChoice::TextStart {
1490            additional_params: Some(additional_params),
1491        } = choice
1492        else {
1493            panic!("expected text-start metadata for code_execution_tool_result");
1494        };
1495        assert_eq!(
1496            additional_params[crate::providers::anthropic::completion::ANTHROPIC_RAW_CONTENT_KEY]["type"],
1497            "code_execution_tool_result"
1498        );
1499        assert_eq!(
1500            additional_params[crate::providers::anthropic::completion::ANTHROPIC_RAW_CONTENT_KEY]["content"]
1501                ["stdout"],
1502            "42\n"
1503        );
1504    }
1505
1506    #[tokio::test]
1507    async fn test_streaming_web_search_blocks_are_preserved_on_final_choice() {
1508        let raw_stream = stream! {
1509            let mut tool_call_state = None;
1510            let mut server_tool_uses = HashMap::new();
1511            let mut thinking_state = None;
1512
1513            let server_tool_use_start = super::handle_event(
1514                &StreamingEvent::ContentBlockStart {
1515                    index: 0,
1516                    content_block: Content::ServerToolUse {
1517                        id: "srvtoolu_01".to_string(),
1518                        name: "web_search".to_string(),
1519                        input: serde_json::Value::Null,
1520                    },
1521                },
1522                &mut tool_call_state,
1523                &mut server_tool_uses,
1524                &mut thinking_state,
1525            );
1526            assert!(
1527                server_tool_use_start.is_none(),
1528                "server_tool_use start should be accumulated until its input JSON is complete"
1529            );
1530
1531            let server_tool_use_delta = super::handle_event(
1532                &StreamingEvent::ContentBlockDelta {
1533                    index: 0,
1534                    delta: ContentDelta::InputJsonDelta {
1535                        partial_json: r#"{"query":"claude shannon birth date"}"#.to_string(),
1536                    },
1537                },
1538                &mut tool_call_state,
1539                &mut server_tool_uses,
1540                &mut thinking_state,
1541            );
1542            assert!(
1543                server_tool_use_delta.is_none(),
1544                "server_tool_use input JSON should not be emitted as a Rig tool-call delta"
1545            );
1546
1547            yield super::handle_event(
1548                &StreamingEvent::ContentBlockStop { index: 0 },
1549                &mut tool_call_state,
1550                &mut server_tool_uses,
1551                &mut thinking_state,
1552            )
1553            .expect("server_tool_use stop should produce completed raw metadata");
1554
1555            yield super::handle_event(
1556                &StreamingEvent::ContentBlockStart {
1557                    index: 1,
1558                    content_block: Content::WebSearchToolResult {
1559                        tool_use_id: "srvtoolu_01".to_string(),
1560                        content: serde_json::json!([{
1561                            "type": "web_search_result",
1562                            "url": "https://example.com/shannon",
1563                            "title": "Claude Shannon",
1564                            "encrypted_content": "encrypted-content"
1565                        }]),
1566                    },
1567                },
1568                &mut tool_call_state,
1569                &mut server_tool_uses,
1570                &mut thinking_state,
1571            )
1572            .expect("web_search_tool_result block should produce raw metadata");
1573
1574            yield super::handle_event(
1575                &StreamingEvent::ContentBlockStart {
1576                    index: 2,
1577                    content_block: Content::Text {
1578                        text: String::new(),
1579                        citations: Vec::new(),
1580                        cache_control: None,
1581                    },
1582                },
1583                &mut tool_call_state,
1584                &mut server_tool_uses,
1585                &mut thinking_state,
1586            )
1587            .expect("text block start should produce a raw choice");
1588
1589            yield super::handle_event(
1590                &StreamingEvent::ContentBlockDelta {
1591                    index: 2,
1592                    delta: ContentDelta::TextDelta {
1593                        text: "Claude Shannon was born on April 30, 1916.".to_string(),
1594                    },
1595                },
1596                &mut tool_call_state,
1597                &mut server_tool_uses,
1598                &mut thinking_state,
1599            )
1600            .expect("text delta should produce a raw choice");
1601
1602            yield super::handle_event(
1603                &StreamingEvent::ContentBlockDelta {
1604                    index: 2,
1605                    delta: ContentDelta::CitationsDelta {
1606                        citation: crate::providers::anthropic::completion::Citation::WebSearchResultLocation {
1607                            cited_text: "Claude Shannon was born on April 30, 1916.".to_string(),
1608                            url: "https://example.com/shannon".to_string(),
1609                            title: Some("Claude Shannon".to_string()),
1610                            encrypted_index: "encrypted-index".to_string(),
1611                        },
1612                    },
1613                },
1614                &mut tool_call_state,
1615                &mut server_tool_uses,
1616                &mut thinking_state,
1617            )
1618            .expect("citation delta should produce a raw choice");
1619
1620            yield Ok(RawStreamingChoice::FinalResponse(StreamingCompletionResponse {
1621                usage: PartialUsage::default(),
1622            }));
1623        };
1624
1625        let mut stream =
1626            crate::streaming::StreamingCompletionResponse::stream(to_stream_result(raw_stream));
1627        while stream.next().await.is_some() {}
1628
1629        let choice_items: Vec<crate::message::AssistantContent> =
1630            stream.choice.clone().into_iter().collect();
1631        assert_eq!(choice_items.len(), 3);
1632        assert!(
1633            choice_items
1634                .iter()
1635                .all(|item| !matches!(item, crate::message::AssistantContent::ToolCall(_))),
1636            "provider-owned web-search blocks must not become Rig client tool calls"
1637        );
1638
1639        let Some(crate::message::AssistantContent::Text(server_tool_use)) = choice_items.first()
1640        else {
1641            panic!("expected raw server_tool_use metadata");
1642        };
1643        assert_eq!(
1644            server_tool_use.additional_params.as_ref().unwrap()
1645                [crate::providers::anthropic::completion::ANTHROPIC_RAW_CONTENT_KEY]["type"],
1646            "server_tool_use"
1647        );
1648        assert_eq!(
1649            server_tool_use.additional_params.as_ref().unwrap()
1650                [crate::providers::anthropic::completion::ANTHROPIC_RAW_CONTENT_KEY]["input"]["query"],
1651            "claude shannon birth date"
1652        );
1653
1654        let Some(crate::message::AssistantContent::Text(web_search_result)) = choice_items.get(1)
1655        else {
1656            panic!("expected raw web_search_tool_result metadata");
1657        };
1658        assert_eq!(
1659            web_search_result.additional_params.as_ref().unwrap()
1660                [crate::providers::anthropic::completion::ANTHROPIC_RAW_CONTENT_KEY]["content"][0]
1661                ["encrypted_content"],
1662            "encrypted-content"
1663        );
1664
1665        let Some(crate::message::AssistantContent::Text(answer)) = choice_items.get(2) else {
1666            panic!("expected answer text");
1667        };
1668        assert_eq!(answer.text, "Claude Shannon was born on April 30, 1916.");
1669        let citations = crate::providers::anthropic::completion::anthropic_citations(answer)
1670            .expect("expected preserved citations");
1671        assert!(matches!(
1672            citations.first(),
1673            Some(crate::providers::anthropic::completion::Citation::WebSearchResultLocation {
1674                encrypted_index,
1675                ..
1676            }) if encrypted_index == "encrypted-index"
1677        ));
1678    }
1679
1680    #[test]
1681    fn test_handle_citations_delta_event_preserves_metadata() {
1682        let event = StreamingEvent::ContentBlockDelta {
1683            index: 0,
1684            delta: ContentDelta::CitationsDelta {
1685                citation: crate::providers::anthropic::completion::Citation::CharLocation {
1686                    cited_text: "The grass is green.".to_string(),
1687                    document_index: 0,
1688                    document_title: Some("Example".to_string()),
1689                    start_char_index: 0,
1690                    end_char_index: 20,
1691                },
1692            },
1693        };
1694
1695        let mut tool_call_state = None;
1696        let mut thinking_state = None;
1697        let result = handle_event(&event, &mut tool_call_state, &mut thinking_state);
1698
1699        assert!(result.is_some());
1700        let choice = result.unwrap().unwrap();
1701        let RawStreamingChoice::TextAdditionalParams(additional_params) = choice else {
1702            panic!("expected TextAdditionalParams choice");
1703        };
1704        assert_eq!(additional_params["citations"][0]["type"], "char_location");
1705    }
1706
1707    #[tokio::test]
1708    async fn test_streaming_citation_deltas_are_preserved_on_final_text() {
1709        let citation = crate::providers::anthropic::completion::Citation::CharLocation {
1710            cited_text: "The grass is green.".to_string(),
1711            document_index: 0,
1712            document_title: Some("Example".to_string()),
1713            start_char_index: 0,
1714            end_char_index: 20,
1715        };
1716
1717        let raw_stream = stream! {
1718            let mut tool_call_state = None;
1719            let mut thinking_state = None;
1720
1721            yield handle_event(
1722                &StreamingEvent::ContentBlockStart {
1723                    index: 0,
1724                    content_block: Content::Text {
1725                        text: String::new(),
1726                        citations: Vec::new(),
1727                        cache_control: None,
1728                    },
1729                },
1730                &mut tool_call_state,
1731                &mut thinking_state,
1732            )
1733            .expect("text block start should produce a raw choice");
1734
1735            yield handle_event(
1736                &StreamingEvent::ContentBlockDelta {
1737                    index: 0,
1738                    delta: ContentDelta::TextDelta {
1739                        text: "the grass is green".to_string(),
1740                    },
1741                },
1742                &mut tool_call_state,
1743                &mut thinking_state,
1744            )
1745            .expect("text delta should produce a raw choice");
1746
1747            yield handle_event(
1748                &StreamingEvent::ContentBlockDelta {
1749                    index: 0,
1750                    delta: ContentDelta::CitationsDelta {
1751                        citation: crate::providers::anthropic::completion::Citation::CharLocation {
1752                            cited_text: "The grass is green.".to_string(),
1753                            document_index: 0,
1754                            document_title: Some("Example".to_string()),
1755                            start_char_index: 0,
1756                            end_char_index: 20,
1757                        },
1758                    },
1759                },
1760                &mut tool_call_state,
1761                &mut thinking_state,
1762            )
1763            .expect("citation delta should produce a raw choice");
1764
1765            yield Ok(RawStreamingChoice::FinalResponse(StreamingCompletionResponse {
1766                usage: PartialUsage::default(),
1767            }));
1768        };
1769
1770        let mut stream =
1771            crate::streaming::StreamingCompletionResponse::stream(to_stream_result(raw_stream));
1772        while stream.next().await.is_some() {}
1773
1774        let choice_items: Vec<crate::message::AssistantContent> =
1775            stream.choice.clone().into_iter().collect();
1776        let Some(crate::message::AssistantContent::Text(text)) = choice_items.first() else {
1777            panic!("expected accumulated text item");
1778        };
1779
1780        assert_eq!(text.text, "the grass is green");
1781        let citations = crate::providers::anthropic::completion::anthropic_citations(text).unwrap();
1782        assert_eq!(citations, vec![citation]);
1783    }
1784
1785    #[test]
1786    fn test_unknown_content_delta_falls_back() {
1787        let json = r#"{"type": "something_new_from_anthropic", "field": "x"}"#;
1788        let delta: ContentDelta = serde_json::from_str(json).unwrap();
1789        assert!(matches!(delta, ContentDelta::Unknown));
1790    }
1791}