mirror of
https://github.com/openai/codex.git
synced 2026-11-08 20:40: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:
cooper-oai
@@ -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(())
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user