codex/codex-rs/ext/guardian-v2/src/sampler.rs

use std::collections::HashMap;
use std::collections::VecDeque;
use std::sync::Arc;
use std::sync::Mutex;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
use std::time::Duration;
use std::time::Instant;

use codex_api::ApiError;
use codex_api::Reasoning;
use codex_api::ReasoningContext;
use codex_api::ResponseEvent;
use codex_api::ResponsesApiRequest;
use codex_api::ResponsesWebsocketClient;
use codex_api::ResponsesWebsocketConnection;
use codex_api::ResponsesWsRequest;
use codex_api::TransportError;
use codex_api::build_session_headers;
use codex_api::create_text_param_for_request;
use codex_extension_api::ExtensionMetrics;
use codex_http_client::HttpClientFactory;
use codex_login::AgentIdentityAuthPolicy;
use codex_login::CodexAuth;
use codex_login::UnauthorizedRecovery;
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::ThreadId;
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 codex_protocol::protocol::TokenUsage;
use http::HeaderValue;
use http::StatusCode;
use serde_json::Value;
use serde_json::json;
use thiserror::Error;
use tokio::sync::OwnedSemaphorePermit;
use tokio::sync::Semaphore;
use tokio::sync::oneshot;

pub(crate) const MODEL: &str = "gpt-5.6-luna";
pub(crate) const CLASSIFICATION_TOKEN_USAGE_METRIC: &str =
    "codex.guardian_v2.classification.token_usage";
const MAX_OUTPUT_BYTES: usize = 8 * 1024;
const INITIAL_WEBSOCKET_CONNECTIONS: usize = 2;
const MAX_WEBSOCKET_CONNECTIONS: usize = 16;
const MAX_SAMPLING_RETRIES: usize = 2;
const MAX_WEBSOCKET_AGE: Duration = Duration::from_secs(55 * 60);
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";
const TURN_METADATA_KEY: &str = "x-codex-turn-metadata";

/// 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>,
    /// Luna model's host-resolved encrypted-compaction compatibility hash.
    pub luna_compaction_hash: Option<String>,
    /// Host-provided metrics capability with the owning session's attribution.
    pub metrics: Option<Arc<dyn ExtensionMetrics>>,
}

/// One tool-less structured Luna request over an already-open connection.
pub struct LunaSamplingRequest {
    /// Trusted instructions describing the requested classification.
    pub instructions: String,
    /// Ordered untrusted input entries that the model should classify.
    pub input: Vec<String>,
    /// Optional bounded screenshots accompanying the transcript.
    pub images: Vec<ContentItem>,
    /// Opaque parent compaction to reuse only for compatible model configurations.
    pub parent_compaction: Option<ResponseItem>,
    /// Current parent model's encrypted-compaction compatibility hash.
    pub parent_compaction_hash: Option<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 newer classification replaced this request when the pool was full.
    #[error("Luna request was superseded by a newer classification")]
    Superseded,
}

struct PooledConnection {
    connection: ResponsesWebsocketConnection,
    // The bridge routes by thread ID, so each socket needs its own identity.
    thread_id: String,
    expires_at: Instant,
    auth_changes: Option<tokio::sync::watch::Receiver<u64>>,
}

struct ConnectionLease {
    connection: PooledConnection,
    idle_connections: Arc<Mutex<Vec<PooledConnection>>>,
    _permit: OwnedSemaphorePermit,
}

impl ConnectionLease {
    fn reuse(self) {
        self.idle_connections
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner)
            .push(self.connection);
    }
}

struct ActiveRequest {
    supersede: oneshot::Sender<()>,
    scored: Arc<AtomicBool>,
}

fn record_token_usage(metrics: Option<&dyn ExtensionMetrics>, token_usage: Option<&TokenUsage>) {
    let (Some(metrics), Some(token_usage)) = (metrics, token_usage) else {
        return;
    };

    for (token_type, value) in [
        ("total", token_usage.total_tokens.max(0)),
        ("input", token_usage.input_tokens.max(0)),
        ("cached_input", token_usage.cached_input()),
        (
            "cache_write_input",
            token_usage.cache_write_input_tokens.max(0),
        ),
        ("non_cached_input", token_usage.non_cached_input()),
        ("output", token_usage.output_tokens.max(0)),
        (
            "reasoning_output",
            token_usage.reasoning_output_tokens.max(0),
        ),
    ] {
        metrics.histogram(
            CLASSIFICATION_TOKEN_USAGE_METRIC,
            value,
            &[("token_type", token_type)],
        );
    }
}

