mermaid_cli/providers/model/

mod.rs

//! Model adapters wrapped as `ModelProvider` implementations.
//!
//! Five providers today: Ollama, Anthropic, Gemini, Meta, and OpenAI-
//! compat (covering OpenAI, OpenRouter, Groq, Cerebras, DeepInfra,
//! Together, and user-defined endpoints). Each wraps the
//! corresponding adapter in `mermaid_model::models::adapters`; the adapter
//! owns the wire format and the wrapper owns the trait shape.

pub mod anthropic;
pub mod gemini;
pub mod meta;
pub mod ollama;
pub mod openai_compat;

use std::sync::Arc;

use async_trait::async_trait;

use mermaid_domain::{ChatRequest, TurnId};
use mermaid_model::models::adapters::ModelLimits;
use mermaid_model::models::adapters::learning::{ParamMemory, Rejections};
use mermaid_model::models::adapters::ollama_sizing::NumCtxSource;
use mermaid_model::models::{FinishReason, ModelError, Result, TokenUsage};
use mermaid_runtime::{NewProviderProbe, RuntimeStore};

use super::ctx::{FinalResponse, StreamContext, StreamEvent};
use mermaid_model::models::ModelCapabilities;

/// Resolved context sizing for a turn. For most providers `model_max ==
/// effective` (the static advertised window). For Ollama they differ:
/// `model_max` is the probed architectural window, while `effective` is what we
/// actually enforce as `num_ctx` (auto-fitted to memory, capped, or an
/// override). Compaction and the status bar use `effective`; "model supports up
/// to X" uses `model_max`.
#[derive(Debug, Clone, Copy, Default)]
pub struct ContextSizing {
    pub model_max: Option<usize>,
    pub effective: Option<usize>,
    /// How `effective` was chosen (Ollama only). `None` for static/advertised.
    pub source: Option<NumCtxSource>,
    /// The model's per-response output ceiling, when the provider exposes one
    /// (`/models` metadata, or a documented static table). Rides the same
    /// resolve→reducer pipeline as the window so `provider_capabilities` can be
    /// refreshed live.
    pub max_output: Option<usize>,
    /// Whether the provider compacts the conversation itself when a turn asks
    /// (`ChatRequest::native_compaction`): Anthropic, until the model refuses.
    pub compacts_natively: bool,
}

/// Where a loaded model actually sits in memory, from a post-turn probe (Ollama
/// `/api/ps`). `total_bytes` is weights + KV + buffers; `size_vram_bytes` is the
/// part resident in VRAM. `size_vram_bytes < total_bytes` means the model spilled
/// to CPU/RAM (partial offload → slow). Only Ollama reports this.
#[derive(Debug, Clone, Copy)]
pub struct ModelPlacement {
    pub size_vram_bytes: u64,
    pub total_bytes: u64,
    /// Auto-converge target: when the model spilled, the largest `num_ctx` that
    /// would fit instead — or `None` if it already fits or shrinking can't help
    /// (weights-bound). Computed against the *measured* footprint.
    pub suggested_num_ctx: Option<u32>,
}

/// Provider-facing interface. A `ModelProvider` impl owns whatever
/// HTTP client / state it needs and exposes `chat()` — that's the
/// whole surface.
#[async_trait]
pub trait ModelProvider: Send + Sync {
    /// `ModelCapabilities` the provider advertises. The reducer reads this
    /// when building the outgoing `ChatRequest` (e.g. whether to
    /// attach reasoning controls).
    fn capabilities(&self) -> &ModelCapabilities;

    /// Resolve the *effective* context window for a turn (what the model will
    /// actually enforce). The default returns the static advertised window;
    /// Ollama overrides this to probe the model's real window and auto-fit
    /// `num_ctx` to host memory, honoring the request's per-model
    /// `ollama_num_ctx` override. `None` means "let the backend decide".
    ///
    /// Awaited only on the effect runtime (never the reducer), so a probe never
    /// blocks the UI.
    async fn resolve_context_window(&self, request: &ChatRequest) -> ContextSizing {
        let _ = request;
        let max = self.capabilities().max_context_tokens;
        ContextSizing {
            model_max: max,
            effective: max,
            source: None,
            max_output: self.capabilities().max_output_tokens,
            compacts_natively: false,
        }
    }

