dbx-tools-model-proxy 0.9.17

//! Streaming protocol translation and SSE response encoding.

use std::{
    convert::Infallible,
    io,
    pin::Pin,
    sync::Arc,
    task::{Context, Poll},
    time::{SystemTime, UNIX_EPOCH},
};

use aigw_anthropic::translate::{stream_event_to_anthropic_sse, NativeSseContext};
use aigw_core::{
    model::StreamEvent,
    translate::{ResponseTranslator, StreamParser},
};
use aigw_openai::{OpenAIResponseTranslator, ResponsesResponseTranslator};
use axum::{
    body::{Body, Bytes},
    http::{header, HeaderMap, StatusCode},
    response::{IntoResponse, Response},
};
use eventsource_stream::{EventStream, Eventsource};
use futures_util::{FutureExt, Stream, StreamExt};
use serde_json::{json, Value};
use tokio::sync::mpsc;

use crate::{
    error::ProxyError,
    events::{HttpExchangeEvent, HttpResponseEvent, ProxyFeeds, ResponseHop, SseEvent},
    protocol::{ClientWire, TargetWire},
    request_log::RequestLogContext,
    runtime::RuntimeGeneration,
    throttle::{response_token_usage, ResponseTokenUsage},
};

const MAX_OBSERVED_SSE_EVENT_BYTES: usize = 1024 * 1024;

/// Request metadata emitted when an SSE body completes or is dropped.
pub(crate) struct StreamLogContext {
    /// Protocol presented by the caller.
    pub(crate) client_wire: ClientWire,
    /// Protocol selected for the upstream request.
    pub(crate) target: TargetWire,
    /// Shared request metadata and local token reservation.
    pub(crate) request: RequestLogContext,
    /// Event feeds used to observe raw upstream chunks on demand.
    pub(crate) feeds: ProxyFeeds,
    /// Upstream body event template for this attempt.
    pub(crate) upstream_event: HttpResponseEvent,
    /// Runtime lease retained until the response body completes or is dropped.
    pub(crate) _runtime: Arc<RuntimeGeneration>,
}

/// Bounded input for observing native SSE frames without changing their wire bytes.
struct NativeUsageInput {
    receiver: mpsc::Receiver<Bytes>,
}

impl Stream for NativeUsageInput {
    type Item = Result<Bytes, Infallible>;

    fn poll_next(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        match self.receiver.poll_recv(context) {
            Poll::Ready(Some(chunk)) => Poll::Ready(Some(Ok(chunk))),
            Poll::Ready(None) => Poll::Ready(None),
            Poll::Pending => Poll::Pending,
        }
    }
}

/// Observes usage from complete native SSE events while retaining only parser state.
struct NativeUsageObserver {
    sender: Option<mpsc::Sender<Bytes>>,
    events: EventStream<NativeUsageInput>,
    event_size: SseEventSize,
    usage: ResponseTokenUsage,
    failure_status: Option<StatusCode>,
    sse_events: Vec<SseEvent>,
    observing: bool,
}

#[derive(Clone, Debug, Default, Eq, PartialEq)]
struct NativeObservation {
    usage: ResponseTokenUsage,
    failure_status: Option<StatusCode>,
    sse_events: Vec<SseEvent>,
}

#[derive(Default)]
struct SseEventSize {
    bytes: usize,
    line_has_data: bool,
    previous_was_cr: bool,
}

impl SseEventSize {
    fn accepts(&mut self, chunk: &[u8]) -> bool {
        for &byte in chunk {
            self.bytes = self.bytes.saturating_add(1);
            if self.bytes > MAX_OBSERVED_SSE_EVENT_BYTES {
                return false;
            }
            match byte {
                b'\r' => {
                    if !self.line_has_data {
                        self.bytes = 0;
                    }
                    self.line_has_data = false;
                    self.previous_was_cr = true;
                }
                b'\n' if self.previous_was_cr => {
                    self.previous_was_cr = false;
                }
                b'\n' => {
                    if !self.line_has_data {
                        self.bytes = 0;
                    }
                    self.line_has_data = false;
                }
                _ => {
                    self.line_has_data = true;
                    self.previous_was_cr = false;
                }
            }
        }
        true
    }
}

