From a7b8c074b577f897111c14de3a5e127b91e2a479 Mon Sep 17 00:00:00 2001 From: jif Date: Thu, 13 Aug 2026 12:30:54 +0000 Subject: [PATCH] Add the Guardian V2 Luna sampler (#38368) ## What changed - Add a `LunaSampler` that opens an authenticated Responses WebSocket and reuses it for structured, tool-free `gpt-5.6-luna` requests. - Carry the host's provider, authentication, proxy, attribution, and service-tier configuration into the connection and requests. - Require strict JSON schemas, preserve per-request reasoning effort and turn metadata, and reject missing or oversized output. ## Testing - Add a WebSocket test covering authentication and attribution headers, connection reuse, structured request fields, per-turn metadata, reasoning effort, and returned output. GitOrigin-RevId: aef85fd8ddcfcb7c76a9ede0d3191e3a0321656b --- codex-rs/Cargo.lock | 13 + codex-rs/ext/guardian-v2/Cargo.toml | 17 +- codex-rs/ext/guardian-v2/src/lib.rs | 7 + codex-rs/ext/guardian-v2/src/sampler.rs | 277 ++++++++++++++++++ codex-rs/ext/guardian-v2/src/sampler_tests.rs | 128 ++++++++ 5 files changed, 441 insertions(+), 1 deletion(-) create mode 100644 codex-rs/ext/guardian-v2/src/sampler.rs create mode 100644 codex-rs/ext/guardian-v2/src/sampler_tests.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index ec86bba976..b1c90a2768 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -3300,7 +3300,20 @@ dependencies = [ name = "codex-guardian-v2" version = "0.0.0" dependencies = [ + "anyhow", + "codex-api", "codex-extension-api", + "codex-http-client", + "codex-login", + "codex-model-provider", + "codex-model-provider-info", + "codex-protocol", + "core_test_support", + "http 1.4.0", + "pretty_assertions", + "serde_json", + "thiserror 2.0.18", + "tokio", ] [[package]] diff --git a/codex-rs/ext/guardian-v2/Cargo.toml b/codex-rs/ext/guardian-v2/Cargo.toml index 3ebdc530b4..a570952a58 100644 --- a/codex-rs/ext/guardian-v2/Cargo.toml +++ b/codex-rs/ext/guardian-v2/Cargo.toml @@ -7,11 +7,26 @@ version.workspace = true [lib] name = "codex_guardian_v2" path = "src/lib.rs" -test = false doctest = false [lints] workspace = true [dependencies] +codex-api = { workspace = true } codex-extension-api = { workspace = true } +codex-http-client = { workspace = true } +codex-login = { workspace = true } +codex-model-provider = { workspace = true } +codex-protocol = { workspace = true } +http = { workspace = true } +serde_json = { workspace = true } +thiserror = { workspace = true } +tokio = { workspace = true, features = ["time"] } + +[dev-dependencies] +anyhow = { workspace = true } +codex-model-provider-info = { workspace = true } +core_test_support = { workspace = true } +pretty_assertions = { workspace = true } +tokio = { workspace = true, features = ["macros", "rt-multi-thread"] } diff --git a/codex-rs/ext/guardian-v2/src/lib.rs b/codex-rs/ext/guardian-v2/src/lib.rs index 4752641123..137779f852 100644 --- a/codex-rs/ext/guardian-v2/src/lib.rs +++ b/codex-rs/ext/guardian-v2/src/lib.rs @@ -1,4 +1,11 @@ +mod sampler; + use codex_extension_api::ExtensionRegistryBuilder; +pub use sampler::LunaSampler; +pub use sampler::LunaSamplerConfig; +pub use sampler::LunaSamplerError; +pub use sampler::LunaSamplingRequest; + /// Installs the Guardian V2 extension without registering contributors yet. pub fn install(_registry: &mut ExtensionRegistryBuilder) {} diff --git a/codex-rs/ext/guardian-v2/src/sampler.rs b/codex-rs/ext/guardian-v2/src/sampler.rs new file mode 100644 index 0000000000..4d652077a9 --- /dev/null +++ b/codex-rs/ext/guardian-v2/src/sampler.rs @@ -0,0 +1,277 @@ +use std::collections::HashMap; + +use codex_api::ApiError; +use codex_api::Reasoning; +use codex_api::ResponseEvent; +use codex_api::ResponsesApiRequest; +use codex_api::ResponsesWebsocketClient; +use codex_api::ResponsesWebsocketConnection; +use codex_api::ResponsesWsRequest; +use codex_api::build_session_headers; +use codex_api::create_text_param_for_request; +use codex_http_client::HttpClientFactory; +use codex_login::AgentIdentityAuthPolicy; +use codex_login::CodexAuth; +use codex_login::default_client::add_originator_header; +use codex_login::default_client::default_headers; +use codex_model_provider::AgentIdentitySessionFallback; +use codex_model_provider::ProviderAuthScope; +use codex_model_provider::SharedModelProvider; +use codex_protocol::error::CodexErr; +use codex_protocol::models::ContentItem; +use codex_protocol::models::ResponseItem; +use codex_protocol::openai_models::ReasoningEffort; +use codex_protocol::protocol::SessionSource; +use http::HeaderValue; +use serde_json::Value; +use thiserror::Error; + +const MODEL: &str = "gpt-5.6-luna"; +const MAX_OUTPUT_BYTES: usize = 8 * 1024; +const RESPONSES_WEBSOCKETS_BETA: &str = "responses_websockets=2026-02-06"; +const RESPONSES_LITE_METADATA_KEY: &str = + "ws_request_header_x_openai_internal_codex_responses_lite"; + +/// Host-owned provider, authentication, and attribution for one Luna connection. +pub struct LunaSamplerConfig { + /// Provider and credentials selected for the owning thread. + pub provider: SharedModelProvider, + /// Effective proxy, custom-CA, and cookie configuration. + pub http_client_factory: HttpClientFactory, + /// Agent-identity policy selected for the owning thread. + pub agent_identity_policy: AgentIdentityAuthPolicy, + /// Host-resolved source used to scope agent-identity authentication. + pub session_source: SessionSource, + /// Owning runtime session identifier. + pub session_id: String, + /// Owning thread identifier. + pub thread_id: String, + /// Optional host-resolved request originator. + pub originator: Option, + /// Optional inference service tier. + pub service_tier: Option, +} + +/// One tool-less structured Luna request over an already-open connection. +pub struct LunaSamplingRequest { + /// Trusted instructions describing the requested classification. + pub instructions: String, + /// Untrusted input that the model should classify. + pub input: String, + /// Strict JSON schema constraining the model response. + pub output_schema: Value, + /// Reasoning budget explicitly selected for this request. + pub reasoning_effort: ReasoningEffort, + /// Owning turn identifier used for request attribution. + pub turn_id: String, +} + +/// Failures returned while connecting or sampling the Luna model. +#[derive(Debug, Error)] +pub enum LunaSamplerError { + /// The thread's provider or scoped credentials could not be resolved. + #[error("could not resolve the Luna model provider: {0}")] + Provider(#[source] CodexErr), + /// The Responses WebSocket could not be opened or streamed. + #[error("Luna Responses WebSocket failed: {0}")] + Api(#[source] ApiError), + /// The provider's WebSocket connect deadline elapsed. + #[error("Luna Responses WebSocket connection timed out")] + ConnectionTimeout, + /// The response did not contain an assistant text value. + #[error("Luna response did not contain assistant output")] + MissingOutput, + /// The response exceeded the bounded output limit. + #[error("Luna response exceeded the output limit")] + OutputTooLarge, +} + +/// A persistent, authenticated Responses WebSocket dedicated to Luna sampling. +pub struct LunaSampler { + connection: ResponsesWebsocketConnection, + session_id: String, + thread_id: String, + service_tier: Option, +} + +impl LunaSampler { + /// Opens the WebSocket before any sample is requested. + pub async fn connect(config: LunaSamplerConfig) -> Result { + let provider = config + .provider + .api_provider() + .await + .map_err(LunaSamplerError::Provider)?; + let auth = config + .provider + .api_auth_for_scope(ProviderAuthScope { + agent_identity_policy: config.agent_identity_policy, + session_source: config.session_source, + agent_identity_session_fallback: AgentIdentitySessionFallback::default(), + }) + .await + .map_err(LunaSamplerError::Provider)? + .auth; + let mut headers = build_session_headers( + Some(config.session_id.clone()), + Some(config.thread_id.clone()), + ); + headers.insert( + "openai-beta", + HeaderValue::from_static(RESPONSES_WEBSOCKETS_BETA), + ); + if let Some(originator) = config.originator.as_deref() { + add_originator_header(&mut headers, originator); + } + if let Ok(request_id) = HeaderValue::from_str(&config.thread_id) { + headers.insert("x-client-request-id", request_id); + } + + let provider_info = config.provider.info(); + if 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 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 connect = client.connect( + &config.http_client_factory, + headers, + default_headers(), + /*turn_state*/ None, + /*telemetry*/ None, + ); + let connection = tokio::time::timeout(provider_info.websocket_connect_timeout(), connect) + .await + .map_err(|_| LunaSamplerError::ConnectionTimeout)? + .map_err(LunaSamplerError::Api)?; + + Ok(Self { + connection, + session_id: config.session_id, + thread_id: config.thread_id, + service_tier: config.service_tier, + }) + } + + /// Sends one structured, tool-less request on the existing WebSocket. + pub async fn sample(&self, request: LunaSamplingRequest) -> Result { + let metadata = HashMap::from([ + ("session_id".to_owned(), self.session_id.clone()), + ("thread_id".to_owned(), self.thread_id.clone()), + ("turn_id".to_owned(), request.turn_id), + (RESPONSES_LITE_METADATA_KEY.to_owned(), "true".to_owned()), + ]); + let request = ResponsesApiRequest { + model: MODEL.to_owned(), + instructions: String::new(), + input: vec![ + ResponseItem::AdditionalTools { + id: None, + role: "developer".to_owned(), + tools: Vec::new(), + }, + ResponseItem::Message { + id: None, + role: "developer".to_owned(), + content: vec![ContentItem::InputText { + text: request.instructions, + }], + phase: None, + internal_chat_message_metadata_passthrough: None, + }, + ResponseItem::Message { + id: None, + role: "user".to_owned(), + content: vec![ContentItem::InputText { + text: request.input, + }], + phase: None, + internal_chat_message_metadata_passthrough: None, + }, + ], + tools: None, + tool_choice: "none".to_owned(), + parallel_tool_calls: false, + reasoning: Some(Reasoning { + effort: Some(request.reasoning_effort), + summary: None, + context: None, + }), + store: false, + stream: true, + stream_options: None, + include: Vec::new(), + service_tier: self.service_tier.clone(), + prompt_cache_key: Some(format!("guardian-v2:{}", self.thread_id)), + text: create_text_param_for_request( + /*verbosity*/ None, + &Some(request.output_schema), + /*output_schema_strict*/ true, + ), + client_metadata: Some(metadata), + }; + let mut stream = self + .connection + .stream_request( + ResponsesWsRequest::ResponseCreate((&request).into()), + /*connection_reused*/ true, + /*turn_state*/ None, + ) + .await + .map_err(LunaSamplerError::Api)?; + + let mut output = String::new(); + let mut deltas = String::new(); + while let Some(event) = stream.rx_event.recv().await { + match event.map_err(LunaSamplerError::Api)? { + ResponseEvent::OutputTextDelta(delta) => deltas.push_str(&delta), + ResponseEvent::OutputItemDone(ResponseItem::Message { role, content, .. }) + if role == "assistant" => + { + for item in content { + if let ContentItem::OutputText { text } = item { + output.push_str(&text); + } + } + } + ResponseEvent::Completed { .. } => { + if !output.is_empty() { + return Ok(output); + } + if !deltas.is_empty() { + return Ok(deltas); + } + return Err(LunaSamplerError::MissingOutput); + } + _ => {} + } + if output.len() > MAX_OUTPUT_BYTES || deltas.len() > MAX_OUTPUT_BYTES { + return Err(LunaSamplerError::OutputTooLarge); + } + } + + Err(LunaSamplerError::MissingOutput) + } +} + +#[cfg(test)] +#[path = "sampler_tests.rs"] +mod tests; diff --git a/codex-rs/ext/guardian-v2/src/sampler_tests.rs b/codex-rs/ext/guardian-v2/src/sampler_tests.rs new file mode 100644 index 0000000000..1cdfdffa58 --- /dev/null +++ b/codex-rs/ext/guardian-v2/src/sampler_tests.rs @@ -0,0 +1,128 @@ +use anyhow::Result; +use codex_http_client::HttpClientFactory; +use codex_http_client::OutboundProxyPolicy; +use codex_login::AgentIdentityAuthPolicy; +use codex_login::AuthManager; +use codex_login::CodexAuth; +use codex_model_provider::create_model_provider; +use codex_model_provider_info::ModelProviderInfo; +use codex_protocol::openai_models::ReasoningEffort; +use codex_protocol::protocol::SessionSource; +use core_test_support::responses; +use core_test_support::responses::ev_assistant_message; +use core_test_support::responses::ev_completed; +use core_test_support::skip_if_no_network; +use pretty_assertions::assert_eq; +use serde_json::json; + +use super::LunaSampler; +use super::LunaSamplerConfig; +use super::LunaSamplingRequest; + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn preconnected_sampler_reuses_authenticated_websocket_for_structured_requests() -> Result<()> +{ + skip_if_no_network!(Ok(())); + + let server = responses::start_websocket_server(vec![vec![ + vec![ + ev_assistant_message("sample-1", r#"{"score":0.25}"#), + ev_completed("response-1"), + ], + vec![ + ev_assistant_message("sample-2", r#"{"score":0.75}"#), + ev_completed("response-2"), + ], + ]]) + .await; + let base_url = server.uri().replacen("ws://", "http://", 1); + let provider = create_model_provider( + ModelProviderInfo::create_openai_provider(Some(format!("{base_url}/v1"))), + Some(AuthManager::from_auth_for_testing(CodexAuth::from_api_key( + "test-api-key", + ))), + ); + + let sampler = LunaSampler::connect(LunaSamplerConfig { + provider, + http_client_factory: HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + agent_identity_policy: AgentIdentityAuthPolicy::JwtOnly, + session_source: SessionSource::Exec, + session_id: "session-1".to_owned(), + thread_id: "thread-1".to_owned(), + originator: Some("guardian-v2-test".to_owned()), + service_tier: None, + }) + .await?; + + let handshake = server.single_handshake(); + assert!(server.single_connection().is_empty()); + assert_eq!( + handshake.header("authorization"), + Some("Bearer test-api-key".to_owned()) + ); + assert_eq!( + handshake.header("OpenAI-Beta"), + Some("responses_websockets=2026-02-06".to_owned()) + ); + assert_eq!(handshake.header("session-id"), Some("session-1".to_owned())); + assert_eq!(handshake.header("thread-id"), Some("thread-1".to_owned())); + assert_eq!( + handshake.header("originator"), + Some("guardian-v2-test".to_owned()) + ); + + let schema = json!({ + "type": "object", + "properties": { "score": { "type": "number" } }, + "required": ["score"], + "additionalProperties": false + }); + let first = sampler + .sample(LunaSamplingRequest { + instructions: "Return a risk score.".to_owned(), + input: "The user requested a README summary.".to_owned(), + output_schema: schema.clone(), + reasoning_effort: ReasoningEffort::None, + turn_id: "turn-1".to_owned(), + }) + .await?; + let second = sampler + .sample(LunaSamplingRequest { + instructions: "Return a risk score.".to_owned(), + input: "The user requested a source review.".to_owned(), + output_schema: schema, + reasoning_effort: ReasoningEffort::Medium, + turn_id: "turn-2".to_owned(), + }) + .await?; + + assert_eq!(first, r#"{"score":0.25}"#); + assert_eq!(second, r#"{"score":0.75}"#); + let requests = server.single_connection(); + assert_eq!(requests.len(), 2); + for (index, request) in requests.iter().enumerate() { + let request = request.body_json(); + assert_eq!(request["type"], "response.create"); + assert_eq!(request["model"], "gpt-5.6-luna"); + assert_eq!(request["input"][0]["tools"], json!([])); + assert_eq!(request["tool_choice"], "none"); + assert_eq!(request["text"]["format"]["strict"], true); + assert_eq!(request["prompt_cache_key"], "guardian-v2:thread-1"); + assert_eq!(request["client_metadata"]["session_id"], "session-1"); + assert_eq!(request["client_metadata"]["thread_id"], "thread-1"); + assert_eq!( + request["client_metadata"]["ws_request_header_x_openai_internal_codex_responses_lite"], + "true" + ); + assert!(request.get("tools").is_none()); + let effort = if index == 0 { "none" } else { "medium" }; + assert_eq!(request["reasoning"]["effort"], effort); + assert_eq!( + request["client_metadata"]["turn_id"], + format!("turn-{}", index + 1) + ); + } + + Ok(()) +}