mermaid_cli/providers/model/

ollama.rs

//! Ollama provider wrapping `OllamaAdapter`.
//!
//! The adapter owns the wire format (NDJSON framing, gpt-oss
//! reasoning dispatch, truncation marker, retry). The wrapper
//! translates `ChatRequest` ↔ `ModelConfig` and bridges the
//! adapter's legacy `StreamCallback` to the typed `StreamEvent`
//! sink. Adapter-internals stay where they are; the architecture
//! boundary is at `ModelProvider::chat`.

use std::sync::Arc;

use async_trait::async_trait;

use mermaid_domain::ChatRequest;
use mermaid_model::models::adapters::ollama::{OllamaAdapter, OllamaModelInfo};
use mermaid_model::models::adapters::ollama_sizing::{
    NumCtxInputs, converge_num_ctx, default_ollama_num_predict, kv_bytes_per_token,
    resolve_ollama_num_ctx,
};
use mermaid_model::models::{BackendConfig, Model, ModelConfig, ModelError, Result};
use mermaid_runtime::{NewProviderProbe, RuntimeStore};

use super::super::ctx::{FinalResponse, StreamContext, StreamEvent};
use super::{
    ContextSizing, ModelPlacement, ModelProvider, learn_output_cap, load_limits_from_db,
    output_cap_from_error, probe_is_stale, retry_cap,
};
use mermaid_model::models::ModelCapabilities;

/// Ollama adapter fronted by `ModelProvider`.
pub struct OllamaProvider {
    adapter: OllamaAdapter,
    capabilities: ModelCapabilities,
    /// Shared app `Config` so `build_model_config` can read Ollama
    /// hardware options (`num_ctx`, `num_gpu`, `num_thread`, `numa`) at
    /// call time. Before F11 these were silently dropped because the
    /// wrapper built `ModelConfig` only from `ChatRequest` fields.
    config: Arc<mermaid_domain::Config>,
    /// Cached `/api/show` probe (context window + dims + weight). Filled once per
    /// process per model — the provider itself is cached by `ProviderFactory`, so
    /// this fires a single network probe per model. Backed by the cross-session
    /// `provider_probes` table. Probe *failures* are not cached (left empty for a
    /// cheap retry next turn).
    ctx_cell: tokio::sync::OnceCell<OllamaModelInfo>,
}

impl OllamaProvider {
    /// Backward-compatible constructor that uses a default app config.
    /// Call `with_app_config` instead when you have one available so
    /// Ollama hardware options actually reach the adapter.
    ///
    /// # Errors
    ///
    /// Only [`Self::with_app_config`]'s.
    pub async fn new(model_name: &str, backend: Arc<BackendConfig>) -> Result<Self> {
        Self::with_app_config(
            model_name,
            backend,
            Arc::new(mermaid_domain::Config::default()),
        )
        .await
    }

    /// Construct with an explicit `domain::Config` reference. Used by
    /// `ProviderFactory::build_provider` so `config.ollama.{num_gpu,
    /// num_ctx, num_thread, numa}` make it into the Ollama request's
    /// `options` block.
    ///
    /// # Errors
    ///
    /// Only [`OllamaAdapter::new`]'s — the HTTP client build. The server is
    /// not contacted here: the model probe is lazy, so a stopped Ollama still
    /// constructs, and with `ollama_autostart` set the recovery hook revives
    /// it on the first request instead.
    pub async fn with_app_config(
        model_name: &str,
        backend: Arc<BackendConfig>,
        config: Arc<mermaid_domain::Config>,
    ) -> Result<Self> {
        let autostart = backend.ollama_autostart;
        let adapter = OllamaAdapter::new(model_name, backend).await?;
        // The chat path is an *intent* path: a dead local server is mermaid's
        // problem, not the user's, so hand the adapter the means to revive one.
        // Enumeration verbs deliberately construct the adapter without this —
        // see `LocalServerRecovery`.
        let adapter = if autostart {
            adapter.with_recovery(Arc::new(crate::ollama::OllamaAutostart))
        } else {
            adapter
        };
        let capabilities = adapter.capabilities().clone();
        Ok(Self {
            adapter,
            capabilities,
            config,
            ctx_cell: tokio::sync::OnceCell::new(),
        })
    }

    /// Probe the model's capabilities, cache-first. In-process `OnceCell` backed
    /// by the cross-session `provider_probes` table; failures aren't cached.
    async fn probe(&self) -> Option<OllamaModelInfo> {
        self.ctx_cell
            .get_or_try_init(|| async { self.load_probe().await.ok_or(()) })
            .await
            .ok()
            .cloned()
    }