impl Default for NativeUsageObserver {
    fn default() -> Self {
        let (sender, receiver) = mpsc::channel(1);
        Self {
            sender: Some(sender),
            events: EventStream::new(NativeUsageInput { receiver }),
            event_size: SseEventSize::default(),
            usage: ResponseTokenUsage::default(),
            failure_status: None,
            sse_events: Vec::new(),
            observing: true,
        }
    }
}

impl NativeUsageObserver {
    /// Observe one upstream chunk and return the exact bytes supplied by the caller.
    async fn observe_chunk(&mut self, chunk: Bytes) -> Bytes {
        if !self.observing {
            return chunk;
        }
        if !self.event_size.accepts(&chunk) {
            self.stop();
            return chunk;
        }
        let Some(sender) = self.sender.as_ref() else {
            return chunk;
        };
        if sender.send(chunk.clone()).await.is_err() {
            self.stop();
            return chunk;
        }
        self.drain_ready();
        chunk
    }

    /// Close the framed input and consume every complete event still buffered by the parser.
    async fn finish(mut self) -> NativeObservation {
        self.sender.take();
        if self.observing {
            while let Some(event) = self.events.next().await {
                match event {
                    Ok(event) => self.observe(&event.event, &event.data),
                    Err(_) => break,
                }
            }
        }
        NativeObservation {
            usage: self.usage,
            failure_status: self.failure_status,
            sse_events: self.sse_events,
        }
    }

    fn drain_ready(&mut self) {
        loop {
            match self.events.next().now_or_never() {
                Some(Some(Ok(event))) => self.observe(&event.event, &event.data),
                Some(Some(Err(_))) | Some(None) => {
                    self.stop();
                    break;
                }
                None => break,
            }
        }
    }

    fn take_sse_events(&mut self) -> Vec<SseEvent> {
        std::mem::take(&mut self.sse_events)
    }

    fn observe(&mut self, event: &str, data: &str) {
        self.sse_events
            .push(SseEvent::new(event.to_owned(), data.to_owned(), None));
        let Ok(payload) = serde_json::from_str::<Value>(data) else {
            return;
        };
        if self.failure_status.is_none() {
            self.failure_status = native_failure_status(event, &payload);
        }
        let usage = response_token_usage(&payload);
        if usage.reported {
            self.usage = usage;
        }
    }

    fn stop(&mut self) {
        self.observing = false;
        self.sender.take();
    }
}

fn native_failure_status(event: &str, payload: &Value) -> Option<StatusCode> {
    let payload_type = payload.get("type").and_then(Value::as_str);
    let failed = matches!(event, "response.failed" | "error")
        || matches!(payload_type, Some("response.failed" | "error"))
        || payload.pointer("/response/status").and_then(Value::as_str) == Some("failed");
    if !failed {
        return None;
    }
    let failure = payload
        .pointer("/response/error")
        .or_else(|| payload.get("error"))
        .unwrap_or(payload)
        .to_string()
        .to_ascii_lowercase();
    Some(
        if failure.contains("rate_limit")
            || failure.contains("rate limit")
            || failure.contains("too many requests")
            || failure.contains("exceeded workspace")
        {
            StatusCode::TOO_MANY_REQUESTS
        } else {
            StatusCode::BAD_GATEWAY
        },
    )
}

struct StreamCompletion {
    context: StreamLogContext,
    response_bytes: u64,
    usage: ResponseTokenUsage,
    status: StatusCode,
    finished: bool,
}

impl StreamCompletion {
    fn new(context: StreamLogContext) -> Self {
        Self {
            context,
            response_bytes: 0,
            usage: ResponseTokenUsage::default(),
            status: StatusCode::OK,
            finished: false,
        }
    }

    fn record_bytes(&mut self, bytes: usize) {
        self.response_bytes = self.response_bytes.saturating_add(bytes as u64);
    }

    fn observe(&mut self, event: &StreamEvent) {
        if let StreamEvent::Usage(usage) = event {
            let input = usage.prompt_tokens.unwrap_or_default();
            let output = usage.completion_tokens.unwrap_or_default();
            self.usage = ResponseTokenUsage {
                reported: true,
                input,
                output,
                total: usage
                    .total_tokens
                    .unwrap_or_else(|| input.saturating_add(output)),
            };
        }
    }

