agent_sdk_providers/

provider.rs

//! LLM provider trait and streaming helpers.
//!
//! This module defines the [`LlmProvider`] trait that all LLM backends implement,
//! as well as the [`collect_stream`] helper for consuming a streaming response.

#[cfg(feature = "anthropic")]
use agent_sdk_foundation::llm::ToolChoice;
use agent_sdk_foundation::llm::{
    ChatOutcome, ChatRequest, ChatResponse, ContentBlock, ThinkingConfig, ThinkingMode, Usage,
};
use anyhow::Result;
use async_trait::async_trait;
use futures::StreamExt;
use serde::{Deserialize, Serialize};

use crate::model_capabilities::{
    ModelCapabilities, default_max_output_tokens, get_model_capabilities,
};
use crate::model_features::{ModelFeatures, get_model_features};
use crate::streaming::{StreamAccumulator, StreamBox, StreamDelta, StreamErrorKind};

/// How a provider satisfies a [`ResponseFormat`](agent_sdk_foundation::llm::ResponseFormat)
/// structured-output request.
///
/// The structured-output runner consults this to decide how to shape the
/// request and where to read the final structured value from the response.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum StructuredOutputSupport {
    /// The provider applies the schema natively (JSON-mode /
    /// structured-outputs) when it sees `request.response_format`. The final
    /// structured value is the JSON in the assistant's text output.
    Native,
    /// The provider has no native JSON-schema mode. The runner injects a
    /// single forced "respond" tool whose `input_schema` is the output schema,
    /// and reads the structured value from that tool call's input.
    ToolForcing,
}

/// A single model entry returned by a provider's live model-listing endpoint.
///
/// This is the *dynamic* counterpart to the static
/// [`ModelCapabilities`] table: it
/// is populated from the provider's own `/models` API at runtime, so newly
/// shipped models appear without an SDK code change. Fields beyond `id` are
/// optional because not every provider's listing endpoint reports them.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ModelInfo {
    /// The model identifier as the provider's chat endpoint expects it.
    pub id: String,
    /// Human-friendly display name, when the listing endpoint provides one.
    pub display_name: Option<String>,
    /// Maximum total context window in tokens, when reported.
    pub context_window: Option<u32>,
    /// Maximum output tokens per response, when reported.
    pub max_output_tokens: Option<u32>,
}

#[async_trait]
pub trait LlmProvider: Send + Sync {
    /// Non-streaming chat completion.
    async fn chat(&self, request: ChatRequest) -> Result<ChatOutcome>;

    /// List the models the provider currently exposes, queried live from the
    /// provider's own model-listing endpoint.
    ///
    /// The default implementation returns an error: model discovery is an
    /// additive capability, so a provider that has not implemented it stays
    /// source-compatible while reporting that the operation is unsupported.
    ///
    /// # Errors
    ///
    /// Returns an error when the provider does not support live model listing,
    /// or when the underlying HTTP request or response parsing fails.
    async fn list_models(&self) -> Result<Vec<ModelInfo>> {
        Err(anyhow::anyhow!(
            "list_models is not supported for provider {}",
            self.provider()
        ))
    }

    /// Cheaply check whether the provider's endpoint is currently reachable,
    /// without dispatching a billable request.
    ///
    /// Durable runtimes poll this to wait out connectivity loss: after a
    /// transport-classified stream failure
    /// ([`StreamErrorKind::is_connectivity`]) they probe until reachability
    /// returns, then re-dispatch the real call. Only an unreachable answer
    /// keeps that wait free of charge, so implementations should answer from
    /// a cheap endpoint — [`probe_http_reachability`] against the API host —
    /// never a model call.
    ///
    /// The default returns `true` ("reachable, or unknown"): a provider that
    /// cannot observe reachability keeps its callers on the bounded retry
    /// path instead of parking them in an offline wait it can never end.
    async fn probe_connectivity(&self) -> bool {
        true
    }

