rust_llm 2.1.1

//! Port of `lib/ruby_llm/protocol.rb`, `protocol/streaming.rb`, and `protocol/stream_accumulator.rb`.
//!
//! RubyLLM builds a protocol class per provider by mixing overrides into a base wire format
//! (`class Mistral::ChatCompletions < Protocols::ChatCompletions; include Mistral::Chat`). Here each
//! wire format is a module of functions that consult the provider for its dialect differences.

pub mod anthropic;
pub mod chat_completions;
pub mod gemini;
pub mod interactions;
pub mod mistral;
pub mod responses;

use serde_json::{Map, Value};

use crate::error::{Error, Result};
use crate::message::{
    Citation, FinishReason, Message, RawResponse, ServerToolCall, Thinking, ToolArguments,
    ToolCall, indexmap_lite::IndexMap,
};
use crate::model::Model;
use crate::providers::{ProtocolName, Provider};
use crate::thinking::ThinkingConfig;
use crate::tokens::Tokens;
use crate::tool::SharedTool;

/// `tool_prefs`: `with_tool_options(choice:, calls:)`.
#[derive(Debug, Clone, Default, PartialEq)]
pub struct ToolPrefs {
    pub choice: Option<ToolChoice>,
    pub calls: Option<ToolCalls>,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ToolChoice {
    Auto,
    None,
    Required,
    Tool(String),
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ToolCalls {
    One,
    Many,
}

/// The normalized `with_schema` payload: `{ name:, schema:, strict: }`.
#[derive(Debug, Clone, PartialEq)]
pub struct Schema {
    pub name: String,
    pub schema: Value,
    pub strict: Option<bool>,
    pub description: Option<String>,
}

/// Everything `Protocol#render_payload` receives.
pub struct Request<'a> {
    pub provider: Provider,
    pub config: &'a crate::Config,
    pub model: &'a Model,
    pub messages: &'a [Message],
    pub tools: &'a [SharedTool],
    pub tool_prefs: &'a ToolPrefs,
    pub temperature: Option<f64>,
    pub max_output_tokens: Option<i64>,
    pub schema: Option<&'a Schema>,
    pub thinking: Option<&'a ThinkingConfig>,
    /// `with_citations`.
    pub citations: bool,
    /// `with_caching`: `None` when not configured.
    pub caching: Option<&'a Caching>,
    pub stream: bool,
}

/// `with_caching(options)` / `with_caching(false)`, the chat's `@caching`.
#[derive(Debug, Clone, PartialEq)]
pub enum Caching {
    /// `with_caching(false)`: send no cache controls and render no cache boundaries.
    Off,
    /// `with_caching(key:, ttl:, mode:, id:, ...)`; empty for the provider's default behavior.
    On(Map<String, Value>),
}

impl Caching {
    /// `caching != false`: whether marked cache boundaries render.
    pub(crate) fn boundaries(caching: Option<&Caching>) -> bool {
        caching != Some(&Caching::Off)
    }

    /// The options of an enabled cache, or `None` when unset or off.
    pub(crate) fn options(caching: Option<&Caching>) -> Option<&Map<String, Value>> {
        match caching {
            Some(Caching::On(options)) => Some(options),
            _ => None,
        }
    }

    /// `prompt_cache_options`: rejects keys outside `allowed`, naming them as Ruby symbols.
    pub(crate) fn checked<'a>(
        caching: Option<&'a Caching>,
        allowed: &[&str],
        label: &str,
    ) -> Result<Option<&'a Map<String, Value>>> {
        let Some(options) = Caching::options(caching) else {
            return Ok(None);
        };
        let unsupported: Vec<String> = options
            .keys()
            .filter(|k| !allowed.contains(&k.as_str()))
            .map(|k| format!(":{k}"))
            .collect();
        if unsupported.is_empty() {
            return Ok(Some(options));
        }
        Err(Error::Argument(format!(
            "{label} prompt caching accepts {}, got {}",
            symbols(allowed),
            unsupported.join(", ")
        )))
    }
}

