From 62fb56ee569e8ed6a9ad7ec575b0e78c25df4b17 Mon Sep 17 00:00:00 2001 From: jif Date: Wed, 26 Aug 2026 14:25:28 +0000 Subject: [PATCH] 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 --- codex-rs/codex-api/src/endpoint/mod.rs | 1 + codex-rs/codex-api/src/endpoint/responses.rs | 38 ++++- .../src/endpoint/responses_websocket.rs | 18 ++- codex-rs/codex-api/src/lib.rs | 1 + codex-rs/codex-api/tests/clients.rs | 21 +++ codex-rs/core/config.schema.json | 4 + codex-rs/core/src/client.rs | 130 +++++++++++++----- codex-rs/core/src/client_tests.rs | 75 ++++++++++ codex-rs/core/src/config/mod.rs | 11 ++ codex-rs/core/src/session/session.rs | 1 + codex-rs/core/tests/common/responses.rs | 2 +- codex-rs/core/tests/common/test_codex.rs | 2 - codex-rs/core/tests/suite/auto_review.rs | 1 + .../core/tests/suite/client_websockets.rs | 24 +++- codex-rs/core/tests/suite/guardian_review.rs | 78 +++++++++-- .../guardian-v2/src/async_scorer/extension.rs | 1 + .../guardian-v2/src/async_scorer/sampler.rs | 62 +++++---- .../src/async_scorer/sampler_tests.rs | 78 +++++++++++ codex-rs/features/src/feature_configs.rs | 3 + codex-rs/features/src/tests.rs | 2 + codex-rs/model-provider-info/src/lib.rs | 9 ++ .../src/model_provider_info_tests.rs | 13 ++ 22 files changed, 488 insertions(+), 87 deletions(-) diff --git a/codex-rs/codex-api/src/endpoint/mod.rs b/codex-rs/codex-api/src/endpoint/mod.rs index ba8ed96e6c..e9f08d2753 100644 --- a/codex-rs/codex-api/src/endpoint/mod.rs +++ b/codex-rs/codex-api/src/endpoint/mod.rs @@ -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; diff --git a/codex-rs/codex-api/src/endpoint/responses.rs b/codex-rs/codex-api/src/endpoint/responses.rs index 804f0027ff..21d3deb928 100644 --- a/codex-rs/codex-api/src/endpoint/responses.rs +++ b/codex-rs/codex-api/src/endpoint/responses.rs @@ -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 { session: EndpointSession, sse_telemetry: Option>, + endpoint: ResponsesEndpoint, } #[derive(Default)] @@ -43,9 +67,16 @@ impl ResponsesClient { 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>, @@ -54,6 +85,7 @@ impl ResponsesClient { Self { session: self.session.with_request_telemetry(request), sse_telemetry: sse, + endpoint: self.endpoint, } } @@ -97,10 +129,6 @@ impl ResponsesClient { .await } - fn path() -> &'static str { - "responses" - } - #[instrument( name = "responses.stream", level = "info", @@ -141,7 +169,7 @@ impl ResponsesClient { .session .stream_encoded_json_with( Method::POST, - Self::path(), + self.endpoint.path(), extra_headers, Some(body), |req| { diff --git a/codex-rs/codex-api/src/endpoint/responses_websocket.rs b/codex-rs/codex-api/src/endpoint/responses_websocket.rs index a85c993d03..b1e009f908 100644 --- a/codex-rs/codex-api/src/endpoint/responses_websocket.rs +++ b/codex-rs/codex-api/src/endpoint/responses_websocket.rs @@ -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 { 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 { 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 = diff --git a/codex-rs/codex-api/src/lib.rs b/codex-rs/codex-api/src/lib.rs index a31a0e2fba..17e3ceb69b 100644 --- a/codex-rs/codex-api/src/lib.rs +++ b/codex-rs/codex-api/src/lib.rs @@ -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; diff --git a/codex-rs/codex-api/tests/clients.rs b/codex-rs/codex-api/tests/clients.rs index 1ad1161848..af43022417 100644 --- a/codex-rs/codex-api/tests/clients.rs +++ b/codex-rs/codex-api/tests/clients.rs @@ -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(); diff --git a/codex-rs/core/config.schema.json b/codex-rs/core/config.schema.json index cc80ed561a..8fd0997d06 100644 --- a/codex-rs/core/config.schema.json +++ b/codex-rs/core/config.schema.json @@ -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, diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index a715d7bb14..c1768066d4 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -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, agent_identity_policy: AgentIdentityAuthPolicy, prompt_cache_key_override: Option, + free_guardian_enabled: bool, http_client_factory: HttpClientFactory, } @@ -302,6 +303,7 @@ struct LastResponse { #[derive(Debug, Default)] struct WebsocketSession { connection: Option, + endpoint: Option, last_request: Option, last_response_rx: Option>, 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, @@ -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 { - 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 { 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( diff --git a/codex-rs/core/src/client_tests.rs b/codex-rs/core/src/client_tests.rs index b580a4dfff..4389d83cee 100644 --- a/codex-rs/core/src/client_tests.rs +++ b/codex-rs/core/src/client_tests.rs @@ -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) = diff --git a/codex-rs/core/src/config/mod.rs b/codex-rs/core/src/config/mod.rs index 252e9ffa22..0ecedf1952 100644 --- a/codex-rs/core/src/config/mod.rs +++ b/codex-rs/core/src/config/mod.rs @@ -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, diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 71ad2afef0..59689978aa 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -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, diff --git a/codex-rs/core/tests/common/responses.rs b/codex-rs/core/tests/common/responses.rs index defa0bffa9..33d4284ecc 100644 --- a/codex-rs/core/tests/common/responses.rs +++ b/codex-rs/core/tests/common/responses.rs @@ -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) } diff --git a/codex-rs/core/tests/common/test_codex.rs b/codex-rs/core/tests/common/test_codex.rs index 7c63fc90af..6e7b9d6e77 100644 --- a/codex-rs/core/tests/common/test_codex.rs +++ b/codex-rs/core/tests/common/test_codex.rs @@ -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; diff --git a/codex-rs/core/tests/suite/auto_review.rs b/codex-rs/core/tests/suite/auto_review.rs index ff454ac429..3754baa3ae 100644 --- a/codex-rs/core/tests/suite/auto_review.rs +++ b/codex-rs/core/tests/suite/auto_review.rs @@ -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??; diff --git a/codex-rs/core/tests/suite/client_websockets.rs b/codex-rs/core/tests/suite/client_websockets.rs index 748cc07728..9d1c184e16 100755 --- a/codex-rs/core/tests/suite/client_websockets.rs +++ b/codex-rs/core/tests/suite/client_websockets.rs @@ -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"); diff --git a/codex-rs/core/tests/suite/guardian_review.rs b/codex-rs/core/tests/suite/guardian_review.rs index 1a0cf9269b..d030669a76 100644 --- a/codex-rs/core/tests/suite/guardian_review.rs +++ b/codex-rs/core/tests/suite/guardian_review.rs @@ -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::>(); 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(()) } diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/extension.rs b/codex-rs/ext/guardian-v2/src/async_scorer/extension.rs index 68a9e8cd09..3d47429f0b 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/extension.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/extension.rs @@ -338,6 +338,7 @@ impl ThreadLifecycleContributor for GuardianV2Extension { .thread_store .get::() .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(), diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/sampler.rs b/codex-rs/ext/guardian-v2/src/async_scorer/sampler.rs index 7db9ca5e05..7dc331c255 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/sampler.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/sampler.rs @@ -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, + /// Whether this thread may use the unmetered Guardian classifier endpoint. + pub free_guardian: bool, /// Optional inference service tier. pub service_tier: Option, /// 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 { 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, diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/sampler_tests.rs b/codex-rs/ext/guardian-v2/src/async_scorer/sampler_tests.rs index d5a99ed84c..903a0d1947 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/sampler_tests.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/sampler_tests.rs @@ -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, diff --git a/codex-rs/features/src/feature_configs.rs b/codex-rs/features/src/feature_configs.rs index 43e6a4b506..eb18554e36 100644 --- a/codex-rs/features/src/feature_configs.rs +++ b/codex-rs/features/src/feature_configs.rs @@ -126,6 +126,9 @@ pub struct GuardianV2ReviewScopeConfigToml { pub struct GuardianV2ConfigToml { #[serde(skip_serializing_if = "Option::is_none")] pub enabled: Option, + /// Route Guardian review and classification through the unmetered Codex endpoints. + #[serde(skip_serializing_if = "Option::is_none")] + pub free_guardian: Option, #[serde(skip_serializing_if = "Option::is_none")] pub classifier_instructions: Option, #[serde(skip_serializing_if = "Option::is_none")] diff --git a/codex-rs/features/src/tests.rs b/codex-rs/features/src/tests.rs index 4725a54ad1..52a5d1f548 100644 --- a/codex-rs/features/src/tests.rs +++ b/codex-rs/features/src/tests.rs @@ -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), diff --git a/codex-rs/model-provider-info/src/lib.rs b/codex-rs/model-provider-info/src/lib.rs index f631e7c6f1..65e2e4cbf7 100644 --- a/codex-rs/model-provider-info/src/lib.rs +++ b/codex-rs/model-provider-info/src/lib.rs @@ -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| { diff --git a/codex-rs/model-provider-info/src/model_provider_info_tests.rs b/codex-rs/model-provider-info/src/model_provider_info_tests.rs index c70d59b39a..d08af0ffb5 100644 --- a/codex-rs/model-provider-info/src/model_provider_info_tests.rs +++ b/codex-rs/model-provider-info/src/model_provider_info_tests.rs @@ -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 {