feat: support map of alternative providers like in TypeScript CLI

This commit is contained in:
Michael Bolin
2025-05-07 12:51:05 -07:00
parent 0360b4d0d7
commit 5743c170e5
16 changed files with 226 additions and 20 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"
tokio = { version = "1", features = [
"io-std",

View File

@@ -26,10 +26,8 @@ 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::models::ResponseItem;
use crate::util::backoff;
@@ -141,13 +139,22 @@ static DEFAULT_TOOLS: LazyLock<Vec<ResponsesApiTool>> = LazyLock::new(|| {
pub struct ModelClient {
model: String,
client: reqwest::Client,
provider_key: String,
provider: crate::model_provider_info::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_key: impl ToString,
provider: crate::model_provider_info::ModelProviderInfo,
) -> Self {
Self {
model: model.to_string(),
client: reqwest::Client::new(),
provider_key: provider_key.to_string(),
provider,
}
}
pub async fn stream(&mut self, prompt: &Prompt) -> Result<ResponseStream> {
@@ -188,7 +195,9 @@ impl ModelClient {
stream: true,
};
let url = format!("{}/v1/responses", *OPENAI_API_BASE);
let base_url = self.provider.api_base(&self.provider_key);
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)?);
@@ -196,10 +205,14 @@ impl ModelClient {
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

@@ -540,6 +540,7 @@ async fn submission_loop(
sess.abort();
}
Op::ConfigureSession {
provider,
model,
instructions,
approval_policy,
@@ -548,7 +549,7 @@ async fn submission_loop(
notify,
cwd,
} => {
info!(model, "Configuring session");
info!(model, provider, "Configuring session");
if !cwd.is_absolute() {
let message = format!("cwd is not absolute: {cwd:?}");
error!(message);
@@ -562,7 +563,25 @@ async fn submission_loop(
return;
}
let client = ModelClient::new(model.clone());
// Load config to resolve provider information & MCP servers.
let config = match Config::load_with_overrides(ConfigOverrides::default()) {
Ok(cfg) => cfg,
Err(e) => {
error!("Failed to load config: {e:#}");
Config::load_default_config_for_test()
}
};
let provider_map = crate::model_provider_info::provider_map(&config);
let provider_info = provider_map
.get(&provider.to_lowercase())
.cloned()
.unwrap_or_else(|| {
crate::model_provider_info::default_providers()["openai"].clone()
});
let client =
ModelClient::new(model.clone(), provider.clone(), provider_info.clone());
// abort any current running session and clone its state
let state = match sess.take() {

View File

@@ -20,6 +20,7 @@ pub async fn init_codex(config: Config) -> anyhow::Result<(CodexWrapper, Event,
let codex = CodexWrapper::new(Codex::spawn(ctrl_c.clone())?);
let init_id = codex
.submit(Op::ConfigureSession {
provider: config.provider.clone(),
model: config.model.clone(),
instructions: config.instructions.clone(),
approval_policy: config.approval_policy,

View File

@@ -19,6 +19,9 @@ pub struct Config {
/// Optional override of model selection.
pub model: String,
/// Selected provider ("openai", "gemini", …)
pub provider: String,
/// Approval policy for executing commands.
pub approval_policy: AskForApproval,
@@ -61,6 +64,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, crate::model_provider_info::ModelProviderInfo>,
}
/// Base config deserialized from ~/.codex/config.toml.
@@ -69,6 +75,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 +102,11 @@ 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
/// (`codex-cli/src/utils/providers.ts`).
#[serde(default)]
pub providers: HashMap<String, crate::model_provider_info::ModelProviderInfo>,
}
impl ConfigToml {
@@ -152,6 +166,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 +192,7 @@ impl Config {
approval_policy,
sandbox_policy,
disable_response_storage,
provider,
} = overrides;
let sandbox_policy = match sandbox_policy {
@@ -195,6 +212,9 @@ impl Config {
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 +242,7 @@ impl Config {
notify: cfg.notify,
instructions,
mcp_servers: cfg.mcp_servers,
providers: cfg.providers,
}
}

View File

@@ -2,28 +2,26 @@ 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";
// Retained for backward compatibility (now includes /v1).
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)
};
pub OPENAI_REQUEST_MAX_RETRIES: u64 = 4;
pub OPENAI_STREAM_MAX_RETRIES: u64 = 10;
// We generally don't want to disconnect; this updates the timeout to be five minutes
// which matches the upstream typescript codex impl.
// We generally don't want to disconnect; this matches the upstream TS CLI.
pub OPENAI_STREAM_IDLE_TIMEOUT_MS: Duration = Duration::from_millis(300_000), |value| {
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

@@ -18,6 +18,7 @@ pub mod linux;
mod mcp_connection_manager;
pub mod mcp_server_config;
mod mcp_tool_call;
mod model_provider_info;
mod models;
pub mod protocol;
mod safety;

View File

@@ -0,0 +1,136 @@
//! 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 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)]
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 default_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()
}
/// Merge built-in defaults with user-defined overrides from the supplied
/// [`crate::config::Config`]. When the same provider key appears in both maps
/// the user-defined entry wins.
pub fn provider_map(cfg: &crate::config::Config) -> HashMap<String, ModelProviderInfo> {
let mut map = default_providers();
map.extend(cfg.providers.clone());
// Normalise keys to lower-case for case-insensitive look-ups.
map.into_iter()
.map(|(k, v)| (k.to_lowercase(), v))
.collect()
}

View File

@@ -20,6 +20,10 @@ pub struct Submission {
pub op: Op,
}
fn default_provider() -> String {
"openai".to_string()
}
/// Submission operation
#[derive(Debug, Clone, Deserialize, Serialize)]
#[serde(tag = "type", rename_all = "snake_case")]
@@ -27,6 +31,11 @@ pub struct Submission {
pub enum Op {
/// Configure the model session.
ConfigureSession {
/// Provider identifier ("openai", "gemini", …). Defaults to
/// "openai" when omitted so that older clients continue to work.
#[serde(default = "default_provider")]
provider: String,
/// 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: "openai".to_string(),
model: config.model,
instructions: None,
approval_policy: config.approval_policy,

View File

@@ -96,6 +96,7 @@ async fn keeps_previous_response_id_between_tasks() {
.submit(Submission {
id: "init".into(),
op: Op::ConfigureSession {
provider: "openai".to_string(),
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: "openai".to_string(),
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) {