mirror of
https://github.com/openai/codex.git
synced 2026-10-18 17:18:52 +00:00
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
This commit is contained in:
jif
Generated
+13
@@ -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]]
|
||||
|
||||
@@ -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"] }
|
||||
|
||||
@@ -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<C: Sync>(_registry: &mut ExtensionRegistryBuilder<C>) {}
|
||||
|
||||
@@ -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<String>,
|
||||
/// Optional inference service tier.
|
||||
pub service_tier: Option<String>,
|
||||
}
|
||||
|
||||
/// 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<String>,
|
||||
}
|
||||
|
||||
impl LunaSampler {
|
||||
/// Opens the WebSocket before any sample is requested.
|
||||
pub async fn connect(config: LunaSamplerConfig) -> Result<Self, LunaSamplerError> {
|
||||
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<String, LunaSamplerError> {
|
||||
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;
|
||||
@@ -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(())
|
||||
}
|
||||
Reference in New Issue
Block a user