    async fn load_probe(&self) -> Option<OllamaModelInfo> {
        let model = self.adapter.name().to_string();
        if let Some(info) = load_probe_from_db(model.clone()).await {
            return Some(info);
        }
        let info = self.adapter.show_model_info().await?;
        save_probe_to_db(model, info.clone()).await;
        Some(info)
    }

    /// Assemble the sizing inputs from the probe, host memory, config, and the
    /// request's per-model override. Built here and in the effect layer from the
    /// same (cached) sources so the effective window can't diverge between the
    /// request and compaction.
    async fn num_ctx_inputs(
        &self,
        info: &OllamaModelInfo,
        override_num_ctx: Option<u32>,
        override_offload: Option<bool>,
    ) -> NumCtxInputs {
        // Live `/context offload` toggle wins; otherwise the persisted default.
        // The provider's `config` is frozen at startup (the factory never
        // rebuilds), so the toggle rides on the request instead of `config`.
        let allow_ram_offload = override_offload.unwrap_or(self.config.ollama.allow_ram_offload);
        // Only fetch the budget we'll actually use: VRAM when keeping the model
        // on the GPU (default), system RAM when offload is allowed.
        let (vram_bytes, system_ram_bytes) = if allow_ram_offload {
            (None, mermaid_model::utils::system_ram_bytes())
        } else {
            (mermaid_model::utils::gpu_vram_bytes().await, None)
        };
        NumCtxInputs {
            model_max: info.context_length,
            dims: info.dims,
            model_weight_bytes: info.weight_bytes,
            per_model_override: override_num_ctx,
            global_num_ctx: self.config.ollama.num_ctx,
            allow_ram_offload,
            vram_bytes,
            system_ram_bytes,
            max_auto_cap: self.config.ollama.max_auto_num_ctx,
            is_cloud: crate::ollama::is_cloud_model(self.adapter.name()),
        }
    }
}

#[async_trait]
impl ModelProvider for OllamaProvider {
    fn capabilities(&self) -> &ModelCapabilities {
        &self.capabilities
    }

    async fn resolve_context_window(&self, request: &ChatRequest) -> ContextSizing {
        let info = self.probe().await.unwrap_or_default();
        let inputs = self
            .num_ctx_inputs(
                &info,
                request.ollama_num_ctx,
                request.ollama_allow_ram_offload,
            )
            .await;
        let model_max = inputs.model_max;
        // Local Ollama imposes no per-response ceiling (num_predict is ours),
        // but Ollama Cloud maps num_predict → max_tokens and 400s above the
        // model's real cap. A cap learned from such a 400 lives in the limits
        // cache; feed it forward so sizing computes min(window room, cap).
        let max_output =
            load_limits_from_db("ollama".to_string(), Model::name(&self.adapter).to_string())
                .await
                .and_then(|l| l.max_output_tokens);
        match resolve_ollama_num_ctx(&inputs) {
            Some(r) => ContextSizing {
                model_max,
                effective: Some(r.value),
                source: Some(r.source),
                max_output,
            },
            // No model_max and nothing configured → omit num_ctx (Ollama default).
            None => ContextSizing {
                model_max,
                effective: None,
                source: None,
                max_output,
            },
        }
    }

    async fn verify_placement(&self, current_num_ctx: Option<usize>) -> Option<ModelPlacement> {
        let (vram, total) = self.adapter.model_placement().await?;
        // A zero total means Ollama reported the model but not its footprint —
        // can't judge placement, so leave it unknown rather than guess.
        if total == 0 {
            return None;
        }
        // On a spill, compute the largest num_ctx that would fit, from the model's
        // KV cost (probe dims) and the *measured* overflow — the auto-converge
        // target. `None` if it already fits or shrinking can't help.
        let suggested_num_ctx = if vram < total {
            let info = self.probe().await.unwrap_or_default();
            current_num_ctx
                .zip(info.dims)
                .and_then(|(current, dims)| {
                    let kv = kv_bytes_per_token(&dims)?;
                    converge_num_ctx(current, vram, total, kv)
                })
                .map(|n| n as u32)
        } else {
            None
        };
        Some(ModelPlacement {
            size_vram_bytes: vram,
            total_bytes: total,
            suggested_num_ctx,
        })
    }

    async fn supports_vision(&self) -> Option<bool> {
        Some(self.adapter.vision_supported().await)
    }