/// A bounded pool of authenticated Responses WebSockets dedicated to Luna sampling.
pub struct LunaSampler {
    config: LunaSamplerConfig,
    idle_connections: Arc<Mutex<Vec<PooledConnection>>>,
    capacity: Arc<Semaphore>,
    active_requests: Mutex<VecDeque<ActiveRequest>>,
}

impl LunaSampler {
    /// Opens the initial WebSockets before any sample is requested.
    pub async fn connect(config: LunaSamplerConfig) -> Result<Self, LunaSamplerError> {
        let sampler = Self {
            config,
            idle_connections: Arc::new(Mutex::new(Vec::with_capacity(MAX_WEBSOCKET_CONNECTIONS))),
            capacity: Arc::new(Semaphore::new(MAX_WEBSOCKET_CONNECTIONS)),
            active_requests: Mutex::new(VecDeque::with_capacity(MAX_WEBSOCKET_CONNECTIONS)),
        };
        for _ in 0..INITIAL_WEBSOCKET_CONNECTIONS {
            let connection = match sampler.open_connection().await {
                Ok(connection) => connection,
                Err(_) => break,
            };
            sampler
                .idle_connections
                .lock()
                .unwrap_or_else(std::sync::PoisonError::into_inner)
                .push(connection);
        }
        Ok(sampler)
    }

    async fn open_connection(&self) -> Result<PooledConnection, LunaSamplerError> {
        let provider = self
            .config
            .provider
            .api_provider()
            .await
            .map_err(LunaSamplerError::Provider)?;
        let auth_manager = self.config.provider.auth_manager();
        let auth_changes = auth_manager.map(|manager| manager.auth_change_receiver());
        let auth = self
            .config
            .provider
            .api_auth_for_scope(ProviderAuthScope {
                agent_identity_policy: self.config.agent_identity_policy,
                session_source: self.config.session_source.clone(),
                agent_identity_session_fallback: AgentIdentitySessionFallback::default(),
            })
            .await
            .map_err(LunaSamplerError::Provider)?
            .auth;
        let thread_id = ThreadId::new().to_string();
        let mut headers = build_session_headers(
            Some(self.config.session_id.clone()),
            Some(thread_id.clone()),
        );
        headers.insert("x-openai-subagent", HeaderValue::from_static("guardian"));
        headers.insert(
            "x-codex-window-id",
            HeaderValue::from_str(&format!("{thread_id}:0")).map_err(|error| {
                LunaSamplerError::Api(ApiError::Stream(format!(
                    "invalid classifier window ID: {error}"
                )))
            })?,
        );
        headers.insert(
            "openai-beta",
            HeaderValue::from_static(RESPONSES_WEBSOCKETS_BETA),
        );
        headers.insert(
            "x-openai-internal-codex-responses-lite",
            HeaderValue::from_static("true"),
        );
        if let Some(originator) = self.config.originator.as_deref() {
            add_originator_header(&mut headers, originator);
        }
        if let Ok(request_id) = HeaderValue::from_str(&thread_id) {
            headers.insert("x-client-request-id", request_id);
        }

        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 connect = client.connect(
            &self.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)?;
        if auth_changes
            .as_ref()
            .is_some_and(|auth| auth.has_changed().unwrap_or(true))
        {
            return Err(LunaSamplerError::Api(ApiError::Stream(
                "authentication changed while connecting".into(),
            )));
        }

        Ok(PooledConnection {
            connection,
            thread_id,
            expires_at: Instant::now() + MAX_WEBSOCKET_AGE,
            auth_changes,
        })
    }

    async fn lease_connection(&self) -> Result<ConnectionLease, LunaSamplerError> {
        let permit = Arc::clone(&self.capacity)
            .acquire_owned()
            .await
            .map_err(|_| LunaSamplerError::ConnectionTimeout)?;
        let connection = loop {
            let idle = self
                .idle_connections
                .lock()
                .unwrap_or_else(std::sync::PoisonError::into_inner)
                .pop();
            match idle {
                Some(connection)
                    if connection
                        .auth_changes
                        .as_ref()
                        .is_none_or(|auth| !auth.has_changed().unwrap_or(true))
                        && Instant::now() < connection.expires_at
                        && !connection.connection.is_closed().await =>
                {
                    break connection;
                }
                Some(_) => {}
                None => break self.open_connection().await?,
            }
        };
        Ok(ConnectionLease {
            connection,
            idle_connections: Arc::clone(&self.idle_connections),
            _permit: permit,
        })
    }

