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
23fn 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 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 map.insert("stream".to_string(), Value::Bool(true));
58
59 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 #[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 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 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 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 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 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 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 _ => 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 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 assert_eq!(streaming_body["stream"], serde_json::Value::Bool(true));
705
706 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 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 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 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 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 assert!(result.is_none());
996
997 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 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 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 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 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 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 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 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 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 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 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 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}