1use std::future::Future;
7
8use async_stream::stream;
9use futures_util::StreamExt;
10use infinity_provider_protocol::{
11 AssistantContent, CompletionError, CompletionRequest, ErrorClass, FinalResponse, Message,
12 ModelStream, Reasoning, StreamChunk, ToolCall, ToolCallDeltaContent, Usage, UserContent,
13};
14use reqwest::Client;
15use serde_json::{Value, json};
16
17use crate::convert::{classify_http_status, error_message_from_body};
18use crate::oauth::{self, OAuthSessionHandle};
19
20const REQUEST_INITIATION_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(60);
24const STREAM_IDLE_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(120);
27const ANTHROPIC_VERSION: &str = "2023-06-01";
28
29pub struct Backend<'a> {
31 pub client: &'a Client,
33 pub base_url: &'a str,
35 pub api_key: &'a str,
37 pub extra_headers: &'a [(String, String)],
39 pub oauth: Option<&'a OAuthSessionHandle>,
41}
42
43pub async fn invoke(
52 backend: &Backend<'_>,
53 upstream_model: &str,
54 default_max_tokens: Option<u64>,
55 request: CompletionRequest,
56 supports_images: bool,
57) -> Result<ModelStream, CompletionError> {
58 let body = request_body(
59 upstream_model,
60 default_max_tokens,
61 &request,
62 supports_images,
63 );
64
65 let oauth_token: Option<String> = match backend.oauth {
68 Some(handle) => Some(ensure_bearer(backend.client, handle, false).await?),
69 None => None,
70 };
71 let mut resp = with_initiation_timeout(
77 post_messages(
78 backend.client,
79 backend.base_url,
80 backend.api_key,
81 oauth_token.as_deref(),
82 backend.extra_headers,
83 &body,
84 ),
85 backend.base_url,
86 )
87 .await?;
88 if resp.status() == reqwest::StatusCode::UNAUTHORIZED {
89 if let Some(handle) = backend.oauth {
90 let token = ensure_bearer(backend.client, handle, true).await?;
91 resp = with_initiation_timeout(
92 post_messages(
93 backend.client,
94 backend.base_url,
95 backend.api_key,
96 Some(token.as_str()),
97 backend.extra_headers,
98 &body,
99 ),
100 backend.base_url,
101 )
102 .await?;
103 }
104 }
105
106 if !resp.status().is_success() {
107 let status = resp.status().as_u16();
108 let text = resp.text().await.unwrap_or_default();
109 let class = classify_http_status(status, &text);
110 return Err(CompletionError::provider(
111 class,
112 format!(
113 "{} -> {status}: {}",
114 backend.base_url,
115 error_message_from_body(&text)
116 ),
117 ));
118 }
119
120 Ok(adapt_stream(resp.bytes_stream(), STREAM_IDLE_TIMEOUT))
121}
122
123fn adapt_stream<S, B, E>(mut byte_stream: S, idle_timeout: std::time::Duration) -> ModelStream
130where
131 S: futures_util::Stream<Item = Result<B, E>> + Unpin + Send + 'static,
132 B: AsRef<[u8]> + Send + 'static,
133 E: std::error::Error + Send + Sync + 'static,
134{
135 Box::pin(stream! {
136 let mut buf = String::new();
137 let mut pending: Vec<u8> = Vec::new();
138 let mut state = StreamState::new();
139 let mut failed = false;
140
141 'chunks: loop {
142 let chunk = tokio::time::timeout(idle_timeout, byte_stream.next()).await;
143 let chunk = match chunk {
144 Err(_) => {
145 yield Err(CompletionError::ResponseError(
146 "stream stalled: no chunks received".to_owned(),
147 ));
148 failed = true;
149 break;
150 }
151 Ok(None) => break,
152 Ok(Some(chunk)) => chunk,
153 };
154 let bytes = match chunk {
155 Ok(b) => b,
156 Err(e) => {
157 yield Err(CompletionError::ResponseError(format!("stream read failed: {e}")));
158 failed = true;
159 break;
160 }
161 };
162 crate::convert::push_bytes(&mut buf, &mut pending, bytes.as_ref());
163 for frame in crate::convert::take_frames(&mut buf) {
164 let Some((event, data)) = split_event(&frame) else {
165 continue;
166 };
167 if data.is_empty() || event == "ping" {
168 continue;
169 }
170 let Ok(v) = serde_json::from_str::<Value>(&data) else {
171 tracing::debug!(
172 "anthropic: skipping malformed frame: {}",
173 data.chars().take(200).collect::<String>()
174 );
175 continue;
176 };
177 let is_terminal_error = event == "error";
178 for item in state.handle_event(&event, &v) {
179 yield item;
180 }
181 if is_terminal_error {
182 failed = true;
183 break 'chunks;
184 }
185 }
186 }
187 if !failed {
188 for item in state.finish() {
189 yield item;
190 }
191 let usage = (state.usage != Usage::default()).then_some(state.usage);
192 yield Ok(StreamChunk::Final(FinalResponse { usage }));
193 }
194 })
195}
196
197fn split_event(frame: &str) -> Option<(String, String)> {
199 let mut event = String::new();
200 let mut data = String::new();
201 for line in frame.lines() {
202 let line = line.trim();
203 if let Some(e) = line.strip_prefix("event:") {
204 e.trim().clone_into(&mut event);
205 } else if let Some(d) = line.strip_prefix("data:") {
206 if !data.is_empty() {
207 data.push('\n');
208 }
209 data.push_str(d.trim());
210 }
211 }
212 if event.is_empty() && data.is_empty() {
213 return None;
214 }
215 Some((event, data))
216}
217
218#[derive(Clone, Copy, PartialEq, Eq)]
220enum BlockKind {
221 Tool,
222 Thinking,
223 Other,
224}
225
226struct StreamState {
228 usage: Usage,
229 tool_index: usize,
230 tool_id: String,
231 tool_name: String,
232 tool_args: String,
233 tool_name_sent: bool,
234 open_block: Option<BlockKind>,
235 thinking: String,
236 thinking_signature: Option<String>,
237}
238
239impl StreamState {
240 fn new() -> Self {
241 Self {
242 usage: Usage::default(),
243 tool_index: 0,
244 tool_id: String::new(),
245 tool_name: String::new(),
246 tool_args: String::new(),
247 tool_name_sent: false,
248 open_block: None,
249 thinking: String::new(),
250 thinking_signature: None,
251 }
252 }
253
254 fn finish(&mut self) -> Vec<Result<StreamChunk, CompletionError>> {
257 let mut out = self.flush_tool();
258 out.extend(self.flush_reasoning());
259 out
260 }
261
262 fn flush_reasoning(&mut self) -> Vec<Result<StreamChunk, CompletionError>> {
269 if self.thinking.is_empty() {
270 self.thinking_signature = None;
271 return Vec::new();
272 }
273 let reasoning = Reasoning::new_with_signature(
274 std::mem::take(&mut self.thinking),
275 self.thinking_signature.take(),
276 );
277 vec![Ok(StreamChunk::Reasoning(reasoning))]
278 }
279
280 fn flush_tool(&mut self) -> Vec<Result<StreamChunk, CompletionError>> {
282 if self.tool_name.is_empty() && self.tool_args.is_empty() {
283 self.tool_id.clear();
284 self.tool_name_sent = false;
285 return Vec::new();
286 }
287 let args: Value = if self.tool_args.trim().is_empty() {
288 json!({})
289 } else {
290 serde_json::from_str(&self.tool_args)
291 .unwrap_or_else(|_| json!({"_raw": self.tool_args}))
292 };
293 let id = if self.tool_id.is_empty() {
294 format!("tool-{}", self.tool_index)
295 } else {
296 self.tool_id.clone()
297 };
298 self.tool_index = self.tool_index.saturating_add(1);
299 let call = StreamChunk::ToolCall(ToolCall::new(id, self.tool_name.clone(), args));
300 self.tool_id.clear();
301 self.tool_name.clear();
302 self.tool_args.clear();
303 self.tool_name_sent = false;
304 vec![Ok(call)]
305 }
306
307 fn handle_event(
308 &mut self,
309 event: &str,
310 v: &Value,
311 ) -> Vec<Result<StreamChunk, CompletionError>> {
312 match event {
313 "message_start" => {
314 if let Some(u) = v.pointer("/message/usage").map(parse_anthropic_usage) {
315 self.usage = u;
316 }
317 Vec::new()
318 }
319 "content_block_start" => self.on_block_start(v),
320 "content_block_delta" => self.on_block_delta(v),
321 "content_block_stop" => self.on_block_stop(),
322 "message_delta" => {
323 if let Some(u) = v.get("usage").map(parse_anthropic_usage) {
324 self.usage.output_tokens = self.usage.output_tokens.max(u.output_tokens);
325 self.usage.total_tokens = self
326 .usage
327 .input_tokens
328 .saturating_add(self.usage.output_tokens);
329 if u.cached_input_tokens > 0 {
330 self.usage.cached_input_tokens = u.cached_input_tokens;
331 }
332 }
333 Vec::new()
334 }
335 "message_stop" => {
336 if let Some(u) = v.get("usage").map(parse_anthropic_usage) {
337 self.usage = u;
338 }
339 Vec::new()
340 }
341 "error" => {
342 let msg = v
343 .pointer("/error/message")
344 .and_then(Value::as_str)
345 .or_else(|| v.get("message").and_then(Value::as_str))
346 .unwrap_or("anthropic stream error");
347 vec![Err(CompletionError::provider(
348 classify_http_status(400, msg),
349 msg.to_owned(),
350 ))]
351 }
352 _ => Vec::new(),
353 }
354 }
355
356 fn on_block_start(&mut self, v: &Value) -> Vec<Result<StreamChunk, CompletionError>> {
357 let block = v.get("content_block").unwrap_or(&Value::Null);
358 match block.get("type").and_then(Value::as_str) {
359 Some("tool_use") => {
360 self.open_block = Some(BlockKind::Tool);
361 }
362 Some("thinking" | "redacted_thinking") => {
365 self.open_block = Some(BlockKind::Thinking);
366 return Vec::new();
367 }
368 _ => {
369 self.open_block = Some(BlockKind::Other);
370 return Vec::new();
371 }
372 }
373 block
374 .get("id")
375 .and_then(Value::as_str)
376 .unwrap_or("")
377 .clone_into(&mut self.tool_id);
378 block
379 .get("name")
380 .and_then(Value::as_str)
381 .unwrap_or("")
382 .clone_into(&mut self.tool_name);
383 self.tool_args.clear();
384 self.tool_name_sent = false;
385 if self.tool_id.is_empty() {
386 self.tool_id = format!("tool-{}", self.tool_index);
387 }
388 if self.tool_name.is_empty() {
389 return Vec::new();
390 }
391 self.tool_name_sent = true;
392 vec![Ok(StreamChunk::ToolCallDelta {
393 id: self.tool_id.clone(),
394 content: ToolCallDeltaContent::Name(self.tool_name.clone()),
395 })]
396 }
397
398 fn on_block_delta(&mut self, v: &Value) -> Vec<Result<StreamChunk, CompletionError>> {
399 let delta = v.get("delta").unwrap_or(&Value::Null);
400 match delta.get("type").and_then(Value::as_str) {
401 Some("text_delta") => {
402 if let Some(t) = delta.get("text").and_then(Value::as_str) {
403 if !t.is_empty() {
404 return vec![Ok(StreamChunk::Text(t.to_owned()))];
405 }
406 }
407 Vec::new()
408 }
409 Some("thinking_delta") => {
410 if let Some(t) = delta.get("thinking").and_then(Value::as_str) {
411 if !t.is_empty() {
412 self.thinking.push_str(t);
413 return vec![Ok(StreamChunk::ReasoningDelta {
414 id: None,
415 text: t.to_owned(),
416 })];
417 }
418 }
419 Vec::new()
420 }
421 Some("signature_delta") => {
422 if let Some(signature) = delta.get("signature").and_then(Value::as_str) {
423 if !signature.is_empty() {
424 self.thinking_signature = Some(signature.to_owned());
425 }
426 }
427 Vec::new()
428 }
429 Some("input_json_delta") => self.on_tool_fragment(delta),
430 _ => Vec::new(),
431 }
432 }
433
434 fn on_tool_fragment(&mut self, delta: &Value) -> Vec<Result<StreamChunk, CompletionError>> {
435 let Some(frag) = delta.get("partial_json").and_then(Value::as_str) else {
436 return Vec::new();
437 };
438 self.tool_args.push_str(frag);
439 let mut out = Vec::new();
440 if !self.tool_name_sent && !self.tool_name.is_empty() {
441 self.tool_name_sent = true;
442 out.push(Ok(StreamChunk::ToolCallDelta {
443 id: self.tool_id.clone(),
444 content: ToolCallDeltaContent::Name(self.tool_name.clone()),
445 }));
446 }
447 if !frag.is_empty() {
448 out.push(Ok(StreamChunk::ToolCallDelta {
449 id: self.tool_id.clone(),
450 content: ToolCallDeltaContent::Delta(frag.to_owned()),
451 }));
452 }
453 out
454 }
455
456 fn on_block_stop(&mut self) -> Vec<Result<StreamChunk, CompletionError>> {
457 let kind = self.open_block.take();
458 match kind {
459 Some(BlockKind::Thinking) => self.flush_reasoning(),
460 _ => self.flush_tool(),
465 }
466 }
467}
468
469async fn with_initiation_timeout(
474 send: impl Future<Output = Result<reqwest::Response, CompletionError>> + Send,
475 base_url: &str,
476) -> Result<reqwest::Response, CompletionError> {
477 tokio::time::timeout(REQUEST_INITIATION_TIMEOUT, send)
480 .await
481 .map_err(|_| {
482 CompletionError::provider(
483 ErrorClass::Transient,
484 format!("{base_url}: request initiation timed out"),
485 )
486 })?
487}
488
489async fn ensure_bearer(
492 client: &Client,
493 handle: &OAuthSessionHandle,
494 force: bool,
495) -> Result<String, CompletionError> {
496 let refreshed = if force {
497 oauth::force_refresh(client, handle).await
498 } else {
499 oauth::ensure_fresh_token(client, handle).await
500 };
501 refreshed.map_err(|e| {
502 let class = if e.retryable {
503 ErrorClass::Transient
504 } else {
505 ErrorClass::Fatal
506 };
507 CompletionError::provider(class, format!("oauth: {}", e.message))
508 })
509}
510
511fn is_oauth_token(api_key: &str) -> bool {
513 api_key.starts_with("sk-ant-oauth")
514}
515
516fn with_oauth_headers(req: reqwest::RequestBuilder, token: &str) -> reqwest::RequestBuilder {
520 req.bearer_auth(token)
521 .header("anthropic-beta", oauth::OAUTH_BETA)
522 .header("anthropic-dangerous-direct-browser-access", "true")
523 .header("User-Agent", oauth::CLAUDE_CODE_USER_AGENT)
524 .header("x-app", "cli")
525}
526
527async fn post_messages(
529 client: &Client,
530 base_url: &str,
531 api_key: &str,
532 oauth_token: Option<&str>,
533 extra_headers: &[(String, String)],
534 body: &Value,
535) -> Result<reqwest::Response, CompletionError> {
536 let url = format!("{}/messages", base_url.trim_end_matches('/'));
537 let mut req = client
538 .post(&url)
539 .header("Content-Type", "application/json")
540 .header("Accept", "text/event-stream")
541 .header("anthropic-version", ANTHROPIC_VERSION);
542 if let Some(token) = oauth_token {
543 req = with_oauth_headers(req, token);
544 } else if is_oauth_token(api_key) {
545 req = with_oauth_headers(req, api_key);
547 } else {
548 req = req
549 .header("x-api-key", api_key)
550 .header("User-Agent", crate::USER_AGENT);
551 }
552 for (k, v) in extra_headers {
553 req = req.header(k.as_str(), v.as_str());
554 }
555 req.json(body).send().await.map_err(transport_error)
556}
557
558fn request_body(
559 upstream_model: &str,
560 default_max_tokens: Option<u64>,
561 request: &CompletionRequest,
562 supports_images: bool,
563) -> Value {
564 let (system, messages) = to_anthropic_messages(
565 request.preamble.as_deref(),
566 &request.chat_history,
567 supports_images,
568 );
569 let mut map = serde_json::Map::new();
570 map.insert("model".to_owned(), json!(upstream_model));
571 map.insert("messages".to_owned(), Value::Array(messages));
572 map.insert("stream".to_owned(), Value::Bool(true));
573 if let Some(system) = system {
574 map.insert("system".to_owned(), system);
575 }
576 let max = request.max_tokens.or(default_max_tokens).unwrap_or(16_384);
577 map.insert("max_tokens".to_owned(), json!(max));
578 if !request.tools.is_empty() {
579 let tools: Vec<Value> = request
580 .tools
581 .iter()
582 .map(|t| {
583 json!({
584 "name": t.name,
585 "description": t.description,
586 "input_schema": t.parameters,
587 })
588 })
589 .collect();
590 map.insert("tools".to_owned(), Value::Array(tools));
591 }
592 let mut body = Value::Object(map);
593 if let Some(extra) = request.additional_params.as_ref() {
594 merge_value(&mut body, extra);
595 }
596 body
597}
598
599fn parse_anthropic_usage(u: &Value) -> Usage {
600 let input = u.get("input_tokens").and_then(Value::as_u64).unwrap_or(0);
601 let output = u.get("output_tokens").and_then(Value::as_u64).unwrap_or(0);
602 Usage {
603 input_tokens: input,
604 output_tokens: output,
605 total_tokens: input.saturating_add(output),
606 cached_input_tokens: u
607 .get("cache_read_input_tokens")
608 .and_then(Value::as_u64)
609 .unwrap_or(0),
610 }
611}
612
613fn to_anthropic_messages(
614 preamble: Option<&str>,
615 history: &[Message],
616 supports_images: bool,
617) -> (Option<Value>, Vec<Value>) {
618 let system = preamble.map(|s| json!([{"type": "text", "text": s}]));
619 let mut messages = Vec::new();
620 for msg in history {
621 match msg {
622 Message::User { content } => {
623 let mut blocks = Vec::new();
624 for c in content {
625 match c {
626 UserContent::Text(t) => {
627 blocks.push(json!({"type": "text", "text": t.text}));
628 }
629 UserContent::Image(img) => {
630 if supports_images {
631 blocks.push(image_block(img));
632 } else {
633 blocks.push(json!({"type": "text", "text": "[image omitted]"}));
634 }
635 }
636 UserContent::ToolResult(r) => {
637 blocks.push(tool_result_block(r, supports_images));
638 }
639 }
640 }
641 if !blocks.is_empty() {
642 messages.push(json!({"role": "user", "content": blocks}));
643 }
644 }
645 Message::Assistant { content } => {
646 let mut blocks = Vec::new();
647 for c in content {
648 match c {
649 AssistantContent::Text(t) => {
650 blocks.push(json!({"type": "text", "text": t.text}));
651 }
652 AssistantContent::ToolCall(tc) => blocks.push(json!({
653 "type": "tool_use",
654 "id": tc.id,
655 "name": tc.function.name,
656 "input": tc.function.arguments,
657 })),
658 AssistantContent::Reasoning(r) => {
659 blocks.push(reasoning_block(r));
660 }
661 }
662 }
663 if !blocks.is_empty() {
664 messages.push(json!({"role": "assistant", "content": blocks}));
665 }
666 }
667 }
668 }
669 (system, messages)
670}
671
672fn tool_result_block(
676 result: &infinity_provider_protocol::ToolResult,
677 supports_images: bool,
678) -> Value {
679 let mut text = String::new();
680 let mut images = Vec::new();
681 for block in &result.content {
682 match block {
683 infinity_provider_protocol::ToolResultContent::Text(t) => {
684 if !text.is_empty() {
685 text.push('\n');
686 }
687 text.push_str(&t.text);
688 }
689 infinity_provider_protocol::ToolResultContent::Image(img) => {
690 if supports_images {
691 images.push(image_block(img));
692 } else {
693 if !text.is_empty() {
694 text.push('\n');
695 }
696 text.push_str("[image omitted]");
697 }
698 }
699 }
700 }
701 if images.is_empty() {
702 return json!({
703 "type": "tool_result",
704 "tool_use_id": result.id,
705 "content": text,
706 });
707 }
708 let mut blocks = Vec::new();
709 if !text.is_empty() {
710 blocks.push(json!({"type": "text", "text": text}));
711 }
712 blocks.extend(images);
713 json!({
714 "type": "tool_result",
715 "tool_use_id": result.id,
716 "content": blocks,
717 })
718}
719
720fn image_block(img: &infinity_provider_protocol::Image) -> Value {
721 match &img.data {
722 infinity_provider_protocol::ImageSource::Url(u) => json!({
723 "type": "image",
724 "source": {"type": "url", "url": u},
725 }),
726 infinity_provider_protocol::ImageSource::Base64(b) => {
727 let mime = img.media_type.as_ref().map_or(
728 "image/png",
729 infinity_provider_protocol::ImageMediaType::to_mime_type,
730 );
731 json!({
732 "type": "image",
733 "source": {"type": "base64", "media_type": mime, "data": b},
734 })
735 }
736 }
737}
738
739fn reasoning_block(r: &Reasoning) -> Value {
740 let text = r.display_text();
741 let signature = r.first_signature().unwrap_or("");
742 json!({"type": "thinking", "thinking": text, "signature": signature})
743}
744
745fn merge_value(dst: &mut Value, src: &Value) {
746 match (dst, src) {
747 (Value::Object(d), Value::Object(s)) => {
748 for (k, v) in s {
749 merge_value(d.entry(k.clone()).or_insert(Value::Null), v);
750 }
751 }
752 (d, s) => *d = s.clone(),
753 }
754}
755
756fn transport_error(e: reqwest::Error) -> CompletionError {
757 if e.is_timeout() || e.is_connect() || e.is_body() || e.is_decode() {
758 CompletionError::provider(ErrorClass::Transient, format!("transport error: {e}"))
759 } else if e.is_request() {
760 CompletionError::RequestError(Box::new(e))
761 } else {
762 CompletionError::provider(ErrorClass::Transient, format!("transport error: {e}"))
763 }
764}
765
766#[cfg(test)]
767mod tests {
768 use super::*;
769 use futures_util::stream;
770
771 fn frame(event: &str, data: &str) -> Vec<u8> {
772 format!("event: {event}\ndata: {data}\n\n").into_bytes()
773 }
774
775 async fn collect<S>(byte_stream: S) -> Vec<Result<StreamChunk, CompletionError>>
776 where
777 S: futures_util::Stream<Item = Result<Vec<u8>, std::io::Error>> + Unpin + Send + 'static,
778 {
779 adapt_stream(byte_stream, std::time::Duration::from_secs(60))
780 .collect()
781 .await
782 }
783
784 fn text_frame(text: &str) -> Vec<u8> {
785 frame(
786 "content_block_delta",
787 &format!("{{\"delta\":{{\"type\":\"text_delta\",\"text\":{text:?}}}}}"),
788 )
789 }
790
791 #[tokio::test]
792 async fn thinking_block_yields_reasoning_with_signature() {
793 let start: Vec<u8> = b"event: message_start\ndata: {\"message\":{\"usage\":{\"input_tokens\":10,\"output_tokens\":0}}}\n\n".to_vec();
794 let thinking_start: Vec<u8> =
795 b"event: content_block_start\ndata: {\"content_block\":{\"type\":\"thinking\"}}\n\n"
796 .to_vec();
797 let thinking: Vec<u8> = b"event: content_block_delta\ndata: {\"delta\":{\"type\":\"thinking_delta\",\"thinking\":\"let me think\"}}\n\n".to_vec();
798 let signature: Vec<u8> = b"event: content_block_delta\ndata: {\"delta\":{\"type\":\"signature_delta\",\"signature\":\"sig123\"}}\n\n".to_vec();
799 let stop: Vec<u8> = b"event: content_block_stop\ndata: {}\n\n".to_vec();
800 let delta: Vec<u8> =
801 b"event: message_delta\ndata: {\"usage\":{\"output_tokens\":5}}\n\n".to_vec();
802 let end: Vec<u8> = b"event: message_stop\ndata: {}\n\n".to_vec();
803 let chunks: Vec<Result<Vec<u8>, std::io::Error>> =
804 vec![start, thinking_start, thinking, signature, stop, delta, end]
805 .into_iter()
806 .map(Ok)
807 .collect();
808 let out = collect(stream::iter(chunks)).await;
809 let reasoning = out.iter().find_map(|item| match item {
810 Ok(StreamChunk::Reasoning(r)) => Some(r),
811 _ => None,
812 });
813 let reasoning = reasoning.expect("expected a Reasoning chunk");
814 assert_eq!(reasoning.first_text(), Some("let me think"));
815 assert_eq!(reasoning.first_signature(), Some("sig123"));
816 assert!(
817 out.iter()
818 .any(|item| matches!(item, Ok(StreamChunk::ReasoningDelta { .. }))),
819 "expected streamed ReasoningDelta, got: {out:?}"
820 );
821 assert!(
822 out.iter().any(|item| matches!(
823 item,
824 Ok(StreamChunk::Final(FinalResponse { usage: Some(_) }))
825 )),
826 "expected Final with usage, got: {out:?}"
827 );
828 }
829
830 #[tokio::test]
831 async fn stream_error_terminates_without_final() {
832 let start: Vec<u8> =
833 b"event: content_block_start\ndata: {\"content_block\":{\"type\":\"text\"}}\n\n"
834 .to_vec();
835 let text: Vec<u8> = b"event: content_block_delta\ndata: {\"delta\":{\"type\":\"text_delta\",\"text\":\"hi\"}}\n\n".to_vec();
836 let error: Vec<u8> =
837 b"event: error\ndata: {\"message\":\"overloaded\",\"type\":\"overloaded_error\"}\n\n"
838 .to_vec();
839 let unreached: Vec<u8> = b"event: content_block_delta\ndata: {\"delta\":{\"type\":\"text_delta\",\"text\":\"unreached\"}}\n\n".to_vec();
840 let chunks: Vec<Result<Vec<u8>, std::io::Error>> = vec![start, text, error, unreached]
841 .into_iter()
842 .map(Ok)
843 .collect();
844 let out = collect(stream::iter(chunks)).await;
845 assert_eq!(out.len(), 2, "error must close the stream: {out:?}");
846 assert!(matches!(&out[0], Ok(StreamChunk::Text(t)) if t == "hi"));
847 match &out[1] {
848 Err(CompletionError::ProviderError { class, .. }) => {
849 assert_eq!(*class, ErrorClass::Transient);
850 }
851 other => panic!("expected transient provider error, got {other:?}"),
852 }
853 }
854
855 #[tokio::test]
856 async fn no_usage_means_no_usage_in_final() {
857 let chunks: Vec<Result<Vec<u8>, std::io::Error>> = vec![Ok(text_frame("hi"))];
858 let out = collect(stream::iter(chunks)).await;
859 assert!(
860 out.iter()
861 .any(|item| matches!(item, Ok(StreamChunk::Final(FinalResponse { usage: None })))),
862 "unreported usage must stay None, got: {out:?}"
863 );
864 }
865}