    async fn retry_after_failure(
        &self,
        error: &LunaSamplerError,
        auth_recovery: &mut Option<UnauthorizedRecovery>,
        retries: &mut usize,
    ) -> bool {
        let retryable = match error {
            LunaSamplerError::ConnectionTimeout
            | LunaSamplerError::Api(
                ApiError::Retryable { .. } | ApiError::Stream(_) | ApiError::ServerOverloaded,
            )
            | LunaSamplerError::Api(ApiError::Transport(
                TransportError::RetryLimit
                | TransportError::Timeout
                | TransportError::Connection(_)
                | TransportError::Network(_),
            )) => true,
            LunaSamplerError::Api(ApiError::Transport(TransportError::Http { status, .. }))
            | LunaSamplerError::Api(ApiError::Api { status, .. }) => {
                if *status == StatusCode::UNAUTHORIZED {
                    let Some(recovery) = auth_recovery.as_mut() else {
                        return false;
                    };
                    if !recovery.has_next() || recovery.next().await.is_err() {
                        return false;
                    }
                    self.idle_connections
                        .lock()
                        .unwrap_or_else(std::sync::PoisonError::into_inner)
                        .clear();
                    return true;
                } else {
                    status.is_server_error() || *status == StatusCode::TOO_MANY_REQUESTS
                }
            }
            LunaSamplerError::Provider(_)
            | LunaSamplerError::MissingOutput
            | LunaSamplerError::OutputTooLarge
            | LunaSamplerError::Superseded
            | LunaSamplerError::Api(
                ApiError::Transport(TransportError::Build(_))
                | ApiError::ContextWindowExceeded
                | ApiError::QuotaExceeded
                | ApiError::UsageNotIncluded
                | ApiError::RateLimit(_)
                | ApiError::InvalidRequest { .. }
                | ApiError::MisalignmentPolicyViolation { .. }
                | ApiError::CyberPolicy { .. },
            ) => false,
        };
        if retryable && *retries < MAX_SAMPLING_RETRIES {
            *retries += 1;
            return true;
        }
        false
    }