/// `:key, :ttl, and :mode` for the Responses/Chat Completions messages, `:ttl` for one key.
fn symbols(keys: &[&str]) -> String {
    let visible: Vec<String> = keys
        .iter()
        .filter(|k| **k != "retention")
        .map(|k| format!(":{k}"))
        .collect();
    match visible.len() {
        0 | 1 => visible.join(""),
        n => format!("{}, and {}", visible[..n - 1].join(", "), visible[n - 1]),
    }
}

/// `apply_end_user`: each protocol's safety identifier field; the rest drop it.
pub(crate) fn apply_end_user(
    protocol: ProtocolName,
    provider: Provider,
    payload: &mut Value,
    identifier: &str,
) {
    let field = match (protocol, provider) {
        (ProtocolName::Anthropic, _) => {
            deep_merge(
                payload,
                &serde_json::json!({ "metadata": { "user_id": identifier } }),
            );
            return;
        }
        (ProtocolName::ChatCompletions | ProtocolName::Responses, Provider::OpenAI) => {
            "safety_identifier"
        }
        (ProtocolName::ChatCompletions | ProtocolName::Responses, Provider::DeepSeek) => "user_id",
        (ProtocolName::ChatCompletions, Provider::OpenRouter) => "user",
        _ => {
            tracing::debug!(
                "{} has no safety identifier parameter, dropping {identifier}",
                provider.display()
            );
            return;
        }
    };
    if let Some(p) = payload.as_object_mut() {
        p.insert(field.into(), identifier.into());
    }
}

/// `apply_compaction`: Anthropic's `context_management` edit, OpenAI's Responses
/// `compact_threshold` entry, or OpenRouter's context-compression plugin; dropped elsewhere.
pub(crate) fn apply_compaction(
    protocol: ProtocolName,
    provider: Provider,
    payload: &mut Value,
    compaction: &Map<String, Value>,
) {
    match (protocol, provider) {
        (ProtocolName::Anthropic, _) => anthropic::apply_compaction(payload, compaction),
        (ProtocolName::Responses, Provider::OpenAI) => {
            let mut entry = serde_json::json!({ "type": "compaction" });
            if let Some(at) = compaction
                .get("at")
                .filter(|v| !v.is_null() && **v != Value::Bool(false))
            {
                entry["compact_threshold"] = at.clone();
            }
            if let Some(p) = payload.as_object_mut() {
                p.insert("context_management".into(), serde_json::json!([entry]));
            }
        }
        (ProtocolName::ChatCompletions, Provider::OpenRouter) => {
            if !compaction.is_empty() {
                tracing::debug!(
                    "OpenRouter compresses context at the model's own limit, dropping {compaction:?}"
                );
            }
            if let Some(p) = payload.as_object_mut() {
                let mut plugins = p
                    .get("plugins")
                    .and_then(Value::as_array)
                    .cloned()
                    .unwrap_or_default();
                plugins.push(serde_json::json!({ "id": "context-compression" }));
                p.insert("plugins".into(), Value::Array(plugins));
            }
        }
        _ => tracing::debug!(
            "{} has no context compaction parameter, dropping {compaction:?}",
            provider.display()
        ),
    }
}

/// Where and how to send a rendered request.
pub struct Endpoint {
    pub path: String,
    pub headers: Vec<(String, String)>,
}

