Route Guardian inference through dedicated endpoints (#40892)

## What changed

- Add the opt-in `features.guardianv2.free_guardian` setting to route eligible
  Guardian reviews to `/guardian` and asynchronous classifications to
  `/guardian-classifier`.
- Support endpoint selection for both HTTP and WebSocket Responses clients,
  including endpoint-aware connection reuse.
- Keep `/responses` for API-key authentication, custom provider URLs,
  non-Guardian models, and configurations where the option is disabled.
- Omit routing hints and `service_tier` from requests sent to the dedicated
  Guardian endpoints.

## Testing

- Cover route selection across authentication modes, provider URLs, HTTP
  fallback, WebSocket review sessions, and classifier sampling.

GitOrigin-RevId: 9c88673c710d009e0239559b9245a4e6c4f4befc
This commit is contained in:
jif
2026-08-26 14:36:48 +00:00
committed by copyberry
parent 10d5a603ae
commit 62fb56ee56
22 changed files with 488 additions and 87 deletions
+1
View File
@@ -27,6 +27,7 @@ pub use realtime_websocket::RealtimeWebsocketEvents;
pub use realtime_websocket::RealtimeWebsocketWriter;
pub use realtime_websocket::session_update_session_json;
pub use responses::ResponsesClient;
pub use responses::ResponsesEndpoint;
pub use responses::ResponsesOptions;
pub use responses_websocket::ResponsesWebsocketClient;
pub use responses_websocket::ResponsesWebsocketClose;
+33 -5
View File
@@ -23,9 +23,33 @@ use std::sync::Arc;
use std::sync::OnceLock;
use tracing::instrument;
/// Responses-compatible inference routes supported by Codex backend.
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
pub enum ResponsesEndpoint {
/// Regular user-owned model inference.
#[default]
Responses,
/// Full Guardian approval-review agent inference.
Guardian,
/// Lightweight asynchronous Guardian risk classification.
GuardianClassifier,
}
impl ResponsesEndpoint {
/// Returns the provider-relative path for this inference surface.
pub const fn path(self) -> &'static str {
match self {
Self::Responses => "/responses",
Self::Guardian => "/guardian",
Self::GuardianClassifier => "/guardian-classifier",
}
}
}
pub struct ResponsesClient<T: HttpTransport> {
session: EndpointSession<T>,
sse_telemetry: Option<Arc<dyn SseTelemetry>>,
endpoint: ResponsesEndpoint,
}
#[derive(Default)]
@@ -43,9 +67,16 @@ impl<T: HttpTransport> ResponsesClient<T> {
Self {
session: EndpointSession::new(transport, provider, auth),
sse_telemetry: None,
endpoint: ResponsesEndpoint::Responses,
}
}
/// Selects a Responses-compatible backend route for subsequent requests.
pub fn with_endpoint(mut self, endpoint: ResponsesEndpoint) -> Self {
self.endpoint = endpoint;
self
}
pub fn with_telemetry(
self,
request: Option<Arc<dyn RequestTelemetry>>,
@@ -54,6 +85,7 @@ impl<T: HttpTransport> ResponsesClient<T> {
Self {
session: self.session.with_request_telemetry(request),
sse_telemetry: sse,
endpoint: self.endpoint,
}
}
@@ -97,10 +129,6 @@ impl<T: HttpTransport> ResponsesClient<T> {
.await
}
fn path() -> &'static str {
"responses"
}
#[instrument(
name = "responses.stream",
level = "info",
@@ -141,7 +169,7 @@ impl<T: HttpTransport> ResponsesClient<T> {
.session
.stream_encoded_json_with(
Method::POST,
Self::path(),
self.endpoint.path(),
extra_headers,
Some(body),
|req| {
@@ -4,6 +4,7 @@ use crate::common::ResponseStream;
use crate::common::ResponsesWsRequest;
use crate::common::SafetyBufferingTreatment;
use crate::common::WS_REQUEST_HEADER_TRACEPARENT_CLIENT_METADATA_KEY;
use crate::endpoint::responses::ResponsesEndpoint;
use crate::error::ApiError;
use crate::provider::Provider;
use crate::rate_limits::parse_rate_limit_event;
@@ -340,6 +341,7 @@ impl ResponsesWebsocketConnection {
pub struct ResponsesWebsocketClient {
provider: Provider,
auth: SharedAuthProvider,
endpoint: ResponsesEndpoint,
}
/// Close frame information captured by a handshake probe.
@@ -369,7 +371,17 @@ pub struct ResponsesWebsocketProbe {
impl ResponsesWebsocketClient {
/// Creates a Responses WebSocket client for an already-resolved provider and auth source.
pub fn new(provider: Provider, auth: SharedAuthProvider) -> Self {
Self { provider, auth }
Self {
provider,
auth,
endpoint: ResponsesEndpoint::Responses,
}
}
/// Selects a Responses-compatible backend route for subsequent connections.
pub fn with_endpoint(mut self, endpoint: ResponsesEndpoint) -> Self {
self.endpoint = endpoint;
self
}
#[instrument(
@@ -388,7 +400,7 @@ impl ResponsesWebsocketClient {
) -> Result<ResponsesWebsocketConnection, ApiError> {
let ws_url = self
.provider
.websocket_url_for_path("responses")
.websocket_url_for_path(self.endpoint.path())
.map_err(|err| ApiError::Stream(format!("failed to build websocket URL: {err}")))?;
let mut headers =
@@ -422,7 +434,7 @@ impl ResponsesWebsocketClient {
) -> Result<ResponsesWebsocketProbe, ApiError> {
let ws_url = self
.provider
.websocket_url_for_path("responses")
.websocket_url_for_path(self.endpoint.path())
.map_err(|err| ApiError::Stream(format!("failed to build websocket URL: {err}")))?;
let mut headers =
+1
View File
@@ -66,6 +66,7 @@ pub use crate::endpoint::RealtimeWebsocketConnection;
pub use crate::endpoint::RealtimeWebsocketEvents;
pub use crate::endpoint::RealtimeWebsocketWriter;
pub use crate::endpoint::ResponsesClient;
pub use crate::endpoint::ResponsesEndpoint;
pub use crate::endpoint::ResponsesOptions;
pub use crate::endpoint::ResponsesWebsocketClient;
pub use crate::endpoint::ResponsesWebsocketClose;
+21
View File
@@ -12,6 +12,7 @@ use codex_api::Compression;
use codex_api::Provider;
use codex_api::ResponsesApiRequest;
use codex_api::ResponsesClient;
use codex_api::ResponsesEndpoint;
use codex_api::ResponsesOptions;
use codex_client::HttpTransport;
use codex_client::Request;
@@ -306,6 +307,26 @@ async fn responses_client_uses_responses_path() -> Result<()> {
Ok(())
}
#[tokio::test]
async fn responses_client_uses_guardian_path() -> Result<()> {
let state = RecordingState::default();
let transport = RecordingTransport::new(state.clone());
let client = ResponsesClient::new(transport, provider("openai"), Arc::new(NoAuth))
.with_endpoint(ResponsesEndpoint::Guardian);
let _stream = client
.stream(
serde_json::json!({ "echo": true }),
HeaderMap::new(),
Compression::None,
/*turn_state*/ None,
)
.await?;
assert_path_ends_with(&state.take_stream_requests(), "/guardian");
Ok(())
}
#[tokio::test]
async fn responses_client_stream_request_preserves_item_ids() -> Result<()> {
let state = RecordingState::default();
+4
View File
@@ -1414,6 +1414,10 @@
"enabled": {
"type": "boolean"
},
"free_guardian": {
"description": "Route Guardian review and classification through the unmetered Codex endpoints.",
"type": "boolean"
},
"max_action_tokens": {
"format": "uint",
"maximum": 100000.0,
+93 -37
View File
@@ -50,6 +50,7 @@ use codex_api::ReqwestTransport;
use codex_api::ResponseCreateWsRequest;
use codex_api::ResponsesApiRequest;
use codex_api::ResponsesClient as ApiResponsesClient;
use codex_api::ResponsesEndpoint;
use codex_api::ResponsesOptions as ApiResponsesOptions;
use codex_api::ResponsesWebsocketClient as ApiWebSocketResponsesClient;
use codex_api::ResponsesWebsocketConnection as ApiWebSocketConnection;
@@ -161,7 +162,6 @@ const RESPONSES_WEBSOCKETS_V2_BETA_HEADER_VALUE: &str = "responses_websockets=20
const X_OPENAI_INTERNAL_CODEX_RESPONSES_LITE_HEADER: &str =
"x-openai-internal-codex-responses-lite";
const REALTIME_CALLS_ENDPOINT: &str = "/realtime/calls";
const RESPONSES_ENDPOINT: &str = "/responses";
const RESPONSES_COMPACT_ENDPOINT: &str = "/responses/compact";
// `/responses/compact` is unary, so the timeout covers the full response rather than one idle
// period between stream events.
@@ -261,6 +261,7 @@ pub struct ModelClient {
state: Arc<ModelClientState>,
agent_identity_policy: AgentIdentityAuthPolicy,
prompt_cache_key_override: Option<String>,
free_guardian_enabled: bool,
http_client_factory: HttpClientFactory,
}
@@ -302,6 +303,7 @@ struct LastResponse {
#[derive(Debug, Default)]
struct WebsocketSession {
connection: Option<ApiWebSocketConnection>,
endpoint: Option<ResponsesEndpoint>,
last_request: Option<ResponsesApiRequest>,
last_response_rx: Option<oneshot::Receiver<LastResponse>>,
last_response_from_untraced_warmup: bool,
@@ -480,10 +482,16 @@ impl ModelClient {
}),
agent_identity_policy,
prompt_cache_key_override: None,
free_guardian_enabled: false,
http_client_factory,
}
}
pub(crate) fn with_free_guardian_enabled(mut self, free_guardian_enabled: bool) -> Self {
self.free_guardian_enabled = free_guardian_enabled;
self
}
pub(crate) fn with_prompt_cache_key_override(
mut self,
prompt_cache_key_override: Option<String>,
@@ -1020,21 +1028,37 @@ impl ModelClient {
})
}
fn responses_endpoint(&self, auth: Option<&CodexAuth>, model: &str) -> ResponsesEndpoint {
if self.free_guardian_enabled
&& crate::guardian::is_basic_session_source(&self.state.session_source)
&& self.uses_codex_backend(auth)
&& self.state.provider.info().supports_codex_backend_routes()
&& model == self.state.provider.approval_review_preferred_model()
{
ResponsesEndpoint::Guardian
} else {
ResponsesEndpoint::Responses
}
}
fn uses_codex_backend(&self, auth: Option<&CodexAuth>) -> bool {
let provider = self.state.provider.info();
auth.is_some_and(CodexAuth::uses_codex_backend)
&& provider.is_openai()
&& provider.requires_openai_auth
&& provider.env_key.is_none()
&& provider.experimental_bearer_token.is_none()
&& provider.auth.is_none()
&& provider.aws.is_none()
}
fn build_routing_hint_header(
&self,
auth: Option<&CodexAuth>,
model: &str,
service_tier: Option<&str>,
) -> Option<HeaderValue> {
let provider = self.state.provider.info();
if !auth.is_some_and(CodexAuth::uses_codex_backend)
|| !provider.is_openai()
|| !provider.requires_openai_auth
|| provider.env_key.is_some()
|| provider.experimental_bearer_token.is_some()
|| provider.auth.is_some()
|| provider.aws.is_some()
{
if !self.uses_codex_backend(auth) {
return None;
}
@@ -1077,6 +1101,7 @@ impl ModelClient {
responses_metadata: &CodexResponsesMetadata,
auth_context: AuthRequestTelemetryContext,
request_route_telemetry: RequestRouteTelemetry,
endpoint: ResponsesEndpoint,
) -> std::result::Result<ApiWebSocketConnection, ApiError> {
let headers = self.build_websocket_headers(responses_metadata).await;
let websocket_telemetry = ModelClientSession::build_websocket_telemetry(
@@ -1089,13 +1114,15 @@ impl ModelClient {
let start = Instant::now();
let result = match tokio::time::timeout(
websocket_connect_timeout,
ApiWebSocketResponsesClient::new(api_provider, api_auth).connect(
&self.http_client_factory,
headers,
codex_login::default_client::default_headers(),
/*turn_state*/ None,
Some(websocket_telemetry),
),
ApiWebSocketResponsesClient::new(api_provider, api_auth)
.with_endpoint(endpoint)
.connect(
&self.http_client_factory,
headers,
codex_login::default_client::default_headers(),
/*turn_state*/ None,
Some(websocket_telemetry),
),
)
.await
{
@@ -1206,6 +1233,7 @@ impl ModelClientSession {
fn reset_websocket_session(&mut self) {
self.websocket_session.connection = None;
self.websocket_session.endpoint = None;
self.websocket_session.last_request = None;
self.websocket_session.last_response_rx = None;
self.websocket_session.last_response_from_untraced_warmup = false;
@@ -1336,6 +1364,7 @@ impl ModelClientSession {
/// This performs only connection setup; it never sends prompt payloads.
pub async fn preconnect_websocket(
&mut self,
model_info: &ModelInfo,
session_telemetry: &SessionTelemetry,
responses_metadata: &CodexResponsesMetadata,
) -> std::result::Result<(), ApiError> {
@@ -1357,6 +1386,9 @@ impl ModelClientSession {
client_setup.agent_identity_telemetry.clone(),
PendingUnauthorizedRetry::default(),
);
let endpoint = self
.client
.responses_endpoint(client_setup.auth.as_ref(), &model_info.slug);
let connection = self
.client
.connect_websocket(
@@ -1365,10 +1397,12 @@ impl ModelClientSession {
client_setup.api_auth,
responses_metadata,
auth_context,
RequestRouteTelemetry::for_endpoint(RESPONSES_ENDPOINT),
RequestRouteTelemetry::for_endpoint(endpoint.path()),
endpoint,
)
.await?;
self.websocket_session.connection = Some(connection);
self.websocket_session.endpoint = Some(endpoint);
self.websocket_session
.set_connection_reused(/*connection_reused*/ false);
Ok(())
@@ -1397,16 +1431,17 @@ impl ModelClientSession {
responses_metadata,
auth_context,
request_route_telemetry,
endpoint,
} = params;
let needs_new = match self.websocket_session.connection.as_ref() {
Some(conn) => conn.is_closed().await,
Some(conn) => {
self.websocket_session.endpoint != Some(endpoint) || conn.is_closed().await
}
None => true,
};
if needs_new {
self.websocket_session.last_request = None;
self.websocket_session.last_response_rx = None;
self.websocket_session.last_response_from_untraced_warmup = false;
self.reset_websocket_session();
let new_conn = match self
.client
.connect_websocket(
@@ -1416,6 +1451,7 @@ impl ModelClientSession {
responses_metadata,
auth_context,
request_route_telemetry,
endpoint,
)
.await
{
@@ -1428,6 +1464,7 @@ impl ModelClientSession {
}
};
self.websocket_session.connection = Some(new_conn);
self.websocket_session.endpoint = Some(endpoint);
self.websocket_session
.set_connection_reused(/*connection_reused*/ false);
} else {
@@ -1490,9 +1527,12 @@ impl ModelClientSession {
let mut pending_retry = PendingUnauthorizedRetry::default();
loop {
let client_setup = self.client.current_client_setup().await?;
let endpoint = self
.client
.responses_endpoint(client_setup.auth.as_ref(), &model_info.slug);
let transport = self
.client
.build_api_transport(&client_setup.api_provider, RESPONSES_ENDPOINT)?;
.build_api_transport(&client_setup.api_provider, endpoint.path())?;
let request_auth_context = AuthRequestTelemetryContext::new(
client_setup.auth.as_ref().map(CodexAuth::auth_mode),
client_setup.api_auth.as_ref(),
@@ -1502,7 +1542,7 @@ impl ModelClientSession {
let (request_telemetry, sse_telemetry) = Self::build_streaming_telemetry(
session_telemetry,
request_auth_context,
RequestRouteTelemetry::for_endpoint(RESPONSES_ENDPOINT),
RequestRouteTelemetry::for_endpoint(endpoint.path()),
self.client.state.auth_env_telemetry.clone(),
);
let compression = self.responses_request_compression(client_setup.auth.as_ref());
@@ -1522,11 +1562,16 @@ impl ModelClientSession {
service_tier.clone(),
responses_metadata,
)?;
if let Some(header_value) = self.client.build_routing_hint_header(
client_setup.auth.as_ref(),
&request.model,
request.service_tier.as_deref(),
) {
if endpoint == ResponsesEndpoint::Guardian {
request.service_tier = None;
}
if endpoint == ResponsesEndpoint::Responses
&& let Some(header_value) = self.client.build_routing_hint_header(
client_setup.auth.as_ref(),
&request.model,
request.service_tier.as_deref(),
)
{
options
.extra_headers
.insert(X_CODEX_ROUTING_HINT_HEADER, header_value);
@@ -1547,6 +1592,7 @@ impl ModelClientSession {
client_setup.api_provider,
client_setup.api_auth,
)
.with_endpoint(endpoint)
.with_telemetry(Some(request_telemetry), Some(sse_telemetry));
let stream_result = client.stream_request(request, options).await;
@@ -1639,6 +1685,9 @@ impl ModelClientSession {
let mut pending_retry = PendingUnauthorizedRetry::default();
loop {
let client_setup = self.client.current_client_setup().await?;
let endpoint = self
.client
.responses_endpoint(client_setup.auth.as_ref(), &model_info.slug);
let request_auth_context = AuthRequestTelemetryContext::new(
client_setup.auth.as_ref().map(CodexAuth::auth_mode),
client_setup.api_auth.as_ref(),
@@ -1653,16 +1702,23 @@ impl ModelClientSession {
service_tier.clone(),
responses_metadata,
)?;
if endpoint == ResponsesEndpoint::Guardian {
request.service_tier = None;
}
request.access_programs = cyber_access_program::for_auth(
client_setup.auth.as_ref(),
prompt.cyber_access_program,
);
let mut websocket_metadata = responses_metadata.clone();
websocket_metadata.routing_hint = self.client.build_routing_hint_header(
client_setup.auth.as_ref(),
&request.model,
request.service_tier.as_deref(),
);
websocket_metadata.routing_hint = if endpoint == ResponsesEndpoint::Responses {
self.client.build_routing_hint_header(
client_setup.auth.as_ref(),
&request.model,
request.service_tier.as_deref(),
)
} else {
None
};
let request_session_telemetry = if warmup {
// `generate=false` prewarm is connection setup, not an inference request.
session_telemetry.clone()
@@ -1682,9 +1738,8 @@ impl ModelClientSession {
api_auth: client_setup.api_auth,
responses_metadata: &websocket_metadata,
auth_context: request_auth_context,
request_route_telemetry: RequestRouteTelemetry::for_endpoint(
RESPONSES_ENDPOINT,
),
request_route_telemetry: RequestRouteTelemetry::for_endpoint(endpoint.path()),
endpoint,
})
.await
{
@@ -2274,6 +2329,7 @@ struct WebsocketConnectParams<'a> {
responses_metadata: &'a CodexResponsesMetadata,
auth_context: AuthRequestTelemetryContext,
request_route_telemetry: RequestRouteTelemetry,
endpoint: ResponsesEndpoint,
}
async fn handle_unauthorized(
+75
View File
@@ -18,6 +18,7 @@ use crate::test_support::responses_metadata as test_responses_metadata;
use codex_api::AgentIdentityTelemetry;
use codex_api::ApiError;
use codex_api::ResponseEvent;
use codex_api::ResponsesEndpoint;
use codex_api::TransportError;
use codex_http_client::HttpClientFactory;
use codex_http_client::OutboundProxyPolicy;
@@ -975,6 +976,80 @@ fn model_client_with_counting_attestation(
(model_client, attestation_calls)
}
#[test]
fn guardian_reviewer_uses_dedicated_endpoint_only_with_codex_backend_auth() {
let (mut model_client, _) =
model_client_with_counting_attestation(/*include_attestation*/ true);
Arc::get_mut(&mut model_client.state)
.expect("test client should have unique session state")
.session_source = SessionSource::SubAgent(SubAgentSource::Other("guardian".to_owned()));
assert_eq!(
model_client.responses_endpoint(
Some(&CodexAuth::create_dummy_chatgpt_auth_for_testing()),
"codex-auto-review",
),
ResponsesEndpoint::Responses
);
model_client = model_client.with_free_guardian_enabled(/*free_guardian_enabled*/ true);
assert_eq!(
model_client.responses_endpoint(
Some(&CodexAuth::create_dummy_chatgpt_auth_for_testing()),
"codex-auto-review",
),
ResponsesEndpoint::Guardian
);
assert_eq!(
model_client.responses_endpoint(
Some(&CodexAuth::create_dummy_chatgpt_auth_for_testing()),
"required-reviewer-model",
),
ResponsesEndpoint::Responses
);
assert_eq!(
model_client.responses_endpoint(
Some(&CodexAuth::create_dummy_chatgpt_auth_for_testing()),
"parent-fallback-model",
),
ResponsesEndpoint::Responses
);
assert_eq!(
model_client.responses_endpoint(
Some(&CodexAuth::from_api_key("test-api-key")),
"codex-auto-review",
),
ResponsesEndpoint::Responses
);
Arc::get_mut(&mut model_client.state)
.expect("test client should have unique session state")
.provider = create_model_provider(
ModelProviderInfo::create_openai_provider(Some("https://proxy.example.com/v1".to_owned())),
Some(AuthManager::from_auth_for_testing(
CodexAuth::create_dummy_chatgpt_auth_for_testing(),
)),
);
assert_eq!(
model_client.responses_endpoint(
Some(&CodexAuth::create_dummy_chatgpt_auth_for_testing()),
"codex-auto-review",
),
ResponsesEndpoint::Responses
);
Arc::get_mut(&mut model_client.state)
.expect("test client should have unique session state")
.session_source = SessionSource::Exec;
assert_eq!(
model_client.responses_endpoint(
Some(&CodexAuth::create_dummy_chatgpt_auth_for_testing()),
"codex-auto-review",
),
ResponsesEndpoint::Responses
);
}
#[tokio::test]
async fn websocket_handshake_includes_attestation_for_chatgpt_codex_responses() {
let (model_client, attestation_calls) =
+11
View File
@@ -1490,6 +1490,17 @@ impl Config {
&self.sqlite
}
/// Whether Guardian may use the unmetered Codex inference endpoints.
pub fn free_guardian_enabled(&self) -> bool {
self.config_layer_stack
.effective_config()
.get("features")
.and_then(|features| features.get("guardianv2"))
.and_then(|guardian| guardian.get("free_guardian"))
.and_then(toml::Value::as_bool)
.unwrap_or(false)
}
/// Resolves the configured, reviewer-catalog, or bundled Guardian policy.
pub fn resolve_guardian_policy<'a>(
&'a self,
+1
View File
@@ -1383,6 +1383,7 @@ impl Session {
attestation_provider,
config.http_client_factory(),
)
.with_free_guardian_enabled(config.free_guardian_enabled())
.with_prompt_cache_key_override(
crate::guardian::prompt_cache_key_override_for_review_session(
&session_configuration.session_source,
+1 -1
View File
@@ -1078,7 +1078,7 @@ where
fn base_mock() -> (MockBuilder, ResponseMock) {
let response_mock = ResponseMock::new();
let mock = Mock::given(method("POST"))
.and(path_regex(".*/responses$"))
.and(path_regex(".*/(responses|guardian)$"))
.and(response_mock.clone());
(mock, response_mock)
}
-2
View File
@@ -562,9 +562,7 @@ impl TestCodexBuilder {
Some(home) => home,
None => Arc::new(TempDir::new()?),
};
let base_url_clone = base_url.clone();
self.config_mutators.push(Box::new(move |config| {
config.model_provider.base_url = Some(base_url_clone);
config.model_provider.supports_websockets = true;
config.experimental_realtime_ws_model = Some("realtime-test-model".to_string());
config.realtime.version = RealtimeWsVersion::V1;
+1
View File
@@ -599,6 +599,7 @@ async fn remote_model_override_uses_catalog_model_for_strict_auto_review() -> Re
guardian_request.body_json()["model"].as_str(),
Some(review_model)
);
assert_eq!(guardian_request.path(), "/v1/responses");
timeout(Duration::from_secs(10), codex.shutdown_and_wait()).await??;
+20 -4
View File
@@ -447,7 +447,11 @@ async fn responses_websocket_preconnect_does_not_replace_turn_trace_payload() {
let mut client_session = harness.client.new_session();
let responses_metadata = websocket_connection_metadata(&harness);
client_session
.preconnect_websocket(&harness.session_telemetry, &responses_metadata)
.preconnect_websocket(
&harness.model_info,
&harness.session_telemetry,
&responses_metadata,
)
.await
.expect("websocket preconnect failed");
let prompt = prompt_with_input(vec![message_item("hello")]);
@@ -484,7 +488,11 @@ async fn responses_websocket_preconnect_reuses_connection() {
let mut client_session = harness.client.new_session();
let responses_metadata = websocket_connection_metadata(&harness);
client_session
.preconnect_websocket(&harness.session_telemetry, &responses_metadata)
.preconnect_websocket(
&harness.model_info,
&harness.session_telemetry,
&responses_metadata,
)
.await
.expect("websocket preconnect failed");
let prompt = prompt_with_input(vec![message_item("hello")]);
@@ -856,7 +864,11 @@ async fn responses_websocket_preconnect_is_reused_even_with_header_changes() {
let mut client_session = harness.client.new_session();
let preconnect_metadata = websocket_connection_metadata(&harness);
client_session
.preconnect_websocket(&harness.session_telemetry, &preconnect_metadata)
.preconnect_websocket(
&harness.model_info,
&harness.session_telemetry,
&preconnect_metadata,
)
.await
.expect("websocket preconnect failed");
let prompt = prompt_with_input(vec![message_item("hello")]);
@@ -1080,7 +1092,11 @@ async fn responses_websocket_preconnect_runs_when_only_v2_feature_enabled() {
let mut client_session = harness.client.new_session();
let responses_metadata = websocket_connection_metadata(&harness);
client_session
.preconnect_websocket(&harness.session_telemetry, &responses_metadata)
.preconnect_websocket(
&harness.model_info,
&harness.session_telemetry,
&responses_metadata,
)
.await
.expect("websocket preconnect failed");
+68 -10
View File
@@ -128,7 +128,16 @@ impl TimeProvider for RecordingTimeProvider {
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn guardian_session_inherits_parent_http_fallback() -> Result<()> {
#[test_case(CodexAuth::from_api_key("test-api-key"), "/v1", true, "/v1/responses"; "api_key_uses_responses")]
#[test_case(CodexAuth::create_dummy_chatgpt_auth_for_testing(), "/backend-api/codex", false, "/backend-api/codex/responses"; "chatgpt_uses_responses_by_default")]
#[test_case(CodexAuth::create_dummy_chatgpt_auth_for_testing(), "/backend-api/codex", true, "/backend-api/codex/guardian"; "chatgpt_uses_guardian_when_enabled")]
#[test_case(CodexAuth::create_dummy_chatgpt_auth_for_testing(), "/v1", true, "/v1/responses"; "custom_openai_url_uses_responses")]
async fn guardian_session_inherits_parent_http_fallback(
auth: CodexAuth,
base_path: &str,
free_guardian: bool,
expected_guardian_path: &str,
) -> Result<()> {
skip_if_no_network!(Ok(()));
skip_if_wine_exec!(
Ok(()),
@@ -162,11 +171,22 @@ async fn guardian_session_inherits_parent_http_fallback() -> Result<()> {
)
.await;
let mut builder = test_codex().with_config(|config| {
config.model_provider.supports_websockets = true;
config.permissions.approval_policy = Constrained::allow_any(AskForApproval::OnRequest);
config.approvals_reviewer = ApprovalsReviewer::User;
});
let base_url = format!("{}{base_path}", server.uri());
let mut builder = test_codex()
.with_auth(auth)
.with_pre_build_hook(move |home| {
fs::write(
home.join("config.toml"),
format!("[features.guardianv2]\nfree_guardian = {free_guardian}\n"),
)
.expect("Guardian endpoint configuration should be written");
})
.with_config(move |config| {
config.model_provider.base_url = Some(base_url);
config.model_provider.supports_websockets = true;
config.permissions.approval_policy = Constrained::allow_any(AskForApproval::OnRequest);
config.approvals_reviewer = ApprovalsReviewer::User;
});
let test = builder.build_with_auto_env(&server).await?;
test.codex
@@ -189,12 +209,19 @@ async fn guardian_session_inherits_parent_http_fallback() -> Result<()> {
let requests = server.received_requests().await.unwrap_or_default();
let websocket_attempts = requests
.iter()
.filter(|request| request.method == Method::GET)
.filter(|request| {
request.method == Method::GET && request.url.path().ends_with("/responses")
})
.count();
assert_eq!(websocket_attempts, 1);
assert!(responses.requests().iter().any(|request| {
request.body_json()["client_metadata"]["x-openai-subagent"].as_str() == Some("guardian")
}));
let guardian_request = responses
.requests()
.into_iter()
.find(|request| {
request.body_json()["client_metadata"]["x-openai-subagent"].as_str() == Some("guardian")
})
.expect("Guardian reviewer inference request");
assert_eq!(guardian_request.path(), expected_guardian_path);
Ok(())
}
@@ -208,6 +235,7 @@ async fn guardian_session_prewarms_and_is_reused_for_first_review(
) -> Result<()> {
skip_if_no_network!(Ok(()));
let uses_codex_backend = auth.uses_codex_backend();
let bundled_models = codex_models_manager::bundled_models_response()?.models;
let catalog_auto_review = bundled_models
.iter()
@@ -275,9 +303,20 @@ async fn guardian_session_prewarms_and_is_reused_for_first_review(
let time_provider = Arc::new(RecordingTimeProvider {
thread_ids: Mutex::new(Vec::new()),
});
let backend_base_url = format!("{}/backend-api/codex", server.uri());
let mut builder = test_codex()
.with_auth(auth)
.with_pre_build_hook(|home| {
fs::write(
home.join("config.toml"),
"[features.guardianv2]\nfree_guardian = true\n",
)
.expect("Guardian endpoint configuration should be written");
})
.with_config(move |config| {
if uses_codex_backend {
config.model_provider.base_url = Some(backend_base_url);
}
let rules_dir = config.codex_home.join("rules");
fs::create_dir_all(&rules_dir).expect("create execution policy directory");
let policy_justification = format!(
@@ -298,6 +337,7 @@ async fn guardian_session_prewarms_and_is_reused_for_first_review(
});
config.model_context_window = Some(900_000);
config.model_auto_compact_token_limit = Some(600_000);
config.service_tier = Some("priority".to_owned());
config.permissions.approval_policy = Constrained::allow_any(AskForApproval::OnRequest);
config.approvals_reviewer = ApprovalsReviewer::AutoReview;
config
@@ -402,6 +442,7 @@ async fn guardian_session_prewarms_and_is_reused_for_first_review(
);
assert_eq!(guardian_review["model"].as_str(), Some(expected_model));
for request in [guardian_prewarm, &guardian_review] {
assert_eq!(request.get("service_tier"), None);
let metadata: serde_json::Value = serde_json::from_str(
request["client_metadata"]["x-codex-turn-metadata"]
.as_str()
@@ -472,6 +513,23 @@ async fn guardian_session_prewarms_and_is_reused_for_first_review(
})
.collect::<Vec<_>>();
assert_eq!(guardian_context_windows, vec![Some(258_400)]);
for handshake in server.handshakes() {
let is_guardian = handshake.header("x-openai-subagent").as_deref() == Some("guardian");
let uses_guardian_endpoint = uses_codex_backend && is_guardian;
assert_eq!(
handshake.uri(),
if uses_guardian_endpoint {
"/backend-api/codex/guardian"
} else if uses_codex_backend {
"/backend-api/codex/responses"
} else {
"/v1/responses"
}
);
if uses_guardian_endpoint {
assert_eq!(handshake.header("x-codex-routing-hint"), None);
}
}
server.shutdown().await;
Ok(())
}
@@ -338,6 +338,7 @@ impl ThreadLifecycleContributor<Config> for GuardianV2Extension {
.thread_store
.get::<ThreadOriginator>()
.map(|originator| originator.0.clone()),
free_guardian: input.config.free_guardian_enabled(),
service_tier: input.config.service_tier.clone(),
luna_compaction_hash,
metrics: input.extension_metrics.clone(),
@@ -12,6 +12,7 @@ use codex_api::Reasoning;
use codex_api::ReasoningContext;
use codex_api::ResponseEvent;
use codex_api::ResponsesApiRequest;
use codex_api::ResponsesEndpoint;
use codex_api::ResponsesWebsocketClient;
use codex_api::ResponsesWebsocketConnection;
use codex_api::ResponsesWsRequest;
@@ -71,6 +72,8 @@ pub struct LunaSamplerConfig {
pub thread_id: String,
/// Optional host-resolved request originator.
pub originator: Option<String>,
/// Whether this thread may use the unmetered Guardian classifier endpoint.
pub free_guardian: bool,
/// Optional inference service tier.
pub service_tier: Option<String>,
/// Luna model's host-resolved encrypted-compaction compatibility hash.
@@ -124,6 +127,7 @@ pub enum LunaSamplerError {
struct PooledConnection {
connection: ResponsesWebsocketConnection,
endpoint: ResponsesEndpoint,
// The bridge routes by thread ID, so each socket needs its own identity.
thread_id: String,
expires_at: Instant,
@@ -209,6 +213,29 @@ impl LunaSampler {
Ok(sampler)
}
async fn responses_endpoint(&self) -> ResponsesEndpoint {
let provider = self.config.provider.info();
if self.config.free_guardian
&& self
.config
.provider
.auth()
.await
.as_ref()
.is_some_and(CodexAuth::uses_codex_backend)
&& provider.supports_codex_backend_routes()
&& provider.requires_openai_auth
&& provider.env_key.is_none()
&& provider.experimental_bearer_token.is_none()
&& provider.auth.is_none()
&& provider.aws.is_none()
{
ResponsesEndpoint::GuardianClassifier
} else {
ResponsesEndpoint::Responses
}
}
async fn open_connection(&self) -> Result<PooledConnection, LunaSamplerError> {
let provider = self
.config
@@ -259,30 +286,8 @@ impl LunaSampler {
}
let provider_info = self.config.provider.info();
if self
.config
.provider
.auth()
.await
.as_ref()
.is_some_and(CodexAuth::uses_codex_backend)
&& provider_info.is_openai()
&& provider_info.requires_openai_auth
&& provider_info.env_key.is_none()
&& provider_info.experimental_bearer_token.is_none()
&& provider_info.auth.is_none()
&& provider_info.aws.is_none()
{
let routing_hint = match self.config.service_tier.as_deref() {
Some(tier) => format!("model={MODEL};tier={tier}"),
None => format!("model={MODEL}"),
};
if let Ok(value) = HeaderValue::from_str(&routing_hint) {
headers.insert("x-codex-routing-hint", value);
}
}
let client = ResponsesWebsocketClient::new(provider, auth);
let endpoint = self.responses_endpoint().await;
let client = ResponsesWebsocketClient::new(provider, auth).with_endpoint(endpoint);
let connect = client.connect(
&self.config.http_client_factory,
headers,
@@ -305,6 +310,7 @@ impl LunaSampler {
Ok(PooledConnection {
connection,
endpoint,
thread_id,
expires_at: Instant::now() + MAX_WEBSOCKET_AGE,
auth_changes,
@@ -485,7 +491,7 @@ impl LunaSampler {
stream: true,
stream_options: None,
include: Vec::new(),
service_tier: self.config.service_tier.clone(),
service_tier: None,
prompt_cache_key: Some(format!("guardian-v2:{}", self.config.thread_id)),
text: None,
client_metadata: None,
@@ -536,6 +542,12 @@ impl LunaSampler {
return Err(error);
}
};
request.service_tier =
if lease.connection.endpoint == ResponsesEndpoint::GuardianClassifier {
None
} else {
self.config.service_tier.clone()
};
let thread_id = &lease.connection.thread_id;
let turn_metadata = json!({
"session_id": self.config.session_id,
@@ -147,6 +147,7 @@ fn sampler_config(base_url: String) -> LunaSamplerConfig {
session_id: "session-1".to_owned(),
thread_id: "thread-1".to_owned(),
originator: Some("guardian-v2-test".to_owned()),
free_guardian: false,
service_tier: None,
luna_compaction_hash: None,
metrics: None,
@@ -244,6 +245,81 @@ impl ExternalAuth for RefreshableAuth {
}
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn classifier_uses_free_endpoint_only_with_codex_backend_auth() -> Result<()> {
skip_if_no_network!(Ok(()));
for (auth, base_path, free_guardian, expected_path, expected_service_tier) in [
(
CodexAuth::create_dummy_chatgpt_auth_for_testing(),
"/backend-api/codex",
false,
"/backend-api/codex/responses",
Some("priority"),
),
(
CodexAuth::create_dummy_chatgpt_auth_for_testing(),
"/backend-api/codex",
true,
"/backend-api/codex/guardian-classifier",
None,
),
(
CodexAuth::create_dummy_chatgpt_auth_for_testing(),
"/v1",
true,
"/v1/responses",
Some("priority"),
),
(
CodexAuth::from_api_key("test-api-key"),
"/v1",
true,
"/v1/responses",
Some("priority"),
),
] {
let events = vec![
ev_assistant_message("classification", "low"),
ev_completed("response-1"),
];
let mut connections = vec![Vec::new(); INITIAL_WEBSOCKET_CONNECTIONS - 1];
connections.push(vec![events]);
let server = responses::start_websocket_server(connections).await;
let base_url = format!(
"http://{}{base_path}",
server.uri().trim_start_matches("ws://")
);
let mut config = sampler_config(base_url.clone());
config.provider = create_model_provider(
ModelProviderInfo::create_openai_provider(Some(base_url)),
Some(AuthManager::from_auth_for_testing(auth)),
);
config.free_guardian = free_guardian;
config.service_tier = Some("priority".to_owned());
let sampler = LunaSampler::connect(config).await?;
assert_eq!(sampler.sample(sample_request("turn-1")).await?, "low");
for handshake in server.handshakes() {
assert_eq!(handshake.uri(), expected_path);
assert_eq!(handshake.header("x-codex-routing-hint"), None);
}
let request = server
.wait_for_request(
/*connection_index*/ INITIAL_WEBSOCKET_CONNECTIONS - 1,
/*request_index*/ 0,
)
.await
.body_json();
assert_eq!(request["service_tier"].as_str(), expected_service_tier);
drop(sampler);
server.shutdown().await;
}
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn preconnected_sampler_reuses_authenticated_websocket_for_classifications() -> Result<()> {
skip_if_no_network!(Ok(()));
@@ -290,6 +366,7 @@ async fn preconnected_sampler_reuses_authenticated_websocket_for_classifications
session_id: "session-1".to_owned(),
thread_id: "thread-1".to_owned(),
originator: Some("guardian-v2-test".to_owned()),
free_guardian: false,
service_tier: None,
luna_compaction_hash: None,
metrics: None,
@@ -509,6 +586,7 @@ async fn sampler_returns_classification_token_before_terminal_response_events()
session_id: "session-1".to_owned(),
thread_id: "thread-1".to_owned(),
originator: None,
free_guardian: false,
service_tier: None,
luna_compaction_hash: None,
metrics: None,
+3
View File
@@ -126,6 +126,9 @@ pub struct GuardianV2ReviewScopeConfigToml {
pub struct GuardianV2ConfigToml {
#[serde(skip_serializing_if = "Option::is_none")]
pub enabled: Option<bool>,
/// Route Guardian review and classification through the unmetered Codex endpoints.
#[serde(skip_serializing_if = "Option::is_none")]
pub free_guardian: Option<bool>,
#[serde(skip_serializing_if = "Option::is_none")]
pub classifier_instructions: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
+2
View File
@@ -163,6 +163,7 @@ fn guardian_v2_feature_config_deserializes_classifier_and_transcript_settings()
r#"
[guardianv2]
enabled = true
free_guardian = true
classifier_instructions = "Review this action"
review_threshold = 0.65
max_tool_call_lag = 2
@@ -192,6 +193,7 @@ max_recent_non_user_entries = 12
features.guardianv2,
Some(FeatureToml::Config(crate::GuardianV2ConfigToml {
enabled: Some(true),
free_guardian: Some(true),
classifier_instructions: Some("Review this action".to_owned()),
review_threshold: Some(0.65),
max_tool_call_lag: Some(2),
+9
View File
@@ -468,6 +468,15 @@ impl ModelProviderInfo {
self.name == OPENAI_PROVIDER_NAME
}
pub fn supports_codex_backend_routes(&self) -> bool {
self.is_openai()
&& self.base_url.as_deref().is_none_or(|base_url| {
base_url
.trim_end_matches('/')
.ends_with("/backend-api/codex")
})
}
pub fn uses_openai_actor_authorization(&self) -> bool {
!self.requires_openai_auth
&& self.http_headers.as_ref().is_some_and(|headers| {
@@ -154,6 +154,19 @@ fn test_header_auth_uses_chatgpt_codex_base_url() {
assert_eq!(api_provider.base_url, CHATGPT_CODEX_BASE_URL);
}
#[test]
fn codex_backend_routes_require_codex_base_url() {
for (base_url, expected) in [
(None, true),
(Some(CHATGPT_CODEX_BASE_URL), true),
(Some("https://chatgpt-staging.com/backend-api/codex/"), true),
(Some("https://proxy.example.com/v1"), false),
] {
let provider = ModelProviderInfo::create_openai_provider(base_url.map(str::to_owned));
assert_eq!(provider.supports_codex_backend_routes(), expected);
}
}
#[test]
fn test_uses_openai_actor_authorization() {
let mut provider = ModelProviderInfo {