    /// Streaming chat completion.
    ///
    /// Returns a stream of [`StreamDelta`] events. The default implementation
    /// calls [`chat()`](Self::chat) and converts the result to a single-chunk stream.
    ///
    /// Providers should override this method to provide true streaming support.
    fn chat_stream(&self, request: ChatRequest) -> StreamBox<'_> {
        Box::pin(async_stream::stream! {
            match self.chat(request).await {
                Ok(outcome) => match outcome {
                    ChatOutcome::Success(response) => {
                        // Emit content as deltas
                        for (idx, block) in response.content.iter().enumerate() {
                            match block {
                                ContentBlock::Text { text } => {
                                    yield Ok(StreamDelta::TextDelta {
                                        delta: text.clone(),
                                        block_index: idx,
                                    });
                                }
                                ContentBlock::Thinking { thinking, .. } => {
                                    yield Ok(StreamDelta::ThinkingDelta {
                                        delta: thinking.clone(),
                                        block_index: idx,
                                    });
                                }
                                ContentBlock::RedactedThinking { .. }
                                | ContentBlock::ToolResult { .. }
                                | ContentBlock::Image { .. }
                                | ContentBlock::Document { .. } => {
                                    // Not streamed in the default implementation
                                }
                                ContentBlock::OpaqueReasoning { provider, data } => {
                                    yield Ok(StreamDelta::OpaqueReasoning {
                                        provider: provider.clone(),
                                        data: data.clone(),
                                        block_index: idx,
                                    });
                                }
                                ContentBlock::ToolUse { id, name, input, thought_signature } => {
                                    yield Ok(StreamDelta::ToolUseStart {
                                        id: id.clone(),
                                        name: name.clone(),
                                        block_index: idx,
                                        thought_signature: thought_signature.clone(),
                                    });
                                    yield Ok(StreamDelta::ToolInputDelta {
                                        id: id.clone(),
                                        delta: serde_json::to_string(input).unwrap_or_default(),
                                        block_index: idx,
                                    });
                                }
                                // `ContentBlock` is `#[non_exhaustive]`; a future
                                // block kind we cannot stream is skipped rather than
                                // panicking the default fallback.
                                _ => {
                                    log::warn!(
                                        "chat_stream fallback skipping unrecognized content block at index {idx}"
                                    );
                                }
                            }
                        }
                        yield Ok(StreamDelta::Usage(response.usage));
                        yield Ok(StreamDelta::Done {
                            stop_reason: response.stop_reason,
                        });
                    }
                    ChatOutcome::RateLimited(retry_after) => {
                        yield Ok(StreamDelta::Error {
                            message: "Rate limited".to_string(),
                            kind: StreamErrorKind::RateLimited(retry_after),
                        });
                    }
                    ChatOutcome::InvalidRequest(msg) => {
                        yield Ok(StreamDelta::Error {
                            message: msg,
                            kind: StreamErrorKind::InvalidRequest,
                        });
                    }
                    ChatOutcome::ServerError(msg) => {
                        yield Ok(StreamDelta::Error {
                            message: msg,
                            kind: StreamErrorKind::ServerError,
                        });
                    }
                    // `ChatOutcome` is `#[non_exhaustive]`; an outcome this SDK
                    // version does not model is surfaced as an unclassified
                    // (non-recoverable) stream error rather than dropped.
                    _ => {
                        yield Ok(StreamDelta::Error {
                            message: "Unrecognized chat outcome".to_string(),
                            kind: StreamErrorKind::Unknown,
                        });
                    }
                },
                Err(e) => yield Err(e),
            }
        })
    }

    fn model(&self) -> &str;
    fn provider(&self) -> &'static str;

    /// The endpoint that actually serves this provider's calls.
    ///
    /// [`Self::provider`] names the *API shape* a request is built for, and
    /// several distinct endpoints share one shape: `OpenRouter`, `Baseten` and
    /// `Fireworks` are all reached through [`OpenAIProvider`] pointed at a
    /// different `base_url`, so all three answer `"openai"`. Durable audit rows
    /// need to tell *model-X via a gateway* from *model-X native*, so this
    /// reports the serving route instead — the two coincide only for providers
    /// that front a single endpoint, which is why the default delegates.
    ///
    /// [`OpenAIProvider`]: crate::impls::openai::OpenAIProvider
    fn route(&self) -> &str {
        self.provider()
    }

    /// Provider-owned thinking configuration, if any.
    fn configured_thinking(&self) -> Option<&ThinkingConfig> {
        None
    }

    /// Canonical capability metadata for this provider/model, if known.
    fn capabilities(&self) -> Option<&'static ModelCapabilities> {
        get_model_capabilities(self.provider(), self.model()).or_else(|| match self.provider() {
            "openai-responses" | "openai-codex" => get_model_capabilities("openai", self.model()),
            "vertex" if self.model().starts_with("claude-") => {
                get_model_capabilities("anthropic", self.model())
            }
            "vertex" => get_model_capabilities("gemini", self.model()),
            _ => None,
        })
    }

    /// API-shape feature metadata for this provider/model, when known.
    ///
    /// This complements [`Self::capabilities`] with endpoint, reasoning,
    /// caching, and tool-control information that does not fit the legacy
    /// context/pricing record.
    fn model_features(&self) -> Option<&'static ModelFeatures> {
        match self.provider() {
            "openai" | "openai-responses" | "openai-codex" => get_model_features(self.model()),
            _ => None,
        }
    }

    /// Validate a thinking configuration against the provider/model capabilities.
    ///
    /// # Errors
    ///
    /// Returns an error when the requested thinking mode is not supported by
    /// the active provider/model capability set.
    fn validate_thinking_config(&self, thinking: Option<&ThinkingConfig>) -> Result<()> {
        let Some(thinking) = thinking else {
            return Ok(());
        };

        if self
            .capabilities()
            .is_some_and(|caps| !caps.supports_thinking)
        {
            return Err(anyhow::anyhow!(
                "thinking is not supported for provider={} model={}",
                self.provider(),
                self.model()
            ));
        }

        if matches!(thinking.mode, ThinkingMode::Adaptive)
            && !self
                .capabilities()
                .is_some_and(|caps| caps.supports_adaptive_thinking)
        {
            return Err(anyhow::anyhow!(
                "adaptive thinking is not supported for provider={} model={}",
                self.provider(),
                self.model()
            ));
        }

        Ok(())
    }

    /// Resolve the effective thinking configuration for a request.
    ///
    /// Request-level thinking overrides provider-owned defaults when present.
    ///
    /// # Errors
    ///
    /// Returns an error when the resolved thinking configuration is not
    /// supported by the active provider/model capability set.
    fn resolve_thinking_config(
        &self,
        request_thinking: Option<&ThinkingConfig>,
    ) -> Result<Option<ThinkingConfig>> {
        let thinking = request_thinking.or_else(|| self.configured_thinking());
        self.validate_thinking_config(thinking)?;
        Ok(thinking.cloned())
    }

    /// Default maximum output tokens for this provider/model when the caller
    /// does not explicitly override `AgentConfig.max_tokens`.
    fn default_max_tokens(&self) -> u32 {
        self.capabilities()
            .and_then(|caps| caps.max_output_tokens)
            .or_else(|| default_max_output_tokens(self.provider(), self.model()))
            .unwrap_or(4096)
    }

    /// How this provider satisfies a structured-output
    /// ([`ResponseFormat`](agent_sdk_foundation::llm::ResponseFormat)) request.
    ///
    /// Providers with a native JSON-schema / JSON-mode wire field
    /// (OpenAI-family, Gemini, Vertex) report
    /// [`StructuredOutputSupport::Native`] and consume
    /// `request.response_format` directly. Providers without one (Anthropic)
    /// report [`StructuredOutputSupport::ToolForcing`] so the runner forces a
    /// single "respond" tool whose schema is the output schema. The default
    /// is the conservative [`StructuredOutputSupport::ToolForcing`], which
    /// works for any tool-capable provider.
    fn structured_output_support(&self) -> StructuredOutputSupport {
        match self.provider() {
            "openai" | "openai-responses" | "openai-codex" | "gemini" => {
                StructuredOutputSupport::Native
            }
            // Vertex multiplexes Anthropic and Gemini models. Only the Gemini
            // side has a native structured-output field; Claude-on-Vertex uses
            // the Messages API shape, which has no `response_format`.
            "vertex" if !self.model().starts_with("claude-") => StructuredOutputSupport::Native,
            _ => StructuredOutputSupport::ToolForcing,
        }
    }
}