/// `Protocol#supported_attachment?`: whether `protocol`, as `provider` speaks it, can render the
/// attachment. Mirrors each media module's predicate (`protocols/*/media.rb`, and the provider
/// dialects in `providers/*/chat.rb` and `media.rb`).
pub fn supported_attachment(
    protocol: ProtocolName,
    provider: Provider,
    attachment: &crate::attachment::Attachment,
) -> bool {
    use crate::attachment::AttachmentType as T;
    let kind = attachment.kind();
    let provider_file = attachment.is_provider_file();
    // `ChatCompletions::Media.supported_attachment?(document_attachments:, audio_attachments:)`.
    let chat_completions = |documents: Option<bool>, audio: bool| {
        if provider_file {
            return documents.is_some();
        }
        match kind {
            T::Image | T::Text => true,
            T::Audio => audio,
            T::Pdf => documents.is_some(),
            T::Document => documents == Some(true),
            _ => false,
        }
    };
    match protocol {
        ProtocolName::Anthropic => provider_file || matches!(kind, T::Image | T::Pdf | T::Text),
        ProtocolName::Gemini => {
            matches!(kind, T::Image | T::Audio | T::Video | T::Pdf | T::Text)
        }
        ProtocolName::Interactions => {
            matches!(kind, T::Image | T::Audio | T::Video | T::Pdf | T::Text)
        }
        // `Perplexity::Agent`: provider files, images, and text.
        ProtocolName::Responses if provider == Provider::Perplexity => {
            provider_file || matches!(kind, T::Image | T::Text)
        }
        ProtocolName::Responses if provider == Provider::DeepSeek => {
            if provider_file {
                kind == T::Image
            } else {
                matches!(kind, T::Image | T::Text)
            }
        }
        ProtocolName::Responses => {
            provider_file || matches!(kind, T::Image | T::Pdf | T::Document | T::Text)
        }
        ProtocolName::Conversations => {
            matches!(kind, T::Image | T::Audio | T::Pdf | T::Document | T::Text)
        }
        ProtocolName::RouterChatCompletions => chat_completions(Some(false), true),
        ProtocolName::ChatCompletions => match provider {
            Provider::Mistral => {
                matches!(kind, T::Image | T::Audio | T::Pdf | T::Document | T::Text)
            }
            Provider::Ollama | Provider::OllamaCloud => {
                matches!(kind, T::Image | T::Audio | T::Text)
            }
            Provider::GPUStack => matches!(kind, T::Image | T::Audio | T::Video | T::Text),
            Provider::Perplexity => matches!(kind, T::Image | T::Pdf | T::Document | T::Text),
            Provider::OpenRouter => kind == T::Video || chat_completions(Some(false), true),
            Provider::DeepSeek | Provider::XAI | Provider::Hetzner => chat_completions(None, false),
            _ => chat_completions(Some(false), true),
        },
    }
}

/// `Protocol#supported_message_attachment?`: [`supported_attachment`] for every message role.
/// (Only Cohere, which this port leaves out, accepts fewer attachments in tool results.)
pub(crate) fn supported_message_attachment(
    protocol: ProtocolName,
    provider: Provider,
    _role: crate::message::Role,
    attachment: &crate::attachment::Attachment,
) -> bool {
    supported_attachment(protocol, provider, attachment)
}

pub fn render(protocol: ProtocolName, req: &Request) -> Result<Value> {
    if req.citations && !req.model.supports("citations") {
        warn_unsupported_citations(protocol, req.model);
    }
    match protocol {
        ProtocolName::ChatCompletions => chat_completions::render_payload(req),
        ProtocolName::Responses => responses::render_payload(req),
        ProtocolName::Anthropic => anthropic::render_payload(req),
        ProtocolName::Gemini => gemini::render_payload(req),
        ProtocolName::Interactions => interactions::render_payload(req),
        ProtocolName::Conversations => mistral::render_payload(req),
        ProtocolName::RouterChatCompletions => router_render_payload(req),
    }
}

/// `Protocols::Perplexity::Router`: the base Chat Completions wire format (not Perplexity's
/// Sonar dialect). Schemas default to strict (`{ strict: true }.merge(schema)`); everything else is
/// left for Router to judge (cfd85579). OpenAI speaks the base format, so it stands in for the
/// dialect; its one OpenAI-only field is renamed back.
fn router_render_payload(req: &Request) -> Result<Value> {
    let strict = req.schema.map(|s| Schema {
        strict: s.strict.or(Some(true)),
        ..s.clone()
    });
    let base = Request {
        provider: Provider::OpenAI,
        schema: strict.as_ref(),
        ..*req
    };
    let mut payload = chat_completions::render_payload(&base)?;
    if let Some(p) = payload.as_object_mut()
        && let Some(max) = p.remove("max_completion_tokens")
    {
        p.insert("max_tokens".into(), max);
    }
    Ok(payload)
}

