From d4a341bd52cc47ff635761334bde8a972d6de943 Mon Sep 17 00:00:00 2001 From: cooper-oai Date: Sat, 7 Nov 2026 03:56:09 +0000 Subject: [PATCH] Reauthenticate runtime requests before reusing WebSocket access (#60055) ## Why An open WebSocket can remain usable after its runtime is paused. Runtime-authenticated requests need a fresh handshake to revalidate access before submitting more work. ## What changed - Reconnect before each runtime-authenticated Responses request and skip speculative preconnection. Preserve turn routing when the runtime identity is unchanged. - Open a fresh connection for each runtime-authenticated Guardian classification instead of prewarming or reusing pooled sockets. - Preserve HTTP fallback for unsupported upgrades and connection failures, while propagating authentication, authorization, rate-limit, and policy errors. ## Testing Add regression coverage for pause and unpause behavior, credential changes, fresh assertions, and Guardian HTTP and WebSocket transports in both runtime authentication modes. GitOrigin-RevId: 9f650147260da5cbbf5bb23dfff0af095fff3955 --- codex-rs/Cargo.lock | 2 + codex-rs/core/src/client.rs | 12 +- .../suite/client_websockets_runtime_auth.rs | 174 ++++++++++-- codex-rs/ext/guardian-v2/Cargo.toml | 2 + .../async_scorer/sampler/connection_pool.rs | 118 ++++++-- .../sampler/connection_pool_tests.rs | 252 ++++++++++++++++++ .../model-provider/src/workspace_routing.rs | 2 +- 7 files changed, 513 insertions(+), 49 deletions(-) diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 5f5ffdf3b9..4078f03af2 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -3764,6 +3764,7 @@ name = "codex-guardian-v2" version = "0.0.0" dependencies = [ "anyhow", + "base64 0.22.1", "codex-analytics", "codex-api", "codex-config", @@ -3781,6 +3782,7 @@ dependencies = [ "codex-prompts", "codex-protocol", "codex-protocol-test-support", + "codex-runtime-agent-instance", "codex-utils-cargo-bin", "codex-websocket-client", "core_test_support", diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index a68d77ddda..de72f5e1a9 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -316,7 +316,8 @@ pub struct ModelClient { /// A turn-scoped streaming session created from a [`ModelClient`]. /// /// The session establishes a Responses WebSocket connection lazily and reuses it across multiple -/// requests within the turn. It also caches per-turn state: +/// requests within the turn. Runtime-authenticated requests reconnect between requests to +/// revalidate the registered credential source. It also caches per-turn state: /// /// - The last full request, so subsequent calls can reuse incremental websocket request payloads /// only when the current request is an incremental extension of the previous one. @@ -1745,7 +1746,10 @@ impl ModelClientSession { session_telemetry: &SessionTelemetry, responses_metadata: &CodexResponsesMetadata, ) -> Result<()> { - if !self.client.responses_websocket_enabled() { + // Runtime admission belongs to the request that will use the connection. + if !self.client.responses_websocket_enabled() + || self.client.state.provider.runtime_auth().is_some() + { return Ok(()); } let provider = Arc::clone(&self.client.state.provider); @@ -1860,8 +1864,10 @@ impl ModelClientSession { }; let connection_key = ResponsesConnectionKey::new(&api_provider, auth_revision); let reset_reason = match self.websocket_session.connection.as_ref() { + // Reauthenticate before each runtime request; an open socket cannot prove liveness. Some(_) - if self.websocket_session.responses_headers != *responses_headers + if runtime_id.is_some() + || self.websocket_session.responses_headers != *responses_headers || self.websocket_session.connection_key.as_ref() != Some(&connection_key) || self.websocket_session.runtime_id != runtime_id || self.websocket_session.parent_routing_affinity diff --git a/codex-rs/core/tests/suite/client_websockets_runtime_auth.rs b/codex-rs/core/tests/suite/client_websockets_runtime_auth.rs index 92a66b966c..4bd2193fd3 100644 --- a/codex-rs/core/tests/suite/client_websockets_runtime_auth.rs +++ b/codex-rs/core/tests/suite/client_websockets_runtime_auth.rs @@ -308,8 +308,9 @@ async fn websocket_reconnects_for_replacement_and_account_switch_within_a_turn( ]; let server = start_websocket_server(vec![ // Leave one unread response on each old connection so the server does not close it. - // Reconnection must follow runtime replacement, not a peer-initiated close. - vec![response.clone(), response.clone(), response.clone()], + // Runtime admission must happen again even when the peer keeps the socket open. + vec![response.clone(), response.clone()], + vec![response.clone(), response.clone()], vec![response.clone(), response.clone()], vec![response], ]) @@ -341,7 +342,7 @@ async fn websocket_reconnects_for_replacement_and_account_switch_within_a_turn( stream_until_complete(&mut session, &harness, &prompt_with_input(input.clone())).await; input.push(message_item("second")); stream_until_complete(&mut session, &harness, &prompt_with_input(input.clone())).await; - assert_eq!(server.handshakes().len(), 1); + assert_eq!(server.handshakes().len(), 2); *credentials.0.write().unwrap() = CodexAuth::from_api_key("key-b"); session @@ -353,18 +354,23 @@ async fn websocket_reconnects_for_replacement_and_account_switch_within_a_turn( ) .await .unwrap(); + assert_eq!( + server.handshakes().len(), + 2, + "preconnect must not admit work" + ); input.push(message_item("third")); stream_until_complete(&mut session, &harness, &prompt_with_input(input.clone())).await; - assert_eq!(server.handshakes().len(), 2); + assert_eq!(server.handshakes().len(), 3); *credentials.0.write().unwrap() = CodexAuth::from_api_key("key-other-user"); input.push(message_item("fourth")); stream_until_complete(&mut session, &harness, &prompt_with_input(input)).await; let handshakes = server.handshakes(); - assert_eq!(handshakes.len(), 3); + assert_eq!(handshakes.len(), 4); let runtime_ids: Vec<_> = handshakes .iter() - .zip(["key-a", "key-b", "key-other-user"]) + .zip(["key-a", "key-a", "key-b", "key-other-user"]) .map(|(h, credential)| { assert_eq!( h.header("cookie").as_deref(), @@ -384,30 +390,31 @@ async fn websocket_reconnects_for_replacement_and_account_switch_within_a_turn( .clone() }) .collect(); - assert_ne!(runtime_ids[0], runtime_ids[1]); + assert_eq!(runtime_ids[0], runtime_ids[1]); assert_ne!(runtime_ids[1], runtime_ids[2]); + assert_ne!(runtime_ids[2], runtime_ids[3]); assert_eq!( server .connections() .iter() .map(Vec::len) .collect::>(), - vec![2, 1, 1] + vec![1, 1, 1, 1] ); - let connections = server.connections(); - assert_eq!( - connections[0][1].body_json()["previous_response_id"], - "resp-1" - ); - assert_eq!( - connections[0][1].body_json()["client_metadata"]["x-codex-turn-state"], - "runtime-state" - ); - for (index, requests) in connections.iter().enumerate().skip(1) { + for (index, requests) in server.connections().iter().enumerate() { let body = requests[0].body_json(); assert_eq!(body["previous_response_id"], Value::Null); - assert_eq!(body["client_metadata"]["x-codex-turn-state"], Value::Null); - assert_eq!(body["input"].as_array().unwrap().len(), index + 2); + // Reconnecting with the same actor preserves turn routing. + let expected_turn_state = if index == 1 { + json!("runtime-state") + } else { + Value::Null + }; + assert_eq!( + body["client_metadata"]["x-codex-turn-state"], + expected_turn_state + ); + assert_eq!(body["input"].as_array().unwrap().len(), index + 1); } for request in registration.received_requests().await.unwrap() { let body: Value = request.body_json().unwrap(); @@ -563,3 +570,130 @@ async fn model_runtime_auth_strips_credentials_on_cross_origin_redirects(mode: R .contains_key("x-openai-agent-assertion") ); } + +#[test_case(RuntimeAuthMode::DualHeader; "dual_header")] +#[test_case(RuntimeAuthMode::AssertionOnly; "assertion_only")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn paused_runtime_cannot_continue_on_an_admitted_websocket(mode: RuntimeAuthMode) { + use futures::SinkExt; + use std::sync::atomic::AtomicUsize; + use std::sync::atomic::Ordering; + use tokio_tungstenite::accept_hdr_async_with_config; + use tokio_tungstenite::tungstenite::Message; + use tokio_tungstenite::tungstenite::extensions::ExtensionsConfig; + use tokio_tungstenite::tungstenite::extensions::compression::deflate::DeflateConfig; + use tokio_tungstenite::tungstenite::handshake::server::Request as HandshakeRequest; + use tokio_tungstenite::tungstenite::handshake::server::Response as HandshakeResponse; + use tokio_tungstenite::tungstenite::protocol::WebSocketConfig; + + skip_if_no_network!(); + let registration = MockServer::start().await; + Mock::given(path("/register")) + .respond_with(registered) + .expect(1) + .mount(®istration) + .await; + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let mut provider = ModelProviderInfo::create_openai_provider(Some(format!( + "ws://{}/v1", + listener.local_addr().unwrap() + ))); + provider.request_max_retries = Some(0); + let (harness, _) = runtime_harness(®istration, provider, mode).await; + let requests = Arc::new(AtomicUsize::new(0)); + let proofs = Arc::new(std::sync::Mutex::new(Vec::new())); + let server_requests = Arc::clone(&requests); + let server_proofs = Arc::clone(&proofs); + let server = tokio::spawn(async move { + let mut connections = tokio::task::JoinSet::new(); + for paused in [false, true, true, false] { + let (stream, _) = listener.accept().await.unwrap(); + let proofs = Arc::clone(&server_proofs); + let requests = Arc::clone(&server_requests); + connections.spawn(async move { + let handshake = move |request: &HandshakeRequest, response: HandshakeResponse| { + proofs.lock().unwrap().push(assertion_payload( + request.headers()["authorization"].to_str().unwrap(), + request + .headers() + .get("x-openai-agent-assertion") + .map(|value| value.to_str().unwrap().to_string()), + mode, + "key-a", + )); + if paused { + let body = json!({"error": { + "message": "The runtime agent instance is paused.", + "code": "runtime_instance_paused", "type": "runtime_agent_error" + }}) + .to_string(); + return Err(HandshakeResponse::builder() + .status(403) + .header("content-type", "application/json") + .header("content-length", body.len()) + .body(Some(body)) + .unwrap()); + } + Ok(response) + }; + let mut extensions = ExtensionsConfig::default(); + extensions.permessage_deflate = Some(DeflateConfig::default()); + let mut config = WebSocketConfig::default(); + config.extensions = extensions; + let Ok(mut socket) = + accept_hdr_async_with_config(stream, handshake, Some(config)).await + else { + return; + }; + // Keep admitted sockets usable: only another handshake observes the pause. + while let Some(Ok(message)) = socket.next().await { + if message.is_text() { + requests.fetch_add(/*val*/ 1, Ordering::SeqCst); + for event in [ev_response_created("resp-1"), ev_completed("resp-1")] { + socket + .send(Message::Text(event.to_string().into())) + .await + .unwrap(); + } + } + } + }); + } + while connections.join_next().await.is_some() {} + }); + + let assert_paused = |error: &codex_protocol::error::CodexErr| { + let codex_protocol::error::CodexErrorDetails::UnexpectedStatus(response) = error.details() + else { + panic!("expected a paused-runtime HTTP response: {error:?}"); + }; + assert_eq!(response.status.as_u16(), 403); + let body: Value = serde_json::from_str(&response.body).unwrap(); + assert_eq!(body["error"]["code"], "runtime_instance_paused"); + }; + let mut session = harness.client.new_session(); + complete(&mut session, &harness).await.unwrap(); + let error = complete(&mut session, &harness).await.unwrap_err(); + assert_paused(&error); + drop(session); + let mut next_turn = harness.client.new_session(); + let error = complete(&mut next_turn, &harness).await.unwrap_err(); + assert_paused(&error); + // An authorized unpause permits the original registration to run again. + complete(&mut next_turn, &harness).await.unwrap(); + assert_eq!(requests.load(Ordering::SeqCst), 2); + let proofs = proofs.lock().unwrap(); + assert_eq!(proofs.len(), 4); + for proof in proofs.iter().skip(1) { + assert_eq!( + proof["runtime_agent_instance"], + proofs[0]["runtime_agent_instance"] + ); + } + let distinct_proofs = proofs + .iter() + .map(|proof| proof["id"].as_str().expect("assertion proof ID")) + .collect::>(); + assert_eq!(distinct_proofs.len(), proofs.len()); + server.abort(); +} diff --git a/codex-rs/ext/guardian-v2/Cargo.toml b/codex-rs/ext/guardian-v2/Cargo.toml index efbe4b9b42..35208d2d3a 100644 --- a/codex-rs/ext/guardian-v2/Cargo.toml +++ b/codex-rs/ext/guardian-v2/Cargo.toml @@ -39,6 +39,8 @@ tracing = { workspace = true } uuid = { workspace = true } [dev-dependencies] +base64 = { workspace = true } +codex-runtime-agent-instance = { workspace = true } codex-model-provider-info = { workspace = true } codex-protocol-test-support = { workspace = true } codex-utils-cargo-bin = { workspace = true } diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/sampler/connection_pool.rs b/codex-rs/ext/guardian-v2/src/async_scorer/sampler/connection_pool.rs index 2470a1c017..bd77e49b5a 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/sampler/connection_pool.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/sampler/connection_pool.rs @@ -97,6 +97,7 @@ pub(super) struct PooledConnection { _permit: OwnedSemaphorePermit, } +#[derive(Clone)] struct ClientSetup { provider: Provider, http_client_factory: HttpClientFactory, @@ -159,6 +160,10 @@ impl ConnectionPool { } fn start_refill(self: &Arc, guard: OwnedMutexGuard<()>) -> Option> { + // Runtime sources must be reauthenticated when a classification is admitted. + if self.config.provider.runtime_auth().is_some() { + return None; + } let mut retry_after = self .retry_after .lock() @@ -226,6 +231,8 @@ impl ConnectionPool { .await .map_err(|error| LunaSamplerError::Api(ApiError::Stream(error.to_string())))?; let admitted = Instant::now(); + let setup = self.client_setup().await?; + let fresh_connection = self.runtime_connection(&setup).await?; let ClientSetup { mut provider, auth, @@ -236,30 +243,34 @@ impl ConnectionPool { redirect_policy, key, auth_changes, - } = self.client_setup().await?; + } = setup; let setup_finished = Instant::now(); - let connection = loop { - let idle = self - .idle_connections - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .pop(); - match idle { - Some(connection) - if connection.key == key - && connection.runtime_id == runtime_id - && connection.auth_owner == auth_owner - && connection - .auth_changes - .as_ref() - .is_none_or(|auth| !auth.has_changed().unwrap_or(true)) - && Instant::now() < connection.expires_at - && !connection.connection.is_closed().await => - { - break Some(connection); + let connection = if runtime_id.is_some() { + fresh_connection + } else { + loop { + let idle = self + .idle_connections + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .pop(); + match idle { + Some(connection) + if connection.key == key + && connection.runtime_id == runtime_id + && connection.auth_owner == auth_owner + && connection + .auth_changes + .as_ref() + .is_none_or(|auth| !auth.has_changed().unwrap_or(true)) + && Instant::now() < connection.expires_at + && !connection.connection.is_closed().await => + { + break Some(connection); + } + Some(_) => {} + None => break None, } - Some(_) => {} - None => break None, } }; let (connection, thread_id, request_kind) = match connection { @@ -458,9 +469,65 @@ impl ConnectionPool { } } + async fn runtime_connection( + &self, + setup: &ClientSetup, + ) -> Result, LunaSamplerError> { + if setup.runtime_id.is_none() || !self.config.provider.info().supports_websockets { + return Ok(None); + } + if self + .retry_after + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .is_some_and(|deadline| Instant::now() < deadline) + { + return Ok(None); + } + let Ok(permit) = Arc::clone(&self.sockets).try_acquire_owned() else { + return Ok(None); + }; + match self.open_connection_with_setup(permit, setup.clone()).await { + Ok(connection) => Ok(Some(connection)), + Err(LunaSamplerError::ConnectionTimeout) => { + *self + .retry_after + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) = + Some(Instant::now() + CONNECT_COOLDOWN); + Ok(None) + } + Err( + error @ LunaSamplerError::Api(ApiError::Transport( + TransportError::Policy(_) | TransportError::Build(_), + )), + ) => Err(error), + Err( + error @ LunaSamplerError::Api(ApiError::Transport(TransportError::Http { + status: + http::StatusCode::UNAUTHORIZED + | http::StatusCode::FORBIDDEN + | http::StatusCode::TOO_MANY_REQUESTS, + .. + })), + ) => Err(error), + // Unsupported upgrades and connection failures can still use HTTP. + Err(_) => Ok(None), + } + } + async fn open_connection( &self, permit: OwnedSemaphorePermit, + ) -> Result { + self.open_connection_with_setup(permit, self.client_setup().await?) + .await + } + + async fn open_connection_with_setup( + &self, + permit: OwnedSemaphorePermit, + setup: ClientSetup, ) -> Result { let ClientSetup { provider, @@ -472,7 +539,7 @@ impl ConnectionPool { redirect_policy: _, key, auth_changes, - } = self.client_setup().await?; + } = setup; let thread_id = ThreadId::new().to_string(); let mut headers = self.headers(&thread_id, request_kind)?; headers.insert( @@ -537,8 +604,8 @@ impl ConnectionPool { .as_ref() .is_some_and(|auth| auth.has_changed().unwrap_or(true)) { - return Err(LunaSamplerError::Api(ApiError::Stream( - "authentication changed while connecting".into(), + return Err(LunaSamplerError::Api(ApiError::Transport( + TransportError::Build("authentication changed while connecting".into()), ))); } @@ -597,6 +664,7 @@ impl ConnectionLease { pub(super) async fn reuse(self) { if let Connection::Websocket(connection) = self.connection + && connection.runtime_id.is_none() && !connection.connection.is_closed().await { self.pool diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/sampler/connection_pool_tests.rs b/codex-rs/ext/guardian-v2/src/async_scorer/sampler/connection_pool_tests.rs index 1121b668ef..c764ab3fe1 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/sampler/connection_pool_tests.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/sampler/connection_pool_tests.rs @@ -653,3 +653,255 @@ async fn http_cache_preserves_no_constraint_redirect_policy_and_configured_heade ); Ok(()) } + +#[tokio::test] +async fn runtime_transports_authenticate_each_classification() -> Result<()> { + use base64::Engine; + use codex_features::RuntimeAuthMode; + use codex_login::CodexAuth; + use codex_runtime_agent_instance::AgentHostAuthorization; + use codex_runtime_agent_instance::RuntimeAuthManager; + use codex_runtime_agent_instance::RuntimeRegistrationMetadata; + use codex_runtime_agent_instance::RuntimeRole; + use serde_json::json; + use wiremock::Mock; + use wiremock::MockServer; + use wiremock::ResponseTemplate; + use wiremock::matchers::method; + use wiremock::matchers::path; + + struct MutableAuth(Mutex); + impl codex_login::ExternalAuth for MutableAuth { + fn resolve(&self) -> codex_login::ExternalAuthFuture<'_, CodexAuth> { + Box::pin(async { Ok(self.0.lock().unwrap().clone()) }) + } + + fn refresh( + &self, + _context: codex_login::ExternalAuthRefreshContext, + ) -> codex_login::ExternalAuthFuture<'_, CodexAuth> { + self.resolve() + } + } + + skip_if_no_network!(Ok(())); + for mode in [RuntimeAuthMode::DualHeader, RuntimeAuthMode::AssertionOnly] { + for transport in [ + "http", + "websocket", + "upgrade_required", + "legacy_endpoint", + "changed_auth", + "denied", + ] { + let websocket = transport == "websocket"; + let denied = transport == "denied"; + let changed_auth = transport == "changed_auth"; + let source = Arc::new(MutableAuth(Mutex::new(CodexAuth::from_api_key( + "guardian-key", + )))); + let registration = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/register")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "runtime_agent_instance_id": "runtime-guardian", + "agent_host_id": "host-guardian", + "actor": {"kind": "platform", "user_id": "actor", "organization_id": "org", "project_id": null}, + }))) + .expect(1) + .mount(®istration) + .await; + let response = vec![ + responses::ev_output_text_delta("low"), + responses::ev_completed("score"), + ]; + let ws = if websocket { + // Each peer leaves another response available, so closure cannot force the retry. + Some( + responses::start_websocket_server(vec![ + vec![response.clone(), response.clone()], + vec![response.clone(), response.clone()], + ]) + .await, + ) + } else { + None + }; + let http = MockServer::start().await; + let upgrade_response = if denied { + ResponseTemplate::new(403).set_body_json(json!({"error": { + "code": "runtime_instance_paused", "type": "runtime_agent_error", + "message": "The runtime agent instance is paused." + }})) + } else if transport == "legacy_endpoint" { + ResponseTemplate::new(404) + } else { + ResponseTemplate::new(426) + }; + let handshake_source = Arc::clone(&source); + Mock::given(method("GET")) + .and(path("/responses")) + .respond_with(move |_request: &wiremock::Request| { + if changed_auth { + *handshake_source.0.lock().unwrap() = + CodexAuth::from_api_key("replacement-key"); + } + upgrade_response.clone() + }) + .expect(match transport { + "upgrade_required" | "legacy_endpoint" => 2, + "changed_auth" | "denied" => 1, + _ => 0, + }) + .mount(&http) + .await; + Mock::given(method("POST")) + .and(path("/responses")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(responses::sse(response)), + ) + .expect(if websocket || denied { + 0 + } else if changed_auth { + 1 + } else { + 2 + }) + .mount(&http) + .await; + let url = ws.as_ref().map_or_else( + || http.uri(), + |ws| ws.uri().replacen("ws://", "http://", /*count*/ 1), + ); + let mut config = sampler_config(url); + let credentials = + AuthManager::from_auth_for_testing(CodexAuth::from_api_key("guardian-key")); + if changed_auth { + credentials.set_external_auth(source.clone()).await?; + } + let runtime = Arc::new(RuntimeAuthManager::new( + Arc::clone(&credentials), + AgentHostAuthorization::InstallationId { + installation_id: uuid::Uuid::nil(), + }, + RuntimeRegistrationMetadata { + harness: Some("codex".into()), + harness_version: None, + role: RuntimeRole::Orchestrator, + }, + format!("{}/register", registration.uri()).parse()?, + )); + let mut info = config.provider.info().clone(); + info.supports_websockets = transport != "http"; + config.provider = codex_model_provider::test_support::runtime_model_provider( + info, + Arc::clone(&credentials), + (runtime, mode), + ); + let sampler = LunaSampler::new(config); + sampler.prewarm().await; + if let Some(ws) = &ws { + assert!( + ws.handshakes().is_empty(), + "prewarming must not admit runtime work" + ); + } + if denied { + let error = sampler.sample(sample_request("denied")).await.unwrap_err(); + assert!( + matches!( + error, + LunaSamplerError::Api(ApiError::Transport(TransportError::Http { + status: http::StatusCode::FORBIDDEN, + .. + })) + ), + "{error:?}" + ); + } else { + assert_eq!(sampler.sample(sample_request("first")).await?, "low"); + if changed_auth { + assert_eq!( + credentials.auth().await.unwrap().get_token()?, + "replacement-key" + ); + } else { + assert_eq!(sampler.sample(sample_request("second")).await?, "low"); + } + } + let headers = if let Some(ws) = &ws { + assert_eq!( + ws.connections().iter().map(Vec::len).collect::>(), + vec![1, 1] + ); + ws.handshakes() + .iter() + .map(|h| { + ( + h.header("authorization").unwrap(), + h.header("x-openai-agent-assertion"), + ) + }) + .collect::>() + } else { + http.received_requests() + .await + .unwrap() + .iter() + .map(|request| { + ( + request.headers["authorization"] + .to_str() + .unwrap() + .to_owned(), + request + .headers + .get("x-openai-agent-assertion") + .map(|value| value.to_str().unwrap().to_owned()), + ) + }) + .collect::>() + }; + let expected_attempts = match transport { + "upgrade_required" | "legacy_endpoint" => 4, + "denied" => 1, + _ => 2, + }; + assert_eq!(headers.len(), expected_attempts); + let mut proofs = std::collections::BTreeSet::new(); + for (authorization, supplemental) in headers { + let proof = match mode { + RuntimeAuthMode::DualHeader => { + assert_eq!(authorization, "Bearer guardian-key"); + supplemental.unwrap() + } + RuntimeAuthMode::AssertionOnly => { + assert!(supplemental.is_none()); + authorization + .strip_prefix("AgentAssertion ") + .unwrap() + .to_owned() + } + }; + let payload: serde_json::Value = serde_json::from_slice( + &base64::engine::general_purpose::URL_SAFE_NO_PAD + .decode(proof.split('.').nth(/*n*/ 2).unwrap())?, + )?; + assert_eq!(payload["runtime_agent_instance"], "runtime-guardian"); + assert!( + proofs.insert(proof), + "each attempt must authenticate afresh" + ); + } + registration.verify().await; + http.verify().await; + drop(sampler); + if let Some(ws) = ws { + ws.shutdown().await; + } + } + } + Ok(()) +} diff --git a/codex-rs/model-provider/src/workspace_routing.rs b/codex-rs/model-provider/src/workspace_routing.rs index ae47503da6..18b48cce4c 100644 --- a/codex-rs/model-provider/src/workspace_routing.rs +++ b/codex-rs/model-provider/src/workspace_routing.rs @@ -85,7 +85,7 @@ pub struct ResolvedResponsesProvider { } /// Effective destination and credentials that permit reusing a Responses socket. -#[derive(Debug, PartialEq, Eq)] +#[derive(Clone, Debug, PartialEq, Eq)] pub struct ResponsesConnectionKey { base_url: String, routing_header: Option,