    /// Best-effort: where the loaded model currently sits in memory. The default
    /// returns `None` (unknown / not applicable); Ollama overrides it to probe
    /// `/api/ps`. Awaited only on the effect runtime, *after* a turn (when the
    /// model is resident), so it never blocks the UI.
    async fn verify_placement(&self, current_num_ctx: Option<usize>) -> Option<ModelPlacement> {
        let _ = current_num_ctx;
        None
    }

    /// Best-effort: whether the active model can actually see images. `None`
    /// means "unknown / not applicable" — the default for providers that don't
    /// probe, and for cloud providers whose vision support is already known
    /// good; `Some(false)` is what drives the no-vision-model warning. Awaited
    /// only on the effect runtime, so a probe never blocks the UI.
    async fn supports_vision(&self) -> Option<bool> {
        None
    }

    /// Stream a chat turn. Typed events flow through
    /// `ctx.sink`; the returned `FinalResponse` carries token usage
    /// and the Anthropic thinking-signature (opaque blob required to
    /// continue extended thinking across turns).
    ///
    /// Cancellation: the provider MUST select! on `ctx.token.
    /// cancelled()` inside any await that could block for more than
    /// a few hundred ms. This is the contract that replaces the old
    /// `check_interrupt` polling pattern.
    async fn chat(&self, request: ChatRequest, ctx: StreamContext) -> Result<FinalResponse>;
}

/// Output collected from a one-shot, non-interactive streaming model call.
#[derive(Debug, Clone, Default)]
pub struct CollectedText {
    pub text: String,
    pub usage: Option<TokenUsage>,
    pub stop_reason: Option<FinishReason>,
    pub reasoning: Option<String>,
}

/// Run a one-shot, non-interactive model call and collect its streamed text
/// into a `CollectedText`. For internal calls whose output is NOT shown to the user
/// as it streams — context compaction, memory consolidation, and the Auto-mode safety classifier.
/// Drains a private event channel, capturing plain text, reasoning trace, final token
/// usage, and stop reasons (e.g. length limits / content filters). The `token` lets the
/// caller cancel the call (e.g. on Ctrl+C) like any other turn work.
pub(crate) async fn collect_text(
    provider: Arc<dyn ModelProvider>,
    turn: TurnId,
    request: ChatRequest,
    token: tokio_util::sync::CancellationToken,
) -> Result<CollectedText> {
    let (stream_tx, mut stream_rx) = tokio::sync::mpsc::channel::<StreamEvent>(128);
    let ctx = StreamContext::new(token, stream_tx, turn);
    let collector = tokio::task::spawn(async move {
        let mut text = String::new();
        let mut reasoning = String::new();
        let mut usage = None;
        let mut stop_reason = None;
        while let Some(event) = stream_rx.recv().await {
            match event {
                StreamEvent::Text(chunk) => text.push_str(&chunk),
                StreamEvent::Reasoning(chunk) => reasoning.push_str(&chunk.text),
                StreamEvent::Done {
                    usage: done_usage,
                    stop_reason: done_stop,
                    ..
                } => {
                    usage = done_usage;
                    stop_reason = done_stop;
                },
                // Status is a user-facing plumbing notice, not content —
                // a text collector has nowhere to surface it.
                StreamEvent::ToolCall(_) | StreamEvent::Status(_) => {},
            }
        }
        (
            text,
            (!reasoning.is_empty()).then_some(reasoning),
            usage,
            stop_reason,
        )
    });

    let response = provider.chat(request, ctx).await;
    let (text, reasoning, stream_usage, stream_stop_reason) = collector
        .await
        .map_err(|err| ModelError::StreamError(format!("collect_text collector failed: {err}")))?;
    match response {
        Ok(final_response) => Ok(CollectedText {
            text,
            usage: final_response.usage.or(stream_usage),
            stop_reason: final_response.stop_reason.or(stream_stop_reason),
            reasoning,
        }),
        Err(err) => Err(err),
    }
}

