mirror of
https://github.com/openai/codex.git
synced 2026-10-31 19:28:49 +00:00
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:
+52
-76
@@ -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,
|
||||
|
||||
@@ -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())
|
||||
|
||||
Reference in New Issue
Block a user