    fn finish(&mut self, failure_status: Option<StatusCode>) {
        self.finished = true;
        if let Some(status) = failure_status {
            self.status = status;
        }
    }

    async fn reconcile(&self) {
        self.context.request.reconcile(self.usage).await;
    }
}

impl Drop for StreamCompletion {
    fn drop(&mut self) {
        self.context.request.stream_completed(
            self.context.client_wire,
            self.context.target,
            self.response_bytes,
            self.usage,
            self.status,
            self.finished,
        );
    }
}

pub(crate) fn stream_response(
    client_wire: ClientWire,
    target: TargetWire,
    upstream: reqwest::Response,
    model: String,
    response_headers: HeaderMap,
    log_context: StreamLogContext,
) -> Result<Response, ProxyError> {
    let event_feeds = log_context.feeds.clone();
    let event_template = log_context.upstream_event.clone();
    let _ = event_feeds
        .http
        .publish(HttpExchangeEvent::from_response(event_template.clone()));
    // Preserve native SSE framing when no protocol translation is required.
    if matches!(
        (client_wire, target),
        (ClientWire::Chat, TargetWire::Chat) | (ClientWire::Responses, TargetWire::Responses)
    ) {
        let mut upstream = upstream.bytes_stream();
        let stream = async_stream::stream! {
            let mut completion = StreamCompletion::new(log_context);
            let mut usage = NativeUsageObserver::default();
            while let Some(chunk) = upstream.next().await {
                match chunk {
                    Ok(chunk) => {
                        let chunk = usage.observe_chunk(chunk).await;
                        for event in usage.take_sse_events() {
                            publish_sse_event(
                                &event_feeds,
                                &event_template,
                                client_wire,
                                event,
                                true,
                            );
                        }
                        completion.record_bytes(chunk.len());
                        yield Ok::<Bytes, io::Error>(chunk);
                    }
                    Err(error) => {
                        completion.finish(Some(StatusCode::BAD_GATEWAY));
                        yield Err(io::Error::other(error.to_string()));
                        return;
                    }
                }
            }
            let observation = usage.finish().await;
            for event in observation.sse_events.iter().cloned() {
                publish_sse_event(
                    &event_feeds,
                    &event_template,
                    client_wire,
                    event,
                    true,
                );
            }
            completion.usage = observation.usage;
            completion.reconcile().await;
            completion.finish(observation.failure_status);
        };
        return Ok(sse_response(Body::from_stream(stream), response_headers));
    }

    let mut parser: Box<dyn StreamParser> = match target {
        TargetWire::Chat => OpenAIResponseTranslator.stream_parser(),
        TargetWire::Responses => ResponsesResponseTranslator.stream_parser(),
        TargetWire::Auto => unreachable!("auto target is resolved before streaming"),
    };
    let mut events = upstream.bytes_stream().eventsource();
    let stream = async_stream::stream! {
        let mut completion = StreamCompletion::new(log_context);
        let mut anthropic = NativeSseContext::with_pinned_model(model.clone());
        let mut chat = ChatSseContext::new(model);
        let mut failed = false;

        while let Some(event) = events.next().await {
            let event = match event {
                Ok(event) => event,
                Err(error) => {
                    let frame = stream_error(client_wire, &error.to_string());
                    completion.record_bytes(frame.len());
                    yield Ok::<Bytes, io::Error>(frame);
                    failed = true;
                    break;
                }
            };
            publish_sse_event(
                &event_feeds,
                &event_template,
                client_wire,
                SseEvent::new(
                    event.event.clone(),
                    event.data.clone(),
                    (!event.id.is_empty()).then_some(event.id.clone()),
                ),
                false,
            );
            let parsed = match parser.parse_event(&event.event, &event.data) {
                Ok(parsed) => parsed,
                Err(error) => {
                    let frame = stream_error(client_wire, &error.to_string());
                    publish_client_sse_frame(
                        &event_feeds,
                        &event_template,
                        client_wire,
                        &frame,
                    );
                    completion.record_bytes(frame.len());
                    yield Ok(frame);
                    failed = true;
                    break;
                }
            };
            for canonical in parsed {
                completion.observe(&canonical);
                for frame in encode_stream_event(
                    client_wire,
                    &mut anthropic,
                    &mut chat,
                    canonical,
                ) {
                    publish_client_sse_frame(
                        &event_feeds,
                        &event_template,
                        client_wire,
                        &frame,
                    );
                    completion.record_bytes(frame.len());
                    yield Ok(frame);
                }
            }
        }
        if !failed {
            // Parsers can buffer terminal usage or completion events until EOF.
            match parser.finish() {
                Ok(parsed) => {
                    for canonical in parsed {
                        completion.observe(&canonical);
                        for frame in encode_stream_event(
                            client_wire,
                            &mut anthropic,
                            &mut chat,
                            canonical,
                        ) {
                            publish_client_sse_frame(
                                &event_feeds,
                                &event_template,
                                client_wire,
                                &frame,
                            );
                            completion.record_bytes(frame.len());
                            yield Ok(frame);
                        }
                    }
                }
                Err(error) => {
                    failed = true;
                    let frame = stream_error(client_wire, &error.to_string());
                    completion.record_bytes(frame.len());
                    yield Ok(frame);
                }
            }
        }
        completion.reconcile().await;
        completion.finish(failed.then_some(StatusCode::BAD_GATEWAY));
    };
    Ok(sse_response(Body::from_stream(stream), response_headers))
}