pub use anthropic::AnthropicProvider;
pub use gemini::GeminiProvider;
pub use meta::MetaProvider;
pub use ollama::OllamaProvider;
pub use openai_compat::OpenAICompatProvider;

/// True when a `provider_probes` row is older than the probe TTL (shared by
/// the Ollama context probe and the per-provider limits probes). An
/// unparseable timestamp is treated as stale so it re-probes.
pub(crate) fn probe_is_stale(probed_at: &str) -> bool {
    use chrono::{DateTime, Utc};
    match DateTime::parse_from_rfc3339(probed_at) {
        Ok(t) => {
            Utc::now()
                .signed_duration_since(t.with_timezone(&Utc))
                .num_days()
                >= mermaid_model::constants::PROVIDER_PROBE_TTL_DAYS
        },
        // Unparseable timestamp → treat as stale and re-probe.
        Err(_) => true,
    }
}

/// Limits learned from one live probe of a provider's models endpoint, cached
/// per (provider, model) in `provider_probes`. A successful fetch WITHOUT
/// limit metadata (or a definitive "model not listed") is cached too — as
/// `None`s — so providers that don't expose limits aren't re-fetched every
/// turn. Fetch *failures* are never cached.
#[derive(serde::Serialize, serde::Deserialize)]
pub(crate) struct CachedLimits {
    pub(crate) max_context_tokens: Option<usize>,
    pub(crate) max_output_tokens: Option<usize>,
}

pub(crate) const LIMITS_PROBE_KEY: &str = "limits_probe";

/// Load fresh cached limits, off the async runtime. Best-effort → `None`.
pub(crate) async fn load_limits_from_db(provider: String, model: String) -> Option<CachedLimits> {
    tokio::task::spawn_blocking(move || {
        let rec = mermaid_runtime::with_shared_store(|store| {
            store
                .provider_probes()
                .get(&provider, &model, LIMITS_PROBE_KEY)
        })
        .ok()??;
        if probe_is_stale(&rec.probed_at) {
            return None;
        }
        serde_json::from_str::<CachedLimits>(&rec.capability_value).ok()
    })
    .await
    .ok()
    .flatten()
}

/// Persist probed limits for subsequent sessions. Best-effort.
pub(crate) async fn save_limits_to_db(provider: String, model: String, limits: &CachedLimits) {
    let value = match serde_json::to_string(limits) {
        Ok(v) => v,
        Err(_) => return,
    };
    let _ = tokio::task::spawn_blocking(move || -> Option<()> {
        mermaid_runtime::with_shared_store(|store| {
            store.provider_probes().upsert(NewProviderProbe {
                provider,
                model_id: model,
                capability_key: LIMITS_PROBE_KEY.into(),
                capability_value: value,
                confidence: "probed".into(),
                error: None,
            })
        })
        .ok()?;
        Some(())
    })
    .await;
}

/// Cache-first limit resolution: return fresh cached limits when present,
/// otherwise run `fetch` against the provider's models endpoint. A successful
/// fetch is cached even when all-`None` (definitive "provider exposes
/// nothing"); a failed fetch is NOT cached and resolves to `None` so the next
/// turn retries.
pub(crate) async fn resolve_limits_cached<F, Fut>(
    provider: &str,
    model: &str,
    fetch: F,
) -> Option<CachedLimits>
where
    F: FnOnce() -> Fut,
    Fut: std::future::Future<Output = Result<ModelLimits>>,
{
    if let Some(cached) = load_limits_from_db(provider.to_string(), model.to_string()).await {
        return Some(cached);
    }
    match fetch().await {
        Ok(limits) => {
            let cached = CachedLimits {
                max_context_tokens: limits.max_context_tokens,
                max_output_tokens: limits.max_output_tokens,
            };
            save_limits_to_db(provider.to_string(), model.to_string(), &cached).await;
            Some(cached)
        },
        // Network/parse failure: don't cache; caller falls back to static.
        Err(_) => None,
    }
}

