Aether/apps/aether-gateway/src/handlers/proxy/websocket/responses/frame.rs

//! Parsed OpenAI Responses WebSocket text frames.
//!
//! A relay frame is parsed once and then shared by the protocol adapter, turn
//! accounting, retry safety, and connection lifecycle code.  Keeping the raw
//! text as a borrow avoids copying the websocket payload while the relay is
//! processing it.

use serde_json::Value;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(super) struct ResponsesWebSocketFrameTerminal {
    pub(super) status_code: u16,
    pub(super) cancelled: bool,
}

#[derive(Debug)]
pub(super) struct ParsedResponsesWebSocketFrame<'a> {
    raw_text: &'a str,
    event: Value,
    event_type: Option<String>,
    status: Option<u16>,
    started: bool,
    terminal: Option<ResponsesWebSocketFrameTerminal>,
    terminal_event: Option<Value>,
    chunked: bool,
}

impl<'a> ParsedResponsesWebSocketFrame<'a> {
    pub(super) fn parse(raw_text: &'a str) -> serde_json::Result<Self> {
        let event = serde_json::from_str::<Value>(raw_text)?;
        let events = protocol_events_of(&event);
        let started = events.iter().copied().any(event_is_started);
        // A batch carries at most one terminal in practice.  Taking the first
        // in document order keeps the outcome deterministic if that ever
        // stops being true.
        let terminal_entry = events
            .iter()
            .copied()
            .find_map(|candidate| terminal_for_event(candidate).map(|term| (candidate, term)));
        let terminal = terminal_entry.map(|(_, terminal)| terminal);
        // The terminal event describes the turn's outcome, so it is the one
        // worth naming in logs and recording as the terminal error body.
        let event_type = terminal_entry
            .map(|(candidate, _)| candidate)
            .or_else(|| events.last().copied())
            .and_then(event_type_of)
            .map(str::to_string);
        let terminal_event = terminal_entry.map(|(candidate, _)| candidate.clone());
        let chunked = event.get("chunks").and_then(Value::as_array).is_some();
        let status = terminal.map(|terminal| terminal.status_code);

        Ok(Self {
            raw_text,
            event,
            event_type,
            status,
            started,
            terminal,
            terminal_event,
            chunked,
        })
    }

    /// The protocol events this frame carries.
    ///
    /// Codex batches standard `response.*` events into a `{"chunks":[...]}`
    /// envelope, so one frame can carry several events — and the terminal one
    /// may be buried inside the batch.  Every consumer that interprets event
    /// semantics must walk this rather than the envelope, or a batched
    /// `response.completed` goes unnoticed and wedges the turn.
    pub(super) fn protocol_events(&self) -> Vec<&Value> {
        protocol_events_of(&self.event)
    }

    /// The individual event that ended the turn, unwrapped from its batch.
    pub(super) fn terminal_event(&self) -> Option<&Value> {
        self.terminal_event.as_ref()
    }

    pub(super) fn is_chunked(&self) -> bool {
        self.chunked
    }

    pub(super) fn raw_text(&self) -> &'a str {
        self.raw_text
    }

    pub(super) fn event(&self) -> &Value {
        &self.event
    }

    pub(super) fn event_type(&self) -> Option<&str> {
        self.event_type.as_deref()
    }

    pub(super) fn status(&self) -> Option<u16> {
        self.status
    }

    pub(super) fn is_started(&self) -> bool {
        self.started
    }

    pub(super) fn is_terminal(&self) -> bool {
        self.terminal.is_some()
    }

    pub(super) fn terminal(&self) -> Option<ResponsesWebSocketFrameTerminal> {
        self.terminal
    }

    /// Return a bounded label suitable for structured logs.  Event payloads
    /// are never inserted directly into a log field.
    pub(super) fn event_type_for_log(&self) -> String {
        self.event_type
            .as_deref()
            .map(safe_websocket_event_label)
            .unwrap_or_else(|| "invalid_json".to_string())
    }
}

/// Flattens a frame into the events it carries.  An envelope may name its own
/// `type` *and* batch further events under `chunks`; both are protocol events.
fn protocol_events_of(event: &Value) -> Vec<&Value> {
    let mut events = Vec::new();
    if event_type_of(event).is_some() {
        events.push(event);
    }
    if let Some(chunks) = event.get("chunks").and_then(Value::as_array) {
        events.extend(chunks.iter().filter(|chunk| event_type_of(chunk).is_some()));
    }
    // An unrecognized shape is still relayed and still accounted for, so it
    // must not vanish from the observer's view of the stream.
    if events.is_empty() {
        events.push(event);
    }
    events
}

fn event_type_of(event: &Value) -> Option<&str> {
    event.get("type").and_then(Value::as_str)
}