/// What a protocol's `render` does after `super` (provider options, provider tools, and
/// `before_request` hooks applied): Mistral Conversations deduplicates its tools and refuses
/// hosted confirmations.
pub(crate) fn finish_render(protocol: ProtocolName, payload: &mut Value) -> Result<()> {
    match protocol {
        ProtocolName::Conversations => mistral::finish_render(payload),
        _ => Ok(()),
    }
}

/// Each protocol's `warn_unsupported_citations`, called from `render_payload`.
fn warn_unsupported_citations(protocol: ProtocolName, model: &Model) {
    let hint = if protocol == ProtocolName::Gemini {
        "Gemini citations come from Google Search grounding: with_provider_options(tools: [{ google_search: {} }])."
    } else {
        "with_citations may have no effect."
    };
    tracing::warn!(
        "{} does not support citations according to the model registry. {hint}",
        model.id
    );
}

pub fn endpoint(
    protocol: ProtocolName,
    provider: Provider,
    model: &Model,
    stream: bool,
) -> Endpoint {
    let path = match protocol {
        ProtocolName::Interactions => "interactions".to_string(),
        ProtocolName::Conversations => "conversations".to_string(),
        // The chat replaces it with `Provider::router_url`, which depends on the configured base.
        ProtocolName::RouterChatCompletions => "router/v1/chat/completions".to_string(),
        ProtocolName::ChatCompletions => "chat/completions".to_string(),
        ProtocolName::Responses if provider == Provider::Perplexity => "v1/agent".to_string(),
        ProtocolName::Responses => "responses".to_string(),
        ProtocolName::Anthropic => "v1/messages".to_string(),
        ProtocolName::Gemini if stream => {
            format!("models/{}:streamGenerateContent?alt=sse", model.id)
        }
        ProtocolName::Gemini => format!("models/{}:generateContent", model.id),
    };
    let headers = if protocol == ProtocolName::Anthropic && stream {
        vec![("Accept-Encoding".into(), "identity".into())]
    } else {
        Vec::new()
    };
    Endpoint { path, headers }
}

pub fn parse_completion(
    protocol: ProtocolName,
    provider: Provider,
    model: &Model,
    raw: RawResponse,
) -> Result<Message> {
    let body = raw.body.clone();
    // `Protocol#parse_completion_response`: `body.nil? || body.empty?` (nil, {}, [], '').
    let empty = match &body {
        Value::Null => true,
        Value::Object(m) => m.is_empty(),
        Value::Array(a) => a.is_empty(),
        Value::String(s) => s.is_empty(),
        _ => false,
    };
    if empty {
        return Err(Error::Api(
            "Provider returned an empty response body".into(),
            None,
        ));
    }
    match protocol {
        // `MultiCompletion#parse_completion_body`: Mistral returns every hosted-tool step.
        ProtocolName::ChatCompletions
            if provider == Provider::Mistral
                && body
                    .pointer("/choices/0/messages")
                    .is_some_and(Value::is_array) =>
        {
            mistral::parse_multi_message(&body, Some(raw))?
                .ok_or_else(|| Error::Api("Provider returned no completion message".into(), None))
        }
        ProtocolName::ChatCompletions => {
            chat_completions::parse_completion_body(provider, &body, raw)
        }
        ProtocolName::Responses => responses::parse_completion_body(provider, &body, raw),
        ProtocolName::Anthropic => anthropic::parse_completion_body(&body, raw),
        ProtocolName::Gemini => gemini::parse_completion_body(model, &body, raw),
        ProtocolName::Interactions => interactions::parse_completion_body(model, &body, Some(raw)),
        ProtocolName::Conversations => mistral::parse_completion_body(&model.id, &body, Some(raw)),
        ProtocolName::RouterChatCompletions => {
            chat_completions::parse_completion_body(provider, &body, raw)
        }
    }
}