/// `provider_probes` key for the request items a provider rejected for one
/// model (`mermaid_model::models::adapters::learning`): a JSON array such as
/// `["effort:xhigh","temperature"]`. TTL-bounded like every probe, so a
/// provider that later adds support is re-asked after the TTL.
pub(crate) const REJECTIONS_PROBE_KEY: &str = "rejected_params";

/// The fresh cached rejections for one model, if any.
fn read_rejections(store: &RuntimeStore, provider: &str, model: &str) -> Option<Rejections> {
    let rec = store
        .provider_probes()
        .get(provider, model, REJECTIONS_PROBE_KEY)
        .ok()??;
    if probe_is_stale(&rec.probed_at) {
        return None;
    }
    serde_json::from_str(&rec.capability_value).ok()
}

/// Store `learned` for one model, merged with whatever a concurrent session
/// cached meanwhile.
fn write_rejections(
    store: &RuntimeStore,
    provider: &str,
    model: &str,
    learned: &Rejections,
) -> anyhow::Result<()> {
    let mut merged = read_rejections(store, provider, model).unwrap_or_default();
    merged.extend(learned);
    store.provider_probes().upsert(NewProviderProbe {
        provider: provider.to_string(),
        model_id: model.to_string(),
        capability_key: REJECTIONS_PROBE_KEY.into(),
        capability_value: serde_json::to_string(&merged)?,
        confidence: "learned".into(),
        error: None,
    })?;
    Ok(())
}

/// A provider wrapper's link between its adapter's in-memory rejections and
/// the cross-session cache: seed once before the first turn, persist
/// whenever a turn learned something new. Best-effort both ways — a store
/// that can't be opened costs a relearned rejection, never a failed turn.
#[derive(Default)]
pub(crate) struct RejectionCache {
    seeded: tokio::sync::OnceCell<()>,
    /// What the store holds as far as this instance knows.
    persisted: std::sync::Mutex<Rejections>,
}

impl RejectionCache {
    /// Seed `memory` from the store, once per provider instance.
    pub(crate) async fn seed(&self, provider: &str, model: &str, memory: &ParamMemory) {
        self.seeded
            .get_or_init(|| async {
                let (provider, model) = (provider.to_string(), model.to_string());
                let cached = tokio::task::spawn_blocking(move || {
                    mermaid_runtime::with_shared_store(|store| {
                        Ok(read_rejections(store, &provider, &model))
                    })
                    .ok()
                    .flatten()
                })
                .await
                .ok()
                .flatten();
                if let Some(cached) = cached {
                    memory.seed(&cached);
                    *self.lock() = cached;
                }
            })
            .await;
    }

    /// Persist `memory` if it learned anything the store doesn't have.
    pub(crate) async fn persist(&self, provider: &str, model: &str, memory: &ParamMemory) {
        let learned = memory.snapshot();
        if *self.lock() == learned {
            return;
        }
        *self.lock() = learned.clone();
        let (provider, model) = (provider.to_string(), model.to_string());
        let _ = tokio::task::spawn_blocking(move || {
            mermaid_runtime::with_shared_store(|store| {
                write_rejections(store, &provider, &model, &learned)
            })
        })
        .await;
    }

    fn lock(&self) -> std::sync::MutexGuard<'_, Rejections> {
        self.persisted
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner)
    }
}

/// Extract a model's real per-response output ceiling from a provider's 400
/// rejection body. Fires ONLY on unambiguous output-cap wordings:
///
/// - Ollama Cloud / MiniMax: `max_tokens (521276) exceeds model's maximum
///   output tokens (131072) for model …`
/// - OpenAI-compat: `max_tokens is too large: … This model supports at most
///   16384 completion tokens …`
///
/// Anything else — in particular context-limit wordings ("prompt is too
/// long", "maximum context length") — returns `None`: learning a window as
/// an output cap would poison the cache, and a missed match just means
/// today's behavior (the error surfaces). Values outside a sanity range are
/// rejected as parser noise.
pub(crate) fn parse_output_cap_message(body: &str) -> Option<usize> {
    let cap = if let Some(rest) = text_after(body, "exceeds model's maximum output tokens") {
        leading_integer(rest)
    } else if body.contains("max_tokens is too large") {
        text_after(body, "supports at most").and_then(leading_integer)
    } else {
        None
    }?;
    (1_024..10_000_000).contains(&cap).then_some(cap)
}