fn event_is_started(event: &Value) -> bool {
    matches!(
        event_type_of(event).unwrap_or_default(),
        "response.created" | "response.in_progress" | "response.queued"
    )
}

/// `response.incomplete` 的合法终态 reason 白名单。
///
/// 这些 reason 表示上游按规则正常结束了本轮响应（写满 `max_output_tokens`、
/// 命中内容过滤、按工具调用截断），标准流解析里它们会变成 `length` /
/// `content_filter` / `tool_calls` 这类正常 finish，和
/// `openai_responses_incomplete_finish_reason` 的既有映射保持一致，因此不能
/// 当成 provider failure 记账。
const LEGITIMATE_RESPONSES_INCOMPLETE_REASONS: [&str; 5] = [
    "max_output_tokens",
    "max_tokens",
    "content_filter",
    "tool_calls",
    "function_call",
];

/// 读取 `response.incomplete` 携带的 `incomplete_details.reason`。
///
/// 标准位置是 `response.incomplete_details.reason`；批量封装偶尔把
/// `incomplete_details` 直接放在事件顶层，两处都要看，否则合法终态会被漏判。
fn responses_incomplete_reason(event: &Value) -> Option<&str> {
    [
        event.pointer("/response/incomplete_details/reason"),
        event.pointer("/incomplete_details/reason"),
    ]
    .into_iter()
    .flatten()
    .filter_map(Value::as_str)
    .map(str::trim)
    .find(|reason| !reason.is_empty())
}

/// 判断一个 `response.incomplete` 是否是合法终态。
///
/// reason 缺失或不在白名单内（例如 `error`、`server_error`）时继续按
/// provider failure 处理：这类 incomplete 说明上游确实没能正常收尾，仍应扣
/// 供应商健康分。
fn responses_incomplete_is_legitimate_terminal(event: &Value) -> bool {
    responses_incomplete_reason(event).is_some_and(|reason| {
        LEGITIMATE_RESPONSES_INCOMPLETE_REASONS
            .iter()
            .any(|candidate| reason.eq_ignore_ascii_case(candidate))
    })
}

fn terminal_for_event(event: &Value) -> Option<ResponsesWebSocketFrameTerminal> {
    match event_type_of(event).unwrap_or_default() {
        "response.completed" => Some(ResponsesWebSocketFrameTerminal {
            status_code: websocket_event_status_code(event, 200),
            cancelled: false,
        }),
        // 合法 incomplete（例如写满 max_output_tokens）是正常终态，默认按 200
        // 记账，不再一律当 502 provider failure；reason 缺失或未知时保留原来的
        // 502 默认值。显式 `status_code` 和 error code 映射仍然优先于默认值，
        // 所以带 `rate_limit_exceeded` 的 incomplete 依旧是 429。
        "response.incomplete" => Some(ResponsesWebSocketFrameTerminal {
            status_code: websocket_event_status_code(
                event,
                if responses_incomplete_is_legitimate_terminal(event) {
                    200
                } else {
                    502
                },
            ),
            cancelled: false,
        }),
        "response.cancelled" => Some(ResponsesWebSocketFrameTerminal {
            status_code: 499,
            cancelled: true,
        }),
        "response.failed" => Some(ResponsesWebSocketFrameTerminal {
            status_code: websocket_event_status_code(event, 502),
            cancelled: false,
        }),
        "error" => Some(ResponsesWebSocketFrameTerminal {
            status_code: websocket_event_status_code(event, 502),
            cancelled: false,
        }),
        _ => None,
    }
}

fn websocket_event_status_code(event: &Value, default: u16) -> u16 {
    if let Some(status_code) = event
        .get("status_code")
        .or_else(|| event.get("status"))
        .or_else(|| {
            event
                .get("response")
                .and_then(|response| response.get("status_code"))
        })
        .and_then(Value::as_u64)
        .and_then(|value| u16::try_from(value).ok())
        .filter(|value| *value > 0)
    {
        return status_code;
    }

    let error_code = [
        event.pointer("/error/type"),
        event.pointer("/error/code"),
        event.pointer("/response/error/type"),
        event.pointer("/response/error/code"),
    ]
    .into_iter()
    .flatten()
    .filter_map(Value::as_str)
    .map(str::to_ascii_lowercase)
    .find(|value| !value.trim().is_empty());
    match error_code.as_deref() {
        Some(
            "usage_limit_reached" | "insufficient_quota" | "rate_limit_exceeded" | "quota_exceeded",
        ) => 429,
        Some("invalid_api_key" | "authentication_error") => 401,
        Some("invalid_request_error" | "invalid_request" | "model_not_found") => 400,
        Some("overloaded" | "server_error" | "service_unavailable") => 503,
        _ => default,
    }
}