/// Timeout for [`probe_http_reachability`]: long enough for a slow TLS
/// handshake, short enough that a half-open network cannot stall the
/// connectivity-wait loop it drives.
const PROBE_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10);

/// `true` when an HTTP response arrives from `url` over a fully-established
/// (and, for HTTPS, certificate-validated) connection.
///
/// The status code is deliberately ignored — a 4xx from the provider's edge
/// proves the network path exactly as well as a 200 does, and needs no
/// credentials. Every transport failure (DNS, connect, TLS, timeout) reports
/// unreachable: this feeds offline-wait loops, where the fail-safe direction
/// is to keep waiting. A captive portal that hijacks the hostname fails
/// certificate validation and therefore stays unreachable until the real
/// endpoint is back.
pub async fn probe_http_reachability(client: &reqwest::Client, url: &str) -> bool {
    client.head(url).timeout(PROBE_TIMEOUT).send().await.is_ok()
}

/// Drop extended thinking when the request forces a specific tool.
///
/// Anthropic-family models (native Anthropic and Claude-on-Vertex) reject a
/// request that pairs extended `thinking` with a `tool_choice` that names a
/// tool (HTTP 400). The forced-tool constraint is orthogonal to *where* the
/// thinking came from: clearing `ChatRequest.thinking` upstream is not enough,
/// because [`resolve_thinking_config`](LlmProvider::resolve_thinking_config)
/// falls back to the provider-configured default when the request field is
/// `None` — resurrecting thinking on the wire. The Claude request builders
/// therefore funnel their already-resolved thinking through this at the wire
/// boundary so a forced-tool request (e.g. structured-output's `respond` tool)
/// stays well-formed regardless of the provider's configured thinking default.
///
/// `ToolChoice::Auto` (and the absence of any `tool_choice`) keeps thinking,
/// matching the API's rule that thinking is only incompatible with a
/// tool-forcing choice.
#[must_use]
#[cfg(feature = "anthropic")]
pub(crate) const fn thinking_for_forced_tool(
    thinking: Option<ThinkingConfig>,
    tool_choice: Option<&ToolChoice>,
) -> Option<ThinkingConfig> {
    match tool_choice {
        Some(ToolChoice::Tool(_)) => None,
        Some(ToolChoice::Auto) | None => thinking,
    }
}

