mirror of
https://github.com/openai/codex.git
synced 2026-10-19 17:28:46 +00:00
Refresh the model catalog before turns after auth changes (#46508)
## Why Credential changes can leave the in-memory model catalog associated with a different identity, causing subsequent turns to use bundled model metadata instead of the active credentials' catalog. ## What changed - Refresh mismatched catalogs before user turns and mailbox-triggered wakeups, resolving lazy command credentials before comparing identities. - Reuse a matching cache or fetch models with a five-second deadline covering auth resolution and cache access. Preserve fallback behavior on failure or timeout; static catalogs need no refresh. - Recheck the active turn after discovery so an interrupted wakeup does not continue starting a turn. ## Testing Add coverage for user and mailbox turns after credential rotation, including switching back after another identity replaces the shared catalog. Verify refreshed context windows and request behavior, fallback on discovery failure or auth timeout, and command-auth token rotation. GitOrigin-RevId: 7e4ad25f1c4476e2023225731924af28c1181767
This commit is contained in:
@@ -1060,6 +1060,15 @@ impl Session {
|
||||
if let Some(service_tier) = service_tier_for_turn {
|
||||
Arc::make_mut(&mut configuration.step_settings).service_tier = Some(service_tier);
|
||||
}
|
||||
if !crate::guardian::is_basic_session_source(&configuration.session_source) {
|
||||
self.services
|
||||
.models_manager
|
||||
.refresh_after_auth_change(
|
||||
self.build_effective_session_config(&configuration)
|
||||
.http_client_factory(),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
let turn_environments = self.activate_turn_environments(&configuration).await;
|
||||
let turn_context = self
|
||||
.new_turn_from_configuration(sub_id, configuration, turn_environments, options)
|
||||
|
||||
@@ -457,6 +457,20 @@ impl Session {
|
||||
Arc::clone(&active_turn.turn_state)
|
||||
};
|
||||
|
||||
self.services
|
||||
.models_manager
|
||||
.refresh_after_auth_change(self.get_config().await.http_client_factory())
|
||||
.await;
|
||||
// A completion-triggered wakeup can be interrupted while discovery waits.
|
||||
if self
|
||||
.active_turn
|
||||
.lock()
|
||||
.await
|
||||
.as_ref()
|
||||
.is_none_or(|turn| !Arc::ptr_eq(&turn.turn_state, &turn_state))
|
||||
{
|
||||
return;
|
||||
}
|
||||
let (input, mut start_options) =
|
||||
self.input_queue.get_pending_input(&self.active_turn).await;
|
||||
if !input.iter().any(
|
||||
|
||||
@@ -1,23 +1,244 @@
|
||||
//! Catalogs must follow auth and provider changes without leaking default billing tiers.
|
||||
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::AtomicBool;
|
||||
use std::sync::atomic::Ordering;
|
||||
use std::time::Duration;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
|
||||
use anyhow::Result;
|
||||
use codex_core::TurnInputRequest;
|
||||
use codex_login::AuthCredentialsStoreMode;
|
||||
use codex_login::AuthHeaders;
|
||||
use codex_login::AuthKeyringBackendKind;
|
||||
use codex_login::CodexAuth;
|
||||
use codex_login::ExternalAuth;
|
||||
use codex_login::ExternalAuthFuture;
|
||||
use codex_login::ExternalAuthRefreshContext;
|
||||
use codex_login::login_with_api_key;
|
||||
use codex_models_manager::bundled_models_response;
|
||||
use codex_models_manager::manager::RefreshStrategy;
|
||||
use codex_protocol::AgentPath;
|
||||
use codex_protocol::openai_models::ModelVisibility;
|
||||
use codex_protocol::openai_models::ModelsResponse;
|
||||
use codex_protocol::openai_models::ToolMode;
|
||||
use codex_protocol::protocol::EventMsg;
|
||||
use codex_protocol::protocol::InterAgentCommunication;
|
||||
use codex_protocol::protocol::Op;
|
||||
use codex_protocol::user_input::UserInput;
|
||||
use core_test_support::responses;
|
||||
use core_test_support::responses::ev_completed;
|
||||
use core_test_support::responses::ev_response_created;
|
||||
use core_test_support::responses::sse;
|
||||
use core_test_support::test_codex::test_codex;
|
||||
use core_test_support::wait_for_event;
|
||||
use http::HeaderMap;
|
||||
use http::HeaderValue;
|
||||
use pretty_assertions::assert_eq;
|
||||
use tempfile::TempDir;
|
||||
use test_case::test_case;
|
||||
use wiremock::Mock;
|
||||
use wiremock::MockServer;
|
||||
use wiremock::ResponseTemplate;
|
||||
use wiremock::matchers::header;
|
||||
use wiremock::matchers::method;
|
||||
use wiremock::matchers::path;
|
||||
|
||||
struct SelectedAuth(CodexAuth);
|
||||
|
||||
impl ExternalAuth for SelectedAuth {
|
||||
fn resolve(&self) -> ExternalAuthFuture<'_, CodexAuth> {
|
||||
Box::pin(async { Ok(self.0.clone()) })
|
||||
}
|
||||
|
||||
fn refresh(&self, _context: ExternalAuthRefreshContext) -> ExternalAuthFuture<'_, CodexAuth> {
|
||||
self.resolve()
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct StallingAuth {
|
||||
stall: AtomicBool,
|
||||
released: CancellationToken,
|
||||
}
|
||||
|
||||
impl ExternalAuth for StallingAuth {
|
||||
fn resolve(&self) -> ExternalAuthFuture<'_, CodexAuth> {
|
||||
Box::pin(async {
|
||||
if self.stall.load(Ordering::SeqCst) {
|
||||
self.released.cancelled().await;
|
||||
}
|
||||
Ok(header_auth("Bearer rotated"))
|
||||
})
|
||||
}
|
||||
|
||||
fn refresh(&self, _context: ExternalAuthRefreshContext) -> ExternalAuthFuture<'_, CodexAuth> {
|
||||
self.resolve()
|
||||
}
|
||||
}
|
||||
|
||||
fn header_auth(token: &'static str) -> CodexAuth {
|
||||
CodexAuth::Headers(AuthHeaders::new(HeaderMap::from_iter([
|
||||
(http::header::AUTHORIZATION, HeaderValue::from_static(token)),
|
||||
(
|
||||
http::header::HeaderName::from_static("chatgpt-account-id"),
|
||||
HeaderValue::from_static("same-account"),
|
||||
),
|
||||
])))
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
enum Input {
|
||||
User,
|
||||
Mail,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
enum RefreshOutcome {
|
||||
Success,
|
||||
Failure,
|
||||
AuthTimeout,
|
||||
}
|
||||
|
||||
#[test_case(Input::User, RefreshOutcome::Success; "user refresh")]
|
||||
#[test_case(Input::User, RefreshOutcome::Failure; "user fallback")]
|
||||
#[test_case(Input::Mail, RefreshOutcome::Success; "mail refresh")]
|
||||
#[test_case(Input::Mail, RefreshOutcome::Failure; "mail fallback")]
|
||||
#[test_case(Input::User, RefreshOutcome::AuthTimeout; "auth timeout")]
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn auth_rotation_refreshes_before_turn_with_best_effort(
|
||||
input: Input,
|
||||
outcome: RefreshOutcome,
|
||||
) -> Result<()> {
|
||||
let server = MockServer::start().await;
|
||||
let catalog = bundled_models_response()?;
|
||||
let mut model = catalog
|
||||
.models
|
||||
.iter()
|
||||
.find(|model| model.slug == "gpt-5.5")
|
||||
.cloned()
|
||||
.expect("bundled model");
|
||||
let fallback_context_window = model.usable_context_window();
|
||||
responses::mount_models_once(&server, catalog).await;
|
||||
let test = test_codex()
|
||||
.with_auth(header_auth("Bearer rotated"))
|
||||
.with_model(&model.slug)
|
||||
.with_config(|config| config.model_provider.request_max_retries = Some(0))
|
||||
.build_with_auto_env(&server)
|
||||
.await?;
|
||||
model.visibility = ModelVisibility::List;
|
||||
model.context_window = Some(481_000);
|
||||
model.max_context_window = Some(481_000);
|
||||
model.effective_context_window_percent = 100;
|
||||
model.use_responses_lite = true;
|
||||
model.tool_mode = Some(ToolMode::CodeModeOnly);
|
||||
let succeeds = matches!(outcome, RefreshOutcome::Success);
|
||||
let stalls_auth = matches!(outcome, RefreshOutcome::AuthTimeout);
|
||||
let expected_context_window = if succeeds {
|
||||
model.usable_context_window()
|
||||
} else {
|
||||
fallback_context_window
|
||||
};
|
||||
let template = if succeeds {
|
||||
ResponseTemplate::new(/*s*/ 200).set_body_json(ModelsResponse {
|
||||
models: vec![model],
|
||||
})
|
||||
} else {
|
||||
ResponseTemplate::new(/*s*/ 503).set_body_string("unavailable")
|
||||
};
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/v1/models"))
|
||||
.and(header("authorization", "Bearer rotated"))
|
||||
.respond_with(template)
|
||||
.expect(/*r*/ u64::from(!stalls_auth))
|
||||
.mount(&server)
|
||||
.await;
|
||||
let stalled_auth = Arc::new(StallingAuth::default());
|
||||
// The shared picker can replace the catalog between turns. Switching back
|
||||
// must refresh even when this was the last identity used to start a turn.
|
||||
test.thread_manager
|
||||
.auth_manager()
|
||||
.set_external_auth(Arc::new(SelectedAuth(header_auth("Bearer other"))))
|
||||
.await?;
|
||||
responses::mount_models_once(&server, bundled_models_response()?).await;
|
||||
test.thread_manager
|
||||
.get_models_manager()
|
||||
.raw_model_catalog(
|
||||
RefreshStrategy::Online,
|
||||
codex_core::test_support::default_http_client_factory(),
|
||||
)
|
||||
.await;
|
||||
let auth: Arc<dyn ExternalAuth> = if stalls_auth {
|
||||
stalled_auth.clone()
|
||||
} else {
|
||||
Arc::new(SelectedAuth(header_auth("Bearer rotated")))
|
||||
};
|
||||
test.thread_manager
|
||||
.auth_manager()
|
||||
.set_external_auth(auth)
|
||||
.await?;
|
||||
stalled_auth.stall.store(stalls_auth, Ordering::SeqCst);
|
||||
let response = responses::mount_sse_once(&server, responses::sse_completed("done")).await;
|
||||
match input {
|
||||
Input::User => {
|
||||
tokio::time::timeout(
|
||||
Duration::from_secs(/*secs*/ 7),
|
||||
test.codex
|
||||
.start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text {
|
||||
text: "hello".into(),
|
||||
text_elements: Vec::new(),
|
||||
}])),
|
||||
)
|
||||
.await??;
|
||||
}
|
||||
Input::Mail => {
|
||||
test.codex
|
||||
.submit(Op::InterAgentCommunication {
|
||||
communication: InterAgentCommunication::new(
|
||||
AgentPath::root().join("worker").expect("valid path"),
|
||||
AgentPath::root(),
|
||||
Vec::new(),
|
||||
"hello".into(),
|
||||
/*trigger_turn*/ true,
|
||||
),
|
||||
start_options: Default::default(),
|
||||
})
|
||||
.await?;
|
||||
}
|
||||
}
|
||||
let started = wait_for_event(&test.codex, |event| {
|
||||
matches!(event, EventMsg::TurnStarted(_))
|
||||
})
|
||||
.await;
|
||||
let EventMsg::TurnStarted(started) = started else {
|
||||
unreachable!()
|
||||
};
|
||||
assert_eq!(started.model_context_window, expected_context_window);
|
||||
if stalls_auth {
|
||||
stalled_auth.stall.store(/*val*/ false, Ordering::SeqCst);
|
||||
stalled_auth.released.cancel();
|
||||
}
|
||||
wait_for_event(&test.codex, |event| {
|
||||
matches!(event, EventMsg::TurnComplete(_))
|
||||
})
|
||||
.await;
|
||||
let request = response.single_request();
|
||||
assert!(request.body_contains_text("hello"));
|
||||
assert_eq!(
|
||||
request.header("authorization").as_deref(),
|
||||
Some("Bearer rotated")
|
||||
);
|
||||
let body = request.body_json();
|
||||
let uses_responses_lite = body["input"]
|
||||
.as_array()
|
||||
.expect("request input")
|
||||
.iter()
|
||||
.any(|item| item["type"] == "additional_tools");
|
||||
assert_eq!(uses_responses_lite, succeeds);
|
||||
test.codex.submit(Op::Shutdown).await?;
|
||||
test.codex.wait_until_terminated().await;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn auth_and_provider_switches_do_not_reuse_chatgpt_catalog() -> Result<()> {
|
||||
|
||||
@@ -603,7 +603,7 @@ mod tests {
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(ModelsResponse {
|
||||
models: vec![model.clone()],
|
||||
}))
|
||||
.expect(2)
|
||||
.expect(/*r*/ 3)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let mut provider = provider_info_with_command_auth();
|
||||
@@ -634,6 +634,12 @@ mod tests {
|
||||
.find(|candidate| candidate.slug == model.slug),
|
||||
Some(&model)
|
||||
);
|
||||
// The cached identity still matches here; resolving command auth must
|
||||
// detect the next token before deciding whether a refresh is needed.
|
||||
manager
|
||||
.refresh_after_auth_change(HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault))
|
||||
.await;
|
||||
assert_eq!(manager.get_remote_models().await, catalog.models);
|
||||
auth.auth().await;
|
||||
let bundled = codex_models_manager::bundled_models_response().unwrap();
|
||||
assert_eq!(manager.get_remote_models().await, bundled.models);
|
||||
@@ -662,7 +668,7 @@ mod tests {
|
||||
.iter()
|
||||
.map(|request| request.headers["authorization"].to_str().unwrap())
|
||||
.collect::<Vec<_>>(),
|
||||
vec!["Bearer token-2", "Bearer token-6"]
|
||||
vec!["Bearer token-2", "Bearer token-5", "Bearer token-9"]
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -138,6 +138,15 @@ pub trait ModelsManager: fmt::Debug + Send + Sync {
|
||||
http_client_factory: HttpClientFactory,
|
||||
) -> ModelsManagerFuture<'_, ModelsResponse>;
|
||||
|
||||
/// Best-effort refresh when the in-memory catalog belongs to different credentials.
|
||||
/// Static catalogs need no refresh. Failures leave the existing cache/default fallback.
|
||||
fn refresh_after_auth_change(
|
||||
&self,
|
||||
_http_client_factory: HttpClientFactory,
|
||||
) -> ModelsManagerFuture<'_, ()> {
|
||||
Box::pin(std::future::ready(()))
|
||||
}
|
||||
|
||||
/// Return the current in-memory remote model catalog without refreshing or loading cache state.
|
||||
fn get_remote_models(&self) -> ModelsManagerFuture<'_, Vec<ModelInfo>>;
|
||||
|
||||
@@ -345,6 +354,36 @@ impl ModelsManager for OpenAiModelsManager {
|
||||
))
|
||||
}
|
||||
|
||||
fn refresh_after_auth_change(
|
||||
&self,
|
||||
http_client_factory: HttpClientFactory,
|
||||
) -> ModelsManagerFuture<'_, ()> {
|
||||
Box::pin(async move {
|
||||
let refresh = async {
|
||||
// Resolve lazy command credentials before comparing catalog identities.
|
||||
if !self.should_refresh_models().await {
|
||||
return Ok(());
|
||||
}
|
||||
let identity = self.endpoint_client.identity();
|
||||
if identity.is_some() && self.remote_models.read().await.identity == identity {
|
||||
return Ok(());
|
||||
}
|
||||
self.refresh_available_models(
|
||||
RefreshStrategy::OnlineIfUncached,
|
||||
&http_client_factory,
|
||||
)
|
||||
.await
|
||||
};
|
||||
// Include auth resolution and cache access in the best-effort deadline.
|
||||
if !matches!(
|
||||
tokio::time::timeout(Duration::from_secs(/*secs*/ 5), refresh).await,
|
||||
Ok(Ok(()))
|
||||
) {
|
||||
tracing::warn!("model catalog refresh after auth change failed or timed out");
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
fn get_remote_models(&self) -> ModelsManagerFuture<'_, Vec<ModelInfo>> {
|
||||
Box::pin(async move {
|
||||
let entry = self.remote_models.read().await;
|
||||
|
||||
Reference in New Issue
Block a user