Share runtime authentication with model providers and catalogs (#57607)

## Why

Model catalog refreshes need to use the same registered actor and credential owner as model requests when runtime authentication is enabled.

## What changed

- Add `ProviderAuthContext` to carry credentials and optional runtime authentication into model providers and catalog clients.
- Select runtime authentication independently for inference and catalog destinations. Reject unsupported OpenAI origins while preserving provider authentication for other destinations.
- Include API organization and project selectors in runtime registration. Reuse registered runtimes across projects within an organization and replace them when the organization changes.
- Preserve transport restrictions and account bindings, and strip sensitive headers on cross-origin redirects.
- Apply the catalog timeout to authentication and fetching while allowing credential refresh and registration to finish safely.

## Testing

Add coverage for runtime-authenticated catalogs, provider isolation, resource selection, gateway authentication, cross-origin redirects, and timeout recovery that preserves registration and refreshed tokens.

GitOrigin-RevId: c14f1a02159031ea0af1f9b26ea890f8dca7d6ad
This commit is contained in:
cooper-oai
2026-10-28 00:51:29 +00:00
committed by copyberry
parent 939f5decb2
commit 9cce575134
36 changed files with 2111 additions and 217 deletions
+52 -76
View File
@@ -71,11 +71,8 @@ use codex_api::build_session_headers;
use codex_api::create_text_param_for_request;
use codex_api::response_create_client_metadata;
use codex_features::RuntimeAuthMode;
use codex_http_client::CHATGPT_ACCOUNT_ID_HEADER;
use codex_http_client::ClientRouteClass;
use codex_http_client::HttpClientFactory;
use codex_http_client::OPENAI_ORGANIZATION_HEADER;
use codex_http_client::OPENAI_PROJECT_HEADER;
use codex_login::AuthManager;
use codex_login::CodexAuth;
use codex_login::RefreshTokenError;
@@ -88,7 +85,6 @@ use codex_otel::WEBSOCKET_CONTINUATION_COUNT_METRIC;
use codex_otel::current_span_w3c_trace_context;
use codex_protocol::ResponseItemId;
use codex_protocol::auth::AuthMode;
use codex_runtime_agent_instance::ActorInfo;
use codex_runtime_agent_instance::RuntimeAuthManager;
use codex_protocol::ThreadId;
@@ -249,7 +245,7 @@ struct CurrentClientSetup {
redirect_policy: ClientRedirectPolicy,
api_auth: SharedAuthProvider,
agent_identity_telemetry: Option<AgentIdentityTelemetry>,
runtime: Option<(String, ActorInfo)>,
runtime_id: Option<String>,
}
impl CurrentClientSetup {
@@ -676,19 +672,12 @@ impl ModelClient {
mode: RuntimeAuthMode,
) -> Result<Self> {
let info = self.state.provider.info();
if !info.requires_openai_auth
|| info.env_key.is_some()
|| info.experimental_bearer_token.is_some()
|| info.auth.is_some()
|| info.aws.is_some()
|| !self
.auth_manager()
.is_some_and(|auth| manager.uses_credentials(&auth))
{
return Err(CodexErr::UnsupportedOperation(
"runtime auth requires the model provider's selected OpenAI credential".to_string(),
));
}
codex_model_provider::validate_runtime_destination(info.base_url.as_deref())?;
codex_model_provider::validate_runtime_provider(
info,
self.auth_manager().as_ref(),
&manager,
)?;
self.runtime_auth = Some((manager, mode));
Ok(self)
}
@@ -825,9 +814,7 @@ impl ModelClient {
) -> Result<RealtimeWebrtcCallStart> {
// Create the media call over HTTP first, then retain matching auth so realtime can attach
// the server-side control WebSocket to the call id from that HTTP response.
let mut client_setup = self
.provider_client_setup(ClientRouting::ConfiguredProvider)
.await?;
let mut client_setup = self.realtime_client_setup().await?;
if let Some(header_value) = self.generate_attestation_header_for().await {
extra_headers.insert(X_OAI_ATTESTATION_HEADER, header_value);
}
@@ -856,9 +843,7 @@ impl ModelClient {
&self,
mut extra_headers: ApiHeaderMap,
) -> Result<ApiHeaderMap> {
let client_setup = self
.provider_client_setup(ClientRouting::ConfiguredProvider)
.await?;
let client_setup = self.realtime_client_setup().await?;
if let Some(header_value) = self.generate_attestation_header_for().await {
extra_headers.insert(X_OAI_ATTESTATION_HEADER, header_value);
}
@@ -1214,11 +1199,12 @@ impl ModelClient {
/// lockstep when auth/provider resolution changes.
async fn current_client_setup(&self, routing: ClientRouting) -> Result<CurrentClientSetup> {
if let Some((manager, mode)) = &self.runtime_auth {
let selected = manager.auth().await.map_err(|err| {
self.state
.provider
.map_api_error(codex_api::AuthError::from(err).into())
})?;
let selected = codex_model_provider::select_runtime_auth(
manager,
&self.state.provider.info().build_header_map()?,
)
.await
.map_err(|error| self.state.provider.map_api_error(error.into()))?;
let auth = selected.auth().clone();
let auth_revision = Some(selected.change_state().generation);
let auth_owner_generation = Some(selected.change_state().owner_generation);
@@ -1239,16 +1225,19 @@ impl ModelClient {
ClientRouting::ConfiguredProvider => (api_provider, ClientRedirectPolicy::Default),
};
let runtime = selected.runtime();
let actor = runtime.actor().info().clone();
validate_runtime_selectors(&api_provider.headers, &actor)?;
let runtime = Some((runtime.id().to_string(), actor));
codex_model_provider::validate_runtime_selectors(
&api_provider.headers,
runtime.actor().info(),
)?;
let runtime_id = Some(runtime.id().to_string());
let http_client_factory = selected.http_client_factory().clone();
let resolved_auth = self
.state
.provider
.compose_auth(ResolvedProviderAuth {
auth: selected.into_auth_provider(*mode),
auth: selected.clone().into_auth_provider(*mode),
agent_identity_telemetry: None,
runtime: Some(selected),
})
.await?;
tracing::debug!(runtime_auth_mode = ?mode, "selected runtime authentication");
@@ -1261,42 +1250,54 @@ impl ModelClient {
redirect_policy,
api_auth: resolved_auth.auth,
agent_identity_telemetry: None,
runtime,
runtime_id,
});
}
self.provider_client_setup(routing).await
self.provider_client_setup(routing, &self.state.provider)
.await
}
// Realtime keeps provider auth until its runtime integration lands later in the stack.
async fn provider_client_setup(&self, routing: ClientRouting) -> Result<CurrentClientSetup> {
async fn realtime_client_setup(&self) -> Result<CurrentClientSetup> {
let provider = if self.state.provider.runtime_auth().is_some() {
create_model_provider(
self.state.provider.info().clone(),
self.state.provider.auth_manager(),
)
} else {
Arc::clone(&self.state.provider)
};
self.provider_client_setup(ClientRouting::ConfiguredProvider, &provider)
.await
}
async fn provider_client_setup(
&self,
routing: ClientRouting,
provider: &SharedModelProvider,
) -> Result<CurrentClientSetup> {
// Capture before resolving credentials so an account switch during setup cannot label
// an old connection with the new owner's revision.
let auth_owner_generation = self.auth_owner_generation();
let auth_changes = self
.state
.provider
let auth_changes = provider
.auth_manager()
.map(|manager| manager.auth_change_receiver());
loop {
let revision = auth_changes.as_ref().map(|changes| *changes.borrow());
let auth = self.state.provider.auth().await;
let auth = provider.auth().await;
let (api_provider, redirect_policy) = match routing {
ClientRouting::Workspace => {
let resolved = self
.state
.provider
let resolved = provider
.responses_api_provider(&self.state.workspace_routing)
.await?;
(resolved.provider, resolved.redirect_policy)
}
ClientRouting::ConfiguredProvider => (
self.state.provider.api_provider().await?,
provider.api_provider().await?,
ClientRedirectPolicy::Default,
),
};
let resolved_auth = self
.state
.provider
let resolved_auth = provider
.api_auth_for_scope(ProviderAuthScope {
agent_identity_policy: self.agent_identity_policy,
session_source: self.state.session_source.clone(),
@@ -1308,7 +1309,7 @@ impl ModelClient {
.await?;
// Command-backed bearer refreshes cannot change workspace routing. Keep the captured
// revisions so a refresh during setup still invalidates cached WebSocket state.
if self.state.provider.info().auth.is_none() {
if provider.info().auth.is_none() {
if self.auth_owner_generation() != auth_owner_generation {
return Err(std::io::Error::other(
"account changed while preparing model request",
@@ -1329,7 +1330,7 @@ impl ModelClient {
redirect_policy,
api_auth: resolved_auth.auth,
agent_identity_telemetry: resolved_auth.agent_identity_telemetry,
runtime: None,
runtime_id: None,
});
}
}
@@ -1735,7 +1736,7 @@ impl ModelClientSession {
auth_context,
request_route_telemetry: RequestRouteTelemetry::for_endpoint("/responses"),
responses_headers: &responses_headers,
runtime_id: client_setup.runtime.map(|(id, _)| id),
runtime_id: client_setup.runtime_id,
})
.await
{
@@ -2164,7 +2165,7 @@ impl ModelClientSession {
api_auth: client_setup.api_auth,
http_client_factory: client_setup.http_client_factory,
auth_owner_generation: client_setup.auth_owner_generation,
runtime_id: client_setup.runtime.map(|(id, _)| id),
runtime_id: client_setup.runtime_id,
responses_metadata: &websocket_metadata,
auth_context: request_auth_context,
request_route_telemetry: RequestRouteTelemetry::for_endpoint("/responses"),
@@ -2877,31 +2878,6 @@ impl AuthRequestTelemetryContext {
}
}
fn validate_runtime_selectors(headers: &ApiHeaderMap, actor: &ActorInfo) -> Result<()> {
let (account, organization, project) = match actor {
ActorInfo::Chatgpt { account_id, .. } => (Some(account_id.as_str()), None, None),
ActorInfo::Platform {
organization_id,
project_id,
..
} => (None, Some(organization_id.as_str()), project_id.as_deref()),
};
for (name, expected) in [
(CHATGPT_ACCOUNT_ID_HEADER, account),
(OPENAI_ORGANIZATION_HEADER, organization),
(OPENAI_PROJECT_HEADER, project),
] {
if let Some(value) = headers.get(name)
&& value.to_str().ok() != expected
{
return Err(CodexErr::UnsupportedOperation(format!(
"configured {name} does not match the registered runtime"
)));
}
}
Ok(())
}
struct WebsocketConnectParams<'a> {
runtime_id: Option<String>,
session_telemetry: &'a SessionTelemetry,
+19 -5
View File
@@ -128,11 +128,16 @@ async fn selected_runtime_transport_stays_revoked(mode: RuntimeAuthMode) {
server.verify().await;
}
#[test_case::test_case(RuntimeAuthMode::DualHeader; "dual_header")]
#[test_case::test_case(RuntimeAuthMode::AssertionOnly; "assertion_only")]
#[test_case::test_case(RuntimeAuthMode::DualHeader, false; "dual_header")]
#[test_case::test_case(RuntimeAuthMode::AssertionOnly, false; "assertion_only")]
#[test_case::test_case(RuntimeAuthMode::DualHeader, true; "runtime_provider_dual_header")]
#[test_case::test_case(RuntimeAuthMode::AssertionOnly, true; "runtime_provider_assertion_only")]
#[tokio::test]
async fn realtime_keeps_provider_auth_until_runtime_integration(mode: RuntimeAuthMode) {
let (client, _) = model_client_with_counting_attestation(/*include_attestation*/ true);
async fn realtime_keeps_provider_auth_until_runtime_integration(
mode: RuntimeAuthMode,
runtime_provider: bool,
) {
let (mut client, _) = model_client_with_counting_attestation(/*include_attestation*/ true);
let registration = MockServer::start().await;
let manager = client.auth_manager().unwrap();
let expected_authorization = format!(
@@ -140,7 +145,7 @@ async fn realtime_keeps_provider_auth_until_runtime_integration(mode: RuntimeAut
manager.auth().await.unwrap().get_token().unwrap()
);
let runtime = Arc::new(RuntimeAuthManager::new(
manager,
manager.clone(),
AgentHostAuthorization::InstallationId {
installation_id: uuid::Uuid::nil(),
},
@@ -151,6 +156,15 @@ async fn realtime_keeps_provider_auth_until_runtime_integration(mode: RuntimeAut
},
format!("{}/register", registration.uri()).parse().unwrap(),
));
if runtime_provider {
let info = client.state.provider.info().clone();
Arc::get_mut(&mut client.state).unwrap().provider =
codex_model_provider::ProviderAuthContext {
credentials: manager,
runtime_auth: Some((Arc::clone(&runtime), mode)),
}
.model_provider(info);
}
let client = client.with_runtime_auth(runtime, mode).unwrap();
let sideband = client
.realtime_sideband_headers(http::HeaderMap::new())