/// The slice of `haystack` after the first occurrence of `marker`.
fn text_after<'a>(haystack: &'a str, marker: &str) -> Option<&'a str> {
    haystack.find(marker).map(|i| &haystack[i + marker.len()..])
}

/// The first integer in `s`, required to start within a few characters —
/// both documented wordings put the number right after the marker (`" ("` /
/// `" "`), and a distant number would belong to something else.
fn leading_integer(s: &str) -> Option<usize> {
    let start = s.find(|c: char| c.is_ascii_digit()).filter(|&i| i <= 8)?;
    s[start..]
        .chars()
        .take_while(char::is_ascii_digit)
        .collect::<String>()
        .parse()
        .ok()
}

/// Decide the output cap for a one-shot retry after learning `learned` from
/// a 400. AUTO (`requested == 0`, including "field omitted") retries at the
/// learned cap; an explicit ask above the cap retries clamped to it; an ask
/// already within the cap returns `None` — the 400 was about something else,
/// so retrying the same request would loop.
pub(crate) fn retry_cap(requested: usize, learned: usize) -> Option<usize> {
    (requested == 0 || requested > learned).then_some(learned)
}

/// `parse_output_cap_message` gated to actual HTTP 400s — the only status
/// where the body names a rejected parameter rather than a transient fault.
pub(crate) fn output_cap_from_error(err: &ModelError) -> Option<usize> {
    match err {
        ModelError::Backend(mermaid_model::models::BackendError::HttpError {
            status: 400,
            message,
            ..
        }) => parse_output_cap_message(message),
        _ => None,
    }
}