fn publish_sse_event(
    feeds: &ProxyFeeds,
    template: &HttpResponseEvent,
    client_wire: ClientWire,
    event: SseEvent,
    include_client: bool,
) {
    let mut upstream = template.clone();
    upstream.headers = None;
    upstream.sse = Some(event.clone());
    upstream.body = None;
    let _ = feeds
        .http
        .publish(HttpExchangeEvent::from_response(upstream));
    if include_client {
        let _ = feeds
            .http
            .publish(HttpExchangeEvent::from_response(client_sse_event(
                template,
                client_wire,
                event,
            )));
    }
}

fn publish_client_sse_frame(
    feeds: &ProxyFeeds,
    template: &HttpResponseEvent,
    client_wire: ClientWire,
    frame: &[u8],
) {
    if let Some(event) = parse_sse_frame(frame) {
        let _ = feeds
            .http
            .publish(HttpExchangeEvent::from_response(client_sse_event(
                template,
                client_wire,
                event,
            )));
    }
}

fn client_sse_event(
    template: &HttpResponseEvent,
    client_wire: ClientWire,
    event: SseEvent,
) -> HttpResponseEvent {
    let path = match client_wire {
        ClientWire::Chat => "/v1/chat/completions",
        ClientWire::Responses => "/v1/responses",
        ClientWire::Anthropic => "/v1/messages",
    };
    HttpResponseEvent {
        request_id: template.request_id,
        hop: ResponseHop::ProxyToClient,
        elapsed_ms: template.elapsed_ms,
        duration_ms: None,
        attempt: None,
        method: "POST".to_owned(),
        host: None,
        path: path.to_owned(),
        status: template.status,
        headers: None,
        response_bytes: None,
        client_protocol: Some(client_wire.label().to_owned()),
        target_protocol: None,
        preferred_model: template.preferred_model.clone(),
        actual_model: template.actual_model.clone(),
        streaming: Some(true),
        fallback_step: template.fallback_step,
        transport_error: None,
        body: None,
        sse: Some(event),
    }
}

fn parse_sse_frame(frame: &[u8]) -> Option<SseEvent> {
    let text = std::str::from_utf8(frame).ok()?;
    let mut event = String::new();
    let mut data = Vec::new();
    let mut id = None;
    for line in text.lines() {
        if let Some(value) = line.strip_prefix("event:") {
            event = value.trim_start().to_owned();
        } else if let Some(value) = line.strip_prefix("data:") {
            data.push(value.trim_start());
        } else if let Some(value) = line.strip_prefix("id:") {
            id = Some(value.trim_start().to_owned());
        }
    }
    (!data.is_empty() || !event.is_empty()).then(|| {
        SseEvent::new(
            if event.is_empty() {
                "message".to_owned()
            } else {
                event
            },
            data.join("\n"),
            id,
        )
    })
}

fn sse_response(body: Body, mut headers: HeaderMap) -> Response {
    headers.insert(
        header::CONTENT_TYPE,
        "text/event-stream".parse().expect("valid content type"),
    );
    headers.insert(
        header::CACHE_CONTROL,
        "no-cache".parse().expect("valid cache control"),
    );
    (StatusCode::OK, headers, body).into_response()
}