/// Per-stream state a protocol's `build_chunk` needs (Anthropic block tracking, Responses citation offsets).
#[derive(Default)]
pub struct StreamState {
    pub anthropic: anthropic::StreamBlocks,
    pub gemini_parts: Vec<Value>,
    pub citation_lengths: Vec<((i64, i64), usize)>,
    /// `OpenRouter::Streaming`'s `@raw_reasoning`: the reasoning_details merged so far.
    pub openrouter_reasoning: Option<Vec<Value>>,
    /// `Interactions::Streaming`'s steps and response.
    pub interactions: interactions::StreamState,
    /// `Mistral::Conversations::Streaming`'s output entries.
    pub conversation: mistral::ConversationStream,
    /// `Mistral::MultiCompletion#stream_response`, when hosted tools were sent.
    pub multi: Option<mistral::MultiStream>,
    /// The model the stream answers for (Interactions and Conversations parse whole responses).
    pub model: Option<Model>,
}

impl StreamState {
    /// The state for a stream of `payload`: Mistral Chat Completions with hosted tools streams
    /// several completions (`MultiCompletion#stream_response`).
    pub fn for_payload(
        protocol: ProtocolName,
        provider: Provider,
        model: &Model,
        payload: &Value,
    ) -> StreamState {
        let multi = mistral::is_multi_stream(protocol, provider, payload)
            .then(mistral::MultiStream::default);
        StreamState {
            multi,
            model: Some(model.clone()),
            ..Default::default()
        }
    }

    fn model_id(&self) -> &str {
        self.model.as_ref().map_or("", |m| m.id.as_str())
    }
}

/// The streamed message: most protocols fold their chunks; Interactions, Conversations, and
/// Mistral multi-completion parse the response they accumulated (`stream_response`). The second
/// value is a last chunk to yield after the stream ends (multi-completion's summary chunk).
pub fn finish_stream(
    protocol: ProtocolName,
    state: &mut StreamState,
    acc: StreamAccumulator,
    raw: RawResponse,
) -> Result<(Message, Option<Message>)> {
    if let Some(multi) = &state.multi {
        let message = mistral::finish_multi_stream(multi, raw)?;
        let last = mistral::final_chunk(&message);
        return Ok((message, Some(last)));
    }
    match protocol {
        ProtocolName::Interactions => {
            interactions::finish_stream(&mut state.interactions, raw).map(|m| (m, None))
        }
        ProtocolName::Conversations => {
            mistral::finish_conversation_stream(&mut state.conversation, raw).map(|m| (m, None))
        }
        _ => acc.into_message(raw).map(|m| (m, None)),
    }
}

pub fn build_chunk(
    protocol: ProtocolName,
    provider: Provider,
    state: &mut StreamState,
    data: &Value,
) -> Result<Message> {
    match protocol {
        ProtocolName::ChatCompletions if state.multi.is_some() => Ok(state
            .multi
            .as_mut()
            .map(|multi| mistral::build_multi_chunk(multi, data))
            .unwrap_or_else(Message::chunk)),
        ProtocolName::ChatCompletions => {
            let mut chunk = chat_completions::build_chunk(provider, data);
            if provider == Provider::OpenRouter {
                let details = data.pointer("/choices/0/delta/reasoning_details");
                chunk.raw_reasoning = chat_completions::accumulate_raw_reasoning(
                    &mut state.openrouter_reasoning,
                    details,
                );
            }
            Ok(chunk)
        }
        ProtocolName::Responses => responses::build_chunk(provider, state, data),
        ProtocolName::Anthropic => Ok(anthropic::build_chunk(&mut state.anthropic, data)),
        ProtocolName::Gemini => Ok(gemini::build_chunk(state, data)),
        ProtocolName::Interactions => {
            let model = state
                .model
                .clone()
                .unwrap_or_else(|| Model::default_for("", provider.slug()));
            interactions::build_chunk(&model, &mut state.interactions, data)
        }
        ProtocolName::Conversations => {
            let model_id = state.model_id().to_string();
            mistral::build_conversation_chunk(&model_id, &mut state.conversation, data)
        }
        ProtocolName::RouterChatCompletions => Ok(chat_completions::build_chunk(provider, data)),
    }
}

