From 258ba436f1f784aa058625aec6b6653df83fe9c6 Mon Sep 17 00:00:00 2001 From: Michael Bolin Date: Mon, 30 Mar 2026 10:48:49 -0700 Subject: [PATCH 1/2] codex-tools: extract discoverable tool models (#16254) ## Why `#16193` moved the pure `tool_search` and `tool_suggest` spec builders into `codex-tools`, but `codex-core` still owned the shared discoverable-tool model that those builders and the `tool_suggest` runtime both depend on. This change continues the migration by moving that reusable model boundary out of `codex-core` as well, so the discovery/suggestion stack uses one shared set of types and `core/src/tools` no longer needs its own `discoverable.rs` module. ## What changed - Moved `DiscoverableTool`, `DiscoverablePluginInfo`, and `filter_tool_suggest_discoverable_tools_for_client()` into `codex-rs/tools/src/tool_discovery.rs` alongside the extracted discovery/suggestion spec builders. - Added `codex-app-server-protocol` as a `codex-tools` dependency so the shared discoverable-tool model can own the connector-side `AppInfo` variant directly. - Updated `core/src/tools/handlers/tool_suggest.rs`, `core/src/tools/spec.rs`, `core/src/tools/router.rs`, `core/src/connectors.rs`, and `core/src/codex.rs` to consume the shared `codex-tools` model instead of the old core-local declarations. - Changed `core/src/plugins/discoverable.rs` to return `DiscoverablePluginInfo` directly, moved the pure client-filter coverage into `tool_discovery_tests.rs`, and deleted the old `core/src/tools/discoverable.rs` module. - Updated `codex-rs/tools/README.md` so the crate boundary documents that `codex-tools` now owns the discoverable-tool models in addition to the discovery/suggestion spec builders. ## Test plan - `cargo test -p codex-tools` - `CARGO_TARGET_DIR=/tmp/codex-core-discoverable-model cargo test -p codex-core --lib tools::handlers::tool_suggest::` - `CARGO_TARGET_DIR=/tmp/codex-core-discoverable-model cargo test -p codex-core --lib tools::spec::` - `CARGO_TARGET_DIR=/tmp/codex-core-discoverable-model cargo test -p codex-core --lib plugins::discoverable::` - `just bazel-lock-check` - `just argument-comment-lint` ## References - #16193 - #16154 - #15923 - #15928 - #15944 - #15953 - #16031 - #16047 - #16129 - #16132 - #16138 - #16141 --- codex-rs/Cargo.lock | 1 + codex-rs/core/src/codex.rs | 3 +- codex-rs/core/src/connectors.rs | 4 +- codex-rs/core/src/plugins/discoverable.rs | 27 ++++-- .../core/src/plugins/discoverable_tests.rs | 32 ++----- codex-rs/core/src/tools/discoverable.rs | 94 ------------------- .../core/src/tools/handlers/tool_suggest.rs | 4 +- .../src/tools/handlers/tool_suggest_tests.rs | 52 +--------- codex-rs/core/src/tools/mod.rs | 1 - codex-rs/core/src/tools/router.rs | 2 +- codex-rs/core/src/tools/spec.rs | 2 +- codex-rs/core/src/tools/spec_tests.rs | 3 +- codex-rs/tools/Cargo.toml | 1 + codex-rs/tools/README.md | 2 +- codex-rs/tools/src/lib.rs | 3 + codex-rs/tools/src/tool_discovery.rs | 75 +++++++++++++++ codex-rs/tools/src/tool_discovery_tests.rs | 49 ++++++++++ 17 files changed, 168 insertions(+), 187 deletions(-) delete mode 100644 codex-rs/core/src/tools/discoverable.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 3dd569c6b0..f8d1de44fc 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2613,6 +2613,7 @@ dependencies = [ name = "codex-tools" version = "0.0.0" dependencies = [ + "codex-app-server-protocol", "codex-code-mode", "codex-protocol", "pretty_assertions", diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 9a8cc9cea7..2c70aad7a7 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -121,6 +121,7 @@ use codex_protocol::request_user_input::RequestUserInputResponse; use codex_rmcp_client::ElicitationResponse; use codex_rmcp_client::OAuthCredentialsStoreMode; use codex_terminal_detection::user_agent; +use codex_tools::filter_tool_suggest_discoverable_tools_for_client; use codex_utils_output_truncation::TruncationPolicy; use codex_utils_stream_parser::AssistantTextChunk; use codex_utils_stream_parser::AssistantTextStreamParser; @@ -6546,7 +6547,7 @@ pub(crate) async fn built_tools( ) .await .map(|discoverable_tools| { - crate::tools::discoverable::filter_tool_suggest_discoverable_tools_for_client( + filter_tool_suggest_discoverable_tools_for_client( discoverable_tools, turn_context.app_server_client_name.as_deref(), ) diff --git a/codex-rs/core/src/connectors.rs b/codex-rs/core/src/connectors.rs index e4999b9b41..a15138f6fb 100644 --- a/codex-rs/core/src/connectors.rs +++ b/codex-rs/core/src/connectors.rs @@ -18,6 +18,7 @@ use codex_connectors::AllConnectorsCacheKey; use codex_connectors::DirectoryListResponse; use codex_login::token_data::TokenData; use codex_protocol::protocol::SandboxPolicy; +use codex_tools::DiscoverableTool; use rmcp::model::ToolAnnotations; use serde::Deserialize; use serde::de::DeserializeOwned; @@ -44,8 +45,6 @@ use crate::mcp_connection_manager::codex_apps_tools_cache_key; use crate::plugins::AppConnectorId; use crate::plugins::PluginsManager; use crate::plugins::list_tool_suggest_discoverable_plugins; -use crate::tools::discoverable::DiscoverablePluginInfo; -use crate::tools::discoverable::DiscoverableTool; use codex_features::Feature; pub use codex_connectors::CONNECTORS_CACHE_TTL; @@ -133,7 +132,6 @@ pub(crate) async fn list_tool_suggest_discoverable_tools_with_auth( .map(DiscoverableTool::from); let discoverable_plugins = list_tool_suggest_discoverable_plugins(config)? .into_iter() - .map(DiscoverablePluginInfo::from) .map(DiscoverableTool::from); Ok(discoverable_connectors .chain(discoverable_plugins) diff --git a/codex-rs/core/src/plugins/discoverable.rs b/codex-rs/core/src/plugins/discoverable.rs index 427706235f..f7006185bf 100644 --- a/codex-rs/core/src/plugins/discoverable.rs +++ b/codex-rs/core/src/plugins/discoverable.rs @@ -9,6 +9,7 @@ use super::PluginsManager; use crate::config::Config; use crate::config::types::ToolSuggestDiscoverableType; use codex_features::Feature; +use codex_tools::DiscoverablePluginInfo; const TOOL_SUGGEST_DISCOVERABLE_PLUGIN_ALLOWLIST: &[&str] = &[ "github@openai-curated", @@ -23,7 +24,7 @@ const TOOL_SUGGEST_DISCOVERABLE_PLUGIN_ALLOWLIST: &[&str] = &[ pub(crate) fn list_tool_suggest_discoverable_plugins( config: &Config, -) -> anyhow::Result> { +) -> anyhow::Result> { if !config.features.enabled(Feature::Plugins) { return Ok(Vec::new()); } @@ -47,7 +48,7 @@ pub(crate) fn list_tool_suggest_discoverable_plugins( return Ok(Vec::new()); }; - let mut discoverable_plugins = Vec::::new(); + let mut discoverable_plugins = Vec::::new(); for plugin in curated_marketplace.plugins { if plugin.installed || (!TOOL_SUGGEST_DISCOVERABLE_PLUGIN_ALLOWLIST.contains(&plugin.id.as_str()) @@ -66,14 +67,28 @@ pub(crate) fn list_tool_suggest_discoverable_plugins( marketplace_path: curated_marketplace.path.clone(), }, ) { - Ok(plugin) => discoverable_plugins.push(plugin.plugin.into()), + Ok(plugin) => { + let plugin: PluginCapabilitySummary = plugin.plugin.into(); + discoverable_plugins.push(DiscoverablePluginInfo { + id: plugin.config_name, + name: plugin.display_name, + description: plugin.description, + has_skills: plugin.has_skills, + mcp_server_names: plugin.mcp_server_names, + app_connector_ids: plugin + .app_connector_ids + .into_iter() + .map(|connector_id| connector_id.0) + .collect(), + }); + } Err(err) => warn!("failed to load discoverable plugin suggestion {plugin_id}: {err:#}"), } } discoverable_plugins.sort_by(|left, right| { - left.display_name - .cmp(&right.display_name) - .then_with(|| left.config_name.cmp(&right.config_name)) + left.name + .cmp(&right.name) + .then_with(|| left.id.cmp(&right.id)) }); Ok(discoverable_plugins) } diff --git a/codex-rs/core/src/plugins/discoverable_tests.rs b/codex-rs/core/src/plugins/discoverable_tests.rs index fd9df9ddf6..70ac887cb6 100644 --- a/codex-rs/core/src/plugins/discoverable_tests.rs +++ b/codex-rs/core/src/plugins/discoverable_tests.rs @@ -5,7 +5,7 @@ use crate::plugins::test_support::write_curated_plugin_sha; use crate::plugins::test_support::write_file; use crate::plugins::test_support::write_openai_curated_marketplace; use crate::plugins::test_support::write_plugins_feature_config; -use crate::tools::discoverable::DiscoverablePluginInfo; +use codex_tools::DiscoverablePluginInfo; use codex_utils_absolute_path::AbsolutePathBuf; use pretty_assertions::assert_eq; use tempfile::tempdir; @@ -18,11 +18,7 @@ async fn list_tool_suggest_discoverable_plugins_returns_uninstalled_curated_plug write_plugins_feature_config(codex_home.path()); let config = load_plugins_config(codex_home.path()).await; - let discoverable_plugins = list_tool_suggest_discoverable_plugins(&config) - .unwrap() - .into_iter() - .map(DiscoverablePluginInfo::from) - .collect::>(); + let discoverable_plugins = list_tool_suggest_discoverable_plugins(&config).unwrap(); assert_eq!( discoverable_plugins, @@ -52,11 +48,7 @@ plugins = false ); let config = load_plugins_config(codex_home.path()).await; - let discoverable_plugins = list_tool_suggest_discoverable_plugins(&config) - .unwrap() - .into_iter() - .map(DiscoverablePluginInfo::from) - .collect::>(); + let discoverable_plugins = list_tool_suggest_discoverable_plugins(&config).unwrap(); assert_eq!(discoverable_plugins, Vec::::new()); } @@ -76,11 +68,7 @@ async fn list_tool_suggest_discoverable_plugins_normalizes_description() { ); let config = load_plugins_config(codex_home.path()).await; - let discoverable_plugins = list_tool_suggest_discoverable_plugins(&config) - .unwrap() - .into_iter() - .map(DiscoverablePluginInfo::from) - .collect::>(); + let discoverable_plugins = list_tool_suggest_discoverable_plugins(&config).unwrap(); assert_eq!( discoverable_plugins, @@ -115,11 +103,7 @@ async fn list_tool_suggest_discoverable_plugins_omits_installed_curated_plugins( .expect("plugin should install"); let refreshed_config = load_plugins_config(codex_home.path()).await; - let discoverable_plugins = list_tool_suggest_discoverable_plugins(&refreshed_config) - .unwrap() - .into_iter() - .map(DiscoverablePluginInfo::from) - .collect::>(); + let discoverable_plugins = list_tool_suggest_discoverable_plugins(&refreshed_config).unwrap(); assert_eq!(discoverable_plugins, Vec::::new()); } @@ -140,11 +124,7 @@ discoverables = [{ type = "plugin", id = "sample@openai-curated" }] ); let config = load_plugins_config(codex_home.path()).await; - let discoverable_plugins = list_tool_suggest_discoverable_plugins(&config) - .unwrap() - .into_iter() - .map(DiscoverablePluginInfo::from) - .collect::>(); + let discoverable_plugins = list_tool_suggest_discoverable_plugins(&config).unwrap(); assert_eq!( discoverable_plugins, diff --git a/codex-rs/core/src/tools/discoverable.rs b/codex-rs/core/src/tools/discoverable.rs deleted file mode 100644 index 78d8935aa7..0000000000 --- a/codex-rs/core/src/tools/discoverable.rs +++ /dev/null @@ -1,94 +0,0 @@ -use crate::plugins::PluginCapabilitySummary; -use codex_app_server_protocol::AppInfo; -use codex_tools::DiscoverableToolType; - -const TUI_CLIENT_NAME: &str = "codex-tui"; - -#[derive(Clone, Debug, PartialEq)] -pub(crate) enum DiscoverableTool { - Connector(Box), - Plugin(Box), -} - -impl DiscoverableTool { - pub(crate) fn tool_type(&self) -> DiscoverableToolType { - match self { - Self::Connector(_) => DiscoverableToolType::Connector, - Self::Plugin(_) => DiscoverableToolType::Plugin, - } - } - - pub(crate) fn id(&self) -> &str { - match self { - Self::Connector(connector) => connector.id.as_str(), - Self::Plugin(plugin) => plugin.id.as_str(), - } - } - - pub(crate) fn name(&self) -> &str { - match self { - Self::Connector(connector) => connector.name.as_str(), - Self::Plugin(plugin) => plugin.name.as_str(), - } - } - - pub(crate) fn install_url(&self) -> Option<&str> { - match self { - Self::Connector(connector) => connector.install_url.as_deref(), - Self::Plugin(_) => None, - } - } -} - -impl From for DiscoverableTool { - fn from(value: AppInfo) -> Self { - Self::Connector(Box::new(value)) - } -} - -impl From for DiscoverableTool { - fn from(value: DiscoverablePluginInfo) -> Self { - Self::Plugin(Box::new(value)) - } -} - -pub(crate) fn filter_tool_suggest_discoverable_tools_for_client( - discoverable_tools: Vec, - app_server_client_name: Option<&str>, -) -> Vec { - if app_server_client_name != Some(TUI_CLIENT_NAME) { - return discoverable_tools; - } - - discoverable_tools - .into_iter() - .filter(|tool| !matches!(tool, DiscoverableTool::Plugin(_))) - .collect() -} - -#[derive(Clone, Debug, PartialEq, Eq)] -pub(crate) struct DiscoverablePluginInfo { - pub(crate) id: String, - pub(crate) name: String, - pub(crate) description: Option, - pub(crate) has_skills: bool, - pub(crate) mcp_server_names: Vec, - pub(crate) app_connector_ids: Vec, -} - -impl From for DiscoverablePluginInfo { - fn from(value: PluginCapabilitySummary) -> Self { - Self { - id: value.config_name, - name: value.display_name, - description: value.description, - has_skills: value.has_skills, - mcp_server_names: value.mcp_server_names, - app_connector_ids: value - .app_connector_ids - .into_iter() - .map(|connector_id| connector_id.0) - .collect(), - } - } -} diff --git a/codex-rs/core/src/tools/handlers/tool_suggest.rs b/codex-rs/core/src/tools/handlers/tool_suggest.rs index 478bd2c837..c68ac8589c 100644 --- a/codex-rs/core/src/tools/handlers/tool_suggest.rs +++ b/codex-rs/core/src/tools/handlers/tool_suggest.rs @@ -8,8 +8,10 @@ use codex_app_server_protocol::McpElicitationSchema; use codex_app_server_protocol::McpServerElicitationRequest; use codex_app_server_protocol::McpServerElicitationRequestParams; use codex_rmcp_client::ElicitationAction; +use codex_tools::DiscoverableTool; use codex_tools::DiscoverableToolAction; use codex_tools::DiscoverableToolType; +use codex_tools::filter_tool_suggest_discoverable_tools_for_client; use rmcp::model::RequestId; use serde::Deserialize; use serde::Serialize; @@ -22,8 +24,6 @@ use crate::mcp::CODEX_APPS_MCP_SERVER_NAME; use crate::tools::context::FunctionToolOutput; use crate::tools::context::ToolInvocation; use crate::tools::context::ToolPayload; -use crate::tools::discoverable::DiscoverableTool; -use crate::tools::discoverable::filter_tool_suggest_discoverable_tools_for_client; use crate::tools::handlers::parse_arguments; use crate::tools::registry::ToolHandler; use crate::tools::registry::ToolKind; diff --git a/codex-rs/core/src/tools/handlers/tool_suggest_tests.rs b/codex-rs/core/src/tools/handlers/tool_suggest_tests.rs index c932e03145..56089a1aa9 100644 --- a/codex-rs/core/src/tools/handlers/tool_suggest_tests.rs +++ b/codex-rs/core/src/tools/handlers/tool_suggest_tests.rs @@ -5,9 +5,9 @@ use crate::plugins::test_support::load_plugins_config; use crate::plugins::test_support::write_curated_plugin_sha; use crate::plugins::test_support::write_openai_curated_marketplace; use crate::plugins::test_support::write_plugins_feature_config; -use crate::tools::discoverable::DiscoverablePluginInfo; -use crate::tools::discoverable::filter_tool_suggest_discoverable_tools_for_client; use codex_app_server_protocol::AppInfo; +use codex_tools::DiscoverablePluginInfo; +use codex_tools::DiscoverableTool; use codex_tools::DiscoverableToolAction; use codex_tools::DiscoverableToolType; use codex_utils_absolute_path::AbsolutePathBuf; @@ -161,54 +161,6 @@ fn build_tool_suggestion_meta_uses_expected_shape() { ); } -#[test] -fn filter_tool_suggest_discoverable_tools_for_codex_tui_omits_plugins() { - let discoverable_tools = vec![ - DiscoverableTool::Connector(Box::new(AppInfo { - id: "connector_google_calendar".to_string(), - name: "Google Calendar".to_string(), - description: Some("Plan events and schedules.".to_string()), - logo_url: None, - logo_url_dark: None, - distribution_channel: None, - branding: None, - app_metadata: None, - labels: None, - install_url: Some("https://example.test/google-calendar".to_string()), - is_accessible: false, - is_enabled: true, - plugin_display_names: Vec::new(), - })), - DiscoverableTool::Plugin(Box::new(DiscoverablePluginInfo { - id: "slack@openai-curated".to_string(), - name: "Slack".to_string(), - description: Some("Search Slack messages".to_string()), - has_skills: true, - mcp_server_names: vec!["slack".to_string()], - app_connector_ids: vec!["connector_slack".to_string()], - })), - ]; - - assert_eq!( - filter_tool_suggest_discoverable_tools_for_client(discoverable_tools, Some("codex-tui"),), - vec![DiscoverableTool::Connector(Box::new(AppInfo { - id: "connector_google_calendar".to_string(), - name: "Google Calendar".to_string(), - description: Some("Plan events and schedules.".to_string()), - logo_url: None, - logo_url_dark: None, - distribution_channel: None, - branding: None, - app_metadata: None, - labels: None, - install_url: Some("https://example.test/google-calendar".to_string()), - is_accessible: false, - is_enabled: true, - plugin_display_names: Vec::new(), - }))] - ); -} - #[test] fn verified_connector_suggestion_completed_requires_accessible_connector() { let accessible_connectors = vec![AppInfo { diff --git a/codex-rs/core/src/tools/mod.rs b/codex-rs/core/src/tools/mod.rs index 57859bba42..e866d1f5f7 100644 --- a/codex-rs/core/src/tools/mod.rs +++ b/codex-rs/core/src/tools/mod.rs @@ -1,6 +1,5 @@ pub mod code_mode; pub mod context; -pub(crate) mod discoverable; pub mod events; pub(crate) mod handlers; pub mod js_repl; diff --git a/codex-rs/core/src/tools/router.rs b/codex-rs/core/src/tools/router.rs index 1145e3a8f2..646c1d1387 100644 --- a/codex-rs/core/src/tools/router.rs +++ b/codex-rs/core/src/tools/router.rs @@ -7,7 +7,6 @@ use crate::sandboxing::SandboxPermissions; use crate::tools::context::SharedTurnDiffTracker; use crate::tools::context::ToolInvocation; use crate::tools::context::ToolPayload; -use crate::tools::discoverable::DiscoverableTool; use crate::tools::registry::AnyToolResult; use crate::tools::registry::ToolRegistry; use crate::tools::spec::ToolsConfig; @@ -18,6 +17,7 @@ use codex_protocol::models::ResponseItem; use codex_protocol::models::SearchToolCallParams; use codex_protocol::models::ShellToolCallParams; use codex_tools::ConfiguredToolSpec; +use codex_tools::DiscoverableTool; use rmcp::model::Tool; use std::collections::HashMap; use std::sync::Arc; diff --git a/codex-rs/core/src/tools/spec.rs b/codex-rs/core/src/tools/spec.rs index 6852671576..d0614fc35b 100644 --- a/codex-rs/core/src/tools/spec.rs +++ b/codex-rs/core/src/tools/spec.rs @@ -7,7 +7,6 @@ use crate::shell::Shell; use crate::shell::ShellType; use crate::tools::code_mode::PUBLIC_TOOL_NAME; use crate::tools::code_mode::WAIT_TOOL_NAME; -use crate::tools::discoverable::DiscoverableTool; use crate::tools::handlers::PLAN_TOOL; use crate::tools::handlers::TOOL_SEARCH_DEFAULT_LIMIT; use crate::tools::handlers::TOOL_SEARCH_TOOL_NAME; @@ -38,6 +37,7 @@ use codex_protocol::protocol::SandboxPolicy; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::SubAgentSource; use codex_tools::CommandToolOptions; +use codex_tools::DiscoverableTool; use codex_tools::DiscoverableToolType; use codex_tools::ShellToolOptions; use codex_tools::SpawnAgentToolOptions; diff --git a/codex-rs/core/src/tools/spec_tests.rs b/codex-rs/core/src/tools/spec_tests.rs index 283bcd7161..16f87ca50e 100644 --- a/codex-rs/core/src/tools/spec_tests.rs +++ b/codex-rs/core/src/tools/spec_tests.rs @@ -4,7 +4,6 @@ use crate::models_manager::model_info::with_config_overrides; use crate::shell::Shell; use crate::shell::ShellType; use crate::tools::ToolRouter; -use crate::tools::discoverable::DiscoverablePluginInfo; use crate::tools::router::ToolRouterParams; use codex_app_server_protocol::AppInfo; use codex_protocol::models::VIEW_IMAGE_TOOL_NAME; @@ -14,6 +13,8 @@ use codex_protocol::openai_models::ModelsResponse; use codex_tools::AdditionalProperties; use codex_tools::CommandToolOptions; use codex_tools::ConfiguredToolSpec; +use codex_tools::DiscoverablePluginInfo; +use codex_tools::DiscoverableTool; use codex_tools::FreeformTool; use codex_tools::ResponsesApiTool; use codex_tools::ResponsesApiWebSearchFilters; diff --git a/codex-rs/tools/Cargo.toml b/codex-rs/tools/Cargo.toml index 9e2d23fb58..3d61dab7d8 100644 --- a/codex-rs/tools/Cargo.toml +++ b/codex-rs/tools/Cargo.toml @@ -8,6 +8,7 @@ version.workspace = true workspace = true [dependencies] +codex-app-server-protocol = { workspace = true } codex-code-mode = { workspace = true } codex-protocol = { workspace = true } rmcp = { workspace = true, default-features = false, features = [ diff --git a/codex-rs/tools/README.md b/codex-rs/tools/README.md index 59533f459d..cafb4b89ac 100644 --- a/codex-rs/tools/README.md +++ b/codex-rs/tools/README.md @@ -27,7 +27,7 @@ schema and Responses API tool primitives that no longer need to live in - local host tool spec builders for shell/exec/request-permissions/view-image - collaboration and agent-job `ToolSpec` builders for spawn/send/wait/close, `request_user_input`, and CSV fanout/reporting -- tool discovery and suggestion models / `ToolSpec` builders for +- discoverable-tool models, client filtering, and `ToolSpec` builders for `tool_search` and `tool_suggest` - `parse_tool_input_schema()` - `parse_dynamic_tool()` diff --git a/codex-rs/tools/src/lib.rs b/codex-rs/tools/src/lib.rs index 9dd629f916..9ae6031ccd 100644 --- a/codex-rs/tools/src/lib.rs +++ b/codex-rs/tools/src/lib.rs @@ -67,12 +67,15 @@ pub use responses_api::mcp_tool_to_deferred_responses_api_tool; pub use responses_api::mcp_tool_to_responses_api_tool; pub use responses_api::tool_definition_to_responses_api_tool; pub use tool_definition::ToolDefinition; +pub use tool_discovery::DiscoverablePluginInfo; +pub use tool_discovery::DiscoverableTool; pub use tool_discovery::DiscoverableToolAction; pub use tool_discovery::DiscoverableToolType; pub use tool_discovery::ToolSearchAppInfo; pub use tool_discovery::ToolSuggestEntry; pub use tool_discovery::create_tool_search_tool; pub use tool_discovery::create_tool_suggest_tool; +pub use tool_discovery::filter_tool_suggest_discoverable_tools_for_client; pub use tool_spec::ConfiguredToolSpec; pub use tool_spec::ResponsesApiWebSearchFilters; pub use tool_spec::ResponsesApiWebSearchUserLocation; diff --git a/codex-rs/tools/src/tool_discovery.rs b/codex-rs/tools/src/tool_discovery.rs index f2b9b4f934..0216cc1703 100644 --- a/codex-rs/tools/src/tool_discovery.rs +++ b/codex-rs/tools/src/tool_discovery.rs @@ -1,10 +1,13 @@ use crate::JsonSchema; use crate::ResponsesApiTool; use crate::ToolSpec; +use codex_app_server_protocol::AppInfo; use serde::Deserialize; use serde::Serialize; use std::collections::BTreeMap; +const TUI_CLIENT_NAME: &str = "codex-tui"; + #[derive(Clone, Debug, PartialEq, Eq)] pub struct ToolSearchAppInfo { pub name: String, @@ -34,6 +37,78 @@ pub enum DiscoverableToolAction { Enable, } +#[derive(Clone, Debug, PartialEq)] +pub enum DiscoverableTool { + Connector(Box), + Plugin(Box), +} + +impl DiscoverableTool { + pub fn tool_type(&self) -> DiscoverableToolType { + match self { + Self::Connector(_) => DiscoverableToolType::Connector, + Self::Plugin(_) => DiscoverableToolType::Plugin, + } + } + + pub fn id(&self) -> &str { + match self { + Self::Connector(connector) => connector.id.as_str(), + Self::Plugin(plugin) => plugin.id.as_str(), + } + } + + pub fn name(&self) -> &str { + match self { + Self::Connector(connector) => connector.name.as_str(), + Self::Plugin(plugin) => plugin.name.as_str(), + } + } + + pub fn install_url(&self) -> Option<&str> { + match self { + Self::Connector(connector) => connector.install_url.as_deref(), + Self::Plugin(_) => None, + } + } +} + +impl From for DiscoverableTool { + fn from(value: AppInfo) -> Self { + Self::Connector(Box::new(value)) + } +} + +impl From for DiscoverableTool { + fn from(value: DiscoverablePluginInfo) -> Self { + Self::Plugin(Box::new(value)) + } +} + +pub fn filter_tool_suggest_discoverable_tools_for_client( + discoverable_tools: Vec, + app_server_client_name: Option<&str>, +) -> Vec { + if app_server_client_name != Some(TUI_CLIENT_NAME) { + return discoverable_tools; + } + + discoverable_tools + .into_iter() + .filter(|tool| !matches!(tool, DiscoverableTool::Plugin(_))) + .collect() +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct DiscoverablePluginInfo { + pub id: String, + pub name: String, + pub description: Option, + pub has_skills: bool, + pub mcp_server_names: Vec, + pub app_connector_ids: Vec, +} + #[derive(Clone, Debug, PartialEq, Eq)] pub struct ToolSuggestEntry { pub id: String, diff --git a/codex-rs/tools/src/tool_discovery_tests.rs b/codex-rs/tools/src/tool_discovery_tests.rs index 7be6bedd2c..1b241c5844 100644 --- a/codex-rs/tools/src/tool_discovery_tests.rs +++ b/codex-rs/tools/src/tool_discovery_tests.rs @@ -1,4 +1,5 @@ use super::*; +use codex_app_server_protocol::AppInfo; use pretty_assertions::assert_eq; use serde_json::json; use std::collections::BTreeMap; @@ -147,3 +148,51 @@ fn discoverable_tool_enums_use_expected_wire_names() { }) ); } + +#[test] +fn filter_tool_suggest_discoverable_tools_for_codex_tui_omits_plugins() { + let discoverable_tools = vec![ + DiscoverableTool::Connector(Box::new(AppInfo { + id: "connector_google_calendar".to_string(), + name: "Google Calendar".to_string(), + description: Some("Plan events and schedules.".to_string()), + logo_url: None, + logo_url_dark: None, + distribution_channel: None, + branding: None, + app_metadata: None, + labels: None, + install_url: Some("https://example.test/google-calendar".to_string()), + is_accessible: false, + is_enabled: true, + plugin_display_names: Vec::new(), + })), + DiscoverableTool::Plugin(Box::new(DiscoverablePluginInfo { + id: "slack@openai-curated".to_string(), + name: "Slack".to_string(), + description: Some("Search Slack messages".to_string()), + has_skills: true, + mcp_server_names: vec!["slack".to_string()], + app_connector_ids: vec!["connector_slack".to_string()], + })), + ]; + + assert_eq!( + filter_tool_suggest_discoverable_tools_for_client(discoverable_tools, Some("codex-tui"),), + vec![DiscoverableTool::Connector(Box::new(AppInfo { + id: "connector_google_calendar".to_string(), + name: "Google Calendar".to_string(), + description: Some("Plan events and schedules.".to_string()), + logo_url: None, + logo_url_dark: None, + distribution_channel: None, + branding: None, + app_metadata: None, + labels: None, + install_url: Some("https://example.test/google-calendar".to_string()), + is_accessible: false, + is_enabled: true, + plugin_display_names: Vec::new(), + }))] + ); +} From be08ad0e4bf9dad7c2f43193238dc053cc010539 Mon Sep 17 00:00:00 2001 From: Michael Bolin Date: Mon, 30 Mar 2026 12:13:05 -0700 Subject: [PATCH 2/2] add support for dynamic auth token for model provider --- codex-rs/core/config.schema.json | 44 +++ codex-rs/core/src/api_bridge.rs | 21 +- codex-rs/core/src/auth_env_telemetry.rs | 1 + codex-rs/core/src/client.rs | 129 +++++++- codex-rs/core/src/config/config_tests.rs | 40 +++ codex-rs/core/src/config/mod.rs | 16 +- codex-rs/core/src/lib.rs | 2 + codex-rs/core/src/model_provider_info.rs | 103 +++++++ .../core/src/model_provider_info_tests.rs | 32 ++ codex-rs/core/src/models_manager/manager.rs | 80 +++-- .../core/src/models_manager/manager_tests.rs | 1 + codex-rs/core/src/provider_auth.rs | 276 +++++++++++++++++ codex-rs/core/src/realtime_conversation.rs | 15 +- codex-rs/core/tests/responses_headers.rs | 3 + codex-rs/core/tests/suite/client.rs | 289 ++++++++++++++++++ .../core/tests/suite/client_websockets.rs | 1 + .../suite/stream_error_allows_next_turn.rs | 1 + .../core/tests/suite/stream_no_completed.rs | 1 + 18 files changed, 997 insertions(+), 58 deletions(-) create mode 100644 codex-rs/core/src/provider_auth.rs diff --git a/codex-rs/core/config.schema.json b/codex-rs/core/config.schema.json index e36f495a1b..953c07b4fe 100644 --- a/codex-rs/core/config.schema.json +++ b/codex-rs/core/config.schema.json @@ -816,10 +816,54 @@ }, "type": "object" }, + "ModelProviderAuthInfo": { + "additionalProperties": false, + "description": "Configuration for obtaining a provider bearer token from a command.", + "properties": { + "args": { + "default": [], + "description": "Command arguments.", + "items": { + "type": "string" + }, + "type": "array" + }, + "command": { + "description": "Command to execute. Bare names are resolved via `PATH`; paths are resolved against `cwd`.", + "type": "string" + }, + "refresh_interval_ms": { + "default": 300000, + "description": "Maximum age for the cached token before rerunning the command.", + "format": "uint64", + "minimum": 0.0, + "type": "integer" + }, + "timeout_ms": { + "default": 5000, + "description": "Maximum time to wait for the token command to exit successfully.", + "format": "uint64", + "minimum": 0.0, + "type": "integer" + } + }, + "required": [ + "command" + ], + "type": "object" + }, "ModelProviderInfo": { "additionalProperties": false, "description": "Serializable representation of a provider definition.", "properties": { + "auth": { + "allOf": [ + { + "$ref": "#/definitions/ModelProviderAuthInfo" + } + ], + "description": "Command-backed bearer-token configuration for this provider." + }, "base_url": { "description": "Base URL for the provider's OpenAI-compatible API.", "type": "string" diff --git a/codex-rs/core/src/api_bridge.rs b/codex-rs/core/src/api_bridge.rs index e7826f9ac6..835a39b63e 100644 --- a/codex-rs/core/src/api_bridge.rs +++ b/codex-rs/core/src/api_bridge.rs @@ -164,18 +164,23 @@ fn extract_x_error_json_code(headers: Option<&HeaderMap>) -> Option { .map(str::to_string) } -pub(crate) fn auth_provider_from_auth( - auth: Option, +pub(crate) fn provider_bearer_token( provider: &ModelProviderInfo, -) -> crate::error::Result { + resolved_provider_token: Option, +) -> crate::error::Result> { if let Some(api_key) = provider.api_key()? { - return Ok(CoreAuthProvider { - token: Some(api_key), - account_id: None, - }); + return Ok(Some(api_key)); } - if let Some(token) = provider.experimental_bearer_token.clone() { + Ok(resolved_provider_token.or_else(|| provider.experimental_bearer_token.clone())) +} + +pub(crate) fn auth_provider_from_resolved_provider_token( + auth: Option, + provider: &ModelProviderInfo, + resolved_provider_token: Option, +) -> crate::error::Result { + if let Some(token) = provider_bearer_token(provider, resolved_provider_token)? { return Ok(CoreAuthProvider { token: Some(token), account_id: None, diff --git a/codex-rs/core/src/auth_env_telemetry.rs b/codex-rs/core/src/auth_env_telemetry.rs index cc5ffa1207..583d79de16 100644 --- a/codex-rs/core/src/auth_env_telemetry.rs +++ b/codex-rs/core/src/auth_env_telemetry.rs @@ -64,6 +64,7 @@ mod tests { env_key: Some("sk-should-not-leak".to_string()), env_key_instructions: None, experimental_bearer_token: None, + auth: None, wire_api: crate::model_provider_info::WireApi::Responses, query_params: None, http_headers: None, diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index af5e3007b4..a6c87cefbc 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -31,7 +31,7 @@ use std::sync::atomic::AtomicBool; use std::sync::atomic::Ordering; use crate::api_bridge::CoreAuthProvider; -use crate::api_bridge::auth_provider_from_auth; +use crate::api_bridge::auth_provider_from_resolved_provider_token; use crate::api_bridge::map_api_error; use crate::auth::UnauthorizedRecovery; use crate::auth_env_telemetry::AuthEnvTelemetry; @@ -104,6 +104,7 @@ use crate::error::Result; use crate::flags::CODEX_RS_SSE_FIXTURE; use crate::model_provider_info::ModelProviderInfo; use crate::model_provider_info::WireApi; +use crate::provider_auth::ProviderAuthResolver; use crate::response_debug_context::extract_response_debug_context; use crate::response_debug_context::extract_response_debug_context_from_api_error; use crate::response_debug_context::telemetry_api_error_message; @@ -134,6 +135,7 @@ struct ModelClientState { auth_manager: Option>, conversation_id: ThreadId, provider: ModelProviderInfo, + provider_auth: ProviderAuthResolver, auth_env_telemetry: AuthEnvTelemetry, session_source: SessionSource, model_verbosity: Option, @@ -269,6 +271,7 @@ impl ModelClient { state: Arc::new(ModelClientState { auth_manager, conversation_id, + provider_auth: ProviderAuthResolver::new(&provider), provider, auth_env_telemetry, session_source, @@ -332,6 +335,10 @@ impl ModelClient { activated } + pub(crate) async fn resolve_provider_bearer_token(&self) -> Result> { + self.state.provider_auth.resolve_token().await + } + /// Compacts the current conversation history using the Compact endpoint. /// /// This is a unary call (no streaming) that returns a new list of @@ -528,11 +535,16 @@ impl ModelClient { Some(manager) => manager.auth().await, None => None, }; + let provider_token = self.resolve_provider_bearer_token().await?; let api_provider = self .state .provider .to_api_provider(auth.as_ref().map(CodexAuth::auth_mode))?; - let api_auth = auth_provider_from_auth(auth.clone(), &self.state.provider)?; + let api_auth = auth_provider_from_resolved_provider_token( + auth.clone(), + &self.state.provider, + provider_token, + )?; Ok(CurrentClientSetup { auth, api_provider, @@ -1020,6 +1032,7 @@ impl ModelClientSession { let mut auth_recovery = auth_manager .as_ref() .map(super::auth::AuthManager::unauthorized_recovery); + let mut provider_auth_retry_available = self.client.state.provider_auth.is_configured(); let mut pending_retry = PendingUnauthorizedRetry::default(); loop { let client_setup = self.client.current_client_setup().await?; @@ -1062,14 +1075,26 @@ impl ModelClientSession { Err(ApiError::Transport( unauthorized_transport @ TransportError::Http { status, .. }, )) if status == StatusCode::UNAUTHORIZED => { - pending_retry = PendingUnauthorizedRetry::from_recovery( - handle_unauthorized( - unauthorized_transport, - &mut auth_recovery, - session_telemetry, + pending_retry = if provider_auth_retry_available { + provider_auth_retry_available = false; + PendingUnauthorizedRetry::from_recovery( + handle_provider_unauthorized( + unauthorized_transport, + &self.client.state.provider_auth, + session_telemetry, + ) + .await?, ) - .await?, - ); + } else { + PendingUnauthorizedRetry::from_recovery( + handle_unauthorized( + unauthorized_transport, + &mut auth_recovery, + session_telemetry, + ) + .await?, + ) + }; continue; } Err(err) => return Err(map_api_error(err)), @@ -1109,6 +1134,7 @@ impl ModelClientSession { let mut auth_recovery = auth_manager .as_ref() .map(super::auth::AuthManager::unauthorized_recovery); + let mut provider_auth_retry_available = self.client.state.provider_auth.is_configured(); let mut pending_retry = PendingUnauthorizedRetry::default(); loop { let client_setup = self.client.current_client_setup().await?; @@ -1162,14 +1188,27 @@ impl ModelClientSession { Err(ApiError::Transport( unauthorized_transport @ TransportError::Http { status, .. }, )) if status == StatusCode::UNAUTHORIZED => { - pending_retry = PendingUnauthorizedRetry::from_recovery( - handle_unauthorized( - unauthorized_transport, - &mut auth_recovery, - session_telemetry, + pending_retry = if provider_auth_retry_available { + provider_auth_retry_available = false; + self.reset_websocket_session(); + PendingUnauthorizedRetry::from_recovery( + handle_provider_unauthorized( + unauthorized_transport, + &self.client.state.provider_auth, + session_telemetry, + ) + .await?, ) - .await?, - ); + } else { + PendingUnauthorizedRetry::from_recovery( + handle_unauthorized( + unauthorized_transport, + &mut auth_recovery, + session_telemetry, + ) + .await?, + ) + }; continue; } Err(err) => return Err(map_api_error(err)), @@ -1667,6 +1706,64 @@ async fn handle_unauthorized( Err(map_api_error(ApiError::Transport(transport))) } +async fn handle_provider_unauthorized( + transport: TransportError, + provider_auth: &ProviderAuthResolver, + session_telemetry: &SessionTelemetry, +) -> Result { + let debug = extract_response_debug_context(&transport); + let mode = "provider_exec"; + let phase = "refresh"; + match provider_auth.refresh_after_unauthorized().await { + Ok(auth_state_changed) => { + session_telemetry.record_auth_recovery( + mode, + phase, + "recovery_succeeded", + debug.request_id.as_deref(), + debug.cf_ray.as_deref(), + debug.auth_error.as_deref(), + debug.auth_error_code.as_deref(), + /*recovery_reason*/ None, + auth_state_changed, + ); + emit_feedback_auth_recovery_tags( + mode, + phase, + "recovery_succeeded", + debug.request_id.as_deref(), + debug.cf_ray.as_deref(), + debug.auth_error.as_deref(), + debug.auth_error_code.as_deref(), + ); + Ok(UnauthorizedRecoveryExecution { mode, phase }) + } + Err(err) => { + session_telemetry.record_auth_recovery( + mode, + phase, + "recovery_failed_permanent", + debug.request_id.as_deref(), + debug.cf_ray.as_deref(), + debug.auth_error.as_deref(), + debug.auth_error_code.as_deref(), + /*recovery_reason*/ None, + /*auth_state_changed*/ None, + ); + emit_feedback_auth_recovery_tags( + mode, + phase, + "recovery_failed_permanent", + debug.request_id.as_deref(), + debug.cf_ray.as_deref(), + debug.auth_error.as_deref(), + debug.auth_error_code.as_deref(), + ); + Err(err) + } + } +} + fn api_error_http_status(error: &ApiError) -> Option { match error { ApiError::Transport(TransportError::Http { status, .. }) => Some(status.as_u16()), diff --git a/codex-rs/core/src/config/config_tests.rs b/codex-rs/core/src/config/config_tests.rs index 00227fe2b3..db41ca6ff4 100644 --- a/codex-rs/core/src/config/config_tests.rs +++ b/codex-rs/core/src/config/config_tests.rs @@ -243,6 +243,45 @@ web_search = false ); } +#[test] +fn rejects_provider_auth_with_env_key() { + let err = toml::from_str::( + r#" +[model_providers.corp] +name = "Corp" +env_key = "CORP_TOKEN" + +[model_providers.corp.auth] +command = "print-token" +"#, + ) + .unwrap_err(); + + assert!( + err.to_string() + .contains("model_providers.corp: provider auth cannot be combined with env_key") + ); +} + +#[test] +fn rejects_provider_auth_zero_refresh_interval() { + let err = toml::from_str::( + r#" +[model_providers.corp] +name = "Corp" + +[model_providers.corp.auth] +command = "print-token" +refresh_interval_ms = 0 +"#, + ) + .unwrap_err(); + + assert!(err.to_string().contains( + "model_providers.corp: provider auth.refresh_interval_ms must be greater than zero" + )); +} + #[test] fn config_toml_deserializes_model_availability_nux() { let toml = r#" @@ -4315,6 +4354,7 @@ model_verbosity = "high" wire_api: crate::WireApi::Responses, env_key_instructions: None, experimental_bearer_token: None, + auth: None, query_params: None, http_headers: None, env_http_headers: None, diff --git a/codex-rs/core/src/config/mod.rs b/codex-rs/core/src/config/mod.rs index 1a0722119b..c479f9f1ab 100644 --- a/codex-rs/core/src/config/mod.rs +++ b/codex-rs/core/src/config/mod.rs @@ -1837,6 +1837,18 @@ Built-in providers cannot be overridden. Rename your custom provider (for exampl } } +fn validate_model_providers( + model_providers: &HashMap, +) -> Result<(), String> { + validate_reserved_model_provider_ids(model_providers)?; + for (key, provider) in model_providers { + provider + .validate() + .map_err(|message| format!("model_providers.{key}: {message}"))?; + } + Ok(()) +} + fn deserialize_model_providers<'de, D>( deserializer: D, ) -> Result, D::Error> @@ -1844,7 +1856,7 @@ where D: serde::Deserializer<'de>, { let model_providers = HashMap::::deserialize(deserializer)?; - validate_reserved_model_provider_ids(&model_providers).map_err(serde::de::Error::custom)?; + validate_model_providers(&model_providers).map_err(serde::de::Error::custom)?; Ok(model_providers) } @@ -1969,7 +1981,7 @@ impl Config { codex_home: PathBuf, config_layer_stack: ConfigLayerStack, ) -> std::io::Result { - validate_reserved_model_provider_ids(&cfg.model_providers) + validate_model_providers(&cfg.model_providers) .map_err(|message| std::io::Error::new(std::io::ErrorKind::InvalidInput, message))?; // Ensure that every field of ConfigRequirements is applied to the final // Config. diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index 5276b09de3..e4cbff1ced 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -52,6 +52,7 @@ pub mod models_manager; mod network_policy_decision; pub mod network_proxy_loader; mod original_image_detail; +mod provider_auth; pub use mcp_connection_manager::MCP_SANDBOX_STATE_CAPABILITY; pub use mcp_connection_manager::MCP_SANDBOX_STATE_METHOD; pub use mcp_connection_manager::SandboxState; @@ -107,6 +108,7 @@ pub use client::X_RESPONSESAPI_INCLUDE_TIMING_METRICS_HEADER; pub use model_provider_info::DEFAULT_LMSTUDIO_PORT; pub use model_provider_info::DEFAULT_OLLAMA_PORT; pub use model_provider_info::LMSTUDIO_OSS_PROVIDER_ID; +pub use model_provider_info::ModelProviderAuthInfo; pub use model_provider_info::ModelProviderInfo; pub use model_provider_info::OLLAMA_OSS_PROVIDER_ID; pub use model_provider_info::OPENAI_PROVIDER_ID; diff --git a/codex-rs/core/src/model_provider_info.rs b/codex-rs/core/src/model_provider_info.rs index 737a47780d..a36ef0d052 100644 --- a/codex-rs/core/src/model_provider_info.rs +++ b/codex-rs/core/src/model_provider_info.rs @@ -9,6 +9,7 @@ use crate::auth::AuthMode; use crate::error::EnvVarError; use codex_api::Provider as ApiProvider; use codex_api::provider::RetryConfig as ApiRetryConfig; +use codex_utils_absolute_path::AbsolutePathBuf; use http::HeaderMap; use http::header::HeaderName; use http::header::HeaderValue; @@ -19,6 +20,8 @@ use std::collections::HashMap; use std::fmt; use std::time::Duration; +const DEFAULT_PROVIDER_AUTH_TIMEOUT_MS: u64 = 5_000; +const DEFAULT_PROVIDER_AUTH_REFRESH_INTERVAL_MS: u64 = 300_000; const DEFAULT_STREAM_IDLE_TIMEOUT_MS: u64 = 300_000; const DEFAULT_STREAM_MAX_RETRIES: u64 = 5; const DEFAULT_REQUEST_MAX_RETRIES: u64 = 4; @@ -66,6 +69,61 @@ impl<'de> Deserialize<'de> for WireApi { } } +/// Configuration for obtaining a provider bearer token from a command. +#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, JsonSchema)] +#[schemars(deny_unknown_fields)] +pub struct ModelProviderAuthInfo { + /// Command to execute. Bare names are resolved via `PATH`; paths are resolved against `cwd`. + pub command: String, + + /// Command arguments. + #[serde(default)] + pub args: Vec, + + /// Maximum time to wait for the token command to exit successfully. + #[serde(default = "default_provider_auth_timeout_ms")] + pub timeout_ms: u64, + + /// Maximum age for the cached token before rerunning the command. + #[serde(default = "default_provider_auth_refresh_interval_ms")] + pub refresh_interval_ms: u64, + + /// Working directory used when running the token command. + #[serde(default = "default_provider_auth_cwd")] + #[schemars(skip)] + pub cwd: AbsolutePathBuf, +} + +impl ModelProviderAuthInfo { + pub(crate) fn timeout(&self) -> Duration { + Duration::from_millis(self.timeout_ms) + } + + pub(crate) fn refresh_interval(&self) -> Duration { + Duration::from_millis(self.refresh_interval_ms) + } +} + +fn default_provider_auth_timeout_ms() -> u64 { + DEFAULT_PROVIDER_AUTH_TIMEOUT_MS +} + +fn default_provider_auth_refresh_interval_ms() -> u64 { + DEFAULT_PROVIDER_AUTH_REFRESH_INTERVAL_MS +} + +fn default_provider_auth_cwd() -> AbsolutePathBuf { + let deserializer = serde::de::value::StrDeserializer::::new("."); + if let Ok(cwd) = AbsolutePathBuf::deserialize(deserializer) { + return cwd; + } + + match AbsolutePathBuf::current_dir() { + Ok(cwd) => cwd, + Err(err) => panic!("provider auth cwd must resolve: {err}"), + } +} + /// Serializable representation of a provider definition. #[derive(Debug, Clone, Deserialize, Serialize, PartialEq, JsonSchema)] #[schemars(deny_unknown_fields)] @@ -86,6 +144,9 @@ pub struct ModelProviderInfo { /// this may be necessary when using this programmatically. pub experimental_bearer_token: Option, + /// Command-backed bearer-token configuration for this provider. + pub auth: Option, + /// Which wire protocol this provider expects. #[serde(default)] pub wire_api: WireApi, @@ -130,6 +191,42 @@ pub struct ModelProviderInfo { } impl ModelProviderInfo { + pub(crate) fn validate(&self) -> std::result::Result<(), String> { + let Some(auth) = self.auth.as_ref() else { + return Ok(()); + }; + + if auth.command.trim().is_empty() { + return Err("provider auth.command must not be empty".to_string()); + } + if auth.timeout_ms == 0 { + return Err("provider auth.timeout_ms must be greater than zero".to_string()); + } + if auth.refresh_interval_ms == 0 { + return Err("provider auth.refresh_interval_ms must be greater than zero".to_string()); + } + + let mut conflicts = Vec::new(); + if self.env_key.is_some() { + conflicts.push("env_key"); + } + if self.experimental_bearer_token.is_some() { + conflicts.push("experimental_bearer_token"); + } + if self.requires_openai_auth { + conflicts.push("requires_openai_auth"); + } + + if conflicts.is_empty() { + Ok(()) + } else { + Err(format!( + "provider auth cannot be combined with {}", + conflicts.join(", ") + )) + } + } + fn build_header_map(&self) -> crate::error::Result { let capacity = self.http_headers.as_ref().map_or(0, HashMap::len) + self.env_http_headers.as_ref().map_or(0, HashMap::len); @@ -246,6 +343,7 @@ impl ModelProviderInfo { env_key: None, env_key_instructions: None, experimental_bearer_token: None, + auth: None, wire_api: WireApi::Responses, query_params: None, http_headers: Some( @@ -277,6 +375,10 @@ impl ModelProviderInfo { pub fn is_openai(&self) -> bool { self.name == OPENAI_PROVIDER_NAME } + + pub(crate) fn has_command_auth(&self) -> bool { + self.auth.is_some() + } } pub const DEFAULT_LMSTUDIO_PORT: u16 = 1234; @@ -338,6 +440,7 @@ pub fn create_oss_provider_with_base_url(base_url: &str, wire_api: WireApi) -> M env_key: None, env_key_instructions: None, experimental_bearer_token: None, + auth: None, wire_api, query_params: None, http_headers: None, diff --git a/codex-rs/core/src/model_provider_info_tests.rs b/codex-rs/core/src/model_provider_info_tests.rs index a5309117ae..ef4fdf10ac 100644 --- a/codex-rs/core/src/model_provider_info_tests.rs +++ b/codex-rs/core/src/model_provider_info_tests.rs @@ -1,5 +1,8 @@ use super::*; +use codex_utils_absolute_path::AbsolutePathBuf; +use codex_utils_absolute_path::AbsolutePathBufGuard; use pretty_assertions::assert_eq; +use tempfile::tempdir; #[test] fn test_deserialize_ollama_model_provider_toml() { @@ -13,6 +16,7 @@ base_url = "http://localhost:11434/v1" env_key: None, env_key_instructions: None, experimental_bearer_token: None, + auth: None, wire_api: WireApi::Responses, query_params: None, http_headers: None, @@ -43,6 +47,7 @@ query_params = { api-version = "2025-04-01-preview" } env_key: Some("AZURE_OPENAI_API_KEY".into()), env_key_instructions: None, experimental_bearer_token: None, + auth: None, wire_api: WireApi::Responses, query_params: Some(maplit::hashmap! { "api-version".to_string() => "2025-04-01-preview".to_string(), @@ -76,6 +81,7 @@ env_http_headers = { "X-Example-Env-Header" = "EXAMPLE_ENV_VAR" } env_key: Some("API_KEY".into()), env_key_instructions: None, experimental_bearer_token: None, + auth: None, wire_api: WireApi::Responses, query_params: None, http_headers: Some(maplit::hashmap! { @@ -121,3 +127,29 @@ supports_websockets = true let provider: ModelProviderInfo = toml::from_str(provider_toml).unwrap(); assert_eq!(provider.websocket_connect_timeout_ms, Some(15_000)); } + +#[test] +fn test_deserialize_provider_auth_config_defaults() { + let base_dir = tempdir().unwrap(); + let provider_toml = r#" +name = "Corp" + +[auth] +command = "./scripts/print-token" +args = ["--format=text"] + "#; + + let _guard = AbsolutePathBufGuard::new(base_dir.path()); + let provider: ModelProviderInfo = toml::from_str(provider_toml).unwrap(); + + assert_eq!( + provider.auth, + Some(ModelProviderAuthInfo { + command: "./scripts/print-token".to_string(), + args: vec!["--format=text".to_string()], + timeout_ms: DEFAULT_PROVIDER_AUTH_TIMEOUT_MS, + refresh_interval_ms: DEFAULT_PROVIDER_AUTH_REFRESH_INTERVAL_MS, + cwd: AbsolutePathBuf::resolve_path_against_base(".", base_dir.path()).unwrap(), + }) + ); +} diff --git a/codex-rs/core/src/models_manager/manager.rs b/codex-rs/core/src/models_manager/manager.rs index 29a1a85767..8618efbbb7 100644 --- a/codex-rs/core/src/models_manager/manager.rs +++ b/codex-rs/core/src/models_manager/manager.rs @@ -1,5 +1,5 @@ use super::cache::ModelsCacheManager; -use crate::api_bridge::auth_provider_from_auth; +use crate::api_bridge::auth_provider_from_resolved_provider_token; use crate::api_bridge::map_api_error; use crate::auth::AuthManager; use crate::auth::AuthMode; @@ -14,6 +14,7 @@ use crate::model_provider_info::ModelProviderInfo; use crate::models_manager::collaboration_mode_presets::CollaborationModesConfig; use crate::models_manager::collaboration_mode_presets::builtin_collaboration_mode_presets; use crate::models_manager::model_info; +use crate::provider_auth::ProviderAuthResolver; use crate::response_debug_context::extract_response_debug_context; use crate::response_debug_context::telemetry_transport_error_message; use crate::util::FeedbackRequestTags; @@ -22,6 +23,7 @@ use codex_api::ModelsClient; use codex_api::RequestTelemetry; use codex_api::ReqwestTransport; use codex_api::TransportError; +use codex_api::error::ApiError; use codex_otel::TelemetryAuthMode; use codex_protocol::config_types::CollaborationModeMask; use codex_protocol::openai_models::ModelInfo; @@ -181,6 +183,7 @@ pub struct ModelsManager { etag: RwLock>, cache_manager: ModelsCacheManager, provider: ModelProviderInfo, + provider_auth: ProviderAuthResolver, } impl ModelsManager { @@ -232,6 +235,7 @@ impl ModelsManager { auth_manager, etag: RwLock::new(None), cache_manager, + provider_auth: ProviderAuthResolver::new(&provider), provider, } } @@ -396,7 +400,9 @@ impl ModelsManager { return Ok(()); } - if self.auth_manager.auth_mode() != Some(AuthMode::Chatgpt) { + if self.auth_manager.auth_mode() != Some(AuthMode::Chatgpt) + && !self.provider.has_command_auth() + { if matches!( refresh_strategy, RefreshStrategy::Offline | RefreshStrategy::OnlineIfUncached @@ -434,36 +440,58 @@ impl ModelsManager { let auth = self.auth_manager.auth().await; let auth_mode = auth.as_ref().map(CodexAuth::auth_mode); let api_provider = self.provider.to_api_provider(auth_mode)?; - let api_auth = auth_provider_from_auth(auth.clone(), &self.provider)?; let auth_env = collect_auth_env_telemetry( &self.provider, self.auth_manager.codex_api_key_env_enabled(), ); - let transport = ReqwestTransport::new(build_reqwest_client()); - let request_telemetry: Arc = Arc::new(ModelsRequestTelemetry { - auth_mode: auth_mode.map(|mode| TelemetryAuthMode::from(mode).to_string()), - auth_header_attached: api_auth.auth_header_attached(), - auth_header_name: api_auth.auth_header_name(), - auth_env, - }); - let client = ModelsClient::new(transport, api_provider, api_auth) - .with_telemetry(Some(request_telemetry)); - let client_version = crate::models_manager::client_version_to_whole(); - let (models, etag) = timeout( - MODELS_REFRESH_TIMEOUT, - client.list_models(&client_version, HeaderMap::new()), - ) - .await - .map_err(|_| CodexErr::Timeout)? - .map_err(map_api_error)?; + let mut provider_auth_retry_available = self.provider_auth.is_configured(); - self.apply_remote_models(models.clone()).await; - *self.etag.write().await = etag.clone(); - self.cache_manager - .persist_cache(&models, etag, client_version) - .await; - Ok(()) + loop { + let provider_token = self.provider_auth.resolve_token().await?; + let api_auth = auth_provider_from_resolved_provider_token( + auth.clone(), + &self.provider, + provider_token, + )?; + let transport = ReqwestTransport::new(build_reqwest_client()); + let request_telemetry: Arc = Arc::new(ModelsRequestTelemetry { + auth_mode: auth_mode.map(|mode| TelemetryAuthMode::from(mode).to_string()), + auth_header_attached: api_auth.auth_header_attached(), + auth_header_name: api_auth.auth_header_name(), + auth_env: auth_env.clone(), + }); + let client = ModelsClient::new(transport, api_provider.clone(), api_auth) + .with_telemetry(Some(request_telemetry)); + + match timeout( + MODELS_REFRESH_TIMEOUT, + client.list_models(&client_version, HeaderMap::new()), + ) + .await + .map_err(|_| CodexErr::Timeout)? + { + Ok((models, etag)) => { + self.apply_remote_models(models.clone()).await; + *self.etag.write().await = etag.clone(); + self.cache_manager + .persist_cache(&models, etag, client_version) + .await; + return Ok(()); + } + Err(ApiError::Transport(TransportError::Http { status, .. })) + if status == http::StatusCode::UNAUTHORIZED + && provider_auth_retry_available => + { + provider_auth_retry_available = false; + self.provider_auth + .refresh_after_unauthorized() + .await + .map(|_| ())?; + } + Err(err) => return Err(map_api_error(err)), + } + } } async fn get_etag(&self) -> Option { diff --git a/codex-rs/core/src/models_manager/manager_tests.rs b/codex-rs/core/src/models_manager/manager_tests.rs index 7b4b2be53b..89c867885c 100644 --- a/codex-rs/core/src/models_manager/manager_tests.rs +++ b/codex-rs/core/src/models_manager/manager_tests.rs @@ -79,6 +79,7 @@ fn provider_for(base_url: String) -> ModelProviderInfo { env_key: None, env_key_instructions: None, experimental_bearer_token: None, + auth: None, wire_api: WireApi::Responses, query_params: None, http_headers: None, diff --git a/codex-rs/core/src/provider_auth.rs b/codex-rs/core/src/provider_auth.rs new file mode 100644 index 0000000000..58ff61977e --- /dev/null +++ b/codex-rs/core/src/provider_auth.rs @@ -0,0 +1,276 @@ +use std::fmt; +use std::path::Path; +use std::path::PathBuf; +use std::process::Stdio; +use std::sync::Arc; +use std::time::Instant; + +use tokio::process::Command; +use tokio::sync::Mutex; + +use crate::error::CodexErr; +use crate::error::Result; +use crate::model_provider_info::ModelProviderAuthInfo; +use crate::model_provider_info::ModelProviderInfo; + +#[derive(Clone, Default)] +pub(crate) struct ProviderAuthResolver { + state: Option>, +} + +impl ProviderAuthResolver { + pub(crate) fn new(provider: &ModelProviderInfo) -> Self { + Self { + state: provider + .auth + .clone() + .map(ProviderAuthState::new) + .map(Arc::new), + } + } + + pub(crate) fn is_configured(&self) -> bool { + self.state.is_some() + } + + pub(crate) async fn resolve_token(&self) -> Result> { + let Some(state) = self.state.as_ref() else { + return Ok(None); + }; + + let mut cached = state.cached_token.lock().await; + if let Some(cached_token) = cached.as_ref() + && cached_token.fetched_at.elapsed() < state.config.refresh_interval() + { + return Ok(Some(cached_token.token.clone())); + } + + let token = run_provider_auth_command(&state.config).await?; + *cached = Some(CachedProviderToken { + token: token.clone(), + fetched_at: Instant::now(), + }); + Ok(Some(token)) + } + + pub(crate) async fn refresh_after_unauthorized(&self) -> Result> { + let Some(state) = self.state.as_ref() else { + return Ok(None); + }; + + let mut cached = state.cached_token.lock().await; + let previous_token = cached.as_ref().map(|token| token.token.clone()); + let token = run_provider_auth_command(&state.config).await?; + let auth_state_changed = previous_token + .as_ref() + .map(|previous_token| previous_token != &token); + *cached = Some(CachedProviderToken { + token, + fetched_at: Instant::now(), + }); + Ok(auth_state_changed) + } +} + +impl fmt::Debug for ProviderAuthResolver { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("ProviderAuthResolver") + .field("configured", &self.is_configured()) + .finish() + } +} + +struct ProviderAuthState { + config: ModelProviderAuthInfo, + cached_token: Mutex>, +} + +impl ProviderAuthState { + fn new(config: ModelProviderAuthInfo) -> Self { + Self { + config, + cached_token: Mutex::new(None), + } + } +} + +struct CachedProviderToken { + token: String, + fetched_at: Instant, +} + +async fn run_provider_auth_command(config: &ModelProviderAuthInfo) -> Result { + let program = resolve_provider_auth_program(&config.command, &config.cwd)?; + let mut command = Command::new(&program); + command + .args(&config.args) + .current_dir(config.cwd.as_path()) + .stdin(Stdio::null()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .kill_on_drop(true); + + let output = tokio::time::timeout(config.timeout(), command.output()) + .await + .map_err(|_| { + CodexErr::InvalidRequest(format!( + "provider auth command `{}` timed out after {} ms", + config.command, config.timeout_ms + )) + })? + .map_err(|err| { + CodexErr::InvalidRequest(format!( + "provider auth command `{}` failed to start: {err}", + config.command + )) + })?; + + if !output.status.success() { + let status = output.status; + let stderr = String::from_utf8_lossy(&output.stderr).trim().to_string(); + let stderr_suffix = if stderr.is_empty() { + String::new() + } else { + format!(": {stderr}") + }; + return Err(CodexErr::InvalidRequest(format!( + "provider auth command `{}` exited with status {status}{stderr_suffix}", + config.command + ))); + } + + let stdout = String::from_utf8(output.stdout).map_err(|_| { + CodexErr::InvalidRequest(format!( + "provider auth command `{}` wrote non-UTF-8 data to stdout", + config.command + )) + })?; + let token = stdout.trim().to_string(); + if token.is_empty() { + return Err(CodexErr::InvalidRequest(format!( + "provider auth command `{}` produced an empty token", + config.command + ))); + } + + Ok(token) +} + +fn resolve_provider_auth_program(command: &str, cwd: &Path) -> Result { + let path = Path::new(command); + if path.is_absolute() || path.components().count() > 1 { + return Ok( + codex_utils_absolute_path::AbsolutePathBuf::resolve_path_against_base(path, cwd)? + .into_path_buf(), + ); + } + + Ok(PathBuf::from(command)) +} + +#[cfg(test)] +mod tests { + use super::*; + use pretty_assertions::assert_eq; + use tempfile::TempDir; + + #[tokio::test] + async fn caches_command_output_until_refreshed() { + let script = ProviderAuthScript::new(&["first-token", "second-token"]).unwrap(); + let provider = ModelProviderInfo { + name: "Test".to_string(), + base_url: None, + env_key: None, + env_key_instructions: None, + experimental_bearer_token: None, + auth: Some(script.auth_config()), + wire_api: crate::WireApi::Responses, + query_params: None, + http_headers: None, + env_http_headers: None, + request_max_retries: None, + stream_max_retries: None, + stream_idle_timeout_ms: None, + websocket_connect_timeout_ms: None, + requires_openai_auth: false, + supports_websockets: false, + }; + let resolver = ProviderAuthResolver::new(&provider); + + let first = resolver.resolve_token().await.unwrap(); + let second = resolver.resolve_token().await.unwrap(); + let changed = resolver.refresh_after_unauthorized().await.unwrap(); + let refreshed = resolver.resolve_token().await.unwrap(); + + assert_eq!(first.as_deref(), Some("first-token")); + assert_eq!(second.as_deref(), Some("first-token")); + assert_eq!(changed, Some(true)); + assert_eq!(refreshed.as_deref(), Some("second-token")); + } + + struct ProviderAuthScript { + tempdir: TempDir, + command: String, + args: Vec, + } + + impl ProviderAuthScript { + fn new(tokens: &[&str]) -> Result { + let tempdir = tempfile::tempdir()?; + let token_file = tempdir.path().join("tokens.txt"); + std::fs::write(&token_file, format!("{}\n", tokens.join("\n")))?; + + #[cfg(unix)] + let (command, args) = { + let script_path = tempdir.path().join("print-token.sh"); + std::fs::write( + &script_path, + "#!/bin/sh\nfirst_line=$(sed -n '1p' tokens.txt)\nprintf '%s\\n' \"$first_line\"\ntail -n +2 tokens.txt > tokens.next\nmv tokens.next tokens.txt\n", + )?; + let mut permissions = std::fs::metadata(&script_path)?.permissions(); + { + use std::os::unix::fs::PermissionsExt; + permissions.set_mode(0o755); + } + std::fs::set_permissions(&script_path, permissions)?; + ("./print-token.sh".to_string(), Vec::new()) + }; + + #[cfg(windows)] + let (command, args) = { + let script_path = tempdir.path().join("print-token.ps1"); + std::fs::write( + &script_path, + "$lines = Get-Content -Path tokens.txt\nif ($lines.Count -eq 0) { exit 1 }\nWrite-Output $lines[0]\n$lines | Select-Object -Skip 1 | Set-Content -Path tokens.txt\n", + )?; + ( + "powershell".to_string(), + vec![ + "-NoProfile".to_string(), + "-ExecutionPolicy".to_string(), + "Bypass".to_string(), + "-File".to_string(), + ".\\print-token.ps1".to_string(), + ], + ) + }; + + Ok(Self { + tempdir, + command, + args, + }) + } + + fn auth_config(&self) -> ModelProviderAuthInfo { + ModelProviderAuthInfo { + command: self.command.clone(), + args: self.args.clone(), + timeout_ms: 1_000, + refresh_interval_ms: 60_000, + cwd: codex_utils_absolute_path::AbsolutePathBuf::try_from(self.tempdir.path()) + .expect("tempdir should be absolute"), + } + } + } +} diff --git a/codex-rs/core/src/realtime_conversation.rs b/codex-rs/core/src/realtime_conversation.rs index 1ddd72d0fd..2bcf742c9a 100644 --- a/codex-rs/core/src/realtime_conversation.rs +++ b/codex-rs/core/src/realtime_conversation.rs @@ -1,5 +1,6 @@ use crate::CodexAuth; use crate::api_bridge::map_api_error; +use crate::api_bridge::provider_bearer_token; use crate::auth::read_openai_api_key_from_env; use crate::codex::Session; use crate::config::RealtimeWsMode; @@ -453,7 +454,12 @@ async fn prepare_realtime_start( ) -> CodexResult { let provider = sess.provider().await; let auth = sess.services.auth_manager.auth().await; - let realtime_api_key = realtime_api_key(auth.as_ref(), &provider)?; + let provider_token = sess + .services + .model_client + .resolve_provider_bearer_token() + .await?; + let realtime_api_key = realtime_api_key(auth.as_ref(), &provider, provider_token)?; let mut api_provider = provider.to_api_provider(Some(crate::auth::AuthMode::ApiKey))?; let config = sess.get_config().await; if let Some(realtime_ws_base_url) = &config.experimental_realtime_ws_base_url { @@ -628,15 +634,12 @@ fn realtime_text_from_handoff_request(handoff: &RealtimeHandoffRequested) -> Opt fn realtime_api_key( auth: Option<&CodexAuth>, provider: &crate::ModelProviderInfo, + provider_token: Option, ) -> CodexResult { - if let Some(api_key) = provider.api_key()? { + if let Some(api_key) = provider_bearer_token(provider, provider_token)? { return Ok(api_key); } - if let Some(token) = provider.experimental_bearer_token.clone() { - return Ok(token); - } - if let Some(api_key) = auth.and_then(CodexAuth::api_key) { return Ok(api_key.to_string()); } diff --git a/codex-rs/core/tests/responses_headers.rs b/codex-rs/core/tests/responses_headers.rs index 515d07f204..536e985d19 100644 --- a/codex-rs/core/tests/responses_headers.rs +++ b/codex-rs/core/tests/responses_headers.rs @@ -46,6 +46,7 @@ async fn responses_stream_includes_subagent_header_on_review() { env_key: None, env_key_instructions: None, experimental_bearer_token: None, + auth: None, wire_api: WireApi::Responses, query_params: None, http_headers: None, @@ -158,6 +159,7 @@ async fn responses_stream_includes_subagent_header_on_other() { env_key: None, env_key_instructions: None, experimental_bearer_token: None, + auth: None, wire_api: WireApi::Responses, query_params: None, http_headers: None, @@ -265,6 +267,7 @@ async fn responses_respects_model_info_overrides_from_config() { env_key: None, env_key_instructions: None, experimental_bearer_token: None, + auth: None, wire_api: WireApi::Responses, query_params: None, http_headers: None, diff --git a/codex-rs/core/tests/suite/client.rs b/codex-rs/core/tests/suite/client.rs index 9be02cf97b..3968277437 100644 --- a/codex-rs/core/tests/suite/client.rs +++ b/codex-rs/core/tests/suite/client.rs @@ -1,5 +1,6 @@ use codex_core::CodexAuth; use codex_core::ModelClient; +use codex_core::ModelProviderAuthInfo; use codex_core::ModelProviderInfo; use codex_core::NewThread; use codex_core::Prompt; @@ -143,6 +144,74 @@ fn write_auth_json( fake_jwt } +struct ProviderAuthCommandFixture { + tempdir: TempDir, + command: String, + args: Vec, +} + +impl ProviderAuthCommandFixture { + fn new(tokens: &[&str]) -> std::io::Result { + let tempdir = tempfile::tempdir()?; + let tokens_file = tempdir.path().join("tokens.txt"); + std::fs::write(&tokens_file, format!("{}\n", tokens.join("\n")))?; + + #[cfg(unix)] + let (command, args) = { + let script_path = tempdir.path().join("print-token.sh"); + std::fs::write( + &script_path, + "#!/bin/sh\nfirst_line=$(sed -n '1p' tokens.txt)\nprintf '%s\\n' \"$first_line\"\ntail -n +2 tokens.txt > tokens.next\nmv tokens.next tokens.txt\n", + )?; + let mut permissions = std::fs::metadata(&script_path)?.permissions(); + { + use std::os::unix::fs::PermissionsExt; + permissions.set_mode(0o755); + } + std::fs::set_permissions(&script_path, permissions)?; + ("./print-token.sh".to_string(), Vec::new()) + }; + + #[cfg(windows)] + let (command, args) = { + let script_path = tempdir.path().join("print-token.ps1"); + std::fs::write( + &script_path, + "$lines = Get-Content -Path tokens.txt\nif ($lines.Count -eq 0) { exit 1 }\nWrite-Output $lines[0]\n$lines | Select-Object -Skip 1 | Set-Content -Path tokens.txt\n", + )?; + ( + "powershell".to_string(), + vec![ + "-NoProfile".to_string(), + "-ExecutionPolicy".to_string(), + "Bypass".to_string(), + "-File".to_string(), + ".\\print-token.ps1".to_string(), + ], + ) + }; + + Ok(Self { + tempdir, + command, + args, + }) + } + + fn auth(&self) -> ModelProviderAuthInfo { + ModelProviderAuthInfo { + command: self.command.clone(), + args: self.args.clone(), + timeout_ms: 1_000, + refresh_interval_ms: 60_000, + cwd: match codex_utils_absolute_path::AbsolutePathBuf::try_from(self.tempdir.path()) { + Ok(cwd) => cwd, + Err(err) => panic!("tempdir should be absolute: {err}"), + }, + } + } +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn resume_includes_initial_messages_and_sends_prior_items() { skip_if_no_network!(); @@ -659,6 +728,223 @@ async fn includes_conversation_id_and_model_headers_in_request() { assert_eq!(request_authorization, "Bearer Test API Key"); } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn provider_auth_command_supplies_bearer_token() { + skip_if_no_network!(); + + let server = MockServer::start().await; + let resp_mock = mount_sse_once( + &server, + sse(vec![ev_response_created("resp1"), ev_completed("resp1")]), + ) + .await; + let auth_fixture = ProviderAuthCommandFixture::new(&["command-token"]).unwrap(); + + let provider = ModelProviderInfo { + name: "corp".into(), + base_url: Some(format!("{}/v1", server.uri())), + env_key: None, + env_key_instructions: None, + experimental_bearer_token: None, + auth: Some(auth_fixture.auth()), + wire_api: WireApi::Responses, + query_params: None, + http_headers: None, + env_http_headers: None, + request_max_retries: Some(0), + stream_max_retries: Some(0), + stream_idle_timeout_ms: Some(5_000), + websocket_connect_timeout_ms: None, + requires_openai_auth: false, + supports_websockets: false, + }; + + let codex_home = TempDir::new().unwrap(); + let mut config = load_default_config_for_test(&codex_home).await; + config.model_provider_id = provider.name.clone(); + config.model_provider = provider.clone(); + let effort = config.model_reasoning_effort; + let summary = config.model_reasoning_summary; + let model = codex_core::test_support::get_model_offline(config.model.as_deref()); + config.model = Some(model.clone()); + let config = Arc::new(config); + let model_info = + codex_core::test_support::construct_model_info_offline(model.as_str(), &config); + let conversation_id = ThreadId::new(); + let session_telemetry = SessionTelemetry::new( + conversation_id, + model.as_str(), + model_info.slug.as_str(), + /*account_id*/ None, + Some("test@test.com".to_string()), + /*auth_mode*/ None, + "test_originator".to_string(), + /*log_user_prompts*/ false, + "test".to_string(), + SessionSource::Exec, + ); + let client = ModelClient::new( + /*auth_manager*/ None, + conversation_id, + provider, + SessionSource::Exec, + config.model_verbosity, + /*enable_request_compression*/ false, + /*include_timing_metrics*/ false, + /*beta_features_header*/ None, + ); + let mut client_session = client.new_session(); + let mut prompt = Prompt::default(); + prompt.input.push(ResponseItem::Message { + id: None, + role: "user".to_string(), + content: vec![ContentItem::InputText { + text: "hello".to_string(), + }], + end_turn: None, + phase: None, + }); + + let mut stream = client_session + .stream( + &prompt, + &model_info, + &session_telemetry, + effort, + summary.unwrap_or(ReasoningSummary::Auto), + /*service_tier*/ None, + /*turn_metadata_header*/ None, + ) + .await + .expect("responses stream to start"); + + while let Some(event) = stream.next().await { + if let Ok(ResponseEvent::Completed { .. }) = event { + break; + } + } + + let request = resp_mock.single_request(); + let request_authorization = request + .header("authorization") + .expect("authorization header"); + assert_eq!(request_authorization, "Bearer command-token"); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn provider_auth_command_refreshes_after_401() { + skip_if_no_network!(); + + let server = MockServer::start().await; + let auth_fixture = ProviderAuthCommandFixture::new(&["first-token", "second-token"]).unwrap(); + + Mock::given(method("POST")) + .and(path("/v1/responses")) + .and(header_regex("Authorization", "Bearer first-token")) + .respond_with(ResponseTemplate::new(401).set_body_string("unauthorized")) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/v1/responses")) + .and(header_regex("Authorization", "Bearer second-token")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_raw( + sse(vec![ev_response_created("resp1"), ev_completed("resp1")]), + "text/event-stream", + ), + ) + .expect(1) + .mount(&server) + .await; + + let provider = ModelProviderInfo { + name: "corp".into(), + base_url: Some(format!("{}/v1", server.uri())), + env_key: None, + env_key_instructions: None, + experimental_bearer_token: None, + auth: Some(auth_fixture.auth()), + wire_api: WireApi::Responses, + query_params: None, + http_headers: None, + env_http_headers: None, + request_max_retries: Some(0), + stream_max_retries: Some(0), + stream_idle_timeout_ms: Some(5_000), + websocket_connect_timeout_ms: None, + requires_openai_auth: false, + supports_websockets: false, + }; + + let codex_home = TempDir::new().unwrap(); + let mut config = load_default_config_for_test(&codex_home).await; + config.model_provider_id = provider.name.clone(); + config.model_provider = provider.clone(); + let effort = config.model_reasoning_effort; + let summary = config.model_reasoning_summary; + let model = codex_core::test_support::get_model_offline(config.model.as_deref()); + config.model = Some(model.clone()); + let config = Arc::new(config); + let model_info = + codex_core::test_support::construct_model_info_offline(model.as_str(), &config); + let conversation_id = ThreadId::new(); + let session_telemetry = SessionTelemetry::new( + conversation_id, + model.as_str(), + model_info.slug.as_str(), + /*account_id*/ None, + Some("test@test.com".to_string()), + /*auth_mode*/ None, + "test_originator".to_string(), + /*log_user_prompts*/ false, + "test".to_string(), + SessionSource::Exec, + ); + let client = ModelClient::new( + /*auth_manager*/ None, + conversation_id, + provider, + SessionSource::Exec, + config.model_verbosity, + /*enable_request_compression*/ false, + /*include_timing_metrics*/ false, + /*beta_features_header*/ None, + ); + let mut client_session = client.new_session(); + let mut prompt = Prompt::default(); + prompt.input.push(ResponseItem::Message { + id: None, + role: "user".to_string(), + content: vec![ContentItem::InputText { + text: "hello".to_string(), + }], + end_turn: None, + phase: None, + }); + + let mut stream = client_session + .stream( + &prompt, + &model_info, + &session_telemetry, + effort, + summary.unwrap_or(ReasoningSummary::Auto), + /*service_tier*/ None, + /*turn_metadata_header*/ None, + ) + .await + .expect("responses stream to start"); + + while let Some(event) = stream.next().await { + if let Ok(ResponseEvent::Completed { .. }) = event { + break; + } + } +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn includes_base_instructions_override_in_request() { skip_if_no_network!(); @@ -1796,6 +2082,7 @@ async fn azure_responses_request_includes_store_and_reasoning_ids() { env_key: None, env_key_instructions: None, experimental_bearer_token: None, + auth: None, wire_api: WireApi::Responses, query_params: None, http_headers: None, @@ -2396,6 +2683,7 @@ async fn azure_overrides_assign_properties_used_for_responses_url() { // Reuse the existing environment variable to avoid using unsafe code env_key: Some(existing_env_var_with_random_value.to_string()), experimental_bearer_token: None, + auth: None, query_params: Some(std::collections::HashMap::from([( "api-version".to_string(), "2025-04-01-preview".to_string(), @@ -2486,6 +2774,7 @@ async fn env_var_overrides_loaded_auth() { )])), env_key_instructions: None, experimental_bearer_token: None, + auth: None, wire_api: WireApi::Responses, http_headers: Some(std::collections::HashMap::from([( "Custom-Header".to_string(), diff --git a/codex-rs/core/tests/suite/client_websockets.rs b/codex-rs/core/tests/suite/client_websockets.rs index 1f94330cb2..836399f0c4 100755 --- a/codex-rs/core/tests/suite/client_websockets.rs +++ b/codex-rs/core/tests/suite/client_websockets.rs @@ -1674,6 +1674,7 @@ fn websocket_provider_with_connect_timeout( env_key: None, env_key_instructions: None, experimental_bearer_token: None, + auth: None, wire_api: WireApi::Responses, query_params: None, http_headers: None, diff --git a/codex-rs/core/tests/suite/stream_error_allows_next_turn.rs b/codex-rs/core/tests/suite/stream_error_allows_next_turn.rs index 23ffc4afb0..159db302b2 100644 --- a/codex-rs/core/tests/suite/stream_error_allows_next_turn.rs +++ b/codex-rs/core/tests/suite/stream_error_allows_next_turn.rs @@ -69,6 +69,7 @@ async fn continue_after_stream_error() { env_key: Some("PATH".into()), env_key_instructions: None, experimental_bearer_token: None, + auth: None, wire_api: WireApi::Responses, query_params: None, http_headers: None, diff --git a/codex-rs/core/tests/suite/stream_no_completed.rs b/codex-rs/core/tests/suite/stream_no_completed.rs index 5d1b214811..df711ee302 100644 --- a/codex-rs/core/tests/suite/stream_no_completed.rs +++ b/codex-rs/core/tests/suite/stream_no_completed.rs @@ -53,6 +53,7 @@ async fn retries_on_early_close() { env_key: Some("PATH".into()), env_key_instructions: None, experimental_bearer_token: None, + auth: None, wire_api: WireApi::Responses, query_params: None, http_headers: None,