fn encode_stream_event(
    client_wire: ClientWire,
    anthropic: &mut NativeSseContext,
    chat: &mut ChatSseContext,
    event: StreamEvent,
) -> Vec<Bytes> {
    match client_wire {
        ClientWire::Anthropic => stream_event_to_anthropic_sse(&event, anthropic)
            .into_iter()
            .map(|frame| Bytes::from(frame.to_sse_bytes()))
            .collect(),
        ClientWire::Chat => chat.encode(event).into_iter().map(Bytes::from).collect(),
        ClientWire::Responses => Vec::new(),
    }
}

fn stream_error(client_wire: ClientWire, message: &str) -> Bytes {
    let payload = match client_wire {
        ClientWire::Anthropic => json!({
            "type": "error",
            "error": {"type": "api_error", "message": message}
        }),
        _ => json!({"error": {"type": "proxy_error", "message": message}}),
    };
    Bytes::from(format!("event: error\ndata: {payload}\n\n"))
}

struct ChatSseContext {
    created: u64,
    id: String,
    model: String,
}

impl ChatSseContext {
    fn new(model: String) -> Self {
        let created = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap_or_default()
            .as_secs();
        Self {
            created,
            id: format!("chatcmpl-proxy-{created}"),
            model,
        }
    }

    fn encode(&mut self, event: StreamEvent) -> Vec<String> {
        let frame = match event {
            StreamEvent::ResponseMeta { id, model } => {
                if !id.is_empty() {
                    self.id = id;
                }
                if !model.is_empty() {
                    self.model = model;
                }
                Some(self.chunk(json!({"role": "assistant"}), Value::Null, None))
            }
            StreamEvent::ContentDelta(text) => {
                Some(self.chunk(json!({"content": text}), Value::Null, None))
            }
            StreamEvent::ReasoningDelta(text) => {
                Some(self.chunk(json!({"reasoning_content": text}), Value::Null, None))
            }
            StreamEvent::ToolCallStart {
                index, id, name, ..
            } => Some(self.chunk(
                json!({
                    "tool_calls": [{
                        "index": index,
                        "id": id,
                        "type": "function",
                        "function": {"name": name, "arguments": ""}
                    }]
                }),
                Value::Null,
                None,
            )),
            StreamEvent::ToolCallDelta {
                index, arguments, ..
            } => Some(self.chunk(
                json!({
                    "tool_calls": [{
                        "index": index,
                        "function": {"arguments": arguments}
                    }]
                }),
                Value::Null,
                None,
            )),
            StreamEvent::Finish(reason) => Some(self.chunk(json!({}), json!(reason), None)),
            StreamEvent::Usage(usage) => {
                Some(self.chunk(json!({}), Value::Null, Some(json!(usage))))
            }
            StreamEvent::Done => return vec!["data: [DONE]\n\n".to_owned()],
            _ => None,
        };
        frame.into_iter().collect()
    }

    fn chunk(&self, delta: Value, finish_reason: Value, usage: Option<Value>) -> String {
        let mut payload = json!({
            "id": self.id,
            "object": "chat.completion.chunk",
            "created": self.created,
            "model": self.model,
            "choices": [{
                "index": 0,
                "delta": delta,
                "finish_reason": finish_reason
            }]
        });
        if let Some(usage) = usage {
            payload["usage"] = usage;
        }
        format!("data: {payload}\n\n")
    }
}

#[cfg(test)]
mod tests {
    use aigw_core::model::{FinishReason, Usage};

    use super::*;

    async fn observed_usage(chunks: impl IntoIterator<Item = Bytes>) -> ResponseTokenUsage {
        let mut observer = NativeUsageObserver::default();
        for chunk in chunks {
            observer.observe_chunk(chunk).await;
        }
        observer.finish().await.usage
    }

    #[test]
    fn parses_complete_sse_frames_for_response_events() {
        let event =
            parse_sse_frame(b"id: 42\nevent: response.completed\ndata: {\"ok\":true}\n\n").unwrap();
        assert_eq!(event.id.as_deref(), Some("42"));
        assert_eq!(event.event, "response.completed");
        assert_eq!(event.data.0, serde_json::json!({"ok": true}));
        assert_eq!(event.data_raw, "eyJvayI6dHJ1ZX0=");
    }