/// Maps an in-stream error body to the HTTP status RubyLLM raises it as.
pub fn streaming_error_status(protocol: ProtocolName) -> fn(&str) -> Option<u16> {
    match protocol {
        ProtocolName::Anthropic => |data| {
            let v: Value = serde_json::from_str(data).ok()?;
            // Only an `{"type": "error"}` event carries a status (anthropic/streaming.rb).
            if v.get("type").and_then(Value::as_str) != Some("error") {
                return None;
            }
            // `Anthropic::Streaming::ERROR_STATUSES`, 500 for any other type.
            Some(match v.pointer("/error/type").and_then(Value::as_str) {
                Some("invalid_request_error") => 400,
                Some("authentication_error") => 401,
                Some("billing_error") => 402,
                Some("permission_error") => 403,
                Some("not_found_error") => 404,
                Some("request_too_large") => 413,
                Some("rate_limit_error") => 429,
                Some("timeout_error") => 504,
                Some("overloaded_error") => 529,
                _ => 500,
            })
        },
        ProtocolName::Gemini => |data| {
            let v: Value = serde_json::from_str(data).ok()?;
            let err = v.get("error").unwrap_or(&v);
            err.get("code").and_then(Value::as_u64).map(|c| c as u16)
        },
        // `ChatCompletions::Streaming#parse_streaming_error`: no status unless `error` is an object.
        ProtocolName::ChatCompletions => |data| {
            let v: Value = serde_json::from_str(data).ok()?;
            let error = v.get("error")?.as_object()?;
            // A numeric code in the HTTP error range is the status (OpenRouter, 6e99b6cc).
            if let Some(code) = error
                .get("code")
                .and_then(Value::as_u64)
                .filter(|c| (400..=599).contains(c))
            {
                return Some(code as u16);
            }
            Some(match error.get("type").and_then(Value::as_str) {
                Some("server_error") => 500,
                Some("rate_limit_exceeded" | "insufficient_quota") => 429,
                _ => 400,
            })
        },
        // `Interactions::Streaming#parse_streaming_error`: the error code's documented status.
        ProtocolName::Interactions => |data| {
            let v: Value = serde_json::from_str(data).ok()?;
            let code = v
                .get("error")
                .unwrap_or(&v)
                .get("code")
                .and_then(Value::as_str);
            Some(match code {
                Some("invalid_request" | "failed_precondition") => 400,
                Some("authentication") => 401,
                Some("payment_required") => 402,
                Some("permission_denied") => 403,
                Some("not_found") => 404,
                Some("rate_limit_exceeded" | "quota_exceeded" | "too_many_requests") => 429,
                Some("service_unavailable") => 503,
                Some("deadline_exceeded") => 504,
                _ => 500,
            })
        },
        ProtocolName::Responses => responses::streaming_error_status,
        _ => |data| {
            let v: Value = serde_json::from_str(data).ok()?;
            let kind = v
                .pointer("/error/type")
                .or_else(|| v.get("code"))
                .and_then(Value::as_str)
                .unwrap_or("");
            Some(match kind {
                "server_error" => 500,
                "rate_limit_exceeded" | "insufficient_quota" => 429,
                _ => 400,
            })
        },
    }
}

/// `Protocol::StreamAccumulator`: folds chunks into the final message.
#[derive(Default)]
pub struct StreamAccumulator {
    content: String,
    model: Option<String>,
    citations: Vec<Citation>,
    /// `@citation_keys`: each kept citation's attributes, so a repeated list is read once.
    citation_keys: std::collections::HashSet<String>,
    thinking_text: Option<String>,
    thinking_signature: Option<String>,
    tool_calls: IndexMap<ToolCall>,
    tool_call_ids_by_index: Vec<(String, String)>,
    latest_tool_call_id: Option<String>,
    tokens: Tokens,
    server_tool_calls: Vec<ServerToolCall>,
    raw_content: Option<Value>,
    raw_reasoning: Option<Value>,
    finish_reason: Option<FinishReason>,
}

impl StreamAccumulator {
    /// Usage reported so far, for billing a stream that fails partway.
    pub fn tokens(&self) -> &Tokens {
        &self.tokens
    }

