From 091008729d5572712ac08964a9602eb71cf31428 Mon Sep 17 00:00:00 2001 From: cooper-oai Date: Tue, 20 Oct 2026 00:41:11 +0000 Subject: [PATCH] Share runtime authentication across requests and credential refreshes (#55413) ## What changed Add `RuntimeAuthManager` to share one registered runtime per credential source and return request-scoped `RuntimeAuth` values that pair the runtime with its credentials and HTTP client factory. - Preserve registration keys and idempotency IDs across retryable failures and refreshes that retain the credential source. - Replace the runtime when the credential source changes, preserve existing request snapshots, and reject late registration responses for a previous source. - Use bounded credential recovery for registration `401` responses and cache `403` denials until the credential source changes. - Expose authenticated actor and agent host information on `RuntimeAgentInstance`. ## Testing Add tests for concurrent registration sharing, retry identity preservation, network policy recovery, credential refresh and replacement, bounded recovery, cached denials, and credential changes during registration. GitOrigin-RevId: ea1e6e10a5939e8c3f7c0c317ccf8a83522c635f --- codex-rs/Cargo.lock | 2 + codex-rs/runtime-agent-instance/Cargo.toml | 3 + codex-rs/runtime-agent-instance/src/actor.rs | 29 ++ codex-rs/runtime-agent-instance/src/auth.rs | 44 +++ codex-rs/runtime-agent-instance/src/lib.rs | 8 + .../runtime-agent-instance/src/manager.rs | 253 ++++++++++++++ .../src/manager_tests.rs | 323 ++++++++++++++++++ .../src/registration.rs | 11 +- .../src/registration_tests.rs | 9 +- .../runtime-agent-instance/src/runtime.rs | 39 ++- 10 files changed, 702 insertions(+), 19 deletions(-) create mode 100644 codex-rs/runtime-agent-instance/src/actor.rs create mode 100644 codex-rs/runtime-agent-instance/src/auth.rs create mode 100644 codex-rs/runtime-agent-instance/src/manager.rs create mode 100644 codex-rs/runtime-agent-instance/src/manager_tests.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 1016bb7ec3..958d6ca2cb 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -4581,6 +4581,7 @@ dependencies = [ "base64 0.22.1", "codex-agent-identity", "codex-http-client", + "codex-login", "codex-utils-redacted-string", "ed25519-dalek", "http 1.4.0", @@ -4588,6 +4589,7 @@ dependencies = [ "rand 0.9.3", "serde", "serde_json", + "tempfile", "thiserror 2.0.18", "tokio", "url", diff --git a/codex-rs/runtime-agent-instance/Cargo.toml b/codex-rs/runtime-agent-instance/Cargo.toml index 1fcd6b99a4..747568271f 100644 --- a/codex-rs/runtime-agent-instance/Cargo.toml +++ b/codex-rs/runtime-agent-instance/Cargo.toml @@ -14,6 +14,7 @@ workspace = true base64 = { workspace = true } codex-agent-identity = { workspace = true } codex-http-client = { workspace = true } +codex-login = { workspace = true } codex-utils-redacted-string = { workspace = true } ed25519-dalek = { workspace = true } http = { workspace = true } @@ -21,11 +22,13 @@ rand = { workspace = true } serde = { workspace = true, features = ["derive"] } serde_json = { workspace = true } thiserror = { workspace = true } +tokio = { workspace = true, features = ["sync"] } url = { workspace = true } uuid = { workspace = true, features = ["serde", "v4"] } zeroize = { workspace = true } [dev-dependencies] pretty_assertions = { workspace = true } +tempfile = { workspace = true } tokio = { workspace = true, features = ["macros", "rt-multi-thread"] } wiremock = { workspace = true } diff --git a/codex-rs/runtime-agent-instance/src/actor.rs b/codex-rs/runtime-agent-instance/src/actor.rs new file mode 100644 index 0000000000..4ef0f0f878 --- /dev/null +++ b/codex-rs/runtime-agent-instance/src/actor.rs @@ -0,0 +1,29 @@ +//! Actor and host information returned by composite registration, without another hydration call. + +use crate::ActorInfo; +use crate::AgentHostAuthorization; + +pub struct AuthenticatedActor { + pub(crate) info: ActorInfo, +} + +impl AuthenticatedActor { + pub fn info(&self) -> &ActorInfo { + &self.info + } +} + +pub struct AgentHost { + pub(crate) id: String, + pub(crate) authorization: AgentHostAuthorization, +} + +impl AgentHost { + pub fn id(&self) -> &str { + &self.id + } + + pub fn authorization(&self) -> &AgentHostAuthorization { + &self.authorization + } +} diff --git a/codex-rs/runtime-agent-instance/src/auth.rs b/codex-rs/runtime-agent-instance/src/auth.rs new file mode 100644 index 0000000000..67d15b6f44 --- /dev/null +++ b/codex-rs/runtime-agent-instance/src/auth.rs @@ -0,0 +1,44 @@ +//! Credentials and transport selected with the runtime for one request. + +use crate::RuntimeAgentInstance; +use codex_http_client::HttpClientFactory; +use codex_login::AuthChangeState; +use codex_login::CodexAuth; +use std::sync::Arc; + +/// Request-scoped credentials and transport selected with their matching runtime. +#[derive(Clone)] +pub struct RuntimeAuth { + pub(crate) runtime: Arc, + pub(crate) auth: CodexAuth, + pub(crate) http_client_factory: HttpClientFactory, + pub(crate) change_state: AuthChangeState, +} + +impl std::fmt::Debug for RuntimeAuth { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("RuntimeAuth") + .field("runtime", &self.runtime) + .field("auth_mode", &self.auth.api_auth_mode()) + .field("change_state", &self.change_state) + .finish_non_exhaustive() + } +} + +impl RuntimeAuth { + pub fn runtime(&self) -> &Arc { + &self.runtime + } + + pub fn auth(&self) -> &CodexAuth { + &self.auth + } + + pub fn http_client_factory(&self) -> &HttpClientFactory { + &self.http_client_factory + } + + pub fn change_state(&self) -> AuthChangeState { + self.change_state + } +} diff --git a/codex-rs/runtime-agent-instance/src/lib.rs b/codex-rs/runtime-agent-instance/src/lib.rs index fb583460ef..aaf258b63d 100644 --- a/codex-rs/runtime-agent-instance/src/lib.rs +++ b/codex-rs/runtime-agent-instance/src/lib.rs @@ -1,9 +1,14 @@ //! Ephemeral OpenAI runtimes: composite registration and per-request assertions. +mod actor; +mod auth; mod bootstrap; +mod manager; mod registration; mod runtime; +pub use actor::AgentHost; +pub use actor::AuthenticatedActor; pub use bootstrap::ActorInfo; pub use bootstrap::AgentHostAuthorization; pub use bootstrap::AuthenticatedActorAuthorization; @@ -12,6 +17,9 @@ pub use bootstrap::RuntimeRegistrationRequest; pub use bootstrap::RuntimeRegistrationResponse; pub use bootstrap::RuntimeRole; +pub use auth::RuntimeAuth; +pub use manager::RuntimeAuthError; +pub use manager::RuntimeAuthManager; pub use registration::DEFAULT_REGISTRATION_URL; pub use registration::RegistrationError; pub use registration::RegistrationErrorCode; diff --git a/codex-rs/runtime-agent-instance/src/manager.rs b/codex-rs/runtime-agent-instance/src/manager.rs new file mode 100644 index 0000000000..5740228eb4 --- /dev/null +++ b/codex-rs/runtime-agent-instance/src/manager.rs @@ -0,0 +1,253 @@ +//! Share one runtime for the current credential source, including bootstrap retries and refreshes. + +use crate::AgentHostAuthorization; +use crate::AuthenticatedActorAuthorization; +use crate::RegistrationError; +use crate::RuntimeAgentInstance; +use crate::RuntimeAuth; +use crate::RuntimeRegistration; +use crate::RuntimeRegistrationClient; +use crate::RuntimeRegistrationMetadata; +use codex_http_client::HttpClientFactory; +use codex_http_client::NetworkPolicyDenied; +use codex_login::AuthChangeState; +use codex_login::AuthManager; +use codex_login::CodexAuth; +use codex_login::RefreshTokenError; +use http::StatusCode; +use std::sync::Arc; +use tokio::sync::Mutex; +use url::Url; + +/// Owned alongside an AuthManager, not created for every thread or request. +pub struct RuntimeAuthManager { + credentials: Arc, + host: AgentHostAuthorization, + metadata: RuntimeRegistrationMetadata, + registration_url: Url, + state: Mutex>, +} + +enum RuntimeState { + Pending { + source_generation: u64, + registration: RuntimeRegistration, + }, + Ready { + source_generation: u64, + runtime: Arc, + }, + Denied { + source_generation: u64, + error: RegistrationError, + }, +} + +impl RuntimeAuthManager { + pub fn new( + credentials: Arc, + host: AgentHostAuthorization, + metadata: RuntimeRegistrationMetadata, + registration_url: Url, + ) -> Self { + Self { + credentials, + host, + metadata, + registration_url, + state: Mutex::new(None), + } + } + + #[expect( + clippy::await_holding_invalid_type, + reason = "runtime initialization serializes registration and refresh for one credential source" + )] + pub async fn auth(&self) -> Result { + // This initializer's lock can span IO. No AuthManager lock is held across registration. + let mut state = self.state.lock().await; + let (auth, http_client_factory) = self + .credentials + .auth_with_http_client_factory() + .await + .ok_or(RuntimeAuthError::Unavailable)?; + let mut recovery = self.credentials.unauthorized_recovery(); + let mut result = self.auth_once(&mut state, auth, http_client_factory).await; + while matches!( + &result, + Err(RuntimeAuthError::Registration( + RegistrationError::Rejected { + status: StatusCode::UNAUTHORIZED, + .. + } + )) + ) && recovery.has_next() + { + recovery.next().await.map_err(|error| match error { + RefreshTokenError::Policy(denied) => RuntimeAuthError::Policy(denied), + RefreshTokenError::Permanent(_) | RefreshTokenError::Transient(_) => { + RuntimeAuthError::RefreshFailed + } + })?; + let (auth, http_client_factory) = self + .credentials + .auth_with_http_client_factory() + .await + .ok_or(RuntimeAuthError::Unavailable)?; + result = self.auth_once(&mut state, auth, http_client_factory).await; + } + result + } + + async fn auth_once( + &self, + state: &mut Option, + auth: CodexAuth, + http_client_factory: HttpClientFactory, + ) -> Result { + let actor_auth = actor_authorization(&auth)?; + let change_state = self.auth_state(&auth)?; + let same_source = match state.as_ref() { + Some( + RuntimeState::Pending { + source_generation, .. + } + | RuntimeState::Ready { + source_generation, .. + } + | RuntimeState::Denied { + source_generation, .. + }, + ) => *source_generation == change_state.source_generation, + None => false, + }; + if !same_source { + *state = Some(RuntimeState::Pending { + registration: RuntimeRegistration::new( + actor_auth.clone(), + self.host.clone(), + self.metadata.clone(), + ) + .map_err(|_| RuntimeAuthError::KeyGeneration)?, + source_generation: change_state.source_generation, + }); + } + let registration = match state.as_mut().ok_or(RuntimeAuthError::Unavailable)? { + RuntimeState::Ready { runtime, .. } => { + return Ok(RuntimeAuth { + runtime: runtime.clone(), + auth, + http_client_factory, + change_state, + }); + } + RuntimeState::Denied { error, .. } => return Err(error.clone().into()), + RuntimeState::Pending { registration, .. } => { + registration.request.actor_auth = actor_auth; + registration + } + }; + let client = RuntimeRegistrationClient::new( + http_client_factory.clone(), + self.registration_url.clone(), + )?; + let result = client.register(registration).await; + if self.auth_state(&auth)?.source_generation != change_state.source_generation { + // A late response must not publish an old identity after a credential switch. + *state = None; + return Err(RuntimeAuthError::Changed); + } + match result { + Ok(info) => { + let runtime = Arc::new(registration.finish(info)); + *state = Some(RuntimeState::Ready { + runtime: runtime.clone(), + source_generation: change_state.source_generation, + }); + Ok(RuntimeAuth { + runtime, + auth, + http_client_factory, + change_state, + }) + } + Err( + error @ RegistrationError::Rejected { + status: StatusCode::FORBIDDEN, + .. + }, + ) => { + *state = Some(RuntimeState::Denied { + source_generation: change_state.source_generation, + error: error.clone(), + }); + Err(error.into()) + } + Err(error) => Err(error.into()), + } + } + + fn auth_state(&self, auth: &CodexAuth) -> Result { + let state = self + .credentials + .auth_change_state_for(auth) + .ok_or(RuntimeAuthError::Changed)?; + if self.credentials.refresh_failure_for_auth(auth).is_some() { + return Err(RuntimeAuthError::RefreshFailed); + } + Ok(state) + } +} + +fn actor_authorization( + auth: &CodexAuth, +) -> Result { + let credential = auth + .get_token() + .map_err(|_| RuntimeAuthError::Unsupported)? + .into(); + Ok(match auth { + CodexAuth::ApiKey(_) => AuthenticatedActorAuthorization::PlatformApiKey { credential }, + CodexAuth::Chatgpt(_) | CodexAuth::ChatgptAuthTokens(_) => { + AuthenticatedActorAuthorization::Oauth { credential } + } + CodexAuth::PersonalAccessToken(_) => { + AuthenticatedActorAuthorization::PersonalAccessToken { credential } + } + CodexAuth::Headers(_) + | CodexAuth::AgentIdentity(_) + | CodexAuth::BedrockApiKey(_) + | CodexAuth::BedrockAccessKeys(_) => return Err(RuntimeAuthError::Unsupported), + }) +} + +#[derive(Debug, thiserror::Error)] +pub enum RuntimeAuthError { + #[error("runtime registration requires an actor credential")] + Unavailable, + #[error("this credential cannot register a runtime; use a supported OpenAI sign-in")] + Unsupported, + #[error("credentials changed during runtime registration; try again")] + Changed, + #[error("runtime registration could not refresh actor credentials")] + RefreshFailed, + #[error("failed to generate runtime signing key")] + KeyGeneration, + #[error(transparent)] + Policy(#[from] NetworkPolicyDenied), + #[error(transparent)] + Registration(RegistrationError), +} + +impl From for RuntimeAuthError { + fn from(error: RegistrationError) -> Self { + match error { + RegistrationError::Policy(denied) => Self::Policy(denied), + error => Self::Registration(error), + } + } +} + +#[cfg(test)] +#[path = "manager_tests.rs"] +mod tests; diff --git a/codex-rs/runtime-agent-instance/src/manager_tests.rs b/codex-rs/runtime-agent-instance/src/manager_tests.rs new file mode 100644 index 0000000000..e4b2108b9c --- /dev/null +++ b/codex-rs/runtime-agent-instance/src/manager_tests.rs @@ -0,0 +1,323 @@ +//! Exercise shared bootstrap against HTTP and the existing external credential-provider boundary. + +use super::*; +use crate::RuntimeRole; +use base64::Engine; +use codex_http_client::DestinationPolicy; +use codex_http_client::HttpClientFactory; +use codex_http_client::NetworkPolicyController; +use codex_http_client::OutboundProxyPolicy; +use codex_login::AuthCredentialsStoreMode; +use codex_login::AuthKeyringBackendKind; +use codex_login::AuthRouteConfig; +use codex_login::ExternalAuth; +use codex_login::ExternalAuthFuture; +use codex_login::ExternalAuthRefreshContext; +use pretty_assertions::assert_eq; +use serde_json::Value; +use serde_json::json; +use std::sync::RwLock; +use std::time::Duration; +use tokio::sync::Notify; +use uuid::Uuid; +use wiremock::Mock; +use wiremock::MockServer; +use wiremock::Request; +use wiremock::ResponseTemplate; +use wiremock::matchers::body_partial_json; +use wiremock::matchers::method; + +struct Credentials { + current: RwLock, + refreshed: CodexAuth, +} + +impl ExternalAuth for Credentials { + fn resolve(&self) -> ExternalAuthFuture<'_, CodexAuth> { + Box::pin(async { Ok(self.current.read().unwrap().clone()) }) + } + + fn refresh(&self, _: ExternalAuthRefreshContext) -> ExternalAuthFuture<'_, CodexAuth> { + Box::pin(async { + *self.current.write().unwrap() = self.refreshed.clone(); + Ok(self.refreshed.clone()) + }) + } +} + +async fn setup( + server: &MockServer, + initial: CodexAuth, + refreshed: CodexAuth, +) -> (Arc, Arc, Arc) { + let credentials = AuthManager::from_auth_for_testing(initial.clone()); + let source = Arc::new(Credentials { + current: RwLock::new(initial), + refreshed, + }); + credentials.set_external_auth(source.clone()).await.unwrap(); + let manager = RuntimeAuthManager::new( + credentials.clone(), + AgentHostAuthorization::InstallationId { + installation_id: Uuid::nil(), + }, + RuntimeRegistrationMetadata { + harness: Some("codex".into()), + harness_version: None, + role: RuntimeRole::OrchestratorAndExecutor, + }, + format!("{}/agent/register", server.uri()).parse().unwrap(), + ); + (Arc::new(manager), source, credentials) +} + +fn registered(request: &Request) -> ResponseTemplate { + let body: Value = request.body_json().unwrap(); + let actor = if body["actor_auth"]["type"] == "oauth" { + json!({"kind": "chatgpt", "user_id": "user-a", "account_id": "workspace-a"}) + } else { + json!({"kind": "platform", "user_id": "user-a", "organization_id": "org-a", "project_id": "project-a"}) + }; + ResponseTemplate::new(200).set_body_json(json!({ + "runtime_agent_instance_id": request.headers["idempotency-key"].to_str().unwrap(), + "agent_host_id": "host-a", "actor": actor, + })) +} + +fn oauth(session: Option<&str>, token: &str) -> CodexAuth { + let claims = json!({ + "iss": "https://auth.openai.com/", "session_id": session, "jti": token, + "https://api.openai.com/auth": {"user_id": "user-a", "chatgpt_account_id": "workspace-a"}, + }); + let payload = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(claims.to_string()); + CodexAuth::from_external_chatgpt_tokens( + &format!("header.{payload}.signature"), + "workspace-a", + /*chatgpt_plan_type*/ None, + ) + .unwrap() +} + +#[tokio::test] +async fn concurrent_requests_share_registration_and_replacement_preserves_in_flight_auth() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(503)) + .up_to_n_times(1) + .with_priority(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .respond_with(registered) + .mount(&server) + .await; + let initial = CodexAuth::from_api_key("key-a"); + let (manager, source, _) = setup(&server, initial.clone(), initial).await; + assert!(manager.auth().await.is_err()); + let (first, second) = tokio::join!(manager.auth(), manager.auth()); + let first = first.unwrap(); + let second = second.unwrap(); + assert!(Arc::ptr_eq(first.runtime(), second.runtime())); + *source.current.write().unwrap() = CodexAuth::from_api_key("key-b"); + let replacement = manager.auth().await.unwrap(); + assert_ne!(first.runtime().id(), replacement.runtime().id()); + assert_eq!(first.auth().get_token().unwrap(), "key-a"); + assert!(!format!("{first:?}").contains("key-a")); + assert_eq!(replacement.auth().get_token().unwrap(), "key-b"); + let requests = server.received_requests().await.unwrap(); + assert_eq!(requests.len(), 3); + assert_eq!(requests[0].body, requests[1].body); + assert_eq!( + requests[0].headers["idempotency-key"], + requests[1].headers["idempotency-key"] + ); + assert_ne!( + requests[1].body_json::().unwrap()["public_key"], + requests[2].body_json::().unwrap()["public_key"] + ); +} + +#[tokio::test] +async fn policy_recovery_preserves_pending_registration() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(503)) + .up_to_n_times(1) + .with_priority(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .respond_with(registered) + .mount(&server) + .await; + let controller = NetworkPolicyController::default(); + let policy = controller.policy(); + let home = tempfile::tempdir().unwrap(); + let credentials = Arc::new( + AuthManager::new( + home.path().to_path_buf(), + /*enable_codex_api_key_env*/ false, + AuthCredentialsStoreMode::Ephemeral, + /*forced_chatgpt_workspace_id*/ None, + /*chatgpt_base_url*/ None, + AuthKeyringBackendKind::Direct, + AuthRouteConfig::from_http_client_factory( + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault) + .with_network_policy(policy.clone()), + ), + ) + .await, + ); + let auth = CodexAuth::from_api_key("key-a"); + credentials + .set_external_auth(Arc::new(Credentials { + current: RwLock::new(auth.clone()), + refreshed: auth, + })) + .await + .unwrap(); + assert!(controller.publish(policy.revision(), DestinationPolicy::Unrestricted)); + let manager = RuntimeAuthManager::new( + credentials, + AgentHostAuthorization::InstallationId { + installation_id: Uuid::nil(), + }, + RuntimeRegistrationMetadata { + harness: Some("codex".into()), + harness_version: None, + role: RuntimeRole::OrchestratorAndExecutor, + }, + format!("{}/agent/register", server.uri()).parse().unwrap(), + ); + assert!(matches!( + manager.auth().await, + Err(RuntimeAuthError::Registration( + RegistrationError::Rejected { + status: StatusCode::SERVICE_UNAVAILABLE, + .. + } + )) + )); + + policy.invalidate(); + assert!(matches!( + manager.auth().await, + Err(RuntimeAuthError::Policy(NetworkPolicyDenied::Unavailable)) + )); + assert_eq!(server.received_requests().await.unwrap().len(), 1); + + assert!(controller.publish(policy.revision(), DestinationPolicy::Unrestricted)); + let selected = manager.auth().await.unwrap(); + let requests = server.received_requests().await.unwrap(); + assert_eq!(requests.len(), 2); + assert_eq!(requests[0].body, requests[1].body); + assert_eq!( + requests[0].headers["idempotency-key"], + requests[1].headers["idempotency-key"] + ); + assert_eq!( + selected.runtime().id(), + requests[1].headers["idempotency-key"].to_str().unwrap() + ); +} + +#[tokio::test] +async fn bootstrap_refresh_reuses_known_sources_but_replaces_sessionless_initialization() { + for session in [Some("login-a"), None] { + let server = MockServer::start().await; + let initial = oauth(session, "old"); + Mock::given(method("POST")) + .and(body_partial_json( + json!({"actor_auth": {"credential": initial.get_token().unwrap()}}), + )) + .respond_with(ResponseTemplate::new(401)) + .with_priority(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .respond_with(registered) + .mount(&server) + .await; + let (manager, source, _) = + setup(&server, initial, oauth(Some("login-a"), "refreshed")).await; + let first = manager.auth().await.unwrap(); + *source.current.write().unwrap() = oauth(Some("login-a"), "next-token"); + let next = manager.auth().await.unwrap(); + assert!(Arc::ptr_eq(first.runtime(), next.runtime())); + assert_ne!( + first.auth().get_token().unwrap(), + next.auth().get_token().unwrap() + ); + let requests = server.received_requests().await.unwrap(); + assert_eq!(requests.len(), 2); + assert_eq!( + requests[0].headers["idempotency-key"] == requests[1].headers["idempotency-key"], + session.is_some() + ); + assert_eq!( + requests[0].body_json::().unwrap()["public_key"] + == requests[1].body_json::().unwrap()["public_key"], + session.is_some() + ); + } +} + +#[tokio::test] +async fn recovery_is_bounded_and_revocation_cannot_trigger_reregistration() { + for status in [401, 403] { + let server = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(status)) + .mount(&server) + .await; + let (manager, _, _) = setup( + &server, + oauth(Some("login-a"), "old"), + oauth(Some("login-a"), "refreshed"), + ) + .await; + assert!( + matches!(manager.auth().await, Err(RuntimeAuthError::Registration(RegistrationError::Rejected { status: actual, .. })) if actual.as_u16() == status) + ); + if status == 403 { + assert!(manager.auth().await.is_err()); + } + assert_eq!( + server.received_requests().await.unwrap().len(), + if status == 401 { 2 } else { 1 } + ); + } +} + +#[tokio::test] +async fn a_late_registration_cannot_publish_the_previous_actor() { + let server = MockServer::start().await; + let received = Arc::new(Notify::new()); + let signal = received.clone(); + Mock::given(method("POST")) + .respond_with(move |request: &Request| { + signal.notify_one(); + registered(request).set_delay(Duration::from_secs(1)) + }) + .mount(&server) + .await; + let initial = CodexAuth::from_api_key("key-a"); + let (manager, source, credentials) = setup(&server, initial.clone(), initial).await; + let task_manager = manager.clone(); + let attempt = tokio::spawn(async move { task_manager.auth().await }); + received.notified().await; + *source.current.write().unwrap() = CodexAuth::from_api_key("key-b"); + credentials.auth().await.unwrap(); + assert!(matches!( + attempt.await.unwrap(), + Err(RuntimeAuthError::Changed) + )); + let current = manager.auth().await.unwrap(); + assert_eq!(current.auth().get_token().unwrap(), "key-b"); + let requests = server.received_requests().await.unwrap(); + assert_eq!(requests.len(), 2); + assert_ne!( + requests[0].headers["idempotency-key"], + requests[1].headers["idempotency-key"] + ); +} diff --git a/codex-rs/runtime-agent-instance/src/registration.rs b/codex-rs/runtime-agent-instance/src/registration.rs index 22763e4731..e52cbef38e 100644 --- a/codex-rs/runtime-agent-instance/src/registration.rs +++ b/codex-rs/runtime-agent-instance/src/registration.rs @@ -1,7 +1,7 @@ //! Direct composite bootstrap. Credentials are sent only in the body, never through redirects. -use crate::RuntimeAgentInstance; use crate::RuntimeRegistration; +use crate::RuntimeRegistrationResponse; use codex_http_client::ClientRouteClass; use codex_http_client::HttpClientFactory; use codex_http_client::HttpError; @@ -52,7 +52,7 @@ impl RuntimeRegistrationClient { pub async fn register( &self, registration: &RuntimeRegistration, - ) -> Result { + ) -> Result { let body = serde_json::to_vec(®istration.request) .map_err(|_| RegistrationError::InvalidRequest)?; if body.len() > MAX_BODY_BYTES { @@ -96,15 +96,12 @@ impl RuntimeRegistrationClient { if info.runtime_agent_instance_id.is_empty() || info.agent_host_id.is_empty() { return Err(RegistrationError::InvalidResponse); } - Ok(RuntimeAgentInstance { - info, - signing_key: registration.signing_key.clone(), - }) + Ok(info) } } /// Errors intentionally omit raw transport details and server messages, which can contain secrets. -#[derive(Debug, thiserror::Error)] +#[derive(Clone, Debug, thiserror::Error)] pub enum RegistrationError { #[error( "runtime registration requires a secure endpoint without credentials or query parameters" diff --git a/codex-rs/runtime-agent-instance/src/registration_tests.rs b/codex-rs/runtime-agent-instance/src/registration_tests.rs index 1770d41fb7..ed76e74c9e 100644 --- a/codex-rs/runtime-agent-instance/src/registration_tests.rs +++ b/codex-rs/runtime-agent-instance/src/registration_tests.rs @@ -3,6 +3,7 @@ use super::*; use crate::AgentHostAuthorization; use crate::AuthenticatedActorAuthorization; +use crate::RuntimeAgentInstance; use crate::RuntimeRegistrationMetadata; use crate::RuntimeRole; use base64::Engine; @@ -77,9 +78,11 @@ async fn registration_retries_keep_the_key_and_sign_fresh_assertions() { .await; let pending = pending(); let client = client(&server); - let runtime = client.register(&pending).await.unwrap(); - let retry = client.register(&pending).await.unwrap(); - assert_eq!(runtime.info(), retry.info()); + let info = client.register(&pending).await.unwrap(); + let retry_info = client.register(&pending).await.unwrap(); + assert_eq!(info, retry_info); + let runtime = pending.finish(info); + let retry = pending.finish(retry_info); let requests = server.received_requests().await.unwrap(); assert_eq!(requests.len(), 2); assert_eq!(requests[0].body, requests[1].body); diff --git a/codex-rs/runtime-agent-instance/src/runtime.rs b/codex-rs/runtime-agent-instance/src/runtime.rs index 967302dd5f..0a1562cab1 100644 --- a/codex-rs/runtime-agent-instance/src/runtime.rs +++ b/codex-rs/runtime-agent-instance/src/runtime.rs @@ -1,5 +1,7 @@ //! A pending registration keeps its key and retry ID; a registered runtime signs v1 assertions. +use crate::AgentHost; +use crate::AuthenticatedActor; use crate::RuntimeRegistrationRequest; use crate::RuntimeRegistrationResponse; use base64::Engine; @@ -48,17 +50,39 @@ impl RuntimeRegistration { signing_key: Arc::new(signing_key), }) } + + pub(crate) fn finish(&self, info: RuntimeRegistrationResponse) -> RuntimeAgentInstance { + RuntimeAgentInstance { + id: info.runtime_agent_instance_id, + actor: AuthenticatedActor { info: info.actor }, + host: AgentHost { + id: info.agent_host_id, + authorization: self.request.agent_host_auth.clone(), + }, + signing_key: self.signing_key.clone(), + } + } } /// A server-registered identity and its memory-only signer, shared for this runtime's lifetime. pub struct RuntimeAgentInstance { - pub(crate) info: RuntimeRegistrationResponse, - pub(crate) signing_key: Arc, + id: String, + pub(crate) actor: AuthenticatedActor, + host: AgentHost, + signing_key: Arc, } impl RuntimeAgentInstance { - pub fn info(&self) -> &RuntimeRegistrationResponse { - &self.info + pub fn id(&self) -> &str { + &self.id + } + + pub fn actor(&self) -> &AuthenticatedActor { + &self.actor + } + + pub fn host(&self) -> &AgentHost { + &self.host } /// Make a fresh assertion for each request attempt or connection handshake. @@ -72,7 +96,7 @@ impl RuntimeAgentInstance { let payload = URL_SAFE_NO_PAD.encode( json!({ "version": 1, - "runtime_agent_instance": self.info.runtime_agent_instance_id, + "runtime_agent_instance": self.id, "iat": issued_at, "id": Uuid::new_v4(), }) @@ -93,10 +117,7 @@ impl RuntimeAgentInstance { impl fmt::Debug for RuntimeAgentInstance { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("RuntimeAgentInstance") - .field( - "runtime_agent_instance_id", - &self.info.runtime_agent_instance_id, - ) + .field("runtime_agent_instance_id", &self.id) .finish_non_exhaustive() } }