    /// Sends one structured, tool-less request on an exclusively leased WebSocket.
    pub async fn sample(&self, request: LunaSamplingRequest) -> Result<String, LunaSamplerError> {
        let turn_id = request.turn_id;
        let mut 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,
            },
        ];
        if request
            .parent_compaction_hash
            .as_deref()
            .zip(self.config.luna_compaction_hash.as_deref())
            .is_some_and(|(parent_hash, luna_hash)| {
                !parent_hash.is_empty() && parent_hash == luna_hash
            })
            && let Some(parent_compaction) = request.parent_compaction
        {
            input.push(parent_compaction);
        }
        input.push(ResponseItem::Message {
            id: None,
            role: "user".to_owned(),
            content: request
                .input
                .into_iter()
                .map(|text| ContentItem::InputText { text })
                .chain(request.images.into_iter().map(|mut image| {
                    if let ContentItem::InputImage { detail, .. } = &mut image {
                        *detail = None;
                    }
                    image
                }))
                .collect(),
            phase: None,
            internal_chat_message_metadata_passthrough: None,
        });
        let mut request = ResponsesApiRequest {
            model: MODEL.to_owned(),
            instructions: String::new(),
            input,
            tools: None,
            tool_choice: "none".to_owned(),
            parallel_tool_calls: false,
            reasoning: Some(Reasoning {
                effort: Some(request.reasoning_effort),
                summary: None,
                context: Some(ReasoningContext::AllTurns),
            }),
            store: false,
            stream: true,
            stream_options: None,
            include: Vec::new(),
            service_tier: self.config.service_tier.clone(),
            prompt_cache_key: Some(format!("guardian-v2:{}", self.config.thread_id)),
            text: create_text_param_for_request(
                /*verbosity*/ None,
                &Some(request.output_schema),
                /*output_schema_strict*/ true,
            ),
            client_metadata: None,
        };
        let (supersede, mut superseded) = oneshot::channel();
        let scored = Arc::new(AtomicBool::new(false));
        {
            let mut active_requests = self
                .active_requests
                .lock()
                .unwrap_or_else(std::sync::PoisonError::into_inner);
            active_requests.retain(|request| !request.supersede.is_closed());
            if active_requests.len() == MAX_WEBSOCKET_CONNECTIONS {
                let oldest_scored = active_requests
                    .iter()
                    .position(|request| request.scored.load(Ordering::Relaxed))
                    .unwrap_or(0);
                if let Some(oldest) = active_requests.remove(oldest_scored) {
                    let _ = oldest.supersede.send(());
                }
            }
            active_requests.push_back(ActiveRequest {
                supersede,
                scored: Arc::clone(&scored),
            });
        }
        let mut retries = 0;
        let mut auth_recovery = self
            .config
            .provider
            .auth_manager()
            .map(|manager| manager.unauthorized_recovery());
        'retry: loop {
            let lease = match tokio::select! {
                biased;
                _ = &mut superseded => return Err(LunaSamplerError::Superseded),
                lease = self.lease_connection() => lease,
            } {
                Ok(lease) => lease,
                Err(error) => {
                    if self
                        .retry_after_failure(&error, &mut auth_recovery, &mut retries)
                        .await
                    {
                        continue;
                    }
                    return Err(error);
                }
            };
            let thread_id = &lease.connection.thread_id;
            let turn_metadata = json!({
                "session_id": self.config.session_id,
                "thread_id": thread_id,
                "guardian_classifier_source_thread_id": self.config.thread_id,
                "turn_id": turn_id,
                "request_kind": "guardian_classifier",
                "is_guardian_mode": true,
            })
            .to_string();
            request.client_metadata = Some(HashMap::from([
                ("session_id".to_owned(), self.config.session_id.clone()),
                ("thread_id".to_owned(), thread_id.clone()),
                ("turn_id".to_owned(), turn_id.clone()),
                ("x-openai-subagent".to_owned(), "guardian".to_owned()),
                // Classifier requests do not advance their own context window.
                ("x-codex-window-id".to_owned(), format!("{thread_id}:0")),
                (RESPONSES_LITE_METADATA_KEY.to_owned(), "true".to_owned()),
                (TURN_METADATA_KEY.to_owned(), turn_metadata),
            ]));
            let mut stream = match lease
                .connection
                .connection
                .stream_request(
                    ResponsesWsRequest::ResponseCreate((&request).into()),
                    /*connection_reused*/ true,
                    /*turn_state*/ None,
                )
                .await
            {
                Ok(stream) => stream,
                Err(error) => {
                    let error = LunaSamplerError::Api(error);
                    if self
                        .retry_after_failure(&error, &mut auth_recovery, &mut retries)
                        .await
                    {
                        continue;
                    }
                    return Err(error);
                }
            };

            let mut output = String::new();
            let mut deltas = String::new();
            while let Some(event) = tokio::select! {
                biased;
                _ = &mut superseded => {
                    return if scored.load(Ordering::Relaxed) && !output.is_empty() {
                        Ok(output)
                    } else {
                        Err(LunaSamplerError::Superseded)
                    };
                }
                event = stream.rx_event.recv() => event,
            } {
                let event = match event {
                    Ok(event) => event,
                    Err(error) => {
                        let error = LunaSamplerError::Api(error);
                        if self
                            .retry_after_failure(&error, &mut auth_recovery, &mut retries)
                            .await
                        {
                            continue 'retry;
                        }
                        return Err(error);
                    }
                };
                match event {
                    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 { token_usage, .. } => {
                        record_token_usage(self.config.metrics.as_deref(), token_usage.as_ref());
                        lease.reuse();
                        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);
                }
                if !output.is_empty() {
                    if serde_json::from_str::<serde_json::Map<String, Value>>(&output).is_ok() {
                        scored.store(true, Ordering::Relaxed);
                    }
                    continue;
                }
                if serde_json::from_str::<serde_json::Map<String, Value>>(&deltas).is_ok() {
                    scored.store(true, Ordering::Relaxed);
                    let mut remaining_events = stream.rx_event;
                    let metrics = self.config.metrics.clone();
                    tokio::spawn(async move {
                        while let Some(event) = tokio::select! {
                            biased;
                            _ = &mut superseded => None,
                            event = remaining_events.recv() => event,
                        } {
                            match event {
                                Ok(ResponseEvent::Completed { token_usage, .. }) => {
                                    record_token_usage(metrics.as_deref(), token_usage.as_ref());
                                    lease.reuse();
                                    break;
                                }
                                Err(_) => break,
                                _ => {}
                            }
                        }
                    });
                    return Ok(deltas);
                }
            }
            return Err(LunaSamplerError::MissingOutput);
        }
    }
}

#[cfg(test)]
#[path = "sampler_tests.rs"]
mod tests;