/// Helper function to consume a stream and collect it into a `ChatResponse`.
///
/// This is useful for providers that want to test their streaming implementation
/// or for cases where you need the full response after streaming.
///
/// # Errors
///
/// Returns an error if the stream yields an error result.
pub async fn collect_stream(mut stream: StreamBox<'_>, model: String) -> Result<ChatOutcome> {
    let mut accumulator = StreamAccumulator::new();
    let mut last_error: Option<(String, StreamErrorKind)> = None;

    while let Some(result) = stream.next().await {
        match result {
            Ok(delta) => {
                if let StreamDelta::Error { message, kind } = &delta {
                    last_error = Some((message.clone(), *kind));
                }
                accumulator.apply(&delta);
            }
            Err(e) => return Err(e),
        }
    }

    // If we encountered an error during streaming, map kind directly
    // to the corresponding `ChatOutcome` variant.  No string-matching
    // heuristic is needed because the kind already records the
    // category at the construction site.
    if let Some((message, kind)) = last_error {
        return Ok(match kind {
            // The rate-limit kind carries whatever retry hint the provider gave
            // (header or error body), so the reconstructed outcome honours the
            // server exactly like the non-streaming `chat()` path does.
            StreamErrorKind::RateLimited(retry_after) => ChatOutcome::RateLimited(retry_after),
            StreamErrorKind::InvalidRequest => ChatOutcome::InvalidRequest(message),
            // `StreamErrorKind::ServerError`, plus the `#[non_exhaustive]`
            // catch-all (`Unknown` / future kinds): an unclassified error is
            // treated as a (non-recoverable) server error so the caller still
            // surfaces the failure rather than silently succeeding.
            _ => ChatOutcome::ServerError(message),
        });
    }

    // Extract usage and stop_reason before consuming the accumulator
    let usage = accumulator.take_usage().unwrap_or(Usage {
        input_tokens: 0,
        output_tokens: 0,
        cached_input_tokens: 0,
        cache_creation_input_tokens: 0,
    });
    let stop_reason = accumulator.take_stop_reason();
    let content = accumulator.into_content_blocks();

    // Log accumulated response for debugging
    log::debug!(
        "Collected stream response: model={} stop_reason={:?} usage={{input_tokens={}, output_tokens={}}} content_blocks={}",
        model,
        stop_reason,
        usage.input_tokens,
        usage.output_tokens,
        content.len()
    );
    for (i, block) in content.iter().enumerate() {
        match block {
            ContentBlock::Text { text } => {
                log::debug!("  content_block[{}]: Text (len={})", i, text.len());
            }
            ContentBlock::Thinking { thinking, .. } => {
                log::debug!("  content_block[{}]: Thinking (len={})", i, thinking.len());
            }
            ContentBlock::RedactedThinking { .. } => {
                log::debug!("  content_block[{i}]: RedactedThinking");
            }
            ContentBlock::OpaqueReasoning { provider, .. } => {
                log::debug!(
                    "  content_block[{i}]: OpaqueReasoning provider={provider} payload=<redacted>"
                );
            }
            ContentBlock::ToolUse {
                id, name, input, ..
            } => {
                log::debug!("  content_block[{i}]: ToolUse id={id} name={name} input={input}");
            }
            ContentBlock::ToolResult {
                tool_use_id,
                content: result_content,
                is_error,
            } => {
                log::debug!(
                    "  content_block[{}]: ToolResult tool_use_id={} is_error={:?} content_len={}",
                    i,
                    tool_use_id,
                    is_error,
                    result_content.len()
                );
            }
            ContentBlock::Image { source } => {
                log::debug!(
                    "  content_block[{i}]: Image media_type={}",
                    source.media_type
                );
            }
            ContentBlock::Document { source } => {
                log::debug!(
                    "  content_block[{i}]: Document media_type={}",
                    source.media_type
                );
            }
            // `ContentBlock` is `#[non_exhaustive]`; log unknown future block
            // kinds generically so the debug dump stays exhaustive.
            _ => {
                log::debug!("  content_block[{i}]: <unrecognized block kind>");
            }
        }
    }

    Ok(ChatOutcome::Success(ChatResponse {
        id: String::new(),
        content,
        model,
        stop_reason,
        usage,
    }))
}