/// Persist an output cap learned from a provider's 400 rejection (the error
/// body names the model's real ceiling). Merges into any existing cached
/// row — reading it raw, ignoring the TTL, since a stale window is still
/// better than dropping it — and upserts as "probed". Best-effort.
pub(crate) async fn learn_output_cap(provider: String, model: String, cap: usize) {
    let _ = tokio::task::spawn_blocking(move || -> Option<()> {
        let existing = mermaid_runtime::with_shared_store(|store| {
            store
                .provider_probes()
                .get(&provider, &model, LIMITS_PROBE_KEY)
        })
        .ok()
        .flatten()
        .and_then(|rec| serde_json::from_str::<CachedLimits>(&rec.capability_value).ok());
        let merged = CachedLimits {
            max_context_tokens: existing.and_then(|l| l.max_context_tokens),
            max_output_tokens: Some(cap),
        };
        let value = serde_json::to_string(&merged).ok()?;
        mermaid_runtime::with_shared_store(|store| {
            store.provider_probes().upsert(NewProviderProbe {
                provider,
                model_id: model,
                capability_key: LIMITS_PROBE_KEY.into(),
                capability_value: value,
                confidence: "probed".into(),
                error: None,
            })
        })
        .ok()?;
        Some(())
    })
    .await;
}

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

    fn temp_store(name: &str) -> RuntimeStore {
        let dir =
            std::env::temp_dir().join(format!("mermaid_rejections_{name}_{}", std::process::id()));
        let _ = std::fs::remove_dir_all(&dir);
        std::fs::create_dir_all(&dir).expect("dir");
        RuntimeStore::open(dir.join("runtime.sqlite3")).expect("store")
    }

    #[test]
    fn rejections_round_trip_through_provider_probes() {
        let store = temp_store("round_trip");
        assert_eq!(read_rejections(&store, "anthropic", "claude-nova-7"), None);
        let learned = Rejections::from(["temperature"]);
        write_rejections(&store, "anthropic", "claude-nova-7", &learned).expect("write");
        assert_eq!(
            read_rejections(&store, "anthropic", "claude-nova-7"),
            Some(learned)
        );
        // Keyed per (provider, model): another model learned nothing.
        assert_eq!(read_rejections(&store, "anthropic", "claude-other"), None);
        let rec = store
            .provider_probes()
            .get("anthropic", "claude-nova-7", REJECTIONS_PROBE_KEY)
            .expect("get")
            .expect("row");
        assert_eq!(rec.confidence, "learned");
        assert_eq!(rec.capability_value, r#"["temperature"]"#);
    }

    #[test]
    fn writing_rejections_merges_with_a_concurrent_sessions() {
        let store = temp_store("merge");
        write_rejections(&store, "openai", "m", &Rejections::from(["temperature"])).expect("a");
        write_rejections(&store, "openai", "m", &Rejections::from(["max_tokens"])).expect("b");
        assert_eq!(
            read_rejections(&store, "openai", "m"),
            Some(Rejections::from(["max_tokens", "temperature"]))
        );
    }

    // The incident wording (Ollama Cloud, minimax-m3), raw and JSON-wrapped.
    const MINIMAX_RAW: &str =
        "max_tokens (521276) exceeds model's maximum output tokens (131072) for model minimax-m3";
    const MINIMAX_JSON: &str = r#"{"error":"max_tokens (521276) exceeds model's maximum output tokens (131072) for model minimax-m3 (ref: a05c9ffb-168f)"}"#;
    const OPENAI_STYLE: &str = r#"{"error":{"message":"max_tokens is too large: 200000. This model supports at most 16384 completion tokens, whereas you provided 200000.","type":"invalid_request_error"}}"#;

    #[test]
    fn parse_output_cap_matches_documented_wordings() {
        assert_eq!(parse_output_cap_message(MINIMAX_RAW), Some(131_072));
        assert_eq!(parse_output_cap_message(MINIMAX_JSON), Some(131_072));
        assert_eq!(parse_output_cap_message(OPENAI_STYLE), Some(16_384));
    }

    #[test]
    fn parse_output_cap_never_matches_context_limit_wordings() {
        // Learning a context window as an output cap would poison the cache —
        // these must all be None even though they mention token limits.
        for body in [
            "prompt is too long: 210000 tokens > 200000 maximum",
            "This model's maximum context length is 128000 tokens",
            "input length and max_tokens exceed context limit: 190000 + 20000 > 200000",
            "the request exceeds the maximum context window of 131072 tokens",
            "rate limit exceeded, try again in 20s",
            "",
        ] {
            assert_eq!(parse_output_cap_message(body), None, "matched: {body}");
        }
    }

    #[test]
    fn parse_output_cap_rejects_nonsense_values() {
        // Sub-1024 and absurd values are parser noise, not real ceilings.
        assert_eq!(
            parse_output_cap_message("exceeds model's maximum output tokens (512)"),
            None
        );
        assert_eq!(
            parse_output_cap_message("exceeds model's maximum output tokens (99999999999)"),
            None
        );
        // A number too far from the marker belongs to something else.
        assert_eq!(
            parse_output_cap_message(
                "exceeds model's maximum output tokens for this deployment tier which is 131072"
            ),
            None
        );
    }

    #[test]
    fn retry_cap_triple() {
        // AUTO (0 / omitted) → retry at the learned cap.
        assert_eq!(retry_cap(0, 131_072), Some(131_072));
        // Explicit ask above the cap → clamp.
        assert_eq!(retry_cap(521_276, 131_072), Some(131_072));
        // Ask already within the cap → the 400 was about something else.
        assert_eq!(retry_cap(4_096, 131_072), None);
    }

    #[test]
    fn output_cap_from_error_gates_on_http_400() {
        let err_400 = ModelError::Backend(mermaid_model::models::BackendError::HttpError {
            status: 400,
            message: MINIMAX_JSON.to_string(),
            debug: Default::default(),
        });
        assert_eq!(output_cap_from_error(&err_400), Some(131_072));
        // Same body on a 500 is a transient fault, not a learned limit.
        let err_500 = ModelError::Backend(mermaid_model::models::BackendError::HttpError {
            status: 500,
            message: MINIMAX_JSON.to_string(),
            debug: Default::default(),
        });
        assert_eq!(output_cap_from_error(&err_500), None);
        assert_eq!(output_cap_from_error(&ModelError::Cancelled), None);
    }
}