    /// `attr_reader :tool_calls`: the calls started so far, keyed by id, arguments still partial.
    pub fn tool_calls(&self) -> &IndexMap<ToolCall> {
        &self.tool_calls
    }

    pub fn add(&mut self, chunk: &Message) {
        if self.model.as_deref().unwrap_or("").is_empty() {
            self.model = chunk.model.clone();
        }
        if let Some(calls) = &chunk.tool_calls {
            for (stream_key, call) in calls.iter() {
                if call.starts_call() {
                    self.start_tool_call(stream_key, call);
                } else {
                    self.append_fragment(stream_key, call);
                }
            }
        }
        if let Some(text) = &chunk.content {
            self.content.push_str(text);
        }
        // `accumulate_citations`: providers like Perplexity repeat the whole list on every chunk,
        // so each citation is keyed by its attributes once instead of compared with every kept one.
        for c in &chunk.citations {
            if self
                .citation_keys
                .insert(serde_json::to_string(c).unwrap_or_default())
            {
                self.citations.push(c.clone());
            }
        }
        if let Some(t) = &chunk.thinking {
            if let Some(text) = &t.text {
                self.thinking_text
                    .get_or_insert_with(String::new)
                    .push_str(text);
            }
            if self.thinking_signature.is_none() {
                self.thinking_signature = t.signature.clone();
            }
        }
        for call in &chunk.server_tool_calls {
            match self
                .server_tool_calls
                .iter_mut()
                .find(|e| call.id.is_some() && e.id == call.id && e.kind == call.kind)
            {
                Some(existing) => *existing = call.clone(),
                None => self.server_tool_calls.push(call.clone()),
            }
        }
        if chunk.raw_content.is_some() {
            self.raw_content = chunk.raw_content.clone();
        }
        if chunk.raw_reasoning.is_some() {
            self.raw_reasoning = chunk.raw_reasoning.clone();
        }
        if chunk.finish_reason.is_some() {
            self.finish_reason = chunk.finish_reason.clone();
        }
        self.tokens.merge_latest(&chunk.tokens);
    }

    fn start_tool_call(&mut self, stream_key: &str, call: &ToolCall) {
        let mut call = call.clone();
        if let ToolArguments::Parsed(m) = &call.arguments
            && m.is_empty()
        {
            call.arguments = ToolArguments::Partial(String::new());
        }
        if call.id.is_empty() {
            call.id = uuid::Uuid::new_v4().to_string();
        }
        let id = call.id.clone();
        self.tool_calls.insert(id.clone(), call);
        self.tool_call_ids_by_index.retain(|(k, _)| k != stream_key);
        self.tool_call_ids_by_index
            .push((stream_key.to_string(), id.clone()));
        self.latest_tool_call_id = Some(id);
    }

    fn append_fragment(&mut self, stream_key: &str, call: &ToolCall) {
        let id = self
            .tool_call_ids_by_index
            .iter()
            .find(|(k, _)| k == stream_key)
            .map(|(_, id)| id.clone())
            .or_else(|| {
                self.tool_calls
                    .contains_key(stream_key)
                    .then(|| stream_key.to_string())
            })
            // A keyless fragment continues the latest call; an unknown key is dropped (find_tool_call).
            .or_else(|| {
                stream_key
                    .is_empty()
                    .then(|| self.latest_tool_call_id.clone())
                    .flatten()
            });
        let Some(existing) = id.and_then(|id| self.tool_calls.get_mut(&id)) else {
            return;
        };
        let fragment = match &call.arguments {
            ToolArguments::Partial(s) => s.clone(),
            ToolArguments::Parsed(m) if m.is_empty() => String::new(),
            ToolArguments::Parsed(m) => Value::Object(m.clone()).to_string(),
        };
        match &mut existing.arguments {
            ToolArguments::Partial(s) => s.push_str(&fragment),
            ToolArguments::Parsed(_) => {}
        }
        if existing.thought_signature.is_none() {
            existing.thought_signature = call.thought_signature.clone();
        }
    }