    #[tokio::test]
    async fn native_usage_observer_returns_upstream_chunks_unchanged() {
        let chunks = vec![
            Bytes::from_static(b"event: completion\r\n"),
            Bytes::from_static(b"data: {\"usage\": null}\r\n\r\n"),
        ];
        let mut observer = NativeUsageObserver::default();
        let mut forwarded = Vec::new();
        for chunk in chunks.iter().cloned() {
            forwarded.push(observer.observe_chunk(chunk).await);
        }
        assert_eq!(forwarded, chunks);
    }

    #[tokio::test]
    async fn native_usage_observer_classifies_in_band_response_failures() {
        let mut rate_limit = NativeUsageObserver::default();
        rate_limit
            .observe_chunk(Bytes::from_static(
                br#"event: response.failed
data: {"type":"response.failed","response":{"status":"failed","error":{"code":"rate_limit_exceeded","message":"Too many requests"}}}

"#,
            ))
            .await;
        assert_eq!(
            rate_limit.finish().await.failure_status,
            Some(StatusCode::TOO_MANY_REQUESTS)
        );

        let mut failure = NativeUsageObserver::default();
        failure
            .observe_chunk(Bytes::from_static(
                br#"event: response.failed
data: {"type":"response.failed","response":{"status":"failed","error":{"code":"server_error"}}}

"#,
            ))
            .await;
        assert_eq!(
            failure.finish().await.failure_status,
            Some(StatusCode::BAD_GATEWAY)
        );
    }

    #[tokio::test]
    async fn native_usage_observer_handles_arbitrary_boundaries_crlf_and_json_whitespace() {
        let event = concat!(
            "event: completion\r\n",
            "data: { \"usage\" : { \"prompt_tokens\" : 8, \"completion_tokens\" : 2, ",
            "\"total_tokens\" : 10 } }\r\n",
            "\r\n"
        );
        let usage = observed_usage(
            event
                .as_bytes()
                .chunks(1)
                .map(Bytes::copy_from_slice)
                .collect::<Vec<_>>(),
        )
        .await;
        assert_eq!(
            usage,
            ResponseTokenUsage {
                reported: true,
                input: 8,
                output: 2,
                total: 10,
            }
        );
    }

    #[tokio::test]
    async fn native_usage_observer_reads_multiline_responses_event() {
        let event = concat!(
            "event: response.completed\r\n",
            "data: {\"type\":\"response.completed\",\"response\":\r\n",
            "data: {\"usage\" : {\"input_tokens\":12,\"output_tokens\":3,\"total_tokens\":15}}}\r\n",
            "\r\n"
        );
        let responses = observed_usage(
            event
                .as_bytes()
                .chunks(7)
                .map(Bytes::copy_from_slice)
                .collect::<Vec<_>>(),
        )
        .await;
        assert_eq!(
            responses,
            ResponseTokenUsage {
                reported: true,
                input: 12,
                output: 3,
                total: 15,
            }
        );
    }

    #[tokio::test]
    async fn native_usage_observer_ignores_null_and_nested_unrelated_usage() {
        let ignored = observed_usage([Bytes::from_static(
            br#"data: {"usage":null,"metadata":{"usage":{"prompt_tokens":90,"completion_tokens":9}}}

"#,
        )])
        .await;
        assert_eq!(ignored, ResponseTokenUsage::default());

        let retained = observed_usage([Bytes::from_static(
            br#"data: {"usage":{"prompt_tokens":8,"completion_tokens":2,"total_tokens":10}}

data: {"usage":null,"metadata":{"usage":{"prompt_tokens":90,"completion_tokens":9}}}

"#,
        )])
        .await;
        assert_eq!(
            retained,
            ResponseTokenUsage {
                reported: true,
                input: 8,
                output: 2,
                total: 10,
            }
        );
    }

    #[tokio::test]
    async fn native_usage_observer_reads_usage_from_event_larger_than_old_tail() {
        let event = format!(
            "data: {{\"usage\":{{\"prompt_tokens\":21,\"completion_tokens\":5}},\"padding\":\"{}\"}}\n\n",
            "x".repeat(256 * 1024)
        );
        let usage = observed_usage(
            event
                .as_bytes()
                .chunks(4093)
                .map(Bytes::copy_from_slice)
                .collect::<Vec<_>>(),
        )
        .await;
        assert_eq!(
            usage,
            ResponseTokenUsage {
                reported: true,
                input: 21,
                output: 5,
                total: 26,
            }
        );
    }