    async fn chat(&self, request: ChatRequest, ctx: StreamContext) -> Result<FinalResponse> {
        // Resolve the effective window first (cache-first probe). Idempotent with
        // the effect layer's call — both read the same cached probe + memory, so
        // what we send as num_ctx matches what compaction assumes.
        let sizing = self.resolve_context_window(&request).await;
        let config =
            build_model_config(&request, &self.config, sizing.effective, sizing.max_output);
        // Ordered relay (F2): the adapter's sync callback pushes into an
        // `UnboundedSender` (synchronous, FIFO). A single relay task drains
        // into the bounded sink in order, avoiding the per-event `tokio::
        // spawn` race that could deliver `Done` before prior tool calls.
        let (relay_tx, relay_handle) = super::stream_bridge::ordered_relay(ctx.sink.clone());
        let callback = super::stream_bridge::forward_callback(relay_tx.clone());

        // Race adapter.chat against the cancellation token. When
        // cancelled, the adapter's stream loop observes the sink
        // closing (we drop `callback`) and exits at its next await.
        // This is the crucial structural win vs. the old
        // `check_interrupt` polling: the adapter doesn't need to
        // know anything about turn IDs — the sink either drains or
        // doesn't, and the tokens handle everything else.
        let chat_fut = async {
            match self
                .adapter
                .chat(&request.messages, &config, Some(callback.clone()))
                .await
            {
                Ok(response) => Ok(response),
                Err(err) => {
                    // Learn-from-400: Ollama Cloud rejects a num_predict above
                    // the model's real output cap and names the cap in the
                    // body. Learn it (persisted — later turns size below it up
                    // front), clamp, retry ONCE. A request that 400s streamed
                    // no events, so the relay is untouched and reusable.
                    let Some(cap) = output_cap_from_error(&err) else {
                        return Err(err);
                    };
                    let sent = config
                        .ollama_options()
                        .num_predict
                        .map_or(0, |v| v.max(0) as usize);
                    if retry_cap(sent, cap).is_none() {
                        return Err(err);
                    }
                    let model = Model::name(&self.adapter).to_string();
                    learn_output_cap("ollama".to_string(), model.clone(), cap).await;
                    let _ = relay_tx.send(StreamEvent::Status(format!(
                        "{model} rejected the output budget; learned its {cap}-token cap and retrying"
                    )));
                    let retry_config =
                        build_model_config(&request, &self.config, sizing.effective, Some(cap));
                    self.adapter
                        .chat(&request.messages, &retry_config, Some(callback.clone()))
                        .await
                },
            }
        };

        let response = tokio::select! {
            biased;
            _ = ctx.token.cancelled() => {
                // Terminal event for a cancelled turn comes from the
                // runner's `drop_scope` once the `TurnScope` drains, so
                // we neither emit `StreamEvent::Done` here nor surface
                // an `UpstreamError`. `ModelError::Cancelled` is the
                // sentinel the runner swallows.
                return Err(ModelError::Cancelled);
            },
            r = chat_fut => r?,
        };

        // F3: the wrapper's `Done` is now the sole terminal event —
        // the adapter no longer emits one from the callback. Carrying
        // `provider_continuation` out of `ModelResponse` here is what
        // lets multi-turn extended thinking round-trip.
        let usage = response.usage.clone();
        let provider_continuation = response.provider_continuation.clone();
        let stop_reason = response.stop_reason.clone();
        // Terminal Done through the ordered relay, then drain (see stream_bridge).
        let _ = relay_tx.send(StreamEvent::Done {
            usage: usage.clone(),
            provider_continuation: provider_continuation.clone(),
            stop_reason: stop_reason.clone(),
        });
        drop(relay_tx);
        mermaid_model::utils::join_logged(relay_handle.take(), "stream_relay").await;

        Ok(FinalResponse {
            usage,
            provider_continuation,
            tool_calls: response.tool_calls.unwrap_or_default(),
            stop_reason,
        })
    }
}

// ─── helpers ────────────────────────────────────────────────────────

