mirror of
https://github.com/openai/codex.git
synced 2026-11-07 20:30:33 +00:00
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
This commit is contained in:
Generated
+2
@@ -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",
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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<_>>(),
|
||||
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::<std::collections::BTreeSet<_>>();
|
||||
assert_eq!(distinct_proofs.len(), proofs.len());
|
||||
server.abort();
|
||||
}
|
||||
|
||||
@@ -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 }
|
||||
|
||||
@@ -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<Self>, guard: OwnedMutexGuard<()>) -> Option<JoinHandle<()>> {
|
||||
// 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<Option<PooledConnection>, 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<PooledConnection, LunaSamplerError> {
|
||||
self.open_connection_with_setup(permit, self.client_setup().await?)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn open_connection_with_setup(
|
||||
&self,
|
||||
permit: OwnedSemaphorePermit,
|
||||
setup: ClientSetup,
|
||||
) -> Result<PooledConnection, LunaSamplerError> {
|
||||
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
|
||||
|
||||
@@ -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<CodexAuth>);
|
||||
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<_>>(),
|
||||
vec![1, 1]
|
||||
);
|
||||
ws.handshakes()
|
||||
.iter()
|
||||
.map(|h| {
|
||||
(
|
||||
h.header("authorization").unwrap(),
|
||||
h.header("x-openai-agent-assertion"),
|
||||
)
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
} 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::<Vec<_>>()
|
||||
};
|
||||
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(())
|
||||
}
|
||||
|
||||
@@ -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<HeaderValue>,
|
||||
|
||||
Reference in New Issue
Block a user