    pub fn into_message(self, raw: RawResponse) -> Result<Message> {
        let finish = self.finish_reason.as_ref().map(|f| f.as_str().to_string());
        let mut tool_calls = IndexMap::new();
        for (id, call) in self.tool_calls.0 {
            let arguments = match call.arguments {
                ToolArguments::Partial(s) if s.is_empty() => Map::new(),
                // `parse_tool_call_arguments`: `ToolCallParseError.new(response:, finish_reason:)`
                // raised with the `JSON::ParserError` as its cause.
                ToolArguments::Partial(s) => serde_json::from_str::<Map<String, Value>>(&s)
                    .map_err(|e| {
                        let response = crate::error::ErrorResponse {
                            status: raw.status,
                            body: raw.body.to_string(),
                            ..Default::default()
                        };
                        Error::tool_call_parse_from(finish.as_deref(), response, e)
                    })?,
                ToolArguments::Parsed(m) => m,
            };
            tool_calls.insert(
                id,
                ToolCall {
                    arguments: ToolArguments::Parsed(arguments),
                    ..call
                },
            );
        }
        let content = self.content;
        let citations = self
            .citations
            .into_iter()
            .map(|mut c| {
                if c.text.is_none()
                    && let (Some(s), Some(e)) = (c.start_index, c.end_index)
                {
                    c.text = char_slice(&content, s, e);
                }
                c
            })
            .collect();
        let mut message = Message::chunk();
        message.content = (!content.is_empty()).then_some(content);
        message.citations = citations;
        message.thinking = Thinking::build(self.thinking_text, self.thinking_signature);
        message.tokens = self.tokens;
        message.server_tool_calls = self.server_tool_calls;
        message.raw_content = self.raw_content;
        message.raw_reasoning = self.raw_reasoning;
        message.finish_reason = self.finish_reason;
        message.model = self.model;
        message.tool_calls = (!tool_calls.is_empty()).then_some(tool_calls);
        message.raw = Some(raw);
        Ok(message.normalized())
    }
}

/// Ruby `content[start...end]` on characters.
pub(crate) fn char_slice(s: &str, start: i64, end: i64) -> Option<String> {
    if start < 0 || end < start {
        return None;
    }
    let chars: Vec<char> = s.chars().collect();
    if start as usize > chars.len() {
        return None;
    }
    Some(
        chars[start as usize..(end as usize).min(chars.len())]
            .iter()
            .collect(),
    )
}

pub(crate) fn normalize_finish_reason(
    reason: Option<&str>,
    table: &[(&str, &str)],
) -> Option<FinishReason> {
    let reason = reason?;
    let symbol = table
        .iter()
        .find(|(k, _)| *k == reason)
        .map(|(_, v)| *v)
        .unwrap_or(reason);
    Some(FinishReason::from_symbol(symbol))
}

pub(crate) fn int(v: Option<&Value>) -> Option<i64> {
    v.and_then(|v| v.as_i64().or_else(|| v.as_f64().map(|f| f as i64)))
}

pub(crate) fn str_of(v: Option<&Value>) -> Option<String> {
    v.and_then(Value::as_str).map(str::to_string)
}

/// `Support::Utils.deep_merge` (used for `provider_options` and tool `provider_options`): nested
/// objects merge key by key, and anything else in `overlay` replaces what `base` has.
pub fn deep_merge(base: &mut Value, overlay: &Value) {
    match (base, overlay) {
        (Value::Object(b), Value::Object(o)) => {
            for (k, v) in o {
                match b.get_mut(k) {
                    Some(existing) if existing.is_object() && v.is_object() => {
                        deep_merge(existing, v)
                    }
                    _ => {
                        b.insert(k.clone(), v.clone());
                    }
                }
            }
        }
        (b, o) => *b = o.clone(),
    }
}

pub(crate) fn tool_call_map(calls: Vec<ToolCall>) -> Option<IndexMap<ToolCall>> {
    if calls.is_empty() {
        return None;
    }
    Some(calls.into_iter().map(|c| (c.id.clone(), c)).collect())
}