    #[tokio::test]
    async fn native_usage_observer_bounds_oversized_event_state() {
        let event = format!(
            "data: {{\"usage\":{{\"prompt_tokens\":21,\"completion_tokens\":5}},\"padding\":\"{}\"}}\n\n",
            "x".repeat(MAX_OBSERVED_SSE_EVENT_BYTES)
        );
        let usage = observed_usage(
            event
                .as_bytes()
                .chunks(4093)
                .map(Bytes::copy_from_slice)
                .collect::<Vec<_>>(),
        )
        .await;
        assert_eq!(usage, ResponseTokenUsage::default());
    }

    #[tokio::test]
    async fn native_usage_observer_ignores_truncated_and_malformed_streams() {
        let truncated = observed_usage([Bytes::from_static(
            br#"data: {"usage":{"prompt_tokens":8,"completion_tokens":2}}"#,
        )])
        .await;
        assert_eq!(truncated, ResponseTokenUsage::default());

        let malformed = observed_usage([
            Bytes::from_static(b"data: \xff\n\n"),
            Bytes::from_static(
                br#"data: {"usage":{"prompt_tokens":8,"completion_tokens":2}}

"#,
            ),
        ])
        .await;
        assert_eq!(malformed, ResponseTokenUsage::default());
    }

    #[test]
    fn canonical_events_encode_as_chat_completion_chunks() {
        let mut context = ChatSseContext::new("requested-model".to_owned());
        let frames = [
            StreamEvent::ResponseMeta {
                id: "response-id".to_owned(),
                model: "upstream-model".to_owned(),
            },
            StreamEvent::ContentDelta("Hello".to_owned()),
            StreamEvent::ToolCallStart {
                index: 0,
                id: "call-1".to_owned(),
                name: "lookup".to_owned(),
            },
            StreamEvent::ToolCallDelta {
                index: 0,
                arguments: r#"{"key":"value"}"#.to_owned(),
            },
            StreamEvent::Finish(FinishReason::ToolCalls),
            StreamEvent::Usage(Usage {
                prompt_tokens: Some(3),
                completion_tokens: Some(5),
                total_tokens: Some(8),
                ..Default::default()
            }),
            StreamEvent::Done,
        ]
        .into_iter()
        .flat_map(|event| context.encode(event))
        .collect::<Vec<_>>();

        assert!(frames[0].contains(r#""role":"assistant""#));
        assert!(frames[1].contains(r#""content":"Hello""#));
        assert!(frames[2].contains(r#""name":"lookup""#));
        assert!(frames[3].contains(r#""arguments":"{\"key\":\"value\"}""#));
        assert!(frames[4].contains(r#""finish_reason":"tool_calls""#));
        assert!(frames[5].contains(r#""total_tokens":8"#));
        assert_eq!(frames[6], "data: [DONE]\n\n");
    }

    #[test]
    fn responses_stream_adapts_to_chat_completion_chunks() {
        let mut parser = ResponsesResponseTranslator.stream_parser();
        let mut context = ChatSseContext::new("requested-model".to_owned());
        let frames = [
            r#"{"type":"response.created","response":{"id":"resp-1","model":"gpt-5.4"}}"#,
            r#"{"type":"response.output_text.delta","delta":"Hello"}"#,
            r#"{"type":"response.completed","response":{"status":"completed","usage":{"input_tokens":2,"output_tokens":1,"total_tokens":3}}}"#,
        ]
        .into_iter()
        .flat_map(|data| parser.parse_event("", data).unwrap())
        .flat_map(|event| context.encode(event))
        .collect::<Vec<_>>();

        assert!(frames[0].contains(r#""id":"resp-1""#));
        assert!(frames[1].contains(r#""content":"Hello""#));
        assert!(frames[2].contains(r#""finish_reason":"stop""#));
        assert!(frames[3].contains(r#""total_tokens":3"#));
        assert_eq!(frames[4], "data: [DONE]\n\n");
    }
}