Skip to main content

infinity_provider_models_dev/
anthropic.rs

1//! Anthropic Messages backend with SSE streaming.
2//!
3//! Used for `ANTHROPIC_API_KEY` (and compatible env aliases). The rest of
4//! the presets go through the OpenAI-compatible backend in `openai.rs`.
5
6use 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
20/// How long to wait for the endpoint to accept a request (connect, send,
21/// response headers). This deliberately covers only initiation: once a
22/// stream is live it must never be cut off by a total deadline.
23const REQUEST_INITIATION_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(60);
24/// How long a live stream may go without yielding a chunk before it is
25/// treated as stalled.
26const STREAM_IDLE_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(120);
27const ANTHROPIC_VERSION: &str = "2023-06-01";
28
29/// How to reach and authenticate one `Anthropic` backend.
30pub struct Backend<'a> {
31    /// Shared HTTP client.
32    pub client: &'a Client,
33    /// `https://api.anthropic.com` (or override).
34    pub base_url: &'a str,
35    /// API key; ignored when `oauth` is present.
36    pub api_key: &'a str,
37    /// Extra headers (e.g. `OpenRouter` attribution).
38    pub extra_headers: &'a [(String, String)],
39    /// Subscription session; `None` for API-key auth.
40    pub oauth: Option<&'a OAuthSessionHandle>,
41}
42
43/// Stream one completion from the `Anthropic` Messages endpoint.
44///
45/// # Errors
46///
47/// Returns [`CompletionError::provider`] when the endpoint rejects the
48/// request (with the provider's retry classification) and
49/// [`CompletionError::ResponseError`] when a live stream breaks mid-flight.
50/// Individual stream items may also yield errors.
51pub 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    // Subscription sessions mint a bearer token (refreshing when needed);
66    // plain API keys pass through untouched.
67    let oauth_token: Option<String> = match backend.oauth {
68        Some(handle) => Some(ensure_bearer(backend.client, handle, false).await?),
69        None => None,
70    };
71    // Bound only initiation (connect, send, response headers). The body
72    // stream below has no total deadline; stalled streams trip the idle
73    // watchdog instead.
74    // A `401` means the bearer died between the freshness check and the
75    // request: refresh once and retry before surfacing the failure.
76    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
123/// Adapt a raw SSE byte stream into protocol chunks.
124///
125/// `idle_timeout` bounds silence between chunks, not the whole stream.
126/// Transport failures, stalls, and terminal error frames close the stream
127/// without emitting tool calls, reasoning, or a `Final` that would read
128/// as a successful completion.
129fn 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
197/// Split one SSE frame into its `event:` name and `data:` payload.
198fn 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/// Which kind of content block is currently open, if tracked.
219#[derive(Clone, Copy, PartialEq, Eq)]
220enum BlockKind {
221    Tool,
222    Thinking,
223    Other,
224}
225
226/// Mutable state for one `Anthropic` streaming response.
227struct 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    /// Flush anything still open at a clean end of stream (mirrors the
255    /// per-block flushes below for streams that end without them).
256    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    /// Complete the accumulated thinking block, if it has text.
263    ///
264    /// Mirrors the reference Bedrock provider: deltas stream the text live
265    /// and the complete block (text plus provider signature, when the
266    /// backend supplied one) is what gets recorded in history. Without the
267    /// signature, the next turn cannot resume extended thinking.
268    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    /// Complete the accumulated tool call, if it has a name or arguments.
281    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            // Thinking blocks (plain or redacted) accumulate server-side;
363            // only blocks with streamed text produce a Reasoning chunk.
364            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            // Tool blocks flush here; text blocks never accumulate tool
461            // state, and an unknown block kind still flushes a tool call in
462            // case its start event was missed. `flush_tool` is a no-op
463            // without a name or arguments.
464            _ => self.flush_tool(),
465        }
466    }
467}
468
469/// Bound a send to initiation (connect, send, response headers).
470///
471/// A timeout here means the endpoint never started responding; the same
472/// request may succeed on retry.
473async fn with_initiation_timeout(
474    send: impl Future<Output = Result<reqwest::Response, CompletionError>> + Send,
475    base_url: &str,
476) -> Result<reqwest::Response, CompletionError> {
477    // `post_messages` already maps transport failures; this only adds the
478    // initiation deadline.
479    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
489/// Resolve the subscription bearer token, refreshing when stale (`force`
490/// refreshes unconditionally after a `401`).
491async 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
511/// Whether a key is a subscription bearer rather than an API key.
512fn is_oauth_token(api_key: &str) -> bool {
513    api_key.starts_with("sk-ant-oauth")
514}
515
516/// Subscription auth headers: `Bearer` (no `x-api-key`) plus the `OAuth`
517/// beta and `Claude Code` identity headers, mirroring the community
518/// bridges. Shared by live sessions and pasted static bearer tokens.
519fn 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
527/// Send one `Messages` request with API-key or subscription auth.
528async 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        // Static subscription token via the `anthropic` preset.
546        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
672/// A tool result as an Anthropic `tool_result` block. Images become real
673/// `image` blocks only when the model accepts them; otherwise a text
674/// placeholder keeps the result well-formed without silent data loss.
675fn 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}