Codex Handbook
app-server/src/outgoing_message.rs 1352 lines
use std::collections::HashMap;use std::sync::Arc;use std::sync::atomic::AtomicI64;use std::sync::atomic::Ordering;use std::time::SystemTime;use std::time::UNIX_EPOCH;use codex_analytics::AnalyticsEventsClient;use codex_app_server_protocol::ClientResponsePayload;use codex_app_server_protocol::JSONRPCErrorError;use codex_app_server_protocol::RequestId;use codex_app_server_protocol::Result;use codex_app_server_protocol::ServerNotification;use codex_app_server_protocol::ServerRequest;use codex_app_server_protocol::ServerRequestPayload;use codex_app_server_protocol::ServerResponse;use codex_otel::span_w3c_trace_context;use codex_protocol::ThreadId;use codex_protocol::protocol::W3cTraceContext;use codex_protocol::request_permissions::RequestPermissionsResponse;use tokio::sync::Mutex;use tokio::sync::mpsc;use tokio::sync::oneshot;use tracing::Instrument;use tracing::Span;use tracing::warn;use crate::error_code::internal_error;use crate::server_request_error::TURN_TRANSITION_PENDING_REQUEST_ERROR_REASON;pub(crate) use codex_app_server_transport::ConnectionId;pub(crate) use codex_app_server_transport::OutgoingError;pub(crate) use codex_app_server_transport::OutgoingMessage;pub(crate) use codex_app_server_transport::OutgoingResponse;pub(crate) use codex_app_server_transport::QueuedOutgoingMessage;#[cfg(test)]use codex_protocol::account::PlanType;pub(crate) type ClientRequestResult = std::result::Result<Result, JSONRPCErrorError>;/// Stable identifier for a client request scoped to a transport connection.#[derive(Clone, Debug, Eq, Hash, PartialEq)]pub(crate) struct ConnectionRequestId {    pub(crate) connection_id: ConnectionId,    pub(crate) request_id: RequestId,}/// Trace data we keep for an incoming request until we send its final/// response or error.#[derive(Clone)]pub(crate) struct RequestContext {    request_id: ConnectionRequestId,    span: Span,    parent_trace: Option<W3cTraceContext>,}impl RequestContext {    pub(crate) fn new(        request_id: ConnectionRequestId,        span: Span,        parent_trace: Option<W3cTraceContext>,    ) -> Self {        Self {            request_id,            span,            parent_trace,        }    }    pub(crate) fn request_trace(&self) -> Option<W3cTraceContext> {        span_w3c_trace_context(&self.span).or_else(|| self.parent_trace.clone())    }    pub(crate) fn span(&self) -> Span {        self.span.clone()    }    fn record_turn_id(&self, turn_id: &str) {        self.span.record("turn.id", turn_id);    }}#[derive(Debug)]pub(crate) enum OutgoingEnvelope {    ToConnection {        connection_id: ConnectionId,        message: OutgoingMessage,        write_complete_tx: Option<oneshot::Sender<()>>,    },    Broadcast {        message: OutgoingMessage,    },}/// Sends messages to the client and manages request callbacks.pub(crate) struct OutgoingMessageSender {    next_server_request_id: AtomicI64,    sender: mpsc::Sender<OutgoingEnvelope>,    request_id_to_callback: Mutex<HashMap<RequestId, PendingCallbackEntry>>,    /// Incoming requests that are still waiting on a final response or error.    /// We keep them here because this is where responses, errors, and    /// disconnect cleanup all get handled.    request_contexts: Mutex<HashMap<ConnectionRequestId, RequestContext>>,    analytics_events_client: AnalyticsEventsClient,}#[derive(Clone)]pub(crate) struct ThreadScopedOutgoingMessageSender {    outgoing: Arc<OutgoingMessageSender>,    connection_ids: Arc<Vec<ConnectionId>>,    thread_id: ThreadId,}struct PendingCallbackEntry {    callback: oneshot::Sender<ClientRequestResult>,    thread_id: Option<ThreadId>,    request: ServerRequest,}impl ThreadScopedOutgoingMessageSender {    pub(crate) fn new(        outgoing: Arc<OutgoingMessageSender>,        connection_ids: Vec<ConnectionId>,        thread_id: ThreadId,    ) -> Self {        Self {            outgoing,            connection_ids: Arc::new(connection_ids),            thread_id,        }    }    pub(crate) async fn send_request(        &self,        payload: ServerRequestPayload,    ) -> (RequestId, oneshot::Receiver<ClientRequestResult>) {        self.outgoing            .send_request_to_connections(                Some(self.connection_ids.as_slice()),                payload,                Some(self.thread_id),            )            .await    }    pub(crate) fn track_effective_permissions_approval_response(        &self,        request_id: RequestId,        response: RequestPermissionsResponse,    ) {        self.outgoing            .analytics_events_client            .track_effective_permissions_approval_response(                now_unix_timestamp_ms(),                request_id,                response,            );    }    pub(crate) async fn send_server_notification(&self, notification: ServerNotification) {        self.outgoing            .analytics_events_client            .track_notification(notification.clone());        if self.connection_ids.is_empty() {            return;        }        self.outgoing            .send_server_notification_to_connections(self.connection_ids.as_slice(), notification)            .await;    }    pub(crate) async fn send_global_server_notification(&self, notification: ServerNotification) {        self.outgoing.send_server_notification(notification).await;    }    pub(crate) async fn abort_pending_server_requests(&self) {        self.outgoing            .cancel_requests_for_thread(                self.thread_id,                Some({                    let mut error = internal_error(                        "client request resolved because the turn state was changed",                    );                    error.data = Some(serde_json::json!({                        "reason": TURN_TRANSITION_PENDING_REQUEST_ERROR_REASON,                    }));                    error                }),            )            .await    }    pub(crate) async fn send_response<T>(&self, request_id: ConnectionRequestId, response: T)    where        T: Into<ClientResponsePayload>,    {        self.outgoing.send_response(request_id, response).await;    }    pub(crate) async fn send_error(        &self,        request_id: ConnectionRequestId,        error: impl Into<JSONRPCErrorError>,    ) {        self.outgoing.send_error(request_id, error).await;    }}impl OutgoingMessageSender {    pub(crate) fn new(        sender: mpsc::Sender<OutgoingEnvelope>,        analytics_events_client: AnalyticsEventsClient,    ) -> Self {        Self {            next_server_request_id: AtomicI64::new(0),            sender,            request_id_to_callback: Mutex::new(HashMap::new()),            request_contexts: Mutex::new(HashMap::new()),            analytics_events_client,        }    }    pub(crate) async fn register_request_context(&self, request_context: RequestContext) {        let mut request_contexts = self.request_contexts.lock().await;        if request_contexts            .insert(request_context.request_id.clone(), request_context)            .is_some()        {            warn!("replaced unresolved request context");        }    }    pub(crate) async fn connection_closed(&self, connection_id: ConnectionId) {        let mut request_contexts = self.request_contexts.lock().await;        request_contexts.retain(|request_id, _| request_id.connection_id != connection_id);    }    pub(crate) async fn request_trace_context(        &self,        request_id: &ConnectionRequestId,    ) -> Option<W3cTraceContext> {        let request_contexts = self.request_contexts.lock().await;        request_contexts            .get(request_id)            .and_then(RequestContext::request_trace)    }    pub(crate) async fn record_request_turn_id(        &self,        request_id: &ConnectionRequestId,        turn_id: &str,    ) {        let request_contexts = self.request_contexts.lock().await;        if let Some(request_context) = request_contexts.get(request_id) {            request_context.record_turn_id(turn_id);        }    }    async fn take_request_context(        &self,        request_id: &ConnectionRequestId,    ) -> Option<RequestContext> {        let mut request_contexts = self.request_contexts.lock().await;        request_contexts.remove(request_id)    }    #[cfg(test)]    async fn request_context_count(&self) -> usize {        self.request_contexts.lock().await.len()    }    pub(crate) async fn send_request(        &self,        request: ServerRequestPayload,    ) -> (RequestId, oneshot::Receiver<ClientRequestResult>) {        self.send_request_to_connections(            /*connection_ids*/ None, request, /*thread_id*/ None,        )        .await    }    fn next_request_id(&self) -> RequestId {        RequestId::Integer(self.next_server_request_id.fetch_add(1, Ordering::Relaxed))    }    pub(crate) async fn send_request_to_connections(        &self,        connection_ids: Option<&[ConnectionId]>,        request: ServerRequestPayload,        thread_id: Option<ThreadId>,    ) -> (RequestId, oneshot::Receiver<ClientRequestResult>) {        let id = self.next_request_id();        let outgoing_message_id = id.clone();        let request = request.request_with_id(outgoing_message_id.clone());        let (tx_approve, rx_approve) = oneshot::channel();        {            let mut request_id_to_callback = self.request_id_to_callback.lock().await;            request_id_to_callback.insert(                id,                PendingCallbackEntry {                    callback: tx_approve,                    thread_id,                    request: request.clone(),                },            );        }        let outgoing_message = OutgoingMessage::Request(request.clone());        let send_result = match connection_ids {            None => {                self.sender                    .send(OutgoingEnvelope::Broadcast {                        message: outgoing_message,                    })                    .await            }            Some(connection_ids) => {                let mut send_error = None;                for connection_id in connection_ids {                    if let Err(err) = self                        .sender                        .send(OutgoingEnvelope::ToConnection {                            connection_id: *connection_id,                            message: outgoing_message.clone(),                            write_complete_tx: None,                        })                        .await                    {                        send_error = Some(err);                        break;                    } else {                        self.analytics_events_client                            .track_server_request(connection_id.0, request.clone());                    }                }                match send_error {                    Some(err) => Err(err),                    None => Ok(()),                }            }        };        if let Err(err) = send_result {            warn!("failed to send request {outgoing_message_id:?} to client: {err:?}");            let mut request_id_to_callback = self.request_id_to_callback.lock().await;            request_id_to_callback.remove(&outgoing_message_id);        }        (outgoing_message_id, rx_approve)    }    pub(crate) async fn replay_requests_to_connection_for_thread(        &self,        connection_id: ConnectionId,        thread_id: ThreadId,    ) {        let requests = self.pending_requests_for_thread(thread_id).await;        for request in requests {            if let Err(err) = self                .sender                .send(OutgoingEnvelope::ToConnection {                    connection_id,                    message: OutgoingMessage::Request(request),                    write_complete_tx: None,                })                .await            {                warn!("failed to resend request to client: {err:?}");            }        }    }    pub(crate) async fn notify_client_response(&self, id: RequestId, result: Result) {        let entry = self.take_request_callback(&id).await;        match entry {            Some((id, entry)) => {                let completed_at_ms = now_unix_timestamp_ms();                if let Ok(response) = entry.request.response_from_result(result.clone())                    && !matches!(response, ServerResponse::PermissionsRequestApproval { .. })                {                    self.analytics_events_client                        .track_server_response(completed_at_ms, response);                }                if let Err(err) = entry.callback.send(Ok(result)) {                    warn!("could not notify callback for {id:?} due to: {err:?}");                }            }            None => {                warn!("could not find callback for {id:?}");            }        }    }    pub(crate) async fn notify_client_error(&self, id: RequestId, error: JSONRPCErrorError) {        let entry = self.take_request_callback(&id).await;        match entry {            Some((id, entry)) => {                warn!("client responded with error for {id:?}: {error:?}");                self.analytics_events_client                    .track_server_request_aborted(now_unix_timestamp_ms(), id.clone());                if let Err(err) = entry.callback.send(Err(error)) {                    warn!("could not notify callback for {id:?} due to: {err:?}");                }            }            None => {                warn!("could not find callback for {id:?}");            }        }    }    pub(crate) async fn cancel_request(&self, id: &RequestId) -> bool {        let entry = self.take_request_callback(id).await;        if let Some((request_id, _entry)) = entry {            self.analytics_events_client                .track_server_request_aborted(now_unix_timestamp_ms(), request_id);            true        } else {            false        }    }    pub(crate) async fn cancel_all_requests(&self, error: Option<JSONRPCErrorError>) {        let entries = {            let mut request_id_to_callback = self.request_id_to_callback.lock().await;            request_id_to_callback                .drain()                .map(|(_, entry)| entry)                .collect::<Vec<_>>()        };        for entry in entries {            self.analytics_events_client                .track_server_request_aborted(now_unix_timestamp_ms(), entry.request.id().clone());            if let Some(error) = error.as_ref()                && let Err(err) = entry.callback.send(Err(error.clone()))            {                let request_id = entry.request.id();                warn!("could not notify callback for {request_id:?} due to: {err:?}");            }        }    }    async fn take_request_callback(        &self,        id: &RequestId,    ) -> Option<(RequestId, PendingCallbackEntry)> {        let mut request_id_to_callback = self.request_id_to_callback.lock().await;        request_id_to_callback.remove_entry(id)    }    pub(crate) async fn pending_requests_for_thread(        &self,        thread_id: ThreadId,    ) -> Vec<ServerRequest> {        let request_id_to_callback = self.request_id_to_callback.lock().await;        let mut requests = request_id_to_callback            .values()            .filter_map(|entry| {                (entry.thread_id == Some(thread_id)).then_some(entry.request.clone())            })            .collect::<Vec<_>>();        requests.sort_by(|left, right| left.id().cmp(right.id()));        requests    }    pub(crate) async fn cancel_requests_for_thread(        &self,        thread_id: ThreadId,        error: Option<JSONRPCErrorError>,    ) {        let entries = {            let mut request_id_to_callback = self.request_id_to_callback.lock().await;            let request_ids = request_id_to_callback                .iter()                .filter_map(|(request_id, entry)| {                    (entry.thread_id == Some(thread_id)).then_some(request_id.clone())                })                .collect::<Vec<_>>();            let mut entries = Vec::with_capacity(request_ids.len());            for request_id in request_ids {                if let Some(entry) = request_id_to_callback.remove(&request_id) {                    entries.push(entry);                }            }            entries        };        for entry in entries {            self.analytics_events_client                .track_server_request_aborted(now_unix_timestamp_ms(), entry.request.id().clone());            if let Some(error) = error.as_ref()                && let Err(err) = entry.callback.send(Err(error.clone()))            {                let request_id = entry.request.id();                warn!("could not notify callback for {request_id:?} due to: {err:?}",);            }        }    }    pub(crate) async fn send_response<T>(&self, request_id: ConnectionRequestId, response: T)    where        T: Into<ClientResponsePayload>,    {        self.send_response_as(request_id, response.into()).await;    }    pub(crate) async fn send_response_as(        &self,        request_id: ConnectionRequestId,        response: ClientResponsePayload,    ) {        let connection_id = request_id.connection_id;        let request_id_for_analytics = request_id.request_id.clone();        let serialized_response = response            .into_jsonrpc_parts_and_payload(request_id.request_id.clone())            .map(|(id, result, response)| {                if let Some(response) = response {                    self.analytics_events_client.track_response(                        connection_id.0,                        request_id_for_analytics,                        response,                    );                }                (id, result)            });        let request_context = self.take_request_context(&request_id).await;        match serialized_response {            Ok((id, result)) => {                let outgoing_message = OutgoingMessage::Response(OutgoingResponse { id, result });                self.send_outgoing_message_to_connection(                    request_context,                    connection_id,                    outgoing_message,                    "response",                )                .await;            }            Err(err) => {                self.send_error_inner(                    request_context,                    request_id,                    internal_error(format!("failed to serialize response: {err}")),                )                .await;            }        }    }    pub(crate) async fn send_server_notification(&self, notification: ServerNotification) {        self.send_server_notification_to_connections(&[], notification)            .await;    }    pub(crate) async fn send_server_notification_to_connections(        &self,        connection_ids: &[ConnectionId],        notification: ServerNotification,    ) {        tracing::trace!(            targeted_connections = connection_ids.len(),            "app-server event: {notification}"        );        let outgoing_message = OutgoingMessage::AppServerNotification(notification.clone());        if connection_ids.is_empty() {            if let Err(err) = self                .sender                .send(OutgoingEnvelope::Broadcast {                    message: outgoing_message,                })                .await            {                warn!("failed to send server notification to client: {err:?}");            }            return;        }        for connection_id in connection_ids {            if let Err(err) = self                .sender                .send(OutgoingEnvelope::ToConnection {                    connection_id: *connection_id,                    message: outgoing_message.clone(),                    write_complete_tx: None,                })                .await            {                warn!("failed to send server notification to client: {err:?}");            }        }    }    pub(crate) async fn send_server_notification_to_connection_and_wait(        &self,        connection_id: ConnectionId,        notification: ServerNotification,    ) {        tracing::trace!("app-server event: {notification}");        let outgoing_message = OutgoingMessage::AppServerNotification(notification.clone());        let (write_complete_tx, write_complete_rx) = oneshot::channel();        if let Err(err) = self            .sender            .send(OutgoingEnvelope::ToConnection {                connection_id,                message: outgoing_message,                write_complete_tx: Some(write_complete_tx),            })            .await        {            warn!("failed to send server notification to client: {err:?}");        }        let _ = write_complete_rx.await;    }    pub(crate) async fn send_error(        &self,        request_id: ConnectionRequestId,        error: impl Into<JSONRPCErrorError>,    ) {        let request_context = self.take_request_context(&request_id).await;        self.send_error_inner(request_context, request_id, error.into())            .await;    }    pub(crate) async fn send_result<T, E>(        &self,        request_id: ConnectionRequestId,        result: std::result::Result<T, E>,    ) where        T: Into<ClientResponsePayload>,        E: Into<JSONRPCErrorError>,    {        match result {            Ok(response) => {                self.send_response(request_id, response).await;            }            Err(error) => self.send_error(request_id, error).await,        }    }    async fn send_error_inner(        &self,        request_context: Option<RequestContext>,        request_id: ConnectionRequestId,        error: JSONRPCErrorError,    ) {        let outgoing_message = OutgoingMessage::Error(OutgoingError {            id: request_id.request_id,            error,        });        self.send_outgoing_message_to_connection(            request_context,            request_id.connection_id,            outgoing_message,            "error",        )        .await;    }    async fn send_outgoing_message_to_connection(        &self,        request_context: Option<RequestContext>,        connection_id: ConnectionId,        message: OutgoingMessage,        message_kind: &'static str,    ) {        let send_fut = self.sender.send(OutgoingEnvelope::ToConnection {            connection_id,            message,            write_complete_tx: None,        });        let send_result = if let Some(request_context) = request_context {            send_fut.instrument(request_context.span()).await        } else {            send_fut.await        };        if let Err(err) = send_result {            warn!("failed to send {message_kind} to client: {err:?}");        }    }}fn now_unix_timestamp_ms() -> u64 {    SystemTime::now()        .duration_since(UNIX_EPOCH)        .unwrap_or_default()        .as_millis()        .try_into()        .unwrap_or_default()}#[cfg(test)]mod tests {    use std::time::Duration;    use codex_app_server_protocol::AccountLoginCompletedNotification;    use codex_app_server_protocol::AccountRateLimitsUpdatedNotification;    use codex_app_server_protocol::AccountUpdatedNotification;    use codex_app_server_protocol::ApplyPatchApprovalParams;    use codex_app_server_protocol::AuthMode;    use codex_app_server_protocol::CommandExecutionApprovalDecision;    use codex_app_server_protocol::CommandExecutionRequestApprovalParams;    use codex_app_server_protocol::ConfigWarningNotification;    use codex_app_server_protocol::DynamicToolCallParams;    use codex_app_server_protocol::FileChangeRequestApprovalParams;    use codex_app_server_protocol::GuardianWarningNotification;    use codex_app_server_protocol::ModelRerouteReason;    use codex_app_server_protocol::ModelReroutedNotification;    use codex_app_server_protocol::ModelVerification;    use codex_app_server_protocol::ModelVerificationNotification;    use codex_app_server_protocol::RateLimitSnapshot;    use codex_app_server_protocol::RateLimitWindow;    use codex_app_server_protocol::ServerResponse;    use codex_app_server_protocol::ToolRequestUserInputParams;    use codex_app_server_protocol::TurnModerationMetadataNotification;    use codex_protocol::ThreadId;    use pretty_assertions::assert_eq;    use serde_json::json;    use std::sync::Arc;    use tokio::time::timeout;    use uuid::Uuid;    use super::*;    #[test]    fn verify_server_notification_serialization() {        let notification =            ServerNotification::AccountLoginCompleted(AccountLoginCompletedNotification {                login_id: Some(Uuid::nil().to_string()),                success: true,                error: None,            });        let jsonrpc_notification = OutgoingMessage::AppServerNotification(notification);        assert_eq!(            json!({                "method": "account/login/completed",                "params": {                    "loginId": Uuid::nil().to_string(),                    "success": true,                    "error": null,                },            }),            serde_json::to_value(jsonrpc_notification)                .expect("ensure the strum macros serialize the method field correctly"),            "ensure the strum macros serialize the method field correctly"        );    }    #[test]    fn verify_account_login_completed_notification_serialization() {        let notification =            ServerNotification::AccountLoginCompleted(AccountLoginCompletedNotification {                login_id: Some(Uuid::nil().to_string()),                success: true,                error: None,            });        let jsonrpc_notification = OutgoingMessage::AppServerNotification(notification);        assert_eq!(            json!({                "method": "account/login/completed",                "params": {                    "loginId": Uuid::nil().to_string(),                    "success": true,                    "error": null,                },            }),            serde_json::to_value(jsonrpc_notification)                .expect("ensure the notification serializes correctly"),            "ensure the notification serializes correctly"        );    }    #[test]    fn verify_account_rate_limits_notification_serialization() {        let notification =            ServerNotification::AccountRateLimitsUpdated(AccountRateLimitsUpdatedNotification {                rate_limits: RateLimitSnapshot {                    limit_id: Some("codex".to_string()),                    limit_name: None,                    primary: Some(RateLimitWindow {                        used_percent: 25,                        window_duration_mins: Some(15),                        resets_at: Some(123),                    }),                    secondary: None,                    credits: None,                    individual_limit: None,                    plan_type: Some(PlanType::Plus),                    rate_limit_reached_type: None,                },            });        let jsonrpc_notification = OutgoingMessage::AppServerNotification(notification);        assert_eq!(            json!({                "method": "account/rateLimits/updated",                "params": {                        "rateLimits": {                        "limitId": "codex",                        "limitName": null,                        "primary": {                            "usedPercent": 25,                            "windowDurationMins": 15,                            "resetsAt": 123                        },                        "secondary": null,                        "credits": null,                        "individualLimit": null,                        "planType": "plus",                        "rateLimitReachedType": null                    }                },            }),            serde_json::to_value(jsonrpc_notification)                .expect("ensure the notification serializes correctly"),            "ensure the notification serializes correctly"        );    }    #[test]    fn verify_account_updated_notification_serialization() {        let notification = ServerNotification::AccountUpdated(AccountUpdatedNotification {            auth_mode: Some(AuthMode::ApiKey),            plan_type: None,        });        let jsonrpc_notification = OutgoingMessage::AppServerNotification(notification);        assert_eq!(            json!({                "method": "account/updated",                "params": {                    "authMode": "apikey",                    "planType": null                },            }),            serde_json::to_value(jsonrpc_notification)                .expect("ensure the notification serializes correctly"),            "ensure the notification serializes correctly"        );    }    #[test]    fn verify_config_warning_notification_serialization() {        let notification = ServerNotification::ConfigWarning(ConfigWarningNotification {            summary: "Config error: using defaults".to_string(),            details: Some("error loading config: bad config".to_string()),            path: None,            range: None,        });        let jsonrpc_notification = OutgoingMessage::AppServerNotification(notification);        assert_eq!(            json!( {                "method": "configWarning",                "params": {                    "summary": "Config error: using defaults",                    "details": "error loading config: bad config",                },            }),            serde_json::to_value(jsonrpc_notification)                .expect("ensure the notification serializes correctly"),            "ensure the notification serializes correctly"        );    }    #[test]    fn verify_guardian_warning_notification_serialization() {        let notification = ServerNotification::GuardianWarning(GuardianWarningNotification {            thread_id: "thread-1".to_string(),            message: "Automatic approval review denied the requested action.".to_string(),        });        let jsonrpc_notification = OutgoingMessage::AppServerNotification(notification);        assert_eq!(            json!({                "method": "guardianWarning",                "params": {                    "threadId": "thread-1",                    "message": "Automatic approval review denied the requested action.",                },            }),            serde_json::to_value(jsonrpc_notification)                .expect("ensure the notification serializes correctly"),            "ensure the notification serializes correctly"        );    }    #[test]    fn verify_model_rerouted_notification_serialization() {        let notification = ServerNotification::ModelRerouted(ModelReroutedNotification {            thread_id: "thread-1".to_string(),            turn_id: "turn-1".to_string(),            from_model: "gpt-5.3-codex".to_string(),            to_model: "gpt-5.2".to_string(),            reason: ModelRerouteReason::HighRiskCyberActivity,        });        let jsonrpc_notification = OutgoingMessage::AppServerNotification(notification);        assert_eq!(            json!({                "method": "model/rerouted",                "params": {                    "threadId": "thread-1",                    "turnId": "turn-1",                    "fromModel": "gpt-5.3-codex",                    "toModel": "gpt-5.2",                    "reason": "highRiskCyberActivity",                },            }),            serde_json::to_value(jsonrpc_notification)                .expect("ensure the notification serializes correctly"),            "ensure the notification serializes correctly"        );    }    #[test]    fn verify_model_verification_notification_serialization() {        let notification = ServerNotification::ModelVerification(ModelVerificationNotification {            thread_id: "thread-1".to_string(),            turn_id: "turn-1".to_string(),            verifications: vec![ModelVerification::TrustedAccessForCyber],        });        let jsonrpc_notification = OutgoingMessage::AppServerNotification(notification);        assert_eq!(            json!({                "method": "model/verification",                "params": {                    "threadId": "thread-1",                    "turnId": "turn-1",                    "verifications": ["trustedAccessForCyber"],                },            }),            serde_json::to_value(jsonrpc_notification)                .expect("ensure the notification serializes correctly"),            "ensure the notification serializes correctly"        );    }    #[test]    fn verify_turn_moderation_metadata_notification_serialization() {        let notification =            ServerNotification::TurnModerationMetadata(TurnModerationMetadataNotification {                thread_id: "thread-1".to_string(),                turn_id: "turn-1".to_string(),                metadata: json!({"presentation": "inline"}),            });        let jsonrpc_notification = OutgoingMessage::AppServerNotification(notification);        assert_eq!(            json!({                "method": "turn/moderationMetadata",                "params": {                    "threadId": "thread-1",                    "turnId": "turn-1",                    "metadata": {"presentation": "inline"},                },            }),            serde_json::to_value(jsonrpc_notification)                .expect("ensure the notification serializes correctly"),            "ensure the notification serializes correctly"        );    }    #[test]    fn server_request_response_from_result_decodes_typed_response() {        let request = ServerRequest::CommandExecutionRequestApproval {            request_id: RequestId::Integer(7),            params: CommandExecutionRequestApprovalParams {                thread_id: "thread-1".to_string(),                turn_id: "turn-1".to_string(),                item_id: "item-1".to_string(),                started_at_ms: 0,                approval_id: None,                reason: None,                network_approval_context: None,                command: Some("echo hi".to_string()),                cwd: None,                command_actions: None,                additional_permissions: None,                proposed_execpolicy_amendment: None,                proposed_network_policy_amendments: None,                available_decisions: None,            },        };        let response = request            .response_from_result(json!({                "decision": "acceptForSession",            }))            .expect("decode typed server response");        let ServerResponse::CommandExecutionRequestApproval {            request_id,            response,        } = response        else {            panic!("expected command execution approval response");        };        assert_eq!(request_id, RequestId::Integer(7));        assert_eq!(            response.decision,            CommandExecutionApprovalDecision::AcceptForSession        );    }    #[tokio::test]    async fn send_response_routes_to_target_connection() {        let (tx, mut rx) = mpsc::channel::<OutgoingEnvelope>(4);        let outgoing =            OutgoingMessageSender::new(tx, codex_analytics::AnalyticsEventsClient::disabled());        let request_id = ConnectionRequestId {            connection_id: ConnectionId(42),            request_id: RequestId::Integer(7),        };        outgoing            .send_response(                request_id.clone(),                ClientResponsePayload::ThreadArchive(                    codex_app_server_protocol::ThreadArchiveResponse {},                ),            )            .await;        let envelope = timeout(Duration::from_secs(1), rx.recv())            .await            .expect("should receive envelope before timeout")            .expect("channel should contain one message");        match envelope {            OutgoingEnvelope::ToConnection {                connection_id,                message,                ..            } => {                assert_eq!(connection_id, ConnectionId(42));                let OutgoingMessage::Response(response) = message else {                    panic!("expected response message");                };                assert_eq!(response.id, request_id.request_id);                assert_eq!(response.result, json!({}));            }            other => panic!("expected targeted response envelope, got: {other:?}"),        }    }    #[tokio::test]    async fn send_response_clears_registered_request_context() {        let (tx, _rx) = mpsc::channel::<OutgoingEnvelope>(4);        let outgoing =            OutgoingMessageSender::new(tx, codex_analytics::AnalyticsEventsClient::disabled());        let request_id = ConnectionRequestId {            connection_id: ConnectionId(42),            request_id: RequestId::Integer(7),        };        outgoing            .register_request_context(RequestContext::new(                request_id.clone(),                tracing::info_span!("app_server.request", rpc.method = "thread/start"),                /*parent_trace*/ None,            ))            .await;        assert_eq!(outgoing.request_context_count().await, 1);        outgoing            .send_response(                request_id,                ClientResponsePayload::ThreadArchive(                    codex_app_server_protocol::ThreadArchiveResponse {},                ),            )            .await;        assert_eq!(outgoing.request_context_count().await, 0);    }    #[tokio::test]    async fn send_error_routes_to_target_connection() {        let (tx, mut rx) = mpsc::channel::<OutgoingEnvelope>(4);        let outgoing =            OutgoingMessageSender::new(tx, codex_analytics::AnalyticsEventsClient::disabled());        let request_id = ConnectionRequestId {            connection_id: ConnectionId(9),            request_id: RequestId::Integer(3),        };        let error = internal_error("boom");        outgoing.send_error(request_id.clone(), error.clone()).await;        let envelope = timeout(Duration::from_secs(1), rx.recv())            .await            .expect("should receive envelope before timeout")            .expect("channel should contain one message");        match envelope {            OutgoingEnvelope::ToConnection {                connection_id,                message,                ..            } => {                assert_eq!(connection_id, ConnectionId(9));                let OutgoingMessage::Error(outgoing_error) = message else {                    panic!("expected error message");                };                assert_eq!(outgoing_error.id, RequestId::Integer(3));                assert_eq!(outgoing_error.error, error);            }            other => panic!("expected targeted error envelope, got: {other:?}"),        }    }    #[tokio::test]    async fn send_server_notification_to_connection_and_wait_tracks_write_completion() {        let (tx, mut rx) = mpsc::channel::<OutgoingEnvelope>(4);        let outgoing =            OutgoingMessageSender::new(tx, codex_analytics::AnalyticsEventsClient::disabled());        let send_task = tokio::spawn(async move {            outgoing                .send_server_notification_to_connection_and_wait(                    ConnectionId(42),                    ServerNotification::ModelRerouted(ModelReroutedNotification {                        thread_id: "thread-1".to_string(),                        turn_id: "turn-1".to_string(),                        from_model: "gpt-5.3-codex".to_string(),                        to_model: "gpt-5.2".to_string(),                        reason: ModelRerouteReason::HighRiskCyberActivity,                    }),                )                .await        });        let envelope = timeout(Duration::from_secs(1), rx.recv())            .await            .expect("should receive envelope before timeout")            .expect("channel should contain one message");        let OutgoingEnvelope::ToConnection {            connection_id,            message,            write_complete_tx,        } = envelope        else {            panic!("expected targeted server notification envelope");        };        assert_eq!(connection_id, ConnectionId(42));        assert!(matches!(message, OutgoingMessage::AppServerNotification(_)));        write_complete_tx            .expect("write completion sender should be attached")            .send(())            .expect("receiver should still be waiting");        timeout(Duration::from_secs(1), send_task)            .await            .expect("send task should finish after write completion is signaled")            .expect("send task should not panic");    }    #[tokio::test]    async fn connection_closed_clears_registered_request_contexts() {        let (tx, _rx) = mpsc::channel::<OutgoingEnvelope>(4);        let outgoing =            OutgoingMessageSender::new(tx, codex_analytics::AnalyticsEventsClient::disabled());        let closed_connection_request = ConnectionRequestId {            connection_id: ConnectionId(9),            request_id: RequestId::Integer(3),        };        let open_connection_request = ConnectionRequestId {            connection_id: ConnectionId(10),            request_id: RequestId::Integer(4),        };        outgoing            .register_request_context(RequestContext::new(                closed_connection_request,                tracing::info_span!("app_server.request", rpc.method = "turn/interrupt"),                /*parent_trace*/ None,            ))            .await;        outgoing            .register_request_context(RequestContext::new(                open_connection_request,                tracing::info_span!("app_server.request", rpc.method = "turn/start"),                /*parent_trace*/ None,            ))            .await;        assert_eq!(outgoing.request_context_count().await, 2);        outgoing.connection_closed(ConnectionId(9)).await;        assert_eq!(outgoing.request_context_count().await, 1);    }    #[tokio::test]    async fn notify_client_error_forwards_error_to_waiter() {        let (tx, _rx) = mpsc::channel::<OutgoingEnvelope>(4);        let outgoing =            OutgoingMessageSender::new(tx, codex_analytics::AnalyticsEventsClient::disabled());        let (request_id, wait_for_result) = outgoing            .send_request(ServerRequestPayload::ApplyPatchApproval(                ApplyPatchApprovalParams {                    conversation_id: ThreadId::new(),                    call_id: "call-id".to_string(),                    file_changes: HashMap::new(),                    reason: None,                    grant_root: None,                },            ))            .await;        let error = internal_error("refresh failed");        outgoing            .notify_client_error(request_id, error.clone())            .await;        let result = timeout(Duration::from_secs(1), wait_for_result)            .await            .expect("wait should not time out")            .expect("waiter should receive a callback");        assert_eq!(result, Err(error));    }    #[tokio::test]    async fn pending_requests_for_thread_returns_thread_requests_in_request_id_order() {        let (tx, _rx) = mpsc::channel::<OutgoingEnvelope>(8);        let outgoing = Arc::new(OutgoingMessageSender::new(            tx,            codex_analytics::AnalyticsEventsClient::disabled(),        ));        let thread_id = ThreadId::new();        let thread_outgoing = ThreadScopedOutgoingMessageSender::new(            outgoing.clone(),            vec![ConnectionId(1)],            thread_id,        );        let (dynamic_tool_request_id, _dynamic_tool_waiter) = thread_outgoing            .send_request(ServerRequestPayload::DynamicToolCall(                DynamicToolCallParams {                    thread_id: thread_id.to_string(),                    turn_id: "turn-1".to_string(),                    call_id: "call-0".to_string(),                    namespace: None,                    tool: "tool".to_string(),                    arguments: json!({}),                },            ))            .await;        let (first_request_id, _first_waiter) = thread_outgoing            .send_request(ServerRequestPayload::ToolRequestUserInput(                ToolRequestUserInputParams {                    thread_id: thread_id.to_string(),                    turn_id: "turn-1".to_string(),                    item_id: "call-1".to_string(),                    questions: vec![],                    auto_resolution_ms: None,                },            ))            .await;        let (second_request_id, _second_waiter) = thread_outgoing            .send_request(ServerRequestPayload::FileChangeRequestApproval(                FileChangeRequestApprovalParams {                    thread_id: thread_id.to_string(),                    turn_id: "turn-1".to_string(),                    item_id: "call-2".to_string(),                    started_at_ms: 0,                    reason: None,                    grant_root: None,                },            ))            .await;        let pending_requests = outgoing.pending_requests_for_thread(thread_id).await;        assert_eq!(            pending_requests                .iter()                .map(ServerRequest::id)                .collect::<Vec<_>>(),            vec![                &dynamic_tool_request_id,                &first_request_id,                &second_request_id            ]        );    }    #[tokio::test]    async fn cancel_requests_for_thread_cancels_all_thread_requests() {        let (tx, _rx) = mpsc::channel::<OutgoingEnvelope>(8);        let outgoing = Arc::new(OutgoingMessageSender::new(            tx,            codex_analytics::AnalyticsEventsClient::disabled(),        ));        let thread_id = ThreadId::new();        let thread_outgoing = ThreadScopedOutgoingMessageSender::new(            outgoing.clone(),            vec![ConnectionId(1)],            thread_id,        );        let (_dynamic_tool_request_id, dynamic_tool_waiter) = thread_outgoing            .send_request(ServerRequestPayload::DynamicToolCall(                DynamicToolCallParams {                    thread_id: thread_id.to_string(),                    turn_id: "turn-1".to_string(),                    call_id: "call-0".to_string(),                    namespace: None,                    tool: "tool".to_string(),                    arguments: json!({}),                },            ))            .await;        let (_request_id, user_input_waiter) = thread_outgoing            .send_request(ServerRequestPayload::ToolRequestUserInput(                ToolRequestUserInputParams {                    thread_id: thread_id.to_string(),                    turn_id: "turn-1".to_string(),                    item_id: "call-1".to_string(),                    questions: vec![],                    auto_resolution_ms: None,                },            ))            .await;        let error = internal_error("tracked request cancelled");        outgoing            .cancel_requests_for_thread(thread_id, Some(error.clone()))            .await;        let dynamic_tool_result = timeout(Duration::from_secs(1), dynamic_tool_waiter)            .await            .expect("dynamic tool waiter should resolve")            .expect("dynamic tool waiter should receive a callback");        let user_input_result = timeout(Duration::from_secs(1), user_input_waiter)            .await            .expect("user input waiter should resolve")            .expect("user input waiter should receive a callback");        assert_eq!(dynamic_tool_result, Err(error.clone()));        assert_eq!(user_input_result, Err(error));        assert!(            outgoing                .pending_requests_for_thread(thread_id)                .await                .is_empty()        );    }}