feat: support map of alternative providers like in TypeScript CLI

This commit is contained in:
Michael Bolin
2025-05-07 14:37:16 -07:00
parent 42617f8726
commit e2a5c000a0
16 changed files with 210 additions and 22 deletions

1
codex-rs/Cargo.lock generated
View File

@@ -528,6 +528,7 @@ dependencies = [
"libc",
"mcp-types",
"mime_guess",
"once_cell",
"openssl-sys",
"patch",
"path-absolutize",

View File

@@ -28,6 +28,7 @@ rand = "0.9"
reqwest = { version = "0.12", features = ["json", "stream"] }
serde = { version = "1", features = ["derive"] }
serde_json = "1"
once_cell = "1.19.0"
thiserror = "2.0.12"
time = { version = "0.3", features = ["formatting", "macros"] }
tokio = { version = "1", features = [

View File

@@ -26,10 +26,9 @@ use tracing::warn;
use crate::error::CodexErr;
use crate::error::Result;
use crate::flags::CODEX_RS_SSE_FIXTURE;
use crate::flags::OPENAI_API_BASE;
use crate::flags::OPENAI_REQUEST_MAX_RETRIES;
use crate::flags::OPENAI_STREAM_IDLE_TIMEOUT_MS;
use crate::flags::get_api_key;
use crate::model_provider_info::ModelProviderInfo;
use crate::models::ResponseItem;
use crate::util::backoff;
@@ -141,13 +140,16 @@ static DEFAULT_TOOLS: LazyLock<Vec<ResponsesApiTool>> = LazyLock::new(|| {
pub struct ModelClient {
model: String,
client: reqwest::Client,
provider: ModelProviderInfo,
}
impl ModelClient {
pub fn new(model: impl ToString) -> Self {
let model = model.to_string();
let client = reqwest::Client::new();
Self { model, client }
pub fn new(model: impl ToString, provider: ModelProviderInfo) -> Self {
Self {
model: model.to_string(),
client: reqwest::Client::new(),
provider,
}
}
pub async fn stream(&mut self, prompt: &Prompt) -> Result<ResponseStream> {
@@ -188,18 +190,28 @@ impl ModelClient {
stream: true,
};
let url = format!("{}/v1/responses", *OPENAI_API_BASE);
let base_url = self.provider.api_base(&self.provider.base_url);
let base_url = base_url.trim_end_matches('/');
let url = format!("{}/responses", base_url);
debug!(url, "POST");
trace!("request payload: {}", serde_json::to_string(&payload)?);
println!(
"request {url:?} payload: {:?}",
serde_json::to_string(&payload)
);
let mut attempt = 0;
loop {
attempt += 1;
let api_key = self
.provider
.api_key()
.ok_or_else(|| crate::error::CodexErr::EnvVar("API_KEY"))?;
let res = self
.client
.post(&url)
.bearer_auth(get_api_key()?)
.bearer_auth(api_key)
.header("OpenAI-Beta", "responses=experimental")
.header(reqwest::header::ACCEPT, "text/event-stream")
.json(&payload)

View File

@@ -564,6 +564,7 @@ async fn submission_loop(
sess.abort();
}
Op::ConfigureSession {
provider,
model,
instructions,
approval_policy,
@@ -572,7 +573,7 @@ async fn submission_loop(
notify,
cwd,
} => {
info!(model, "Configuring session");
info!("Configuring session: model={model}; provider={provider:?}");
if !cwd.is_absolute() {
let message = format!("cwd is not absolute: {cwd:?}");
error!(message);
@@ -586,7 +587,7 @@ async fn submission_loop(
return;
}
let client = ModelClient::new(model.clone());
let client = ModelClient::new(model.clone(), provider.clone());
// abort any current running session and clone its state
let state = match sess.take() {

View File

@@ -16,10 +16,16 @@ use tokio::sync::Notify;
/// is received as a response to the initial `ConfigureSession` submission so
/// that callers can surface the information to the UI.
pub async fn init_codex(config: Config) -> anyhow::Result<(CodexWrapper, Event, Arc<Notify>)> {
let provider = config
.providers
.get(&config.provider)
.ok_or_else(|| anyhow::anyhow!("provider {} not found in config", config.provider))?;
let ctrl_c = notify_on_sigint();
let codex = CodexWrapper::new(Codex::spawn(ctrl_c.clone())?);
let init_id = codex
.submit(Op::ConfigureSession {
provider: provider.clone(),
model: config.model.clone(),
instructions: config.instructions.clone(),
approval_policy: config.approval_policy,

View File

@@ -1,5 +1,7 @@
use crate::flags::OPENAI_DEFAULT_MODEL;
use crate::mcp_server_config::McpServerConfig;
use crate::model_provider_info::ModelProviderInfo;
use crate::model_provider_info::built_in_model_providers;
use crate::protocol::AskForApproval;
use crate::protocol::SandboxPermission;
use crate::protocol::SandboxPolicy;
@@ -19,6 +21,9 @@ pub struct Config {
/// Optional override of model selection.
pub model: String,
/// Selected provider ("openai", "openrouter", ...)
pub provider: String,
/// Approval policy for executing commands.
pub approval_policy: AskForApproval,
@@ -61,6 +66,9 @@ pub struct Config {
/// Definition for MCP servers that Codex can reach out to for tool calls.
pub mcp_servers: HashMap<String, McpServerConfig>,
/// Combined provider map (defaults merged with user-defined overrides).
pub providers: HashMap<String, ModelProviderInfo>,
}
/// Base config deserialized from ~/.codex/config.toml.
@@ -69,6 +77,9 @@ pub struct ConfigToml {
/// Optional override of model selection.
pub model: Option<String>,
/// Selected provider
pub provider: Option<String>,
/// Default approval policy for executing commands.
pub approval_policy: Option<AskForApproval>,
@@ -93,6 +104,10 @@ pub struct ConfigToml {
/// Definition for MCP servers that Codex can reach out to for tool calls.
#[serde(default)]
pub mcp_servers: HashMap<String, McpServerConfig>,
/// User-defined provider entries that extend/override the built-in list.
#[serde(default)]
pub providers: HashMap<String, ModelProviderInfo>,
}
impl ConfigToml {
@@ -152,6 +167,8 @@ pub struct ConfigOverrides {
pub approval_policy: Option<AskForApproval>,
pub sandbox_policy: Option<SandboxPolicy>,
pub disable_response_storage: Option<bool>,
pub provider: Option<String>,
}
impl Config {
@@ -176,6 +193,7 @@ impl Config {
approval_policy,
sandbox_policy,
disable_response_storage,
provider,
} = overrides;
let sandbox_policy = match sandbox_policy {
@@ -193,8 +211,17 @@ impl Config {
}
};
let mut model_providers = built_in_model_providers();
// Merge user-defined providers into the built-in list.
for (key, provider) in cfg.providers.into_iter() {
model_providers.entry(key).or_insert(provider);
}
Self {
model: model.or(cfg.model).unwrap_or_else(default_model),
provider: provider
.or(cfg.provider)
.unwrap_or_else(|| "openai".to_string()),
cwd: cwd.map_or_else(
|| {
tracing::info!("cwd not set, using current dir");
@@ -222,6 +249,7 @@ impl Config {
notify: cfg.notify,
instructions,
mcp_servers: cfg.mcp_servers,
providers: model_providers,
}
}

View File

@@ -2,13 +2,13 @@ use std::time::Duration;
use env_flags::env_flags;
use crate::error::CodexErr;
use crate::error::Result;
env_flags! {
pub OPENAI_DEFAULT_MODEL: &str = "o3";
pub OPENAI_API_BASE: &str = "https://api.openai.com";
pub OPENAI_API_BASE: &str = "https://api.openai.com/v1";
/// Fallback when the provider-specific key is not set.
pub OPENAI_API_KEY: Option<&str> = None;
pub OPENAI_TIMEOUT_MS: Duration = Duration::from_millis(300_000), |value| {
value.parse().map(Duration::from_millis)
};
@@ -21,9 +21,6 @@ env_flags! {
value.parse().map(Duration::from_millis)
};
/// Fixture path for offline tests (see client.rs).
pub CODEX_RS_SSE_FIXTURE: Option<&str> = None;
}
pub fn get_api_key() -> Result<&'static str> {
OPENAI_API_KEY.ok_or_else(|| CodexErr::EnvVar("OPENAI_API_KEY"))
}

View File

@@ -7,6 +7,7 @@
mod client;
pub mod codex;
pub use codex::Codex;
pub mod codex_wrapper;
pub mod config;
pub mod error;
@@ -18,6 +19,8 @@ pub mod linux;
mod mcp_connection_manager;
pub mod mcp_server_config;
mod mcp_tool_call;
mod model_provider_info;
pub use model_provider_info::ModelProviderInfo;
mod models;
pub mod protocol;
mod rollout;
@@ -25,5 +28,3 @@ mod safety;
mod user_notification;
pub mod util;
mod zdr_transcript;
pub use codex::Codex;

View File

@@ -0,0 +1,124 @@
//! Registry of model providers supported by Codex.
//!
//! Providers can be defined in two places:
//! 1. Built-in defaults compiled into the binary so Codex works out-of-the-box.
//! 2. User-defined entries inside `~/.codex/config.toml` under the `providers`
//! key. These override or extend the defaults at runtime.
//!
//! The combined mapping is surfaced via [`provider_map()`] and used by helper
//! functions in [`crate::flags`] to resolve API keys and base URLs.
use serde::Deserialize;
use serde::Serialize;
use std::collections::HashMap;
/// Serializable representation of a provider definition.
///
/// All fields are owned `String`s so that user-defined providers loaded from
/// disk can be stored alongside the built-ins without lifetime headaches.
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct ModelProviderInfo {
/// Friendly display name (optional for built-ins).
#[serde(default)]
pub name: String,
/// Base URL for the providers OpenAI-compatible API.
pub base_url: String,
/// Environment variable that stores the users API key for this provider.
pub env_key: String,
}
impl ModelProviderInfo {
/// Returns the API key for this provider if present in the environment.
pub fn api_key(&self) -> Option<String> {
std::env::var(&self.env_key).ok()
}
/// Determines the base URL for API requests, giving precedence to the
/// `{{PROVIDER}}_BASE_URL` environment variable when it is set.
pub fn api_base(&self, provider_key: &str) -> String {
let override_key = format!("{}_BASE_URL", provider_key.to_uppercase());
if let Ok(val) = std::env::var(&override_key) {
if !val.is_empty() {
return val;
}
}
self.base_url.clone()
}
}
/// Built-in default provider list mirrors `codex-cli/src/utils/providers.ts`.
/// Built-in provider registry. Public so callers (e.g. flags.rs) can resolve
/// information without needing a full [`crate::config::Config`].
pub fn built_in_model_providers() -> HashMap<String, ModelProviderInfo> {
use ModelProviderInfo as P;
[
(
"openai",
P {
name: "OpenAI".into(),
base_url: "https://api.openai.com/v1".into(),
env_key: "OPENAI_API_KEY".into(),
},
),
(
"openrouter",
P {
name: "OpenRouter".into(),
base_url: "https://openrouter.ai/api/v1".into(),
env_key: "OPENROUTER_API_KEY".into(),
},
),
(
"gemini",
P {
name: "Gemini".into(),
base_url: "https://generativelanguage.googleapis.com/v1beta/openai".into(),
env_key: "GEMINI_API_KEY".into(),
},
),
(
"ollama",
P {
name: "Ollama".into(),
base_url: "http://localhost:11434/v1".into(),
env_key: "OLLAMA_API_KEY".into(),
},
),
(
"mistral",
P {
name: "Mistral".into(),
base_url: "https://api.mistral.ai/v1".into(),
env_key: "MISTRAL_API_KEY".into(),
},
),
(
"deepseek",
P {
name: "DeepSeek".into(),
base_url: "https://api.deepseek.com".into(),
env_key: "DEEPSEEK_API_KEY".into(),
},
),
(
"xai",
P {
name: "xAI".into(),
base_url: "https://api.x.ai/v1".into(),
env_key: "XAI_API_KEY".into(),
},
),
(
"groq",
P {
name: "Groq".into(),
base_url: "https://api.groq.com/openai/v1".into(),
env_key: "GROQ_API_KEY".into(),
},
),
]
.into_iter()
.map(|(k, v)| (k.to_string(), v))
.collect()
}

View File

@@ -11,6 +11,8 @@ use mcp_types::CallToolResult;
use serde::Deserialize;
use serde::Serialize;
use crate::model_provider_info::ModelProviderInfo;
/// Submission Queue Entry - requests from user
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct Submission {
@@ -27,6 +29,11 @@ pub struct Submission {
pub enum Op {
/// Configure the model session.
ConfigureSession {
/// Provider identifier ("openai", "openrouter", ...). Defaults to
/// "openai" when omitted so that older clients continue to work.
// #[serde(default = "default_provider")]
provider: ModelProviderInfo,
/// If not specified, server will use its default model.
model: String,
/// Model instructions

View File

@@ -61,6 +61,7 @@ async fn spawn_codex() -> Codex {
.submit(Submission {
id: "init".into(),
op: Op::ConfigureSession {
provider: config.providers.get("openai").unwrap().clone(),
model: config.model,
instructions: None,
approval_policy: config.approval_policy,

View File

@@ -1,6 +1,7 @@
use std::time::Duration;
use codex_core::Codex;
use codex_core::ModelProviderInfo;
use codex_core::config::Config;
use codex_core::protocol::InputItem;
use codex_core::protocol::Op;
@@ -82,11 +83,14 @@ async fn keeps_previous_response_id_between_tasks() {
// Update environment `set_var` is `unsafe` starting with the 2024
// edition so we group the calls into a single `unsafe { … }` block.
unsafe {
std::env::set_var("OPENAI_API_KEY", "test-key");
std::env::set_var("OPENAI_API_BASE", server.uri());
std::env::set_var("OPENAI_REQUEST_MAX_RETRIES", "0");
std::env::set_var("OPENAI_STREAM_MAX_RETRIES", "0");
}
let model_provider = ModelProviderInfo {
name: "openai".into(),
base_url: format!("{}/v1", server.uri()),
env_key: "test-key".into(),
};
let codex = Codex::spawn(std::sync::Arc::new(tokio::sync::Notify::new())).unwrap();
@@ -96,6 +100,7 @@ async fn keeps_previous_response_id_between_tasks() {
.submit(Submission {
id: "init".into(),
op: Op::ConfigureSession {
provider: model_provider,
model: config.model,
instructions: None,
approval_policy: config.approval_policy,

View File

@@ -84,6 +84,7 @@ async fn retries_on_early_close() {
.submit(Submission {
id: "init".into(),
op: Op::ConfigureSession {
provider: config.providers.get("openai").unwrap().clone(),
model: config.model,
instructions: None,
approval_policy: config.approval_policy,

View File

@@ -66,6 +66,7 @@ pub async fn run_main(cli: Cli) -> anyhow::Result<()> {
None
},
cwd: cwd.map(|p| p.canonicalize().unwrap_or(p)),
provider: None,
};
let config = Config::load_with_overrides(overrides)?;

View File

@@ -158,6 +158,7 @@ impl CodexToolCallParam {
approval_policy: approval_policy.map(Into::into),
sandbox_policy,
disable_response_storage,
provider: None,
};
let cfg = codex_core::config::Config::load_with_overrides(overrides)?;

View File

@@ -58,6 +58,7 @@ pub fn run_main(cli: Cli) -> std::io::Result<()> {
None
},
cwd: cli.cwd.clone().map(|p| p.canonicalize().unwrap_or(p)),
provider: None,
};
#[allow(clippy::print_stderr)]
match Config::load_with_overrides(overrides) {