fn safe_websocket_event_label(value: &str) -> String {
    let value = value.trim();
    if value.is_empty()
        || value.len() > 80
        || !value
            .bytes()
            .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'.' | b'_' | b'-'))
    {
        return "unknown".to_string();
    }
    value.to_string()
}

#[cfg(test)]
mod tests {
    use super::ParsedResponsesWebSocketFrame;

    #[test]
    fn parses_started_frame_once_with_raw_text_and_event_metadata() {
        let raw = r#"{"type":"response.in_progress","response":{"status":200}}"#;
        let frame = ParsedResponsesWebSocketFrame::parse(raw).expect("valid frame");

        assert_eq!(frame.raw_text(), raw);
        assert_eq!(frame.event_type(), Some("response.in_progress"));
        assert_eq!(frame.status(), None);
        assert!(frame.is_started());
        assert!(!frame.is_terminal());
        assert_eq!(frame.event()["response"]["status"], 200);
        assert_eq!(frame.event_type_for_log(), "response.in_progress");
    }

    #[test]
    fn classifies_terminal_status_and_cancellation() {
        let completed = ParsedResponsesWebSocketFrame::parse(
            r#"{"type":"response.completed","status_code":201}"#,
        )
        .expect("valid frame");
        assert_eq!(completed.status(), Some(201));
        assert_eq!(
            completed
                .terminal()
                .map(|terminal| (terminal.status_code, terminal.cancelled)),
            Some((201, false))
        );

        let cancelled = ParsedResponsesWebSocketFrame::parse(r#"{"type":"response.cancelled"}"#)
            .expect("valid frame");
        assert_eq!(cancelled.status(), Some(499));
        assert_eq!(
            cancelled
                .terminal()
                .map(|terminal| (terminal.status_code, terminal.cancelled)),
            Some((499, true))
        );

        let error = ParsedResponsesWebSocketFrame::parse(
            r#"{"type":"error","status_code":429,"error":{"type":"usage_limit_reached"}}"#,
        )
        .expect("valid frame");
        assert_eq!(error.status(), Some(429));
        assert!(error.is_terminal());

        let failed = ParsedResponsesWebSocketFrame::parse(
            r#"{"type":"response.failed","response":{"error":{"code":"rate_limit_exceeded"}}}"#,
        )
        .expect("valid frame");
        assert_eq!(failed.status(), Some(429));
    }

    #[test]
    fn a_legitimate_incomplete_is_a_terminal_but_not_a_provider_failure() {
        for reason in [
            "max_output_tokens",
            "max_tokens",
            "content_filter",
            "tool_calls",
            "function_call",
            "MAX_OUTPUT_TOKENS",
        ] {
            let raw = format!(
                r#"{{"type":"response.incomplete","response":{{"status":"incomplete","incomplete_details":{{"reason":"{reason}"}}}}}}"#
            );
            let frame = ParsedResponsesWebSocketFrame::parse(&raw).expect("valid frame");

            assert!(frame.is_terminal(), "{reason} should end the turn");
            assert_eq!(
                frame
                    .terminal()
                    .map(|terminal| (terminal.status_code, terminal.cancelled)),
                Some((200, false)),
                "{reason} is a legitimate terminal result, not a 502 provider failure"
            );
        }
    }

    #[test]
    fn a_top_level_incomplete_details_reason_is_also_honored() {
        let frame = ParsedResponsesWebSocketFrame::parse(
            r#"{"type":"response.incomplete","incomplete_details":{"reason":"max_output_tokens"}}"#,
        )
        .expect("valid frame");

        assert_eq!(frame.status(), Some(200));
    }

    #[test]
    fn an_incomplete_without_a_legitimate_reason_stays_a_provider_failure() {
        for raw in [
            r#"{"type":"response.incomplete"}"#,
            r#"{"type":"response.incomplete","response":{"incomplete_details":null}}"#,
            r#"{"type":"response.incomplete","response":{"incomplete_details":{"reason":""}}}"#,
            r#"{"type":"response.incomplete","response":{"incomplete_details":{"reason":"error"}}}"#,
            r#"{"type":"response.incomplete","response":{"incomplete_details":{"reason":"server_error"}}}"#,
        ] {
            let frame = ParsedResponsesWebSocketFrame::parse(raw).expect("valid frame");

            assert_eq!(
                frame.status(),
                Some(502),
                "an incomplete without a known-good reason must stay a provider failure: {raw}"
            );
        }
    }

    #[test]
    fn a_legitimate_incomplete_still_respects_an_explicit_provider_status() {
        let explicit = ParsedResponsesWebSocketFrame::parse(
            r#"{"type":"response.incomplete","status_code":503,"response":{"incomplete_details":{"reason":"max_output_tokens"}}}"#,
        )
        .expect("valid frame");
        assert_eq!(explicit.status(), Some(503));

        let quota = ParsedResponsesWebSocketFrame::parse(
            r#"{"type":"response.incomplete","response":{"error":{"code":"rate_limit_exceeded"},"incomplete_details":{"reason":"max_output_tokens"}}}"#,
        )
        .expect("valid frame");
        assert_eq!(quota.status(), Some(429));
    }

    #[test]
    fn a_legitimate_incomplete_batched_inside_a_chunks_envelope_is_not_a_failure() {
        let frame = ParsedResponsesWebSocketFrame::parse(
            r#"{"chunks":[{"type":"response.output_text.delta","delta":"hi"},{"type":"response.incomplete","response":{"incomplete_details":{"reason":"max_output_tokens"},"usage":{"total_tokens":9}}}]}"#,
        )
        .expect("valid frame");

        assert!(frame.is_chunked());
        assert!(frame.is_terminal());
        assert_eq!(frame.status(), Some(200));
        assert_eq!(frame.event_type(), Some("response.incomplete"));
        assert_eq!(
            frame.terminal_event().and_then(|event| event
                .pointer("/response/usage/total_tokens")
                .and_then(serde_json::Value::as_u64)),
            Some(9)
        );
    }

    #[test]
    fn detects_a_terminal_batched_inside_a_chunks_envelope() {
        let frame = ParsedResponsesWebSocketFrame::parse(
            r#"{"chunks":[{"type":"response.output_text.delta","delta":"hi"},{"type":"response.completed","response":{"usage":{"total_tokens":8}}}]}"#,
        )
        .expect("valid frame");

        assert!(frame.is_chunked());
        assert!(frame.is_terminal());
        assert_eq!(frame.status(), Some(200));
        // The label and the recorded error body must name the event that ended
        // the turn, not the envelope.
        assert_eq!(frame.event_type(), Some("response.completed"));
        assert_eq!(
            frame.terminal_event().and_then(|event| event
                .pointer("/response/usage/total_tokens")
                .and_then(serde_json::Value::as_u64)),
            Some(8)
        );
        assert_eq!(frame.protocol_events().len(), 2);
    }

    #[test]
    fn detects_a_start_event_batched_inside_a_chunks_envelope() {
        let frame = ParsedResponsesWebSocketFrame::parse(
            r#"{"chunks":[{"type":"codex.rate_limits"},{"type":"response.created"}]}"#,
        )
        .expect("valid frame");

        assert!(frame.is_started());
        assert!(!frame.is_terminal());
        assert_eq!(frame.protocol_events().len(), 2);
    }

    #[test]
    fn an_envelope_may_carry_its_own_type_alongside_batched_events() {
        let frame = ParsedResponsesWebSocketFrame::parse(
            r#"{"type":"codex.response.metadata","chunks":[{"type":"response.failed","response":{"error":{"code":"rate_limit_exceeded"}}}]}"#,
        )
        .expect("valid frame");

        assert_eq!(frame.protocol_events().len(), 2);
        assert!(frame.is_terminal());
        assert_eq!(frame.status(), Some(429));
        assert_eq!(frame.event_type(), Some("response.failed"));
    }

    #[test]
    fn a_batch_without_a_terminal_does_not_end_the_turn() {
        let frame = ParsedResponsesWebSocketFrame::parse(
            r#"{"chunks":[{"type":"response.output_text.delta","delta":"a"},{"type":"response.output_text.delta","delta":"b"}]}"#,
        )
        .expect("valid frame");

        assert!(!frame.is_terminal());
        assert!(!frame.is_started());
        assert!(frame.terminal_event().is_none());
    }

    #[test]
    fn an_unrecognized_shape_is_still_surfaced_as_one_event() {
        let frame =
            ParsedResponsesWebSocketFrame::parse(r#"{"unexpected":true}"#).expect("valid frame");

        assert_eq!(frame.protocol_events().len(), 1);
        assert!(!frame.is_chunked());
        assert!(!frame.is_terminal());
        assert_eq!(frame.event_type(), None);
        assert_eq!(frame.event_type_for_log(), "invalid_json");
    }

    #[test]
    fn preserves_safe_log_label_boundaries() {
        let unsafe_label =
            ParsedResponsesWebSocketFrame::parse(r#"{"type":"not safe / contains spaces"}"#)
                .expect("valid frame");
        assert_eq!(unsafe_label.event_type_for_log(), "unknown");

        let missing_label =
            ParsedResponsesWebSocketFrame::parse(r#"{"message":"ok"}"#).expect("valid frame");
        assert_eq!(missing_label.event_type_for_log(), "invalid_json");
    }

    #[test]
    fn rejects_invalid_json() {
        assert!(ParsedResponsesWebSocketFrame::parse("not-json").is_err());
    }
}