From cb00a67569450aea2fdb1b1f31fc042fc0b5416f Mon Sep 17 00:00:00 2001 From: cooper-oai Date: Tue, 20 Oct 2026 17:22:03 +0000 Subject: [PATCH] Add runtime assertion authentication modes and an auth provider (#55636) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## What changed Add `RuntimeAuthProvider` to resolve request authentication through the runtime auth manager. Support `dual_header` mode, which sends a bearer credential alongside `x-openai-agent-assertion`, and `assertion_only` mode, which uses `Authorization: AgentAssertion …`. Replace stale authentication headers and sign a fresh assertion for each request attempt. Propagate registration failures without falling back to credential-only authentication, preserving policy errors and classifying transport failures, rate limits, and server errors as transient. Add the disabled-by-default `features.runtime_agent_instances` configuration with independent `enabled` and `mode` settings, schema support, and structured override merging. Selecting a mode alone does not enable the feature. ## Testing Add coverage for both modes with API key and ChatGPT credentials, stale header replacement, fresh assertions, credential refresh without runtime replacement, registration failures, and configuration parsing and merging. GitOrigin-RevId: 9e93d25762600f0ff592c5248fe4c8c2c8182c8d --- codex-rs/Cargo.lock | 2 + codex-rs/config/src/merge.rs | 6 +- codex-rs/config/src/merge_tests.rs | 57 ++++--- codex-rs/config/src/schema.rs | 9 + codex-rs/core/config.schema.json | 46 ++++++ codex-rs/features/src/lib.rs | 20 +++ codex-rs/features/src/runtime_auth.rs | 31 ++++ codex-rs/features/src/tests.rs | 22 +++ codex-rs/runtime-agent-instance/Cargo.toml | 2 + codex-rs/runtime-agent-instance/src/auth.rs | 133 ++++++++++++++- codex-rs/runtime-agent-instance/src/lib.rs | 1 + .../src/manager_tests.rs | 154 ++++++++++++++++-- 12 files changed, 444 insertions(+), 39 deletions(-) create mode 100644 codex-rs/features/src/runtime_auth.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index dc8951b23d..91f53aa981 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -4581,6 +4581,8 @@ version = "0.0.0" dependencies = [ "base64 0.22.1", "codex-agent-identity", + "codex-api", + "codex-features", "codex-http-client", "codex-login", "codex-utils-redacted-string", diff --git a/codex-rs/config/src/merge.rs b/codex-rs/config/src/merge.rs index ef67c26f72..70855f98f2 100644 --- a/codex-rs/config/src/merge.rs +++ b/codex-rs/config/src/merge.rs @@ -70,7 +70,11 @@ pub fn is_structured_feature_path>(path: &[S]) -> bool { features.as_ref() == "features" && matches!( feature.as_ref(), - "code_mode" | "multi_agent_v2" | "network_proxy" | "sleep_tool" + "code_mode" + | "multi_agent_v2" + | "network_proxy" + | "sleep_tool" + | "runtime_agent_instances" ) } diff --git a/codex-rs/config/src/merge_tests.rs b/codex-rs/config/src/merge_tests.rs index 2c09762751..e5e04a51a7 100644 --- a/codex-rs/config/src/merge_tests.rs +++ b/codex-rs/config/src/merge_tests.rs @@ -295,33 +295,38 @@ fn code_mode_overrides_preserve_enabled_state_and_cell_overhead() { } #[test] -fn sleep_tool_overrides_preserve_disabled_state_and_mode() { - for feature_path in ["features", "profiles.work.features"] { - let disabled = ( - format!("{feature_path}.sleep_tool"), - TomlValue::Boolean(false), - ); - let mode = ( - format!("{feature_path}.sleep_tool.mode"), - TomlValue::String("always_on".to_string()), - ); - let expected = parse_toml(&format!( - "[{feature_path}.sleep_tool]\nenabled = false\nmode = \"always_on\"\n", - )); - for overrides in [vec![disabled.clone(), mode.clone()], vec![mode, disabled]] { - assert_eq!(crate::build_cli_overrides_layer(&overrides), expected); - } +fn structured_mode_overrides_preserve_disabled_state_and_mode() { + for (feature, selected_mode) in [ + ("sleep_tool", "always_on"), + ("runtime_agent_instances", "assertion_only"), + ] { + for feature_path in ["features", "profiles.work.features"] { + let disabled = ( + format!("{feature_path}.{feature}"), + TomlValue::Boolean(false), + ); + let mode = ( + format!("{feature_path}.{feature}.mode"), + TomlValue::String(selected_mode.to_string()), + ); + let expected = parse_toml(&format!( + "[{feature_path}.{feature}]\nenabled = false\nmode = \"{selected_mode}\"\n", + )); + for overrides in [vec![disabled.clone(), mode.clone()], vec![mode, disabled]] { + assert_eq!(crate::build_cli_overrides_layer(&overrides), expected); + } - let boolean_layer = parse_toml(&format!("[{feature_path}]\nsleep_tool = false\n")); - let table_layer = parse_toml(&format!( - "[{feature_path}.sleep_tool]\nmode = \"always_on\"\n", - )); - for (mut base, overlay) in [ - (boolean_layer.clone(), table_layer.clone()), - (table_layer, boolean_layer), - ] { - merge_toml_values(&mut base, &overlay); - assert_eq!(base, expected); + let boolean_layer = parse_toml(&format!("[{feature_path}]\n{feature} = false\n")); + let table_layer = parse_toml(&format!( + "[{feature_path}.{feature}]\nmode = \"{selected_mode}\"\n", + )); + for (mut base, overlay) in [ + (boolean_layer.clone(), table_layer.clone()), + (table_layer, boolean_layer), + ] { + merge_toml_values(&mut base, &overlay); + assert_eq!(base, expected); + } } } } diff --git a/codex-rs/config/src/schema.rs b/codex-rs/config/src/schema.rs index 770e73ee5c..f7d6b6a7d2 100644 --- a/codex-rs/config/src/schema.rs +++ b/codex-rs/config/src/schema.rs @@ -153,6 +153,15 @@ pub fn features_schema(schema_gen: &mut SchemaGenerator) -> Schema { ); continue; } + if feature.id == codex_features::Feature::RuntimeAgentInstances { + validation.properties.insert( + feature.key.to_string(), + schema_gen.subschema_for::>(), + ); + continue; + } if feature.id == codex_features::Feature::AppsMcpPathOverride { validation.properties.insert( feature.key.to_string(), diff --git a/codex-rs/core/config.schema.json b/codex-rs/core/config.schema.json index 79fae8e37b..bc6d091f69 100644 --- a/codex-rs/core/config.schema.json +++ b/codex-rs/core/config.schema.json @@ -1162,6 +1162,9 @@ "rollout_budget": { "$ref": "#/definitions/FeatureToml_for_RolloutBudgetConfigToml" }, + "runtime_agent_instances": { + "$ref": "#/definitions/FeatureToml_for_RuntimeAgentInstancesConfigToml" + }, "runtime_metrics": { "type": "boolean" }, @@ -1719,6 +1722,16 @@ } ] }, + "FeatureToml_for_RuntimeAgentInstancesConfigToml": { + "anyOf": [ + { + "type": "boolean" + }, + { + "$ref": "#/definitions/RuntimeAgentInstancesConfigToml" + } + ] + }, "FeatureToml_for_SleepToolConfigToml": { "anyOf": [ { @@ -4274,6 +4287,36 @@ }, "type": "object" }, + "RuntimeAgentInstancesConfigToml": { + "additionalProperties": false, + "properties": { + "enabled": { + "type": "boolean" + }, + "mode": { + "$ref": "#/definitions/RuntimeAuthMode" + } + }, + "type": "object" + }, + "RuntimeAuthMode": { + "oneOf": [ + { + "description": "Send the actor credential and a runtime assertion together during migration.", + "enum": [ + "dual_header" + ], + "type": "string" + }, + { + "description": "Send only the runtime assertion as the request credential.", + "enum": [ + "assertion_only" + ], + "type": "string" + } + ] + }, "SandboxMode": { "enum": [ "read-only", @@ -7574,6 +7617,9 @@ "rollout_budget": { "$ref": "#/definitions/FeatureToml_for_RolloutBudgetConfigToml" }, + "runtime_agent_instances": { + "$ref": "#/definitions/FeatureToml_for_RuntimeAgentInstancesConfigToml" + }, "runtime_metrics": { "type": "boolean" }, diff --git a/codex-rs/features/src/lib.rs b/codex-rs/features/src/lib.rs index dd79cedce0..d3a1183215 100644 --- a/codex-rs/features/src/lib.rs +++ b/codex-rs/features/src/lib.rs @@ -16,6 +16,7 @@ use toml::Table; mod feature_configs; mod legacy; +mod runtime_auth; pub use feature_configs::AsyncCompactionConfigToml; pub use feature_configs::CodeModeConfigToml; pub use feature_configs::CodeModeHostConfigToml; @@ -42,6 +43,8 @@ pub use feature_configs::TokenBudgetConfigToml; pub use feature_configs::ToolRegistryConfigToml; use legacy::LegacyFeatureToggles; pub use legacy::legacy_feature_keys; +pub use runtime_auth::RuntimeAgentInstancesConfigToml; +pub use runtime_auth::RuntimeAuthMode; /// High-level lifecycle stage for a feature. #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -442,6 +445,8 @@ pub enum Feature { RetainClientToolOutputs, /// Use Agent Identity for ChatGPT-authenticated sessions. UseAgentIdentity, + /// Register an ephemeral runtime for OpenAI request authentication. + RuntimeAgentInstances, /// Enable workspace dependency support. WorkspaceDependencies, @@ -904,6 +909,8 @@ pub struct FeaturesToml { pub current_time_reminder: Option>, #[serde(default, skip_serializing_if = "Option::is_none")] pub sleep_tool: Option>, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub runtime_agent_instances: Option>, #[serde(default, rename = "apps_mcp_path_override", skip_serializing)] #[schemars(skip)] removed_apps_mcp_path_override: Option>, @@ -980,6 +987,13 @@ impl FeaturesToml { if let Some(enabled) = self.sleep_tool.as_ref().and_then(FeatureToml::enabled) { entries.insert(Feature::SleepTool.key().to_string(), enabled); } + if let Some(enabled) = self + .runtime_agent_instances + .as_ref() + .and_then(FeatureToml::enabled) + { + entries.insert(Feature::RuntimeAgentInstances.key().to_string(), enabled); + } entries } } @@ -2132,6 +2146,12 @@ pub const FEATURES: &[FeatureSpec] = &[ stage: Stage::UnderDevelopment, default_enabled: false, }, + FeatureSpec { + id: Feature::RuntimeAgentInstances, + key: "runtime_agent_instances", + stage: Stage::UnderDevelopment, + default_enabled: false, + }, FeatureSpec { id: Feature::WorkspaceDependencies, key: "workspace_dependencies", diff --git a/codex-rs/features/src/runtime_auth.rs b/codex-rs/features/src/runtime_auth.rs new file mode 100644 index 0000000000..be7501489f --- /dev/null +++ b/codex-rs/features/src/runtime_auth.rs @@ -0,0 +1,31 @@ +//! Configuration for runtime assertions; enabling the feature is separate from selecting its mode. + +use crate::FeatureConfig; +use schemars::JsonSchema; +use serde::Deserialize; +use serde::Serialize; + +#[derive(Serialize, Deserialize, Debug, Clone, Copy, Default, PartialEq, Eq, JsonSchema)] +#[serde(rename_all = "snake_case")] +pub enum RuntimeAuthMode { + /// Send the actor credential and a runtime assertion together during migration. + #[default] + DualHeader, + /// Send only the runtime assertion as the request credential. + AssertionOnly, +} + +#[derive(Serialize, Deserialize, Debug, Clone, Default, PartialEq, Eq, JsonSchema)] +#[serde(deny_unknown_fields)] +pub struct RuntimeAgentInstancesConfigToml { + #[serde(skip_serializing_if = "Option::is_none")] + pub enabled: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub mode: Option, +} + +impl FeatureConfig for RuntimeAgentInstancesConfigToml { + fn enabled(&self) -> Option { + self.enabled + } +} diff --git a/codex-rs/features/src/tests.rs b/codex-rs/features/src/tests.rs index 0fdfebc50a..a0fd13decf 100644 --- a/codex-rs/features/src/tests.rs +++ b/codex-rs/features/src/tests.rs @@ -29,6 +29,28 @@ fn sleep_tool_config_rejects_unknown_mode() { assert!(toml::from_str::("[sleep_tool]\nmode = 'off'").is_err()); } +#[test] +fn runtime_auth_config_accepts_modes_without_implicitly_enabling_the_feature() { + for mode in ["dual_header", "assertion_only"] { + let config: FeaturesToml = + toml::from_str(&format!("[runtime_agent_instances]\nmode = '{mode}'")).unwrap(); + let mut features = Features::with_defaults(); + features.apply_toml(&config); + assert!(!features.enabled(Feature::RuntimeAgentInstances)); + + let config: FeaturesToml = toml::from_str(&format!( + "[runtime_agent_instances]\nenabled = true\nmode = '{mode}'" + )) + .unwrap(); + features.apply_toml(&config); + assert!(features.enabled(Feature::RuntimeAgentInstances)); + } + assert!( + toml::from_str::("[runtime_agent_instances]\nmode = 'credential_only'") + .is_err() + ); +} + #[test] fn under_development_features_are_disabled_by_default() { for spec in crate::FEATURES { diff --git a/codex-rs/runtime-agent-instance/Cargo.toml b/codex-rs/runtime-agent-instance/Cargo.toml index 747568271f..1c000678ab 100644 --- a/codex-rs/runtime-agent-instance/Cargo.toml +++ b/codex-rs/runtime-agent-instance/Cargo.toml @@ -13,6 +13,8 @@ workspace = true [dependencies] base64 = { workspace = true } codex-agent-identity = { workspace = true } +codex-api = { workspace = true } +codex-features = { workspace = true } codex-http-client = { workspace = true } codex-login = { workspace = true } codex-utils-redacted-string = { workspace = true } diff --git a/codex-rs/runtime-agent-instance/src/auth.rs b/codex-rs/runtime-agent-instance/src/auth.rs index 67d15b6f44..e3214c6f63 100644 --- a/codex-rs/runtime-agent-instance/src/auth.rs +++ b/codex-rs/runtime-agent-instance/src/auth.rs @@ -1,11 +1,27 @@ -//! Credentials and transport selected with the runtime for one request. +//! Apply runtime auth from the selected request credentials. No credential-only fallback is allowed. +use crate::ActorInfo; +use crate::RegistrationError; use crate::RuntimeAgentInstance; +use crate::RuntimeAuthError; +use crate::RuntimeAuthManager; +use codex_api::AuthError; +use codex_api::AuthHeadersFuture; +use codex_api::AuthProvider; +use codex_api::AuthProviderFuture; +use codex_features::RuntimeAuthMode; use codex_http_client::HttpClientFactory; +use codex_http_client::Request; use codex_login::AuthChangeState; use codex_login::CodexAuth; +use http::HeaderMap; +use http::HeaderValue; +use http::StatusCode; +use http::header::AUTHORIZATION; use std::sync::Arc; +const AGENT_ASSERTION: &str = "x-openai-agent-assertion"; + /// Request-scoped credentials and transport selected with their matching runtime. #[derive(Clone)] pub struct RuntimeAuth { @@ -25,6 +41,50 @@ impl std::fmt::Debug for RuntimeAuth { } } +/// Auth adapter for OpenAI clients. Share its manager across requests and threads; +/// never select this provider for a third-party endpoint. +pub struct RuntimeAuthProvider { + manager: Arc, + mode: RuntimeAuthMode, +} + +impl RuntimeAuthProvider { + pub fn new(manager: Arc, mode: RuntimeAuthMode) -> Self { + Self { manager, mode } + } +} + +impl AuthProvider for RuntimeAuthProvider { + fn add_auth_headers(&self, _headers: &mut HeaderMap) -> Result<(), AuthError> { + Err(AuthError::Build( + "authentication requires async header resolution".to_string(), + )) + } + + fn resolve_auth_headers(&self) -> AuthHeadersFuture<'_> { + Box::pin(async move { + let auth = self.manager.auth().await?; + auth.auth_headers(self.mode) + }) + } + + fn apply_auth(&self, mut request: Request) -> AuthProviderFuture<'_> { + Box::pin(async move { + let headers = self.resolve_auth_headers().await?; + for name in [ + AUTHORIZATION.as_str(), + AGENT_ASSERTION, + "chatgpt-account-id", + "x-openai-fedramp", + ] { + request.headers.remove(name); + } + request.headers.extend(headers); + Ok(request) + }) + } +} + impl RuntimeAuth { pub fn runtime(&self) -> &Arc { &self.runtime @@ -41,4 +101,75 @@ impl RuntimeAuth { pub fn change_state(&self) -> AuthChangeState { self.change_state } + + fn auth_headers(&self, mode: RuntimeAuthMode) -> Result { + let assertion = self + .runtime() + .assertion() + .map_err(|_| AuthError::Build("failed to sign runtime assertion".to_string()))?; + let mut headers = HeaderMap::new(); + match mode { + RuntimeAuthMode::DualHeader => { + let token = self.auth().get_token().map_err(|_| { + AuthError::Build("runtime actor credential is unavailable".to_string()) + })?; + headers.insert(AUTHORIZATION, secret_header(&format!("Bearer {token}"))?); + headers.insert(AGENT_ASSERTION, secret_header(assertion.as_str())?); + } + RuntimeAuthMode::AssertionOnly => { + headers.insert( + AUTHORIZATION, + secret_header(&format!("AgentAssertion {}", assertion.as_str()))?, + ); + } + } + if let ActorInfo::Chatgpt { account_id, .. } = self.runtime().actor().info() { + headers.insert( + "chatgpt-account-id", + HeaderValue::from_str(account_id) + .map_err(|_| AuthError::Build("invalid runtime account header".to_string()))?, + ); + } + if self.auth().is_fedramp_account() { + headers.insert("x-openai-fedramp", HeaderValue::from_static("true")); + } + Ok(headers) + } +} + +fn secret_header(value: &str) -> Result { + let mut header = HeaderValue::from_str(value) + .map_err(|_| AuthError::Build("invalid runtime authentication header".to_string()))?; + header.set_sensitive(true); + Ok(header) +} + +impl From for AuthError { + fn from(error: RuntimeAuthError) -> Self { + let transient = match &error { + RuntimeAuthError::Policy(denied) + | RuntimeAuthError::Registration(RegistrationError::Policy(denied)) => { + return Self::Policy(*denied); + } + RuntimeAuthError::Registration(RegistrationError::Transport) => true, + RuntimeAuthError::Registration(RegistrationError::Rejected { status, .. }) => { + status.is_server_error() || *status == StatusCode::TOO_MANY_REQUESTS + } + RuntimeAuthError::Unavailable + | RuntimeAuthError::Unsupported + | RuntimeAuthError::Changed + | RuntimeAuthError::RefreshFailed + | RuntimeAuthError::KeyGeneration + | RuntimeAuthError::Registration( + RegistrationError::InvalidEndpoint + | RegistrationError::InvalidRequest + | RegistrationError::InvalidResponse, + ) => false, + }; + if transient { + Self::Transient(error.to_string()) + } else { + Self::Build(error.to_string()) + } + } } diff --git a/codex-rs/runtime-agent-instance/src/lib.rs b/codex-rs/runtime-agent-instance/src/lib.rs index aaf258b63d..db25096162 100644 --- a/codex-rs/runtime-agent-instance/src/lib.rs +++ b/codex-rs/runtime-agent-instance/src/lib.rs @@ -9,6 +9,7 @@ mod runtime; pub use actor::AgentHost; pub use actor::AuthenticatedActor; +pub use auth::RuntimeAuthProvider; pub use bootstrap::ActorInfo; pub use bootstrap::AgentHostAuthorization; pub use bootstrap::AuthenticatedActorAuthorization; diff --git a/codex-rs/runtime-agent-instance/src/manager_tests.rs b/codex-rs/runtime-agent-instance/src/manager_tests.rs index e4b2108b9c..e0805f5425 100644 --- a/codex-rs/runtime-agent-instance/src/manager_tests.rs +++ b/codex-rs/runtime-agent-instance/src/manager_tests.rs @@ -1,8 +1,12 @@ //! Exercise shared bootstrap against HTTP and the existing external credential-provider boundary. use super::*; +use crate::RuntimeAuthProvider; use crate::RuntimeRole; use base64::Engine; +use codex_api::AuthError; +use codex_api::AuthProvider; +use codex_features::RuntimeAuthMode; use codex_http_client::DestinationPolicy; use codex_http_client::HttpClientFactory; use codex_http_client::NetworkPolicyController; @@ -13,6 +17,10 @@ use codex_login::AuthRouteConfig; use codex_login::ExternalAuth; use codex_login::ExternalAuthFuture; use codex_login::ExternalAuthRefreshContext; +use http::HeaderMap; +use http::HeaderValue; +use http::Method; +use http::header::AUTHORIZATION; use pretty_assertions::assert_eq; use serde_json::Value; use serde_json::json; @@ -177,7 +185,7 @@ async fn policy_recovery_preserves_pending_registration() { .await .unwrap(); assert!(controller.publish(policy.revision(), DestinationPolicy::Unrestricted)); - let manager = RuntimeAuthManager::new( + let manager = Arc::new(RuntimeAuthManager::new( credentials, AgentHostAuthorization::InstallationId { installation_id: Uuid::nil(), @@ -188,25 +196,24 @@ async fn policy_recovery_preserves_pending_registration() { role: RuntimeRole::OrchestratorAndExecutor, }, format!("{}/agent/register", server.uri()).parse().unwrap(), - ); + )); + let provider = RuntimeAuthProvider::new(manager.clone(), RuntimeAuthMode::DualHeader); assert!(matches!( - manager.auth().await, - Err(RuntimeAuthError::Registration( - RegistrationError::Rejected { - status: StatusCode::SERVICE_UNAVAILABLE, - .. - } - )) + provider.resolve_auth_headers().await, + Err(AuthError::Transient(_)) )); policy.invalidate(); assert!(matches!( - manager.auth().await, - Err(RuntimeAuthError::Policy(NetworkPolicyDenied::Unavailable)) + provider.resolve_auth_headers().await, + Err(AuthError::Policy(NetworkPolicyDenied::Unavailable)) )); assert_eq!(server.received_requests().await.unwrap().len(), 1); assert!(controller.publish(policy.revision(), DestinationPolicy::Unrestricted)); + let headers = provider.resolve_auth_headers().await.unwrap(); + assert_eq!(headers["authorization"], "Bearer key-a"); + assert!(headers.contains_key("x-openai-agent-assertion")); let selected = manager.auth().await.unwrap(); let requests = server.received_requests().await.unwrap(); assert_eq!(requests.len(), 2); @@ -321,3 +328,128 @@ async fn a_late_registration_cannot_publish_the_previous_actor() { requests[1].headers["idempotency-key"] ); } + +#[tokio::test] +async fn request_auth_replaces_stale_headers_and_signs_each_attempt() { + for mode in [RuntimeAuthMode::DualHeader, RuntimeAuthMode::AssertionOnly] { + for initial in [ + CodexAuth::from_api_key("key-a"), + oauth(Some("login-a"), "token-a"), + ] { + let server = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(registered) + .mount(&server) + .await; + let (manager, _, _) = setup(&server, initial.clone(), initial.clone()).await; + let provider = RuntimeAuthProvider::new(manager.clone(), mode); + let mut previous_assertion = None; + for _ in 0..2 { + let mut request = codex_http_client::Request::new( + Method::POST, + format!("{}/responses", server.uri()), + ); + for name in [ + "authorization", + "x-openai-agent-assertion", + "chatgpt-account-id", + "x-openai-fedramp", + ] { + request + .headers + .insert(name, HeaderValue::from_static("stale")); + } + request + .headers + .insert("openai-project", HeaderValue::from_static("project-a")); + let mut headers = provider.apply_auth(request).await.unwrap().headers; + let mut expected = HeaderMap::new(); + expected.insert("openai-project", HeaderValue::from_static("project-a")); + if initial.is_chatgpt_auth() { + expected.insert( + "chatgpt-account-id", + HeaderValue::from_static("workspace-a"), + ); + } + let assertion = match mode { + RuntimeAuthMode::DualHeader => { + let bearer = headers.remove(AUTHORIZATION).unwrap(); + assert!(bearer.is_sensitive()); + assert_eq!(bearer, format!("Bearer {}", initial.get_token().unwrap())); + headers.remove("x-openai-agent-assertion").unwrap() + } + RuntimeAuthMode::AssertionOnly => headers.remove(AUTHORIZATION).unwrap(), + }; + assert!(assertion.is_sensitive()); + assert_eq!(headers, expected); + assert_ne!(previous_assertion.as_ref(), Some(&assertion)); + let assertion_text = assertion.to_str().unwrap(); + let assertion_text = match mode { + RuntimeAuthMode::DualHeader => assertion_text, + RuntimeAuthMode::AssertionOnly => { + assertion_text.strip_prefix("AgentAssertion ").unwrap() + } + }; + let segments: Vec<_> = assertion_text.split('.').collect(); + let claims: Value = serde_json::from_slice( + &base64::engine::general_purpose::URL_SAFE_NO_PAD + .decode(segments[2]) + .unwrap(), + ) + .unwrap(); + assert_eq!( + claims["runtime_agent_instance"], + manager.auth().await.unwrap().runtime().id() + ); + previous_assertion = Some(assertion); + } + assert_eq!(server.received_requests().await.unwrap().len(), 1); + } + } +} + +#[tokio::test] +async fn header_auth_uses_refreshed_credentials_without_replacing_the_runtime() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(registered) + .mount(&server) + .await; + let initial = oauth(Some("login-a"), "old"); + let (manager, source, _) = setup(&server, initial.clone(), initial).await; + let provider = RuntimeAuthProvider::new(manager.clone(), RuntimeAuthMode::DualHeader); + let first = provider.resolve_auth_headers().await.unwrap(); + let runtime = manager.auth().await.unwrap(); + let refreshed = oauth(Some("login-a"), "new"); + *source.current.write().unwrap() = refreshed.clone(); + let next = provider.resolve_auth_headers().await.unwrap(); + assert_eq!( + next[AUTHORIZATION], + format!("Bearer {}", refreshed.get_token().unwrap()) + ); + assert_ne!(first[AUTHORIZATION], next[AUTHORIZATION]); + assert!(Arc::ptr_eq( + runtime.runtime(), + manager.auth().await.unwrap().runtime() + )); + assert_eq!(server.received_requests().await.unwrap().len(), 1); +} + +#[tokio::test] +async fn registration_failure_never_returns_credential_only_auth() { + for status in [403, 404, 429, 503] { + let server = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(status)) + .mount(&server) + .await; + let initial = CodexAuth::from_api_key("key-a"); + let (manager, _, _) = setup(&server, initial.clone(), initial).await; + let provider = RuntimeAuthProvider::new(manager, RuntimeAuthMode::DualHeader); + let error = provider.resolve_auth_headers().await.unwrap_err(); + assert_eq!( + matches!(error, AuthError::Transient(_)), + status == 429 || status == 503 + ); + } +}