From 0301b5865bd9158232e792abebc6233ba9d3a444 Mon Sep 17 00:00:00 2001 From: Cooper Gamble Date: Wed, 3 Jun 2026 06:58:55 +0000 Subject: [PATCH] [codex] attribute model catalog HTTP state [ci changed_files] --- codex-rs/Cargo.lock | 1 + codex-rs/app-server/src/message_processor.rs | 10 +- codex-rs/app-server/src/models.rs | 4 +- .../request_processors/catalog_processor.rs | 15 +- .../external_agent_config_processor.rs | 12 +- .../request_processors/thread_processor.rs | 14 +- .../src/request_processors/turn_processor.rs | 3 +- codex-rs/core/src/client.rs | 9 ++ codex-rs/core/src/codex_delegate.rs | 8 +- codex-rs/core/src/codex_thread.rs | 9 ++ codex-rs/core/src/session/mod.rs | 14 +- codex-rs/core/src/session/review.rs | 6 +- codex-rs/core/src/session/session.rs | 4 + codex-rs/core/src/session/tests.rs | 13 ++ .../core/src/session/tests/guardian_tests.rs | 1 + codex-rs/core/src/session/turn.rs | 9 +- codex-rs/core/src/session/turn_context.rs | 16 +- codex-rs/core/src/thread_manager.rs | 69 +++++++++ codex-rs/core/src/thread_manager_tests.rs | 4 + .../tests/suite/subagent_notifications.rs | 1 + codex-rs/memories/write/src/runtime.rs | 1 + .../model-provider/src/models_endpoint.rs | 14 +- codex-rs/models-manager/Cargo.toml | 1 + codex-rs/models-manager/src/manager.rs | 138 +++++++++++++++++- codex-rs/models-manager/src/manager_tests.rs | 34 +++++ 25 files changed, 381 insertions(+), 29 deletions(-) diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 750b874be0..3d827aad46 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -3357,6 +3357,7 @@ dependencies = [ "chrono", "codex-app-server-protocol", "codex-collaboration-mode-templates", + "codex-http-state", "codex-login", "codex-otel", "codex-protocol", diff --git a/codex-rs/app-server/src/message_processor.rs b/codex-rs/app-server/src/message_processor.rs index 8df998cd31..70c45591f6 100644 --- a/codex-rs/app-server/src/message_processor.rs +++ b/codex-rs/app-server/src/message_processor.rs @@ -895,7 +895,11 @@ impl MessageProcessor { .map(|response| Some(response.into())), ClientRequest::ExternalAgentConfigImport { params, .. } => self .external_agent_config_processor - .import(request_id.clone(), params) + .import( + request_id.clone(), + params, + app_server_client_name.as_deref(), + ) .await .map(|()| None), ClientRequest::ConfigValueWrite { params, .. } => { @@ -1188,7 +1192,9 @@ impl MessageProcessor { self.plugin_processor.plugin_uninstall(params).await } ClientRequest::ModelList { params, .. } => { - self.catalog_processor.model_list(params).await + self.catalog_processor + .model_list(params, app_server_client_name.as_deref()) + .await } ClientRequest::ExperimentalFeatureList { params, .. } => { self.catalog_processor diff --git a/codex-rs/app-server/src/models.rs b/codex-rs/app-server/src/models.rs index e53f13df84..a772007b1e 100644 --- a/codex-rs/app-server/src/models.rs +++ b/codex-rs/app-server/src/models.rs @@ -5,6 +5,7 @@ use codex_app_server_protocol::ModelServiceTier; use codex_app_server_protocol::ModelUpgradeInfo; use codex_app_server_protocol::ReasoningEffortOption; use codex_core::ThreadManager; +use codex_http_state::HttpStateSurface; use codex_models_manager::manager::RefreshStrategy; use codex_protocol::openai_models::ModelPreset; use codex_protocol::openai_models::ReasoningEffortPreset; @@ -12,9 +13,10 @@ use codex_protocol::openai_models::ReasoningEffortPreset; pub async fn supported_models( thread_manager: Arc, include_hidden: bool, + http_state_surface: HttpStateSurface, ) -> Vec { thread_manager - .list_models(RefreshStrategy::OnlineIfUncached) + .list_models_for_surface(RefreshStrategy::OnlineIfUncached, http_state_surface) .await .into_iter() .filter(|preset| include_hidden || preset.show_in_picker) diff --git a/codex-rs/app-server/src/request_processors/catalog_processor.rs b/codex-rs/app-server/src/request_processors/catalog_processor.rs index 6f17e4614f..5f2d681912 100644 --- a/codex-rs/app-server/src/request_processors/catalog_processor.rs +++ b/codex-rs/app-server/src/request_processors/catalog_processor.rs @@ -1,5 +1,6 @@ use super::*; use codex_config::config_toml::ConfigToml; +use codex_http_state::HttpStateSurface; use futures::StreamExt; #[derive(Clone)] @@ -156,8 +157,9 @@ impl CatalogRequestProcessor { pub(crate) async fn model_list( &self, params: ModelListParams, + app_server_client_name: Option<&str>, ) -> Result, JSONRPCErrorError> { - Self::list_models(self.thread_manager.clone(), params) + Self::list_models(self.thread_manager.clone(), params, app_server_client_name) .await .map(|response| Some(response.into())) } @@ -248,13 +250,22 @@ impl CatalogRequestProcessor { async fn list_models( thread_manager: Arc, params: ModelListParams, + app_server_client_name: Option<&str>, ) -> Result { let ModelListParams { limit, cursor, include_hidden, } = params; - let models = supported_models(thread_manager, include_hidden.unwrap_or(false)).await; + let http_state_surface = HttpStateSurface::from_app_server_client_name( + app_server_client_name.unwrap_or_default(), + ); + let models = supported_models( + thread_manager, + include_hidden.unwrap_or(false), + http_state_surface, + ) + .await; let total = models.len(); if total == 0 { diff --git a/codex-rs/app-server/src/request_processors/external_agent_config_processor.rs b/codex-rs/app-server/src/request_processors/external_agent_config_processor.rs index 87956182d0..fcd4c03827 100644 --- a/codex-rs/app-server/src/request_processors/external_agent_config_processor.rs +++ b/codex-rs/app-server/src/request_processors/external_agent_config_processor.rs @@ -34,6 +34,7 @@ use codex_external_agent_sessions::ImportedExternalAgentSession; use codex_external_agent_sessions::PendingSessionImport; use codex_external_agent_sessions::prepare_validated_session_imports; use codex_external_agent_sessions::record_imported_session; +use codex_http_state::HttpStateSurface; use codex_protocol::ThreadId; use codex_protocol::protocol::InitialHistory; use codex_thread_store::ThreadMetadataPatch; @@ -174,7 +175,11 @@ impl ExternalAgentConfigRequestProcessor { &self, request_id: ConnectionRequestId, params: ExternalAgentConfigImportParams, + app_server_client_name: Option<&str>, ) -> Result<(), JSONRPCErrorError> { + let http_state_surface = HttpStateSurface::from_app_server_client_name( + app_server_client_name.unwrap_or_default(), + ); let needs_runtime_refresh = migration_items_need_runtime_refresh(¶ms.migration_items); let has_migration_items = !params.migration_items.is_empty(); let has_plugin_imports = params.migration_items.iter().any(|item| { @@ -223,7 +228,10 @@ impl ExternalAgentConfigRequestProcessor { .prepare_validated_session_imports(pending_session_imports); for pending_session_import in pending_session_imports { match session_processor - .import_external_agent_session(pending_session_import.session) + .import_external_agent_session( + pending_session_import.session, + http_state_surface, + ) .await { Ok(imported_thread_id) => { @@ -277,6 +285,7 @@ impl ExternalAgentConfigRequestProcessor { async fn import_external_agent_session( &self, session: ImportedExternalAgentSession, + http_state_surface: HttpStateSurface, ) -> Result { let ImportedExternalAgentSession { cwd, @@ -312,6 +321,7 @@ impl ExternalAgentConfigRequestProcessor { metrics_service_name: None, parent_trace: None, environments, + http_state_surface: Some(http_state_surface), }) .await .map_err(|err| internal_error(format!("failed to import session: {err}")))?; diff --git a/codex-rs/app-server/src/request_processors/thread_processor.rs b/codex-rs/app-server/src/request_processors/thread_processor.rs index 32dfb1eeb9..3fbda2c9b8 100644 --- a/codex-rs/app-server/src/request_processors/thread_processor.rs +++ b/codex-rs/app-server/src/request_processors/thread_processor.rs @@ -1,5 +1,6 @@ use super::*; use crate::error_code::method_not_found; +use codex_http_state::HttpStateSurface; use codex_protocol::models::BUILT_IN_PERMISSION_PROFILE_DANGER_FULL_ACCESS; use codex_protocol::models::BUILT_IN_PERMISSION_PROFILE_WORKSPACE; @@ -1088,6 +1089,9 @@ impl ThreadRequestProcessor { metrics_service_name: service_name, parent_trace: request_trace, environments, + http_state_surface: Some(HttpStateSurface::from_app_server_client_name( + app_server_client_name.as_deref().unwrap_or_default(), + )), }) .instrument(tracing::info_span!( "app_server.thread_start.create_thread", @@ -2529,11 +2533,14 @@ impl ThreadRequestProcessor { match self .thread_manager - .resume_thread_with_history( + .resume_thread_with_history_for_surface( config.clone(), thread_history, self.auth_manager.clone(), self.request_trace_context(&request_id).await, + Some(HttpStateSurface::from_app_server_client_name( + app_server_client_name.as_deref().unwrap_or_default(), + )), ) .await { @@ -3238,7 +3245,7 @@ impl ThreadRequestProcessor { .. } = self .thread_manager - .fork_thread_from_history( + .fork_thread_from_history_for_surface( ForkSnapshot::Interrupted, config, InitialHistory::Resumed(ResumedHistory { @@ -3248,6 +3255,9 @@ impl ThreadRequestProcessor { }), thread_source.map(Into::into), self.request_trace_context(&request_id).await, + Some(HttpStateSurface::from_app_server_client_name( + app_server_client_name.as_deref().unwrap_or_default(), + )), ) .await .map_err(|err| match err { diff --git a/codex-rs/app-server/src/request_processors/turn_processor.rs b/codex-rs/app-server/src/request_processors/turn_processor.rs index 25beee3cdb..3190e7f1b4 100644 --- a/codex-rs/app-server/src/request_processors/turn_processor.rs +++ b/codex-rs/app-server/src/request_processors/turn_processor.rs @@ -1087,7 +1087,7 @@ impl TurnRequestProcessor { .. } = self .thread_manager - .fork_thread_from_history( + .fork_thread_from_history_for_surface( ForkSnapshot::Interrupted, config.clone(), InitialHistory::Resumed(ResumedHistory { @@ -1097,6 +1097,7 @@ impl TurnRequestProcessor { }), /*thread_source*/ None, self.request_trace_context(request_id).await, + parent_thread.http_state_surface(), ) .await .map_err(|err| { diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index cd7e75f9b0..403010e945 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -1808,6 +1808,15 @@ fn subagent_header_value(session_source: &SessionSource) -> Option { } } +pub(crate) fn http_state_surface_for_session( + session_source: &SessionSource, + app_server_client_name: Option<&str>, +) -> HttpStateSurface { + app_server_client_name + .map(HttpStateSurface::from_app_server_client_name) + .unwrap_or_else(|| http_state_surface_for_session_source(session_source)) +} + fn http_state_surface_for_session_source(session_source: &SessionSource) -> HttpStateSurface { match session_source { SessionSource::Cli => HttpStateSurface::CodexTui, diff --git a/codex-rs/core/src/codex_delegate.rs b/codex-rs/core/src/codex_delegate.rs index fc0f4f2e9e..f92e7c65a5 100644 --- a/codex-rs/core/src/codex_delegate.rs +++ b/codex-rs/core/src/codex_delegate.rs @@ -91,6 +91,7 @@ pub(crate) async fn run_codex_thread_interactive( extensions: Arc::clone(&parent_session.services.extensions), conversation_history, session_source: SessionSource::SubAgent(subagent_source.clone()), + http_state_surface: parent_session.services.model_client.http_state_surface(), forked_from_thread_id, parent_thread_id: Some(parent_session.conversation_id), thread_source: Some(ThreadSource::Subagent), @@ -110,13 +111,6 @@ pub(crate) async fn run_codex_thread_interactive( })) .or_cancel(&cancel_token) .await??; - if let Some(surface) = parent_session.services.model_client.http_state_surface() { - codex - .session - .services - .model_client - .set_http_state_surface(surface); - } let thread_config = codex.thread_config_snapshot().await; let client_metadata = parent_session.app_server_client_metadata().await; emit_subagent_session_started( diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index cd90829c0c..cbc7044637 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/codex-rs/core/src/codex_thread.rs @@ -6,6 +6,7 @@ use crate::session::Codex; use crate::session::SessionSettingsUpdate; use crate::session::SteerInputError; use codex_features::Feature; +use codex_http_state::HttpStateSurface; use codex_otel::SessionTelemetry; use codex_protocol::ThreadId; use codex_protocol::config_types::ApprovalsReviewer; @@ -302,6 +303,14 @@ impl CodexThread { .await } + pub fn http_state_surface(&self) -> Option { + self.codex + .session + .services + .model_client + .http_state_surface() + } + /// Preview persistent thread settings overrides without committing them. pub async fn preview_thread_settings_overrides( &self, diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 46a0ec5f30..12a75077c2 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -57,6 +57,7 @@ use codex_features::Feature; use codex_features::unstable_features_warning_event; use codex_hooks::Hooks; use codex_hooks::HooksConfig; +use codex_http_state::HttpStateSurface; use codex_login::AuthManager; use codex_login::CodexAuth; use codex_login::auth_env_telemetry::collect_auth_env_telemetry; @@ -171,6 +172,7 @@ use tracing::warn; use uuid::Uuid; use crate::client::ModelClient; +use crate::client::http_state_surface_for_session; use crate::codex_thread::ThreadConfigSnapshot; use crate::compact::collect_user_messages; use crate::config::Config; @@ -400,6 +402,7 @@ pub(crate) struct CodexSpawnArgs { pub(crate) extensions: Arc>, pub(crate) conversation_history: InitialHistory, pub(crate) session_source: SessionSource, + pub(crate) http_state_surface: Option, pub(crate) forked_from_thread_id: Option, pub(crate) parent_thread_id: Option, pub(crate) thread_source: Option, @@ -484,6 +487,7 @@ impl Codex { extensions, conversation_history, session_source, + http_state_surface, forked_from_thread_id, parent_thread_id, thread_source, @@ -530,6 +534,9 @@ impl Codex { }; let config = Arc::new(config); + let http_state_surface = http_state_surface.unwrap_or_else(|| { + http_state_surface_for_session(&session_source, /*client*/ None) + }); let refresh_strategy = if session_source.is_non_root_agent() { codex_models_manager::manager::RefreshStrategy::Offline } else { @@ -541,10 +548,12 @@ impl Codex { codex_models_manager::manager::RefreshStrategy::Offline ) { - let _ = models_manager.list_models(refresh_strategy).await; + let _ = models_manager + .list_models_for_surface(refresh_strategy, http_state_surface) + .await; } let model = models_manager - .get_default_model(&config.model, refresh_strategy) + .get_default_model_for_surface(&config.model, refresh_strategy, http_state_surface) .await; // Resolve base instructions for the session. Priority order: @@ -647,6 +656,7 @@ impl Codex { parent_rollout_thread_trace, attestation_provider, multi_agent_version, + http_state_surface, )) .await .map_err(|e| { diff --git a/codex-rs/core/src/session/review.rs b/codex-rs/core/src/session/review.rs index fe78e033ec..16d4b347bd 100644 --- a/codex-rs/core/src/session/review.rs +++ b/codex-rs/core/src/session/review.rs @@ -29,7 +29,10 @@ pub(super) async fn spawn_review_thread( let available_models = sess .services .models_manager - .list_models(RefreshStrategy::OnlineIfUncached) + .list_models_for_surface( + RefreshStrategy::OnlineIfUncached, + parent_turn_context.http_state_surface, + ) .await; let unified_exec_shell_mode = UnifiedExecShellMode::for_session( codex_tools::unified_exec_feature_mode_for_features(review_features.get()), @@ -124,6 +127,7 @@ pub(super) async fn spawn_review_thread( current_date: parent_turn_context.current_date.clone(), timezone: parent_turn_context.timezone.clone(), app_server_client_name: parent_turn_context.app_server_client_name.clone(), + http_state_surface: parent_turn_context.http_state_surface, developer_instructions: None, user_instructions: None, compact_prompt: parent_turn_context.compact_prompt.clone(), diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 05773abb8c..20a432c9bf 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -507,6 +507,7 @@ impl Session { parent_rollout_thread_trace: ThreadTraceContext, attestation_provider: Option>, multi_agent_version: Option, + http_state_surface: HttpStateSurface, ) -> anyhow::Result> { debug!( "Configuring session: model={}; provider={:?}", @@ -1054,6 +1055,9 @@ impl Session { code_mode_service: crate::tools::code_mode::CodeModeService::new(), environment_manager, }; + services + .model_client + .set_http_state_surface(http_state_surface); services .model_client .set_window_generation(window_generation); diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index d4b4942a88..456c405bc1 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -3404,6 +3404,7 @@ async fn includes_timed_out_message() { async fn turn_context_with_model_updates_model_fields() { let (session, mut turn_context) = make_session_and_context().await; turn_context.reasoning_effort = Some(ReasoningEffortConfig::Minimal); + turn_context.http_state_surface = HttpStateSurface::CodexDesktop; let updated = turn_context .with_model("gpt-5.4".to_string(), &session.services.models_manager) .await; @@ -3435,6 +3436,7 @@ async fn turn_context_with_model_updates_model_fields() { updated.truncation_policy, expected_model_info.truncation_policy.into() ); + assert_eq!(updated.http_state_surface, HttpStateSurface::CodexDesktop); } #[test] @@ -4605,6 +4607,7 @@ async fn session_new_fails_when_zsh_fork_enabled_without_packaged_zsh() { codex_rollout_trace::ThreadTraceContext::disabled(), /*attestation_provider*/ None, Some(config.multi_agent_version_from_features()), + HttpStateSurface::CodexExec, ) .await; @@ -4810,6 +4813,10 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { per_turn_config, model_info, &models_manager, + services + .model_client + .http_state_surface() + .expect("model client should have HTTP state"), /*network*/ None, turn_environments, session_configuration.cwd.clone(), @@ -4955,6 +4962,7 @@ async fn make_session_with_config_and_rx( codex_rollout_trace::ThreadTraceContext::disabled(), /*attestation_provider*/ None, Some(config.multi_agent_version_from_features()), + HttpStateSurface::CodexExec, ) .await?; @@ -5067,6 +5075,7 @@ async fn make_session_with_history_source_and_agent_control_and_rx( codex_rollout_trace::ThreadTraceContext::disabled(), /*attestation_provider*/ None, Some(config.multi_agent_version_from_features()), + HttpStateSurface::CodexExec, ) .await?; @@ -6659,6 +6668,10 @@ where per_turn_config, model_info, &models_manager, + services + .model_client + .http_state_surface() + .expect("model client should have HTTP state"), /*network*/ None, turn_environments, session_configuration.cwd.clone(), diff --git a/codex-rs/core/src/session/tests/guardian_tests.rs b/codex-rs/core/src/session/tests/guardian_tests.rs index 2977e6bdf5..6077917a4d 100644 --- a/codex-rs/core/src/session/tests/guardian_tests.rs +++ b/codex-rs/core/src/session/tests/guardian_tests.rs @@ -692,6 +692,7 @@ async fn guardian_subagent_does_not_inherit_parent_exec_policy_rules() { session_source: SessionSource::SubAgent(SubAgentSource::Other( GUARDIAN_REVIEWER_NAME.to_string(), )), + http_state_surface: None, forked_from_thread_id: None, parent_thread_id: None, thread_source: None, diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index d1e9e59f70..a74e7fffbe 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -1995,7 +1995,14 @@ async fn try_run_sampling_request( } ResponseEvent::ModelsEtag(etag) => { // Update internal state with latest models etag - sess.services.models_manager.refresh_if_new_etag(etag).await; + if let Some(http_state_surface) = sess.services.model_client.http_state_surface() { + sess.services + .models_manager + .refresh_if_new_etag_for_surface(etag, http_state_surface) + .await; + } else { + sess.services.models_manager.refresh_if_new_etag(etag).await; + } } ResponseEvent::Completed { response_id, diff --git a/codex-rs/core/src/session/turn_context.rs b/codex-rs/core/src/session/turn_context.rs index d600368137..311e925f27 100644 --- a/codex-rs/core/src/session/turn_context.rs +++ b/codex-rs/core/src/session/turn_context.rs @@ -1,5 +1,6 @@ use super::*; use crate::SkillLoadOutcome; +use crate::client::http_state_surface_for_session; use crate::config::GhostSnapshotConfig; use crate::environment_selection::ResolvedTurnEnvironments; use codex_model_provider::SharedModelProvider; @@ -76,6 +77,7 @@ pub struct TurnContext { pub(crate) current_date: Option, pub(crate) timezone: Option, pub(crate) app_server_client_name: Option, + pub(crate) http_state_surface: HttpStateSurface, pub(crate) developer_instructions: Option, pub(crate) compact_prompt: Option, pub(crate) user_instructions: Option, @@ -224,7 +226,7 @@ impl TurnContext { ); let features = self.features.clone(); let available_models = models_manager - .list_models(RefreshStrategy::OnlineIfUncached) + .list_models_for_surface(RefreshStrategy::OnlineIfUncached, self.http_state_surface) .await; Self { @@ -251,6 +253,7 @@ impl TurnContext { current_date: self.current_date.clone(), timezone: self.timezone.clone(), app_server_client_name: self.app_server_client_name.clone(), + http_state_surface: self.http_state_surface, developer_instructions: self.developer_instructions.clone(), compact_prompt: self.compact_prompt.clone(), user_instructions: self.user_instructions.clone(), @@ -472,6 +475,7 @@ impl Session { per_turn_config: Config, model_info: ModelInfo, models_manager: &SharedModelsManager, + http_state_surface: HttpStateSurface, network: Option, environments: ResolvedTurnEnvironments, cwd: AbsolutePathBuf, @@ -551,6 +555,7 @@ impl Session { current_date: Some(current_date), timezone: Some(timezone), app_server_client_name: session_configuration.app_server_client_name.clone(), + http_state_surface, developer_instructions: session_configuration.developer_instructions.clone(), compact_prompt: session_configuration.compact_prompt.clone(), user_instructions: session_configuration.user_instructions.clone(), @@ -792,6 +797,15 @@ impl Session { per_turn_config, model_info, &self.services.models_manager, + self.services + .model_client + .http_state_surface() + .unwrap_or_else(|| { + http_state_surface_for_session( + &session_configuration.session_source, + session_configuration.app_server_client_name.as_deref(), + ) + }), self.services .network_proxy .load_full() diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index d130a35815..1ec5eaec13 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -24,6 +24,7 @@ use codex_exec_server::EnvironmentManager; use codex_extension_api::ExtensionRegistry; use codex_extension_api::empty_extension_registry; use codex_features::Feature; +use codex_http_state::HttpStateSurface; use codex_login::AuthManager; use codex_login::CodexAuth; use codex_model_provider::create_model_provider; @@ -182,6 +183,7 @@ pub struct StartThreadOptions { pub metrics_service_name: Option, pub parent_trace: Option, pub environments: Vec, + pub http_state_surface: Option, } pub(crate) struct ResumeThreadWithHistoryOptions { @@ -451,6 +453,17 @@ impl ThreadManager { .await } + pub async fn list_models_for_surface( + &self, + refresh_strategy: RefreshStrategy, + http_state_surface: HttpStateSurface, + ) -> Vec { + self.state + .models_manager + .list_models_for_surface(refresh_strategy, http_state_surface) + .await + } + pub fn list_collaboration_modes(&self) -> Vec { self.state.models_manager.list_collaboration_modes() } @@ -576,6 +589,7 @@ impl ThreadManager { metrics_service_name: None, parent_trace: None, environments, + http_state_surface: None, })) .await } @@ -605,6 +619,7 @@ impl ThreadManager { Arc::clone(&self.state.auth_manager), self.agent_control(), session_source, + options.http_state_surface, /*parent_thread_id*/ None, forked_from_thread_id, thread_source, @@ -679,6 +694,24 @@ impl ThreadManager { initial_history: InitialHistory, auth_manager: Arc, parent_trace: Option, + ) -> CodexResult { + self.resume_thread_with_history_for_surface( + config, + initial_history, + auth_manager, + parent_trace, + /*http_state_surface*/ None, + ) + .await + } + + pub async fn resume_thread_with_history_for_surface( + &self, + config: Config, + initial_history: InitialHistory, + auth_manager: Arc, + parent_trace: Option, + http_state_surface: Option, ) -> CodexResult { let environments = default_thread_environment_selections( self.state.environment_manager.as_ref(), @@ -693,6 +726,7 @@ impl ThreadManager { auth_manager, self.agent_control(), session_source, + http_state_surface, /*parent_thread_id*/ None, /*forked_from_thread_id*/ None, thread_source, @@ -721,6 +755,7 @@ impl ThreadManager { InitialHistory::New, Arc::clone(&self.state.auth_manager), self.agent_control(), + /*http_state_surface*/ None, /*parent_thread_id*/ None, /*forked_from_thread_id*/ None, /*thread_source*/ None, @@ -754,6 +789,7 @@ impl ThreadManager { auth_manager, self.agent_control(), session_source, + /*http_state_surface*/ None, /*parent_thread_id*/ None, /*forked_from_thread_id*/ None, thread_source, @@ -883,6 +919,30 @@ impl ThreadManager { history, thread_source, parent_trace, + /*http_state_surface*/ None, + ) + .await + } + + pub async fn fork_thread_from_history_for_surface( + &self, + snapshot: S, + config: Config, + history: InitialHistory, + thread_source: Option, + parent_trace: Option, + http_state_surface: Option, + ) -> CodexResult + where + S: Into, + { + self.fork_thread_with_initial_history( + snapshot.into(), + config, + history, + thread_source, + parent_trace, + http_state_surface, ) .await } @@ -894,6 +954,7 @@ impl ThreadManager { history: InitialHistory, thread_source: Option, parent_trace: Option, + http_state_surface: Option, ) -> CodexResult { // `forked_from_id()` describes this history's existing lineage. When // forking a resumed thread, the child copies the resumed thread itself. @@ -924,6 +985,7 @@ impl ThreadManager { history, Arc::clone(&self.state.auth_manager), self.agent_control(), + http_state_surface, /*parent_thread_id*/ None, forked_from_thread_id, thread_source, @@ -1127,6 +1189,7 @@ impl ThreadManagerState { Arc::clone(&self.auth_manager), agent_control, session_source, + /*http_state_surface*/ None, parent_thread_id, forked_from_thread_id, thread_source, @@ -1163,6 +1226,7 @@ impl ThreadManagerState { Arc::clone(&self.auth_manager), agent_control, session_source, + /*http_state_surface*/ None, parent_thread_id, /*forked_from_thread_id*/ None, thread_source, @@ -1200,6 +1264,7 @@ impl ThreadManagerState { Arc::clone(&self.auth_manager), agent_control, session_source, + /*http_state_surface*/ None, parent_thread_id, forked_from_thread_id, thread_source, @@ -1222,6 +1287,7 @@ impl ThreadManagerState { initial_history: InitialHistory, auth_manager: Arc, agent_control: AgentControl, + http_state_surface: Option, parent_thread_id: Option, forked_from_thread_id: Option, thread_source: Option, @@ -1237,6 +1303,7 @@ impl ThreadManagerState { auth_manager, agent_control, self.session_source.clone(), + http_state_surface, parent_thread_id, forked_from_thread_id, thread_source, @@ -1259,6 +1326,7 @@ impl ThreadManagerState { auth_manager: Arc, agent_control: AgentControl, session_source: SessionSource, + http_state_surface: Option, parent_thread_id: Option, forked_from_thread_id: Option, thread_source: Option, @@ -1320,6 +1388,7 @@ impl ThreadManagerState { extensions: Arc::clone(&self.extensions), conversation_history: initial_history, session_source, + http_state_surface, forked_from_thread_id, parent_thread_id, thread_source, diff --git a/codex-rs/core/src/thread_manager_tests.rs b/codex-rs/core/src/thread_manager_tests.rs index 52a0f8eb3b..02596c56a9 100644 --- a/codex-rs/core/src/thread_manager_tests.rs +++ b/codex-rs/core/src/thread_manager_tests.rs @@ -333,6 +333,7 @@ async fn start_thread_rejects_explicit_local_environment_when_default_provider_i environment_id: "local".to_string(), cwd: config.cwd.clone(), }], + http_state_surface: None, }) .await; let err = match result { @@ -463,6 +464,7 @@ async fn start_thread_keeps_internal_threads_hidden_from_normal_lookups() { metrics_service_name: None, parent_trace: None, environments: Vec::new(), + http_state_surface: None, }) .await .expect("internal thread should start"); @@ -518,6 +520,7 @@ async fn resume_and_fork_do_not_restore_thread_environments_from_rollout() { metrics_service_name: None, parent_trace: None, environments: environments.clone(), + http_state_surface: None, }) .await .expect("start source thread"); @@ -788,6 +791,7 @@ async fn resume_stopped_thread_from_rollout_preserves_thread_source() { metrics_service_name: None, parent_trace: None, environments: Vec::new(), + http_state_surface: None, }) .await .expect("start source thread"); diff --git a/codex-rs/core/tests/suite/subagent_notifications.rs b/codex-rs/core/tests/suite/subagent_notifications.rs index 52fd1a25e7..dad2660f38 100644 --- a/codex-rs/core/tests/suite/subagent_notifications.rs +++ b/codex-rs/core/tests/suite/subagent_notifications.rs @@ -755,6 +755,7 @@ async fn subagent_stop_replaces_stop_and_skips_internal_subagents() -> Result<() metrics_service_name: None, parent_trace: None, environments: Vec::new(), + http_state_surface: None, }) .await?; diff --git a/codex-rs/memories/write/src/runtime.rs b/codex-rs/memories/write/src/runtime.rs index 36099fe2e8..616cd82ed9 100644 --- a/codex-rs/memories/write/src/runtime.rs +++ b/codex-rs/memories/write/src/runtime.rs @@ -251,6 +251,7 @@ impl MemoryStartupContext { metrics_service_name: None, parent_trace: None, environments, + http_state_surface: None, }) .await?; diff --git a/codex-rs/model-provider/src/models_endpoint.rs b/codex-rs/model-provider/src/models_endpoint.rs index 8a72beea70..2348917846 100644 --- a/codex-rs/model-provider/src/models_endpoint.rs +++ b/codex-rs/model-provider/src/models_endpoint.rs @@ -10,6 +10,8 @@ use codex_api::auth_header_telemetry; use codex_api::map_api_error; use codex_feedback::FeedbackRequestTags; use codex_feedback::emit_feedback_request_tags_with_auth_env; +use codex_http_state::HttpStateContext; +use codex_http_state::HttpStateSurface; use codex_login::AuthEnvTelemetry; use codex_login::AuthManager; use codex_login::CodexAuth; @@ -27,6 +29,7 @@ use http::HeaderMap; use tokio::time::timeout; use crate::auth::resolve_provider_auth; +use crate::auth::with_native_integrity_state; const MODELS_REFRESH_TIMEOUT: Duration = Duration::from_secs(5); const MODELS_ENDPOINT: &str = "/models"; @@ -81,13 +84,22 @@ impl ModelsEndpointClient for OpenAiModelsEndpoint { async fn list_models( &self, client_version: &str, + http_state_surface: Option, ) -> CoreResult<(Vec, Option)> { let _timer = codex_otel::start_global_timer("codex.remote_models.fetch_update.duration_ms", &[]); let auth = self.auth().await; let auth_mode = auth.as_ref().map(CodexAuth::auth_mode); let api_provider = self.provider_info.to_api_provider(auth_mode)?; - let api_auth = resolve_provider_auth(auth.as_ref(), &self.provider_info)?; + let api_auth = with_native_integrity_state( + resolve_provider_auth(auth.as_ref(), &self.provider_info)?, + auth.as_ref(), + http_state_surface.and_then(|surface| { + self.auth_manager.as_ref().map(|auth_manager| { + HttpStateContext::new(auth_manager.codex_home().to_path_buf(), surface) + }) + }), + ); let transport = ReqwestTransport::new(build_reqwest_client()); let auth_telemetry = auth_header_telemetry(api_auth.as_ref()); let request_telemetry: Arc = Arc::new(ModelsRequestTelemetry { diff --git a/codex-rs/models-manager/Cargo.toml b/codex-rs/models-manager/Cargo.toml index f46bf2b285..60bd6ea332 100644 --- a/codex-rs/models-manager/Cargo.toml +++ b/codex-rs/models-manager/Cargo.toml @@ -17,6 +17,7 @@ async-trait = { workspace = true } chrono = { workspace = true, features = ["serde"] } codex-app-server-protocol = { workspace = true } codex-collaboration-mode-templates = { workspace = true } +codex-http-state = { workspace = true } codex-login = { workspace = true } codex-otel = { workspace = true } codex-protocol = { workspace = true } diff --git a/codex-rs/models-manager/src/manager.rs b/codex-rs/models-manager/src/manager.rs index ac4044ac20..6aab5df997 100644 --- a/codex-rs/models-manager/src/manager.rs +++ b/codex-rs/models-manager/src/manager.rs @@ -4,6 +4,7 @@ use crate::config::ModelsManagerConfig; use crate::model_info; use async_trait::async_trait; use codex_app_server_protocol::AuthMode; +use codex_http_state::HttpStateSurface; use codex_login::AuthManager; use codex_protocol::config_types::CollaborationModeMask; use codex_protocol::error::Result as CoreResult; @@ -41,6 +42,7 @@ pub trait ModelsEndpointClient: fmt::Debug + Send + Sync { async fn list_models( &self, client_version: &str, + http_state_surface: Option, ) -> CoreResult<(Vec, Option)>; } @@ -91,9 +93,37 @@ pub trait ModelsManager: fmt::Debug + Send + Sync { .await } + /// List all available models, attributing any network refresh to `http_state_surface`. + async fn list_models_for_surface( + &self, + refresh_strategy: RefreshStrategy, + http_state_surface: HttpStateSurface, + ) -> Vec { + async move { + let catalog = self + .raw_model_catalog_for_surface(refresh_strategy, http_state_surface) + .await; + self.build_available_models(catalog.models) + } + .instrument(tracing::info_span!( + "list_models", + refresh_strategy = %refresh_strategy + )) + .await + } + /// Return the active raw model catalog, refreshing according to the specified strategy. async fn raw_model_catalog(&self, refresh_strategy: RefreshStrategy) -> ModelsResponse; + /// Return the active raw model catalog, attributing any network refresh to `http_state_surface`. + async fn raw_model_catalog_for_surface( + &self, + refresh_strategy: RefreshStrategy, + _http_state_surface: HttpStateSurface, + ) -> ModelsResponse { + self.raw_model_catalog(refresh_strategy).await + } + /// Return the current in-memory remote model catalog without refreshing or loading cache state. async fn get_remote_models(&self) -> Vec; @@ -157,6 +187,30 @@ pub trait ModelsManager: fmt::Debug + Send + Sync { .await } + /// Get the model identifier to use, attributing any network refresh to `http_state_surface`. + async fn get_default_model_for_surface( + &self, + model: &Option, + refresh_strategy: RefreshStrategy, + http_state_surface: HttpStateSurface, + ) -> String { + async move { + if let Some(model) = model.as_ref() { + return model.to_string(); + } + default_model_from_available( + self.list_models_for_surface(refresh_strategy, http_state_surface) + .await, + ) + } + .instrument(tracing::info_span!( + "get_default_model", + model.provided = model.is_some(), + refresh_strategy = %refresh_strategy + )) + .await + } + // todo(aibrahim): look if we can tighten it to pub(crate) /// Look up model metadata, applying remote overrides and config adjustments. async fn get_model_info(&self, model: &str, config: &ModelsManagerConfig) -> ModelInfo { @@ -172,6 +226,15 @@ pub trait ModelsManager: fmt::Debug + Send + Sync { /// /// Uses `Online` strategy to fetch latest models when ETags differ. async fn refresh_if_new_etag(&self, etag: String); + + /// Refresh models for a new ETag, attributing any network refresh to `http_state_surface`. + async fn refresh_if_new_etag_for_surface( + &self, + etag: String, + _http_state_surface: HttpStateSurface, + ) { + self.refresh_if_new_etag(etag).await; + } } /// Shared model manager handle used across runtime services. @@ -235,6 +298,22 @@ impl ModelsManager for OpenAiModelsManager { } } + async fn raw_model_catalog_for_surface( + &self, + refresh_strategy: RefreshStrategy, + http_state_surface: HttpStateSurface, + ) -> ModelsResponse { + if let Err(err) = self + .refresh_available_models_for_surface(refresh_strategy, http_state_surface) + .await + { + error!("failed to refresh available models: {err}"); + } + ModelsResponse { + models: self.get_remote_models().await, + } + } + async fn get_remote_models(&self) -> Vec { self.remote_models.read().await.clone() } @@ -252,6 +331,26 @@ impl ModelsManager for OpenAiModelsManager { } async fn refresh_if_new_etag(&self, etag: String) { + self.refresh_if_new_etag_inner(etag, /*http_state_surface*/ None) + .await; + } + + async fn refresh_if_new_etag_for_surface( + &self, + etag: String, + http_state_surface: HttpStateSurface, + ) { + self.refresh_if_new_etag_inner(etag, Some(http_state_surface)) + .await; + } +} + +impl OpenAiModelsManager { + async fn refresh_if_new_etag_inner( + &self, + etag: String, + http_state_surface: Option, + ) { let current_etag = self.get_etag().await; if current_etag.clone().is_some() && current_etag.as_deref() == Some(etag.as_str()) { if let Err(err) = self.cache_manager.renew_cache_ttl().await { @@ -259,15 +358,34 @@ impl ModelsManager for OpenAiModelsManager { } return; } - if let Err(err) = self.refresh_available_models(RefreshStrategy::Online).await { + if let Err(err) = self + .refresh_available_models_inner(RefreshStrategy::Online, http_state_surface) + .await + { error!("failed to refresh available models: {err}"); } } -} -impl OpenAiModelsManager { /// Refresh available models according to the specified strategy. async fn refresh_available_models(&self, refresh_strategy: RefreshStrategy) -> CoreResult<()> { + self.refresh_available_models_inner(refresh_strategy, /*http_state_surface*/ None) + .await + } + + async fn refresh_available_models_for_surface( + &self, + refresh_strategy: RefreshStrategy, + http_state_surface: HttpStateSurface, + ) -> CoreResult<()> { + self.refresh_available_models_inner(refresh_strategy, Some(http_state_surface)) + .await + } + + async fn refresh_available_models_inner( + &self, + refresh_strategy: RefreshStrategy, + http_state_surface: Option, + ) -> CoreResult<()> { if !self.should_refresh_models().await { if matches!( refresh_strategy, @@ -291,18 +409,24 @@ impl OpenAiModelsManager { return Ok(()); } info!("models cache: cache miss, fetching remote models"); - self.fetch_and_update_models().await + self.fetch_and_update_models(http_state_surface).await } RefreshStrategy::Online => { // Always fetch from network - self.fetch_and_update_models().await + self.fetch_and_update_models(http_state_surface).await } } } - async fn fetch_and_update_models(&self) -> CoreResult<()> { + async fn fetch_and_update_models( + &self, + http_state_surface: Option, + ) -> CoreResult<()> { let client_version = crate::client_version_to_whole(); - let (models, etag) = self.endpoint_client.list_models(&client_version).await?; + let (models, etag) = self + .endpoint_client + .list_models(&client_version, http_state_surface) + .await?; self.apply_remote_models(models.clone()).await; *self.etag.write().await = etag.clone(); self.cache_manager diff --git a/codex-rs/models-manager/src/manager_tests.rs b/codex-rs/models-manager/src/manager_tests.rs index ede7cfe795..e3f5b29ee3 100644 --- a/codex-rs/models-manager/src/manager_tests.rs +++ b/codex-rs/models-manager/src/manager_tests.rs @@ -2,6 +2,7 @@ use super::*; use crate::ModelsManagerConfig; use chrono::Utc; use codex_app_server_protocol::AuthMode; +use codex_http_state::HttpStateSurface; use codex_login::AuthCredentialsStoreMode; use codex_login::AuthManager; use codex_login::CodexAuth; @@ -76,6 +77,7 @@ struct TestModelsEndpoint { uses_codex_backend: bool, responses: Mutex>>, fetch_count: AtomicUsize, + surfaces: Mutex>>, } impl TestModelsEndpoint { @@ -85,6 +87,7 @@ impl TestModelsEndpoint { uses_codex_backend: true, responses: Mutex::new(responses.into()), fetch_count: AtomicUsize::new(0), + surfaces: Mutex::new(Vec::new()), }) } @@ -94,12 +97,20 @@ impl TestModelsEndpoint { uses_codex_backend: false, responses: Mutex::new(responses.into()), fetch_count: AtomicUsize::new(0), + surfaces: Mutex::new(Vec::new()), }) } fn fetch_count(&self) -> usize { self.fetch_count.load(Ordering::SeqCst) } + + fn surfaces(&self) -> Vec> { + self.surfaces + .lock() + .expect("surfaces lock should not be poisoned") + .clone() + } } #[derive(Debug)] @@ -157,8 +168,13 @@ impl ModelsEndpointClient for TestModelsEndpoint { async fn list_models( &self, _client_version: &str, + http_state_surface: Option, ) -> CoreResult<(Vec, Option)> { self.fetch_count.fetch_add(1, Ordering::SeqCst); + self.surfaces + .lock() + .expect("surfaces lock should not be poisoned") + .push(http_state_surface); let models = self .responses .lock() @@ -333,6 +349,22 @@ async fn get_model_info_rejects_multi_segment_namespace_suffix_matching() { assert!(model_info.used_fallback_model_metadata); } +#[tokio::test] +async fn list_models_for_surface_attributes_network_refresh() { + let codex_home = tempdir().expect("temp dir"); + let endpoint = TestModelsEndpoint::new(Vec::new()); + let manager = openai_manager_for_tests(codex_home.path().to_path_buf(), endpoint.clone()); + + manager + .list_models_for_surface(RefreshStrategy::Online, HttpStateSurface::CodexDesktop) + .await; + + assert_eq!( + endpoint.surfaces(), + vec![Some(HttpStateSurface::CodexDesktop)] + ); +} + #[tokio::test] async fn refresh_available_models_sorts_by_priority() { let remote_models = vec![ @@ -500,6 +532,7 @@ async fn refresh_available_models_keeps_merging_for_api_auth() { uses_codex_backend: false, responses: Mutex::new(vec![remote_models.clone()].into()), fetch_count: AtomicUsize::new(0), + surfaces: Mutex::new(Vec::new()), }); let manager = openai_manager_for_tests_with_auth( codex_home.path().to_path_buf(), @@ -733,6 +766,7 @@ impl ModelsEndpointClient for TestAuthAwareModelsEndpoint { async fn list_models( &self, _client_version: &str, + _http_state_surface: Option, ) -> CoreResult<(Vec, Option)> { self.fetch_count.fetch_add(1, Ordering::SeqCst); let models = self