diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index d748df8697b7..c1f3e4fdcc4f 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -3870,6 +3870,7 @@ dependencies = [ "codex-home", "codex-http-client", "codex-install-context", + "codex-local-models", "codex-login", "codex-mcp", "codex-memories-write", @@ -3926,6 +3927,18 @@ dependencies = [ [[package]] name = "codex-local-models" version = "0.162.1" +dependencies = [ + "codex-http-client", + "codex-login", + "codex-model-provider", + "codex-model-provider-info", + "codex-models-manager", + "codex-protocol", + "reqwest 0.12.28", + "serde_json", + "tokio", + "wiremock", +] [[package]] name = "codex-local-resume" diff --git a/codex-rs/local/cli/Cargo.toml b/codex-rs/local/cli/Cargo.toml index 99ca1fa41821..d5e70310dd5f 100644 --- a/codex-rs/local/cli/Cargo.toml +++ b/codex-rs/local/cli/Cargo.toml @@ -23,6 +23,7 @@ test = false [dependencies] codex-cli = { path = "../../cli" } +codex-local-models = { path = "../models" } dirs = { workspace = true } toml_edit = { workspace = true } codex-utils-process = { workspace = true } diff --git a/codex-rs/local/cli/src/main.rs b/codex-rs/local/cli/src/main.rs index 20b94d2265d8..9c45c2d09632 100644 --- a/codex-rs/local/cli/src/main.rs +++ b/codex-rs/local/cli/src/main.rs @@ -20,6 +20,7 @@ fn main() -> anyhow::Result<()> { }; // SAFETY: still single-threaded; no runtime or worker threads exist yet. unsafe { std::env::set_var("CODEX_HOME", &codex_home) }; + codex_local_models::register(); let remote_control_disabled = codex_app_server::take_remote_control_disabled_env(); arg0_dispatch_or_else(move |arg0_paths: Arg0DispatchPaths| async move { Box::pin(cli_main(argv, arg0_paths, remote_control_disabled)).await?; diff --git a/codex-rs/local/models/Cargo.toml b/codex-rs/local/models/Cargo.toml index fa0543f14ffd..e9210f3c6c59 100644 --- a/codex-rs/local/models/Cargo.toml +++ b/codex-rs/local/models/Cargo.toml @@ -9,3 +9,17 @@ workspace = true [lib] doctest = false + +[dependencies] +codex-http-client = { workspace = true } +codex-login = { workspace = true } +codex-model-provider = { workspace = true } +codex-model-provider-info = { workspace = true } +codex-models-manager = { workspace = true } +codex-protocol = { workspace = true } +reqwest = { workspace = true, features = ["json"] } +serde_json = { workspace = true } +tokio = { workspace = true, features = ["rt", "macros", "sync"] } + +[dev-dependencies] +wiremock = { workspace = true } diff --git a/codex-rs/local/models/src/lib.rs b/codex-rs/local/models/src/lib.rs index 8b137891791f..894274e72361 100644 --- a/codex-rs/local/models/src/lib.rs +++ b/codex-rs/local/models/src/lib.rs @@ -1 +1,349 @@ +//! Ollama-backed model catalog for `codex-local` (`/model` picker, req 8). +//! +//! Lists models via `GET /api/tags` (PAIR merges the cluster's models) and reads each +//! context window via `POST /api/show`. +use std::fmt; +use std::sync::Arc; +use std::time::Duration; + +use codex_http_client::HttpClientFactory; +use codex_http_client::OutboundProxyPolicy; +use codex_login::AuthManager; +use codex_model_provider_info::ModelProviderInfo; +use codex_models_manager::ModelsManagerConfig; +use codex_models_manager::manager::ModelsEndpointClient; +use codex_models_manager::manager::ModelsEndpointFuture; +use codex_models_manager::manager::ModelsEndpointResponse; +use codex_models_manager::manager::ModelsManager; +use codex_models_manager::manager::ModelsManagerFuture; +use codex_models_manager::manager::OpenAiModelsManager; +use codex_models_manager::manager::RefreshStrategy; +use codex_models_manager::manager::SharedModelsManager; +use codex_models_manager::model_info::model_info_from_slug; +use codex_protocol::config_types::CollaborationModeMask; +use codex_protocol::error::CodexErr; +use codex_protocol::error::Result as CoreResult; +use codex_protocol::openai_models::ApplyPatchToolType; +use codex_protocol::openai_models::ModelInfo; +use codex_protocol::openai_models::ModelPreset; +use codex_protocol::openai_models::ModelVisibility; +use codex_protocol::openai_models::ModelsResponse; +use serde_json::Value; +use tokio::sync::TryLockError; + +const REQUEST_TIMEOUT: Duration = Duration::from_secs(10); +/// Matches upstream's fallback metadata when `/api/show` gives no context length. +const FALLBACK_CONTEXT_WINDOW: i64 = 272_000; + +#[derive(Debug)] +pub struct OllamaModelsEndpoint { + root: String, + freeform_apply_patch: bool, + client: reqwest::Client, +} + +impl OllamaModelsEndpoint { + /// `root` is the Ollama root (no `/v1`). `freeform_apply_patch` defaults to off. + pub fn new(root: impl Into, freeform_apply_patch: bool) -> Self { + Self { + root: root.into().trim_end_matches('/').to_string(), + freeform_apply_patch, + client: reqwest::Client::new(), + } + } + + async fn context_window(&self, name: &str) -> Option { + let resp = self + .client + .post(format!("{}/api/show", self.root)) + .timeout(REQUEST_TIMEOUT) + .json(&serde_json::json!({ "model": name })) + .send() + .await + .ok()? + .error_for_status() + .ok()?; + let body: Value = resp.json().await.ok()?; + body.get("model_info")? + .as_object()? + .iter() + .find(|(k, _)| k.ends_with(".context_length")) + .and_then(|(_, v)| v.as_i64()) + } + + async fn fetch(&self) -> CoreResult> { + let invalid = |e: &dyn fmt::Display| CodexErr::InvalidRequest(format!("ollama: {e}")); + let body: Value = self + .client + .get(format!("{}/api/tags", self.root)) + .timeout(REQUEST_TIMEOUT) + .send() + .await + .and_then(reqwest::Response::error_for_status) + .map_err(|e| invalid(&e))? + .json() + .await + .map_err(|e| invalid(&e))?; + let mut names: Vec = Vec::new(); + for m in body["models"].as_array().into_iter().flatten() { + if let Some(n) = m["name"].as_str().or_else(|| m["model"].as_str()) + && !names.iter().any(|x| x == n) + { + names.push(n.to_string()); + } + } + let mut set = tokio::task::JoinSet::new(); + for (i, n) in names.iter().enumerate() { + let (this, n) = (self.clone_light(), n.clone()); + set.spawn(async move { (i, this.context_window(&n).await) }); + } + let mut windows = vec![None; names.len()]; + while let Some(Ok((i, w))) = set.join_next().await { + windows[i] = w; + } + Ok(names + .into_iter() + .zip(windows) + .enumerate() + .map(|(i, (name, w))| self.model_info(&name, w, i)) + .collect()) + } + + fn clone_light(&self) -> Self { + Self { + root: self.root.clone(), + freeform_apply_patch: self.freeform_apply_patch, + client: self.client.clone(), + } + } + + fn model_info(&self, name: &str, window: Option, index: usize) -> ModelInfo { + let mut info = model_info_from_slug(name); + let window = window.unwrap_or(FALLBACK_CONTEXT_WINDOW); + info.visibility = ModelVisibility::List; + info.priority = i32::try_from(index).unwrap_or(i32::MAX); + info.context_window = Some(window); + info.max_context_window = Some(window); + info.used_fallback_model_metadata = false; + if self.freeform_apply_patch { + info.apply_patch_tool_type = Some(ApplyPatchToolType::Freeform); + } + info + } +} + +impl ModelsEndpointClient for OllamaModelsEndpoint { + fn identity(&self) -> Option { + Some(format!("codex-local:{}", self.root)) + } + + fn has_command_auth(&self) -> bool { + // The manager only fetches a catalog for command-auth or ChatGPT-backend providers; + // Ollama needs no auth but must be fetched, so it opts in via this flag. + true + } + + fn uses_codex_backend(&self) -> ModelsEndpointFuture<'_, bool> { + Box::pin(async { false }) + } + + fn list_models<'a>( + &'a self, + _client_version: &'a str, + _http_client_factory: HttpClientFactory, + ) -> ModelsEndpointFuture<'a, CoreResult> { + Box::pin(async move { + Ok(ModelsEndpointResponse { + models: self.fetch().await?, + etag: None, + identity: format!("codex-local:{}", self.root), + }) + }) + } +} + +fn factory(info: &ModelProviderInfo) -> SharedModelsManager { + let base = info.base_url.as_deref().unwrap_or("http://127.0.0.1:11434"); + let root = base.trim_end_matches('/').trim_end_matches("/v1"); + Arc::new(LoadsOnDemand(Arc::new( + OpenAiModelsManager::new_without_cache( + Arc::new(OllamaModelsEndpoint::new(root, false)), + None, + ) + // Ollama's catalog is authoritative: do not merge the bundled gpt-* models. The initial + // list is therefore empty, which `LoadsOnDemand::ensure_loaded` treats as "not loaded". + .with_provider_catalog(), + ))) +} + +/// Upstream resolves model metadata from the in-memory catalog only, which stays empty +/// when a run never lists models (e.g. `exec`), so every model would get the "Model +/// metadata not found" fallback warning. This wrapper loads the catalog on first lookup. +#[derive(Debug)] +struct LoadsOnDemand(SharedModelsManager); + +impl LoadsOnDemand { + async fn ensure_loaded(&self) { + if self.0.get_remote_models().await.is_empty() { + let factory = HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault); + self.0 + .raw_model_catalog(RefreshStrategy::OnlineIfUncached, factory) + .await; + } + } +} + +impl ModelsManager for LoadsOnDemand { + fn raw_model_catalog( + &self, + refresh_strategy: RefreshStrategy, + http_client_factory: HttpClientFactory, + ) -> ModelsManagerFuture<'_, ModelsResponse> { + self.0 + .raw_model_catalog(refresh_strategy, http_client_factory) + } + + fn list_models( + &self, + refresh_strategy: RefreshStrategy, + http_client_factory: HttpClientFactory, + ) -> ModelsManagerFuture<'_, Vec> { + self.0.list_models(refresh_strategy, http_client_factory) + } + + fn refresh_after_auth_change( + &self, + http_client_factory: HttpClientFactory, + ) -> ModelsManagerFuture<'_, ()> { + self.0.refresh_after_auth_change(http_client_factory) + } + + fn get_remote_models(&self) -> ModelsManagerFuture<'_, Vec> { + self.0.get_remote_models() + } + + fn try_get_remote_models(&self) -> Result, TryLockError> { + self.0.try_get_remote_models() + } + + fn auth_manager(&self) -> Option<&AuthManager> { + self.0.auth_manager() + } + + fn list_collaboration_modes(&self) -> Vec { + self.0.list_collaboration_modes() + } + + fn get_default_model<'a>( + &'a self, + model: &'a Option, + allow_provider_model_fallback: bool, + refresh_strategy: RefreshStrategy, + http_client_factory: HttpClientFactory, + ) -> ModelsManagerFuture<'a, String> { + self.0.get_default_model( + model, + allow_provider_model_fallback, + refresh_strategy, + http_client_factory, + ) + } + + fn get_model_info<'a>( + &'a self, + model: &'a str, + config: &'a ModelsManagerConfig, + ) -> ModelsManagerFuture<'a, ModelInfo> { + Box::pin(async move { + self.ensure_loaded().await; + self.0.get_model_info(model, config).await + }) + } + + fn refresh_if_new_etag( + &self, + etag: String, + http_client_factory: HttpClientFactory, + ) -> ModelsManagerFuture<'_, ()> { + self.0.refresh_if_new_etag(etag, http_client_factory) + } +} + +/// Registers the H2 factory. Call once, before creating the `codex-local` provider. +pub fn register() { + codex_model_provider::register_local_models_factory(factory); +} + +#[cfg(test)] +mod tests { + use super::*; + use wiremock::Mock; + use wiremock::MockServer; + use wiremock::ResponseTemplate; + use wiremock::matchers::body_partial_json; + use wiremock::matchers::method; + use wiremock::matchers::path; + + #[tokio::test] + async fn lists_tags_with_context_windows() { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/api/tags")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "models": [{"name": "a:1"}, {"name": "b:2"}, {"name": "c:3"}, {"name": "a:1"}] + }))) + .mount(&server) + .await; + for (n, ctx) in [("a:1", 1000), ("b:2", 2000), ("c:3", 3000)] { + Mock::given(method("POST")) + .and(path("/api/show")) + .and(body_partial_json(serde_json::json!({"model": n}))) + .respond_with(ResponseTemplate::new(200).set_body_json( + serde_json::json!({"model_info": {"general.architecture": "x", "x.context_length": ctx}}), + )) + .mount(&server) + .await; + } + let ep = OllamaModelsEndpoint::new(server.uri(), true); + let models = ep.fetch().await.unwrap(); + let got: Vec<_> = models + .iter() + .map(|m| (m.slug.as_str(), m.context_window.unwrap())) + .collect(); + assert_eq!(got, vec![("a:1", 1000), ("b:2", 2000), ("c:3", 3000)]); + assert!(models.iter().all(|m| m.apply_patch_tool_type.is_some())); + } + + #[tokio::test] + async fn catalog_is_ollama_only_without_bundled_models() { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/api/tags")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "models": [{"name": "a:1"}, {"name": "b:2"}, {"name": "c:3"}] + }))) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/api/show")) + .respond_with(ResponseTemplate::new(200).set_body_json( + serde_json::json!({"model_info": {"general.architecture": "x", "x.context_length": 1000}}), + )) + .mount(&server) + .await; + let info = ModelProviderInfo { + base_url: Some(format!("{}/v1", server.uri())), + ..ModelProviderInfo::create_openai_provider(None) + }; + let manager = factory(&info); + let factory = HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault); + let presets = manager + .list_models(RefreshStrategy::Online, factory) + .await; + let mut ids: Vec<_> = presets.iter().map(|p| p.model.clone()).collect(); + ids.sort(); + assert_eq!(ids, vec!["a:1", "b:2", "c:3"]); + assert!(ids.iter().all(|i| !i.starts_with("gpt-"))); + } +} diff --git a/codex-rs/model-provider/src/lib.rs b/codex-rs/model-provider/src/lib.rs index 9c818af7dd81..eee2c75fe29c 100644 --- a/codex-rs/model-provider/src/lib.rs +++ b/codex-rs/model-provider/src/lib.rs @@ -31,7 +31,9 @@ pub use codex_model_provider_info::AMAZON_BEDROCK_RUNTIME_PROVIDER_ID; pub use codex_model_provider_info::CHATGPT_CODEX_BASE_URL; pub use codex_model_provider_info::RemoteCompactionSupport; pub use codex_protocol::account::ProviderAccount; +pub use provider::LocalModelsFactory; pub use provider::ModelProvider; +pub use provider::register_local_models_factory; pub use provider::ModelProviderFuture; pub use provider::ProviderAccountError; pub use provider::ProviderAccountResult; diff --git a/codex-rs/model-provider/src/provider.rs b/codex-rs/model-provider/src/provider.rs index f8f0790dde81..d58d16e618c6 100644 --- a/codex-rs/model-provider/src/provider.rs +++ b/codex-rs/model-provider/src/provider.rs @@ -357,6 +357,18 @@ pub fn create_model_provider( )) } +/// codex-local hook H2: factory registered by `codex-local-cli`, used only for `codex-local`. +pub type LocalModelsFactory = fn(&ModelProviderInfo) -> SharedModelsManager; +static LOCAL_MODELS_FACTORY: std::sync::OnceLock = std::sync::OnceLock::new(); + +pub fn register_local_models_factory(factory: LocalModelsFactory) { + let _ = LOCAL_MODELS_FACTORY.set(factory); +} + +fn local_models_manager(info: &ModelProviderInfo) -> Option { + (info.name == "codex-local").then(|| LOCAL_MODELS_FACTORY.get().map(|f| f(info)))? +} + /// Runtime model provider that orchestrates primary and gateway credentials. #[derive(Clone, Debug)] struct ConfiguredModelProvider { @@ -396,6 +408,9 @@ impl ConfiguredModelProvider { model_catalog, )); } + if let Some(manager) = local_models_manager(&self.info) { + return manager; + } let endpoint = Arc::new(OpenAiModelsEndpoint::new( self.info.clone(), self.auth_manager.clone(), @@ -1403,4 +1418,23 @@ printf '%s\n' '{"AccessKeyId":"exported","SecretAccessKey":"secret"}' } } } + + #[test] + fn local_models_factory_is_consulted_only_for_codex_local() { + static CALLS: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0); + fn factory(_: &ModelProviderInfo) -> SharedModelsManager { + CALLS.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + Arc::new(StaticModelsManager::new(None, ModelsResponse::default())) + } + let named = |name: &str| ModelProviderInfo { + name: name.to_string(), + ..Default::default() + }; + assert!(local_models_manager(&named("codex-local")).is_none()); + register_local_models_factory(factory); + assert!(local_models_manager(&named("openai")).is_none()); + assert_eq!(CALLS.load(std::sync::atomic::Ordering::SeqCst), 0); + assert!(local_models_manager(&named("codex-local")).is_some()); + assert_eq!(CALLS.load(std::sync::atomic::Ordering::SeqCst), 1); + } } diff --git a/local/FORK_SURFACE.md b/local/FORK_SURFACE.md index a4a10c659e1d..ef5e8b26f1cd 100644 --- a/local/FORK_SURFACE.md +++ b/local/FORK_SURFACE.md @@ -3,7 +3,7 @@ | Hook | Upstream file | Purpose | Status | |------|---------------|---------|--------| | H1 | `codex-rs/cli/src/main.rs` (+1 line in `codex-rs/cli/build.rs`) | `cli_main` takes argv (`parse_from`) so `codex-local-cli`, which `include!`s this file as its crate root, can call it with its own argv. Stock `main` is gated `#[cfg(not(codex_local))]`. 11 lines added, 2 removed in main.rs (the `Box::pin(cli_main(..))` call grew from 1 to 6 lines via rustfmt); build.rs declares the cfg. | Done | -| H2 | TODO | TODO | TODO | +| H2 | `codex-rs/model-provider/src/provider.rs` (+ 2 `pub use` lines in `lib.rs`) | `OnceLock` `LOCAL_MODELS_FACTORY` + `register_local_models_factory`; `create_models_manager` consults it only when provider `name == "codex-local"` (no provider id exists at this layer). Unregistered: stock behaviour. Lines: 17 hook + 19 unit test. Re-apply: add the static, `local_models_manager()` helper and the 3-line early return after the `config_model_catalog` check in `create_models_manager`; re-export from lib.rs. | done | | H3 | TODO | TODO | TODO | | H4 | TODO | TODO | TODO | | H5 | TODO | TODO | TODO | diff --git a/local/fork-surface.toml b/local/fork-surface.toml index 85347398e41e..0657eff503c8 100644 --- a/local/fork-surface.toml +++ b/local/fork-surface.toml @@ -1,5 +1,5 @@ # Upstream paths the fork is allowed to modify. -allowlist = ["codex-rs/Cargo.toml", "codex-rs/cli/src/main.rs", "codex-rs/cli/build.rs", "codex-rs/core/src/tools/router.rs", "codex-rs/core/src/tools/router_tests.rs"] +allowlist = ["codex-rs/Cargo.toml", "codex-rs/cli/src/main.rs", "codex-rs/cli/build.rs", "codex-rs/core/src/tools/router.rs", "codex-rs/core/src/tools/router_tests.rs", "codex-rs/model-provider/src/provider.rs", "codex-rs/model-provider/src/lib.rs"] generated = ["codex-rs/Cargo.lock"] [budget]