/// `num_ctx` is the resolved effective window (auto-fitted, or override/global);
/// passing it here keeps a single source of truth — the direct `config.ollama.
/// num_ctx` forward is gone because the resolver already considered it.
/// `provider_max_output` is a known per-response ceiling (learned from an
/// Ollama Cloud 400), bounding the derived `num_predict`.
fn build_model_config(
    request: &ChatRequest,
    app_config: &mermaid_domain::Config,
    num_ctx: Option<usize>,
    provider_max_output: Option<usize>,
) -> ModelConfig {
    let mut mc = ModelConfig {
        model: request.model_id.clone(),
        temperature: request.temperature,
        max_tokens: request.max_tokens,
        reasoning: request.reasoning,
        system_prompt: Some(request.system_prompt.clone()),
        dynamic_system_suffix: request.instructions.clone(),
        tools: request.tools.iter().map(|t| t.to_openai_json()).collect(),
        output_schema: request.output_schema.clone(),
        ..Default::default()
    };
    // Effective context window (auto-fitted to memory / override / global).
    if let Some(n) = num_ctx {
        mc.set_backend_option("ollama".into(), "num_ctx".into(), n.to_string());
    }
    // Output cap: AUTO (max_tokens == 0) gets the full room num_ctx leaves
    // after the prompt; an explicit max_tokens is an exact cap bounded by that
    // room. Without a cap Ollama generates unbounded and only stops when the
    // window fills — the truncation bug. `None` = AUTO with an unknown window →
    // omit num_predict so Ollama applies its own default.
    if let Some(num_predict) = default_ollama_num_predict(
        request.max_tokens,
        num_ctx,
        estimate_prompt_tokens(request),
        provider_max_output,
    ) {
        mc.set_backend_option(
            "ollama".into(),
            "num_predict".into(),
            num_predict.to_string(),
        );
    }

    // F11: forward Ollama hardware options from the user's app config (num_ctx is
    // now handled above via the resolver).
    if let Some(v) = app_config.ollama.num_gpu {
        mc.set_backend_option("ollama".into(), "num_gpu".into(), v.to_string());
    }
    if let Some(v) = app_config.ollama.num_thread {
        mc.set_backend_option("ollama".into(), "num_thread".into(), v.to_string());
    }
    if let Some(v) = app_config.ollama.numa {
        mc.set_backend_option("ollama".into(), "numa".into(), v.to_string());
    }
    mc
}

/// Rough prompt-token estimate (≈4 chars/token) for bounding `num_predict`
/// against the remaining room in `num_ctx`. Approximate by design — it only
/// gates the output cap, never the prompt itself.
fn estimate_prompt_tokens(request: &ChatRequest) -> usize {
    let chars = request.system_prompt.len()
        + request.instructions.as_deref().map_or(0, str::len)
        + request
            .messages
            .iter()
            .map(|m| m.content.len())
            .sum::<usize>();
    chars / 4
}

/// Load a cached probe from `provider_probes` (within TTL). Runs the blocking
/// SQLite read off the async runtime. Best-effort: any failure → `None`.
async fn load_probe_from_db(model: String) -> Option<OllamaModelInfo> {
    tokio::task::spawn_blocking(move || {
        let store = RuntimeStore::open_default().ok()?;
        let rec = store
            .provider_probes()
            .get("ollama", &model, "context_probe")
            .ok()??;
        if probe_is_stale(&rec.probed_at) {
            return None;
        }
        serde_json::from_str::<OllamaModelInfo>(&rec.capability_value).ok()
    })
    .await
    .ok()
    .flatten()
}