#[cfg(test)]
mod tests {
    use super::*;
    use anyhow::Result;
    use async_trait::async_trait;

    struct Stub {
        provider: &'static str,
        model: &'static str,
    }

    #[async_trait]
    impl LlmProvider for Stub {
        async fn chat(&self, _request: ChatRequest) -> Result<ChatOutcome> {
            Ok(ChatOutcome::ServerError("unused".to_owned()))
        }

        fn model(&self) -> &str {
            self.model
        }

        fn provider(&self) -> &'static str {
            self.provider
        }
    }

    fn support_for(provider: &'static str, model: &'static str) -> StructuredOutputSupport {
        Stub { provider, model }.structured_output_support()
    }

    #[test]
    fn native_providers_report_native_support() {
        for provider in ["openai", "openai-responses", "openai-codex", "gemini"] {
            assert_eq!(
                support_for(provider, "any-model"),
                StructuredOutputSupport::Native,
                "{provider} should be native"
            );
        }
    }

    #[test]
    fn anthropic_reports_tool_forcing() {
        assert_eq!(
            support_for("anthropic", "claude-sonnet-4-5"),
            StructuredOutputSupport::ToolForcing
        );
    }

    #[test]
    fn vertex_is_native_for_gemini_models_and_tool_forcing_for_claude() {
        assert_eq!(
            support_for("vertex", "gemini-3-flash-preview"),
            StructuredOutputSupport::Native
        );
        assert_eq!(
            support_for("vertex", "claude-sonnet-4-5"),
            StructuredOutputSupport::ToolForcing
        );
    }

    #[test]
    fn unknown_provider_defaults_to_tool_forcing() {
        assert_eq!(
            support_for("some-new-provider", "x"),
            StructuredOutputSupport::ToolForcing
        );
    }

    #[test]
    #[cfg(feature = "anthropic")]
    fn thinking_for_forced_tool_drops_thinking_when_a_tool_is_forced() {
        let cfg = ThinkingConfig::new(10_000);
        let forced = ToolChoice::Tool("respond".to_owned());
        assert!(
            thinking_for_forced_tool(Some(cfg), Some(&forced)).is_none(),
            "thinking must be dropped when tool_choice names a tool"
        );
    }

    #[test]
    #[cfg(feature = "anthropic")]
    fn thinking_for_forced_tool_keeps_thinking_for_auto_and_none() {
        let auto = ToolChoice::Auto;
        assert!(
            thinking_for_forced_tool(Some(ThinkingConfig::new(10_000)), Some(&auto)).is_some(),
            "thinking must survive with ToolChoice::Auto"
        );
        assert!(
            thinking_for_forced_tool(Some(ThinkingConfig::new(10_000)), None).is_some(),
            "thinking must survive with no tool_choice"
        );
    }

    #[test]
    #[cfg(feature = "anthropic")]
    fn thinking_for_forced_tool_is_a_noop_when_thinking_absent() {
        let forced = ToolChoice::Tool("respond".to_owned());
        assert!(thinking_for_forced_tool(None, Some(&forced)).is_none());
    }

    fn stream_of(deltas: Vec<StreamDelta>) -> StreamBox<'static> {
        Box::pin(futures::stream::iter(deltas.into_iter().map(Ok)))
    }

    #[tokio::test]
    async fn collect_stream_preserves_the_rate_limit_retry_hint() -> Result<()> {
        let hint = std::time::Duration::from_secs(30);
        let outcome = collect_stream(
            stream_of(vec![StreamDelta::Error {
                message: "Rate limited".to_owned(),
                kind: StreamErrorKind::RateLimited(Some(hint)),
            }]),
            "test-model".to_owned(),
        )
        .await?;

        assert!(
            matches!(outcome, ChatOutcome::RateLimited(Some(delay)) if delay == hint),
            "streamed Retry-After must survive reconstruction, got {outcome:?}"
        );
        Ok(())
    }

    #[tokio::test]
    async fn collect_stream_reports_no_hint_when_the_provider_gave_none() -> Result<()> {
        let outcome = collect_stream(
            stream_of(vec![StreamDelta::Error {
                message: "Rate limited".to_owned(),
                kind: StreamErrorKind::RateLimited(None),
            }]),
            "test-model".to_owned(),
        )
        .await?;

        assert!(
            matches!(outcome, ChatOutcome::RateLimited(None)),
            "a hintless rate limit must not invent a delay, got {outcome:?}"
        );
        Ok(())
    }

    #[tokio::test]
    async fn default_chat_stream_carries_the_chat_retry_hint_into_the_error_delta() -> Result<()> {
        // A provider that only implements `chat()` gets its streaming from the
        // trait's default adapter; a rate limit's hint must not be lost there.
        struct RateLimitedStub {
            model: &'static str,
        }

        #[async_trait]
        impl LlmProvider for RateLimitedStub {
            async fn chat(&self, _request: ChatRequest) -> Result<ChatOutcome> {
                Ok(ChatOutcome::RateLimited(Some(
                    std::time::Duration::from_secs(12),
                )))
            }
            fn model(&self) -> &str {
                self.model
            }
            fn provider(&self) -> &'static str {
                "stub"
            }
        }

        let stub = RateLimitedStub {
            model: "stub-model",
        };
        let request = ChatRequest::new("", vec![agent_sdk_foundation::llm::Message::user("hi")]);
        let outcome =
            collect_stream(Box::pin(stub.chat_stream(request)), "stub-model".to_owned()).await?;

        assert!(
            matches!(outcome, ChatOutcome::RateLimited(Some(delay))
                if delay == std::time::Duration::from_secs(12)),
            "default chat_stream must forward chat()'s Retry-After, got {outcome:?}"
        );
        Ok(())
    }

    #[tokio::test]
    async fn collect_stream_maps_other_kinds_without_a_hint() -> Result<()> {
        let outcome = collect_stream(
            stream_of(vec![StreamDelta::Error {
                message: "boom".to_owned(),
                kind: StreamErrorKind::ServerError,
            }]),
            "test-model".to_owned(),
        )
        .await?;

        assert!(matches!(outcome, ChatOutcome::ServerError(msg) if msg == "boom"));
        Ok(())
    }
}