/// Persist a probe to `provider_probes` for subsequent sessions. Best-effort.
async fn save_probe_to_db(model: String, info: OllamaModelInfo) {
    let _ = tokio::task::spawn_blocking(move || -> Option<()> {
        let value = serde_json::to_string(&info).ok()?;
        let store = RuntimeStore::open_default().ok()?;
        store
            .provider_probes()
            .upsert(NewProviderProbe {
                provider: "ollama".into(),
                model_id: model,
                capability_key: "context_probe".into(),
                capability_value: value,
                confidence: "probed".into(),
                error: None,
            })
            .ok()?;
        Some(())
    })
    .await;
}

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

    #[test]
    fn build_model_config_maps_request_fields() {
        let req = ChatRequest {
            model_id: "ollama/test".to_string(),
            messages: vec![],
            system_prompt: "sys".to_string(),
            instructions: Some("instructions text".to_string()),
            reasoning: mermaid_model::models::ReasoningLevel::High,
            temperature: 0.3,
            max_tokens: 2048,
            tools: vec![],

            ollama_num_ctx: None,
            ollama_allow_ram_offload: None,
            resolved_context_window: None,
            resolved_max_output: None,
            output_schema: None,
            suppress_auto_compact: false,
            suppressed_builtin_tools: Vec::new(),
        };
        let app_cfg = mermaid_domain::Config::default();
        let cfg = build_model_config(&req, &app_cfg, None, None);
        assert_eq!(cfg.model, "ollama/test");
        assert_eq!(cfg.temperature, 0.3);
        assert_eq!(cfg.max_tokens, 2048);
        assert_eq!(cfg.reasoning, mermaid_model::models::ReasoningLevel::High);
        assert_eq!(cfg.system_prompt.as_deref(), Some("sys"));
        assert_eq!(
            cfg.dynamic_system_suffix.as_deref(),
            Some("instructions text")
        );
    }

    /// F11 regression guard: Ollama hardware options in the user's app config
    /// must land in the `ModelConfig`'s `backend_options` so the adapter's
    /// `build_request_body` emits them under `options`. `num_ctx` now arrives via
    /// the resolver param (not a direct config forward), and `num_predict` is
    /// always derived.
    #[test]
    fn build_model_config_forwards_ollama_hardware_options() {
        let req = ChatRequest {
            model_id: "ollama/test".to_string(),
            messages: vec![],
            system_prompt: "sys".to_string(),
            instructions: None,
            reasoning: mermaid_model::models::ReasoningLevel::Medium,
            temperature: 0.7,
            max_tokens: 4096,
            tools: vec![],

            ollama_num_ctx: None,
            ollama_allow_ram_offload: None,
            resolved_context_window: None,
            resolved_max_output: None,
            output_schema: None,
            suppress_auto_compact: false,
            suppressed_builtin_tools: Vec::new(),
        };
        let mut app_cfg = mermaid_domain::Config::default();
        app_cfg.ollama.num_gpu = Some(10);
        app_cfg.ollama.num_thread = Some(8);
        app_cfg.ollama.numa = Some(true);

        // The effective num_ctx (8192) is passed in by the resolver.
        let cfg = build_model_config(&req, &app_cfg, Some(8192), None);
        let opts = cfg.ollama_options();
        assert_eq!(opts.num_ctx, Some(8192));
        assert_eq!(opts.num_gpu, Some(10));
        assert_eq!(opts.num_thread, Some(8));
        assert_eq!(opts.numa, Some(true));
        assert!(opts.num_predict.is_some(), "num_predict is always derived");
    }

    /// An explicit `max_tokens` is an exact cap, bounded by the room left in
    /// `num_ctx` (no reasoning reserve added — hard means hard).
    #[test]
    fn build_model_config_derives_num_predict() {
        let req = ChatRequest {
            model_id: "ollama/test".to_string(),
            messages: vec![],
            system_prompt: String::new(),
            instructions: None,
            reasoning: mermaid_model::models::ReasoningLevel::Max,
            temperature: 0.7,
            max_tokens: 4096,
            tools: vec![],

            ollama_num_ctx: None,
            ollama_allow_ram_offload: None,
            resolved_context_window: None,
            resolved_max_output: None,
            output_schema: None,
            suppress_auto_compact: false,
            suppressed_builtin_tools: Vec::new(),
        };
        let cfg = build_model_config(
            &req,
            &mermaid_domain::Config::default(),
            Some(131_072),
            None,
        );
        // Exactly the 4096 cap; plenty of room in a 131072 window.
        assert_eq!(cfg.ollama_options().num_predict, Some(4_096));
    }

    /// The minimax incident, recomputed: a learned provider cap bounds the
    /// AUTO `num_predict` below the full window room, so the retry (and every
    /// later turn) sends a value Ollama Cloud accepts.
    #[test]
    fn build_model_config_caps_num_predict_at_learned_ceiling() {
        let req = ChatRequest {
            model_id: "ollama/minimax-m3:cloud".to_string(),
            messages: vec![],
            system_prompt: String::new(),
            instructions: None,
            reasoning: mermaid_model::models::ReasoningLevel::Medium,
            temperature: 0.7,
            max_tokens: 0, // AUTO — the incident's configuration
            tools: vec![],

            ollama_num_ctx: None,
            ollama_allow_ram_offload: None,
            resolved_context_window: None,
            resolved_max_output: None,
            output_schema: None,
            suppress_auto_compact: false,
            suppressed_builtin_tools: Vec::new(),
        };
        let app_cfg = mermaid_domain::Config::default();
        // Without a learned cap, AUTO hands over the full window room —
        // 524_288 minus margin — which Ollama Cloud 400s for minimax-m3.
        let uncapped = build_model_config(&req, &app_cfg, Some(524_288), None);
        assert!(uncapped.ollama_options().num_predict.unwrap() > 131_072);
        // With the learned 131_072 cap, sizing stays at the ceiling.
        let capped = build_model_config(&req, &app_cfg, Some(524_288), Some(131_072));
        assert_eq!(capped.ollama_options().num_predict, Some(131_072));
    }
}