provider_core/

openai_compat_chat.rs

// SPDX-License-Identifier: AGPL-3.0-or-later

//! Transport-independent OpenAI-compatible Chat Completions adapter shared by
//! the follow-on providers (PRD-001s) and the DeepSeek direct API phase
//! (PRD-001y S1): Qwen/DashScope, xAI Grok, Z.ai/GLM, Moonshot/Kimi, and
//! DeepSeek (`api.deepseek.com`).
//!
//! Each provider exposes an OpenAI-compatible chat-completions wire API with
//! SSE streaming (`include_usage: true`, `data: [DONE]` terminal). The
//! providers differ only in base URL, path, and model-id namespace, all of
//! which the route identity carries. This adapter is therefore provider
//! neutral: it validates that the route is one of the follow-on phases and
//! builds the request from the route's own endpoint and model id.
//!
//! The wire behavior is deliberately strict and matches the OpenRouter adapter
//! contract: usage appears at most once, the terminal is required and
//! duplicated once, and every provider error after a visible or ambiguous
//! boundary denies replay.

use crate::ollama_chat::{
    BoundedRetryAfter, HttpMethod, NetworkPolicy, ReplayBoundary, ReplayDecision,
};
use router_protocol::{
    AttemptId, CanonicalRequest, ContentBlock, InputItem, MessageRole, PhaseKind, RouteIdentity,
    RouteIdentityError, ToolChoice, TransportKind,
};
use serde_json::{Map, Value, json};
use std::collections::{BTreeMap, HashSet};
use std::fmt;

pub const DEFAULT_MAX_SSE_EVENT_BYTES: usize = 256 * 1024;

/// The follow-on phase kinds this adapter serves.
pub const fn is_follow_on_phase(phase: PhaseKind) -> bool {
    matches!(
        phase,
        PhaseKind::QwenSubscription
            | PhaseKind::QwenApi
            | PhaseKind::GrokSubscription
            | PhaseKind::GrokApi
            | PhaseKind::ZaiSubscription
            | PhaseKind::ZaiApi
            | PhaseKind::KimiSubscription
            | PhaseKind::KimiApi
            | PhaseKind::DeepseekApi
    )
}

#[derive(Clone, PartialEq, Eq)]
pub struct OpenAiCompatSecretReferenceMarker(String);

impl OpenAiCompatSecretReferenceMarker {
    pub fn as_str(&self) -> &str {
        &self.0
    }
}

impl fmt::Debug for OpenAiCompatSecretReferenceMarker {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter.write_str("OpenAiCompatSecretReferenceMarker([REDACTED])")
    }
}

#[derive(Clone, PartialEq)]
pub struct OpenAiCompatRequestSpec {
    pub method: HttpMethod,
    pub origin: String,
    pub path: String,
    pub content_type: &'static str,
    pub accept: &'static str,
    pub bearer_secret: OpenAiCompatSecretReferenceMarker,
    pub network: NetworkPolicy,
    body: Value,
}

impl OpenAiCompatRequestSpec {
    pub fn body(&self) -> &Value {
        &self.body
    }

    pub fn body_bytes(&self) -> Result<Vec<u8>, serde_json::Error> {
        serde_json::to_vec(&self.body)
    }
}

impl fmt::Debug for OpenAiCompatRequestSpec {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter
            .debug_struct("OpenAiCompatRequestSpec")
            .field("method", &self.method)
            .field("origin", &self.origin)
            .field("path", &self.path)
            .field("content_type", &self.content_type)
            .field("accept", &self.accept)
            .field("bearer_secret", &self.bearer_secret)
            .field("network", &self.network)
            .field("body", &"[REDACTED]")
            .finish()
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum OpenAiCompatRequestErrorKind {
    InvalidRouteIdentity,
    UnsupportedRoute,
    InvalidCanonicalRequest,
    AmbiguousToolName,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct OpenAiCompatRequestError {
    pub kind: OpenAiCompatRequestErrorKind,
}

impl fmt::Display for OpenAiCompatRequestError {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter.write_str("the follow-on chat request could not be prepared")
    }
}

impl std::error::Error for OpenAiCompatRequestError {}

impl From<RouteIdentityError> for OpenAiCompatRequestError {
    fn from(_: RouteIdentityError) -> Self {
        Self {
            kind: OpenAiCompatRequestErrorKind::InvalidRouteIdentity,
        }
    }
}

pub fn build_request(
    route: &RouteIdentity,
    request: &CanonicalRequest,
) -> Result<OpenAiCompatRequestSpec, OpenAiCompatRequestError> {
    route.validate()?;
    if !is_follow_on_phase(route.phase) || route.transport != TransportKind::DirectHttps {
        return Err(OpenAiCompatRequestError {
            kind: OpenAiCompatRequestErrorKind::UnsupportedRoute,
        });
    }
    request.validate().map_err(|_| OpenAiCompatRequestError {
        kind: OpenAiCompatRequestErrorKind::InvalidCanonicalRequest,
    })?;
    let secret_ref = route
        .credential_ref
        .clone()
        .ok_or(OpenAiCompatRequestError {
            kind: OpenAiCompatRequestErrorKind::InvalidRouteIdentity,
        })?;

    Ok(OpenAiCompatRequestSpec {
        method: HttpMethod::Post,
        origin: route.endpoint.origin.clone(),
        path: route.endpoint.path.clone(),
        content_type: "application/json",
        accept: "text/event-stream",
        bearer_secret: OpenAiCompatSecretReferenceMarker(secret_ref),
        network: NetworkPolicy {
            allow_proxy: false,
            allow_dns: true,
            follow_redirects: false,
            require_tls: true,
        },
        body: build_body(route, request)?,
    })
}

fn build_body(
    route: &RouteIdentity,
    request: &CanonicalRequest,
) -> Result<Value, OpenAiCompatRequestError> {
    let mut object = Map::new();
    object.insert("model".to_owned(), Value::String(route.model_id.clone()));
    object.insert("messages".to_owned(), Value::Array(build_messages(request)));
    object.insert("stream".to_owned(), Value::Bool(true));
    object.insert("stream_options".to_owned(), json!({"include_usage": true}));

    // `tool_choice` is only valid alongside a `tools` array for OpenAI-compat
    // providers; emit both or neither.
    let emit_tools = !request.tools.is_empty() && request.tool_choice != ToolChoice::None;
    if emit_tools {
        object.insert("tool_choice".to_owned(), openai_compat_tool_choice(request));
        let mut names = HashSet::new();
        let mut tools = Vec::with_capacity(request.tools.len());
        for tool in &request.tools {
            let name = provider_tool_name(tool.namespace.as_deref(), &tool.name);
            if !names.insert(name.clone()) {
                return Err(OpenAiCompatRequestError {
                    kind: OpenAiCompatRequestErrorKind::AmbiguousToolName,
                });
            }
            tools.push(json!({
                "type": "function",
                "function": {
                    "name": name,
                    "description": tool.description,
                    "parameters": tool.parameters,
                    "strict": tool.strict,
                }
            }));
        }
        object.insert("tools".to_owned(), Value::Array(tools));
    }
    if let Some(maximum) = request.max_output_tokens {
        object.insert("max_tokens".to_owned(), Value::from(maximum));
    }
    if let Some(format) = &request.structured_output {
        object.insert(
            "response_format".to_owned(),
            json!({
                "type": "json_schema",
                "json_schema": {
                    "name": format.name,
                    "strict": format.strict,
                    "schema": format.schema,
                }
            }),
        );
    }
    Ok(Value::Object(object))
}

fn build_messages(request: &CanonicalRequest) -> Vec<Value> {
    request
        .input
        .iter()
        .map(|item| match item {
            InputItem::Message { role, content } => json!({
                "role": role_name(*role),
                "content": openai_compat_content(content),
            }),
            InputItem::FunctionCall {
                call_id,
                namespace,
                name,
                arguments,
            } => json!({
                "role": "assistant",
                "content": null,
                "tool_calls": [{
                    "id": call_id.as_str(),
                    "type": "function",
                    "function": {
                        "name": provider_tool_name(namespace.as_deref(), name),
                        "arguments": arguments.to_string(),
                    }
                }]
            }),
            InputItem::FunctionCallOutput { call_id, output } => json!({
                "role": "tool",
                "tool_call_id": call_id.as_str(),
                "content": output.to_string(),
            }),
        })
        .collect()
}

fn openai_compat_content(content: &[ContentBlock]) -> Value {
    if content
        .iter()
        .all(|block| matches!(block, ContentBlock::Text { .. }))
    {
        let text = content
            .iter()
            .filter_map(|block| match block {
                ContentBlock::Text { text } => Some(text.as_str()),
                ContentBlock::ImageReference { .. } => None,
            })
            .collect::<String>();
        return Value::String(text);
    }
    Value::Array(
        content
            .iter()
            .map(|block| match block {
                ContentBlock::Text { text } => json!({"type": "text", "text": text}),
                ContentBlock::ImageReference { url } => {
                    json!({"type": "image_url", "image_url": {"url": url}})
                }
            })
            .collect(),
    )
}

const fn role_name(role: MessageRole) -> &'static str {
    match role {
        MessageRole::Developer => "developer",
        MessageRole::System => "system",
        MessageRole::User => "user",
        MessageRole::Assistant => "assistant",
    }
}

fn provider_tool_name(namespace: Option<&str>, name: &str) -> String {
    namespace.map_or_else(
        || name.to_owned(),
        |namespace| format!("{namespace}__{name}"),
    )
}

fn openai_compat_tool_choice(request: &CanonicalRequest) -> Value {
    match &request.tool_choice {
        ToolChoice::Auto => Value::String("auto".to_owned()),
        ToolChoice::None => Value::String("none".to_owned()),
        ToolChoice::Required => Value::String("required".to_owned()),
        ToolChoice::Named(name) => {
            let provider_name = request
                .tools
                .iter()
                .find(|tool| tool.name == *name)
                .map_or_else(
                    || name.clone(),
                    |tool| provider_tool_name(tool.namespace.as_deref(), &tool.name),
                );
            json!({"type":"function","function":{"name":provider_name}})
        }
    }
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct OpenAiCompatUsage {
    pub prompt_tokens: u64,
    pub completion_tokens: u64,
    pub total_tokens: u64,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub enum OpenAiCompatStreamEvent {
    ContentDelta {
        text: String,
    },
    ToolCallDelta {
        id: String,
        ordinal: u64,
        index: u64,
        name: Option<String>,
        arguments: Option<String>,
    },
    FinishReason {
        reason: String,
    },
    Usage(OpenAiCompatUsage),
    Done,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum OpenAiCompatStreamErrorKind {
    InvalidUtf8,
    InvalidJson,
    InvalidFrame,
    EventTooLarge,
    ProviderError,
    MissingUsage,
    DuplicateUsage,
    MissingTerminal,
    DuplicateTerminal,
    DataAfterTerminal,
    ParserFailed,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct OpenAiCompatStreamError {
    pub kind: OpenAiCompatStreamErrorKind,
    pub byte_offset: u64,
}

impl fmt::Display for OpenAiCompatStreamError {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter.write_str("the OpenAI-compatible response stream was invalid")
    }
}

impl std::error::Error for OpenAiCompatStreamError {}

#[derive(Clone, Debug)]
struct ToolState {
    id: String,
    ordinal: u64,
}

pub struct OpenAiCompatSseParser {
    attempt_ref: String,
    line_buffer: Vec<u8>,
    data_lines: Vec<String>,
    data_bytes: usize,
    max_event_bytes: usize,
    byte_offset: u64,
    tool_states: BTreeMap<u64, ToolState>,
    next_tool_ordinal: u64,
    usage_seen: bool,
    terminal: bool,
    failed: bool,
    replay_guard: OpenAiCompatReplayGuard,
}

impl fmt::Debug for OpenAiCompatSseParser {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter
            .debug_struct("OpenAiCompatSseParser")
            .field("buffered_line_bytes", &self.line_buffer.len())
            .field("data_bytes", &self.data_bytes)
            .field("max_event_bytes", &self.max_event_bytes)
            .field("byte_offset", &self.byte_offset)
            .field("tool_count", &self.tool_states.len())
            .field("usage_seen", &self.usage_seen)
            .field("terminal", &self.terminal)
            .field("failed", &self.failed)
            .field("replay_decision", &self.replay_guard.decision())
            .finish()
    }
}

impl OpenAiCompatSseParser {
    pub fn new(attempt_id: &AttemptId) -> Self {
        Self::with_max_event_bytes(attempt_id, DEFAULT_MAX_SSE_EVENT_BYTES)
    }

    pub fn with_max_event_bytes(attempt_id: &AttemptId, max_event_bytes: usize) -> Self {
        Self {
            attempt_ref: attempt_id.as_str().to_owned(),
            line_buffer: Vec::new(),
            data_lines: Vec::new(),
            data_bytes: 0,
            max_event_bytes: max_event_bytes.max(1),
            byte_offset: 0,
            tool_states: BTreeMap::new(),
            next_tool_ordinal: 0,
            usage_seen: false,
            terminal: false,
            failed: false,
            replay_guard: OpenAiCompatReplayGuard::default(),
        }
    }

    pub fn mark_body_written(&mut self) {
        self.replay_guard.mark_body_written();
    }

    pub fn mark_forwarded_tool_result(&mut self) {
        self.replay_guard.mark_forwarded_tool_result();
    }

    pub fn mark_transport_outcome_unknown(&mut self) {
        self.replay_guard.mark_transport_outcome_unknown();
    }

    pub fn replay_decision(&self) -> ReplayDecision {
        self.replay_guard.decision()
    }

    pub fn feed(
        &mut self,
        bytes: &[u8],
    ) -> Result<Vec<OpenAiCompatStreamEvent>, OpenAiCompatStreamError> {
        if self.failed {
            return Err(self.error(OpenAiCompatStreamErrorKind::ParserFailed));
        }
        let mut events = Vec::new();
        for &byte in bytes {
            self.byte_offset = self.byte_offset.saturating_add(1);
            if byte == b'\n' {
                let mut line = std::mem::take(&mut self.line_buffer);
                if line.last() == Some(&b'\r') {
                    line.pop();
                }
                if let Err(error) = self.process_line(&line, &mut events) {
                    self.failed = true;
                    return Err(error);
                }
            } else {
                self.line_buffer.push(byte);
                if self.line_buffer.len() > self.max_event_bytes {
                    self.failed = true;
                    return Err(self.error(OpenAiCompatStreamErrorKind::EventTooLarge));
                }
            }
        }
        Ok(events)
    }

    pub fn finish(&mut self) -> Result<Vec<OpenAiCompatStreamEvent>, OpenAiCompatStreamError> {
        if self.failed {
            return Err(self.error(OpenAiCompatStreamErrorKind::ParserFailed));
        }
        let mut events = Vec::new();
        if !self.line_buffer.is_empty() {
            let mut line = std::mem::take(&mut self.line_buffer);
            if line.last() == Some(&b'\r') {
                line.pop();
            }
            if let Err(error) = self.process_line(&line, &mut events) {
                self.failed = true;
                return Err(error);
            }
        }
        if !self.data_lines.is_empty()
            && let Err(error) = self.dispatch_event(&mut events)
        {
            self.failed = true;
            return Err(error);
        }
        if !self.terminal {
            self.failed = true;
            return Err(self.error(OpenAiCompatStreamErrorKind::MissingTerminal));
        }
        Ok(events)
    }

    fn process_line(
        &mut self,
        line: &[u8],
        events: &mut Vec<OpenAiCompatStreamEvent>,
    ) -> Result<(), OpenAiCompatStreamError> {
        let line = std::str::from_utf8(line)
            .map_err(|_| self.error(OpenAiCompatStreamErrorKind::InvalidUtf8))?;
        if line.is_empty() {
            return self.dispatch_event(events);
        }
        if line.starts_with(':') {
            return Ok(());
        }
        let (field, value) = line.split_once(':').unwrap_or((line, ""));
        if field != "data" {
            return Ok(());
        }
        let value = value.strip_prefix(' ').unwrap_or(value);
        let separator = usize::from(!self.data_lines.is_empty());
        self.data_bytes = self
            .data_bytes
            .saturating_add(value.len())
            .saturating_add(separator);
        if self.data_bytes > self.max_event_bytes {
            return Err(self.error(OpenAiCompatStreamErrorKind::EventTooLarge));
        }
        self.data_lines.push(value.to_owned());
        Ok(())
    }

    fn dispatch_event(
        &mut self,
        events: &mut Vec<OpenAiCompatStreamEvent>,
    ) -> Result<(), OpenAiCompatStreamError> {
        if self.data_lines.is_empty() {
            return Ok(());
        }
        let data = self.data_lines.join("\n");
        self.data_lines.clear();
        self.data_bytes = 0;
        if data == "[DONE]" {
            if self.terminal {
                return Err(self.error(OpenAiCompatStreamErrorKind::DuplicateTerminal));
            }
            if !self.usage_seen {
                return Err(self.error(OpenAiCompatStreamErrorKind::MissingUsage));
            }
            self.terminal = true;
            events.push(OpenAiCompatStreamEvent::Done);
            return Ok(());
        }
        if self.terminal {
            return Err(self.error(OpenAiCompatStreamErrorKind::DataAfterTerminal));
        }
        let value: Value = serde_json::from_str(&data)
            .map_err(|_| self.error(OpenAiCompatStreamErrorKind::InvalidJson))?;
        let object = value
            .as_object()
            .ok_or_else(|| self.error(OpenAiCompatStreamErrorKind::InvalidFrame))?;
        if object.contains_key("error") {
            return Err(self.error(OpenAiCompatStreamErrorKind::ProviderError));
        }
        self.parse_usage(object, events)?;
        self.parse_choices(object, events)
    }

    fn parse_usage(
        &mut self,
        object: &Map<String, Value>,
        events: &mut Vec<OpenAiCompatStreamEvent>,
    ) -> Result<(), OpenAiCompatStreamError> {
        let Some(usage) = object.get("usage") else {
            return Ok(());
        };
        if usage.is_null() {
            return Ok(());
        }
        if self.usage_seen {
            return Err(self.error(OpenAiCompatStreamErrorKind::DuplicateUsage));
        }
        let usage = usage
            .as_object()
            .ok_or_else(|| self.error(OpenAiCompatStreamErrorKind::InvalidFrame))?;
        let prompt_tokens = required_u64(usage, "prompt_tokens", self.byte_offset)?;
        let completion_tokens = required_u64(usage, "completion_tokens", self.byte_offset)?;
        let total_tokens = required_u64(usage, "total_tokens", self.byte_offset)?;
        self.usage_seen = true;
        events.push(OpenAiCompatStreamEvent::Usage(OpenAiCompatUsage {
            prompt_tokens,
            completion_tokens,
            total_tokens,
        }));
        Ok(())
    }

    fn parse_choices(
        &mut self,
        object: &Map<String, Value>,
        events: &mut Vec<OpenAiCompatStreamEvent>,
    ) -> Result<(), OpenAiCompatStreamError> {
        let Some(choices) = object.get("choices") else {
            return Ok(());
        };
        let choices = choices
            .as_array()
            .ok_or_else(|| self.error(OpenAiCompatStreamErrorKind::InvalidFrame))?;
        for choice in choices {
            let choice = choice
                .as_object()
                .ok_or_else(|| self.error(OpenAiCompatStreamErrorKind::InvalidFrame))?;
            if choice.contains_key("error") {
                return Err(self.error(OpenAiCompatStreamErrorKind::ProviderError));
            }
            if let Some(delta) = choice.get("delta") {
                let delta = delta
                    .as_object()
                    .ok_or_else(|| self.error(OpenAiCompatStreamErrorKind::InvalidFrame))?;
                if delta.contains_key("error") {
                    return Err(self.error(OpenAiCompatStreamErrorKind::ProviderError));
                }
                if let Some(content) = delta.get("content")
                    && !content.is_null()
                {
                    let text = content
                        .as_str()
                        .ok_or_else(|| self.error(OpenAiCompatStreamErrorKind::InvalidFrame))?;
                    if !text.is_empty() {
                        let event = OpenAiCompatStreamEvent::ContentDelta {
                            text: text.to_owned(),
                        };
                        self.replay_guard.observe(&event);
                        events.push(event);
                    }
                }
                if let Some(tool_calls) = delta.get("tool_calls") {
                    self.parse_tool_calls(tool_calls, events)?;
                }
            }
            if let Some(reason) = choice.get("finish_reason")
                && !reason.is_null()
            {
                let reason = reason
                    .as_str()
                    .ok_or_else(|| self.error(OpenAiCompatStreamErrorKind::InvalidFrame))?;
                if reason == "error" {
                    return Err(self.error(OpenAiCompatStreamErrorKind::ProviderError));
                }
                events.push(OpenAiCompatStreamEvent::FinishReason {
                    reason: reason.to_owned(),
                });
            }
        }
        Ok(())
    }

    fn parse_tool_calls(
        &mut self,
        value: &Value,
        events: &mut Vec<OpenAiCompatStreamEvent>,
    ) -> Result<(), OpenAiCompatStreamError> {
        let calls = value
            .as_array()
            .ok_or_else(|| self.error(OpenAiCompatStreamErrorKind::InvalidFrame))?;
        for call in calls {
            let call = call
                .as_object()
                .ok_or_else(|| self.error(OpenAiCompatStreamErrorKind::InvalidFrame))?;
            let index = call
                .get("index")
                .and_then(Value::as_u64)
                .ok_or_else(|| self.error(OpenAiCompatStreamErrorKind::InvalidFrame))?;
            let incoming_id = call.get("id").and_then(Value::as_str);
            let state = if let Some(state) = self.tool_states.get(&index) {
                if incoming_id.is_some_and(|id| id != state.id) {
                    return Err(self.error(OpenAiCompatStreamErrorKind::InvalidFrame));
                }
                state.clone()
            } else {
                let ordinal = self.next_tool_ordinal;
                self.next_tool_ordinal = self.next_tool_ordinal.saturating_add(1);
                let state = ToolState {
                    id: incoming_id.map_or_else(
                        || synthesized_tool_id(&self.attempt_ref, ordinal),
                        str::to_owned,
                    ),
                    ordinal,
                };
                self.tool_states.insert(index, state.clone());
                state
            };
            let function = call.get("function").and_then(Value::as_object);
            let name = function
                .and_then(|function| function.get("name"))
                .and_then(Value::as_str)
                .map(str::to_owned);
            let arguments = function
                .and_then(|function| function.get("arguments"))
                .and_then(Value::as_str)
                .map(str::to_owned);
            if name.is_some() || arguments.is_some() || incoming_id.is_some() {
                let event = OpenAiCompatStreamEvent::ToolCallDelta {
                    id: state.id,
                    ordinal: state.ordinal,
                    index,
                    name,
                    arguments,
                };
                self.replay_guard.observe(&event);
                events.push(event);
            }
        }
        Ok(())
    }

    fn error(&self, kind: OpenAiCompatStreamErrorKind) -> OpenAiCompatStreamError {
        OpenAiCompatStreamError {
            kind,
            byte_offset: self.byte_offset,
        }
    }
}

fn required_u64(
    object: &Map<String, Value>,
    name: &str,
    byte_offset: u64,
) -> Result<u64, OpenAiCompatStreamError> {
    object
        .get(name)
        .and_then(Value::as_u64)
        .ok_or(OpenAiCompatStreamError {
            kind: OpenAiCompatStreamErrorKind::InvalidFrame,
            byte_offset,
        })
}

fn synthesized_tool_id(attempt_ref: &str, ordinal: u64) -> String {
    const FNV_OFFSET: u64 = 0xcbf2_9ce4_8422_2325;
    const FNV_PRIME: u64 = 0x0000_0100_0000_01b3;
    let mut value = FNV_OFFSET;
    for byte in attempt_ref
        .bytes()
        .chain(*b":")
        .chain(ordinal.to_be_bytes())
    {
        value ^= u64::from(byte);
        value = value.wrapping_mul(FNV_PRIME);
    }
    format!("call_{value:016x}")
}

#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct OpenAiCompatReplayGuard {
    boundary: Option<ReplayBoundary>,
    body_written: bool,
}

impl OpenAiCompatReplayGuard {
    pub fn mark_body_written(&mut self) {
        self.body_written = true;
    }

    pub fn observe(&mut self, event: &OpenAiCompatStreamEvent) {
        match event {
            OpenAiCompatStreamEvent::ContentDelta { text } if !text.is_empty() => {
                self.cross(ReplayBoundary::PublicContent);
            }
            OpenAiCompatStreamEvent::ToolCallDelta { .. } => {
                self.cross(ReplayBoundary::VisibleToolCall);
            }
            OpenAiCompatStreamEvent::ContentDelta { .. }
            | OpenAiCompatStreamEvent::FinishReason { .. }
            | OpenAiCompatStreamEvent::Usage(_)
            | OpenAiCompatStreamEvent::Done => {}
        }
    }

    pub fn mark_forwarded_tool_result(&mut self) {
        self.cross(ReplayBoundary::ForwardedToolResult);
    }

    pub fn mark_transport_outcome_unknown(&mut self) {
        if self.body_written {
            self.cross(ReplayBoundary::AmbiguousDispatch);
        }
    }

    pub fn decision(&self) -> ReplayDecision {
        self.boundary
            .map_or(ReplayDecision::Allowed, ReplayDecision::Denied)
    }

    fn cross(&mut self, boundary: ReplayBoundary) {
        if self.boundary.is_none() {
            self.boundary = Some(boundary);
        }
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum OpenAiCompatFailureClass {
    InvalidRequest,
    Authentication,
    PaymentRequired,
    Authorization,
    NotFound,
    Timeout,
    Unprocessable,
    RateLimited,
    UpstreamServer,
    UnexpectedStatus,
    MismatchedContentType,
    RedirectRejected,
    Cancelled,
    TransportTimeout,
    AmbiguousDispatch,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct OpenAiCompatFailure {
    pub class: OpenAiCompatFailureClass,
    pub status: Option<u16>,
    pub retry_after: Option<BoundedRetryAfter>,
}

/// The only retry authorization exposed by the follow-on adapter.  This is
/// intentionally a policy decision rather than a transport convenience: the
/// caller still owns the budget reservation and must record the result.
pub const AUTHORIZED_RATE_LIMIT_RETRY_AFTER_MAX_SECONDS: u32 = 2;

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum OpenAiCompatRetryDecision {
    RetryOnce { retry_after_seconds: u32 },
    NoRetry,
}

/// Applies the same deliberately narrow 429 policy as the OpenRouter adapter.
pub fn authorized_rate_limit_retry(
    failure: OpenAiCompatFailure,
    replay: ReplayDecision,
    retry_already_used: bool,
) -> OpenAiCompatRetryDecision {
    let retry_after_seconds = failure
        .retry_after
        .filter(|retry_after| !retry_after.capped)
        .map(|retry_after| retry_after.seconds);
    if retry_already_used
        || !matches!(failure.class, OpenAiCompatFailureClass::RateLimited)
        || !matches!(replay, ReplayDecision::Allowed)
        || retry_after_seconds
            .is_none_or(|seconds| seconds > AUTHORIZED_RATE_LIMIT_RETRY_AFTER_MAX_SECONDS)
    {
        return OpenAiCompatRetryDecision::NoRetry;
    }
    match retry_after_seconds {
        Some(retry_after_seconds) => OpenAiCompatRetryDecision::RetryOnce {
            retry_after_seconds,
        },
        None => OpenAiCompatRetryDecision::NoRetry,
    }
}

pub fn normalize_http_response(
    status: u16,
    content_type: Option<&str>,
    retry_after: Option<&str>,
    redirected: bool,
) -> Result<(), OpenAiCompatFailure> {
    if redirected || (300..400).contains(&status) {
        return Err(failure(
            OpenAiCompatFailureClass::RedirectRejected,
            Some(status),
            retry_after,
        ));
    }
    if status != 200 {
        let class = match status {
            400 => OpenAiCompatFailureClass::InvalidRequest,
            401 => OpenAiCompatFailureClass::Authentication,
            402 => OpenAiCompatFailureClass::PaymentRequired,
            403 => OpenAiCompatFailureClass::Authorization,
            404 => OpenAiCompatFailureClass::NotFound,
            408 => OpenAiCompatFailureClass::Timeout,
            422 => OpenAiCompatFailureClass::Unprocessable,
            429 => OpenAiCompatFailureClass::RateLimited,
            500..=599 => OpenAiCompatFailureClass::UpstreamServer,
            _ => OpenAiCompatFailureClass::UnexpectedStatus,
        };
        return Err(failure(class, Some(status), retry_after));
    }
    let media_type = content_type
        .and_then(|value| value.split(';').next())
        .map(str::trim);
    if media_type != Some("text/event-stream") {
        return Err(failure(
            OpenAiCompatFailureClass::MismatchedContentType,
            Some(status),
            None,
        ));
    }
    Ok(())
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum BodyWriteState {
    NotWritten,
    Written,
    Unknown,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum TransportFailureSignal {
    Cancelled,
    Timeout,
}

pub fn normalize_transport_failure(
    signal: TransportFailureSignal,
    body_write: BodyWriteState,
) -> OpenAiCompatFailure {
    let class = match body_write {
        BodyWriteState::Written | BodyWriteState::Unknown => {
            OpenAiCompatFailureClass::AmbiguousDispatch
        }
        BodyWriteState::NotWritten => match signal {
            TransportFailureSignal::Cancelled => OpenAiCompatFailureClass::Cancelled,
            TransportFailureSignal::Timeout => OpenAiCompatFailureClass::TransportTimeout,
        },
    };
    failure(class, None, None)
}

fn failure(
    class: OpenAiCompatFailureClass,
    status: Option<u16>,
    retry_after: Option<&str>,
) -> OpenAiCompatFailure {
    OpenAiCompatFailure {
        class,
        status,
        retry_after: BoundedRetryAfter::parse(retry_after),
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use router_protocol::{
        BillingAuthorityClass, CapabilityRequirements, CatalogSourceKind, CredentialSourceKind,
        EndpointIdentity, ExecutionLocation, FunctionTool, HarnessInstanceId, M1_WIRE_VERSION,
        MessageRole, ModelNamespace, PhaseKind, PhasePlacement, ProviderService, ReasoningIntent,
        SessionId, TurnId,
    };

    fn route() -> RouteIdentity {
        RouteIdentity {
            phase: PhaseKind::KimiApi,
            placement: PhasePlacement::Builtin,
            provider: ProviderService::MoonshotKimi,
            transport: TransportKind::DirectHttps,
            endpoint: EndpointIdentity {
                origin: router_protocol::MOONSHOT_ORIGIN.to_owned(),
                path: router_protocol::MOONSHOT_CHAT_PATH.to_owned(),
            },
            credential_source: CredentialSourceKind::SecretRef,
            credential_ref: Some("env:STEADYMUX_KIMI_KEY".to_owned()),
            credential_refs: vec!["env:STEADYMUX_KIMI_KEY".to_owned()],
            catalog_source: CatalogSourceKind::KimiCatalog,
            catalog_ref: "catalog:kimi:verified".to_owned(),
            billing_authority: BillingAuthorityClass::KimiMeteredApi,
            model_namespace: ModelNamespace::MoonshotKimi,
            model_id: "kimi-k2.6".to_owned(),
            execution_location: ExecutionLocation::RemoteCloud,
            account_ref: "account:kimi".to_owned(),
            capacity_domain_ref: "capacity:kimi".to_owned(),
            local_execution_proof: false,
            approved_local_inventory: Vec::new(),
            reasoning_modes: Vec::new(),
        }
    }

    fn request(prompt: &str) -> CanonicalRequest {
        CanonicalRequest {
            wire_version: M1_WIRE_VERSION.to_owned(),
            model: "router-auto".to_owned(),
            stream: true,
            harness_instance: HarnessInstanceId::parse("fixture-cli").expect("fixture"),
            session_id: SessionId::parse("session-kimi").expect("fixture"),
            turn_id: TurnId::parse("turn-kimi").expect("fixture"),
            input: vec![InputItem::Message {
                role: MessageRole::User,
                content: vec![ContentBlock::Text {
                    text: prompt.to_owned(),
                }],
            }],
            tools: vec![FunctionTool {
                namespace: Some("weather".to_owned()),
                name: "lookup".to_owned(),
                description: "lookup weather".to_owned(),
                parameters: json!({"type":"object"}),
                strict: true,
            }],
            tool_choice: ToolChoice::Auto,
            structured_output: None,
            reasoning: ReasoningIntent::None,
            max_output_tokens: Some(64),
            max_spend_micro_usd: None,
            capabilities: CapabilityRequirements::default(),
        }
    }

    #[test]
    fn exact_request_has_redacted_secret_and_prompt_and_no_openrouter_provider_block() {
        let spec = build_request(&route(), &request("prompt-canary")).expect("request");
        assert_eq!(spec.method, HttpMethod::Post);
        assert_eq!(spec.origin, router_protocol::MOONSHOT_ORIGIN);
        assert_eq!(spec.path, router_protocol::MOONSHOT_CHAT_PATH);
        assert_eq!(spec.accept, "text/event-stream");
        assert!(spec.network.require_tls);
        assert!(!spec.network.follow_redirects);
        assert_eq!(
            spec.body(),
            &json!({
                "model":"kimi-k2.6",
                "messages":[{"role":"user","content":"prompt-canary"}],
                "stream":true,
                "stream_options":{"include_usage":true},
                "tool_choice":"auto",
                "tools":[{"type":"function","function":{
                    "name":"weather__lookup",
                    "description":"lookup weather",
                    "parameters":{"type":"object"},
                    "strict":true
                }}],
                "max_tokens":64
            })
        );
        assert_eq!(spec.bearer_secret.as_str(), "env:STEADYMUX_KIMI_KEY");
        let debug = format!("{spec:?}");
        assert!(!debug.contains("STEADYMUX_KIMI_KEY"));
        assert!(!debug.contains("prompt-canary"));
        assert_eq!(debug.matches("REDACTED").count(), 2);
    }

    #[test]
    fn tool_choice_emitted_only_alongside_tools() {
        // Empty tools with ToolChoice::Required must not emit a bare
        // `tool_choice` (invalid for OpenAI-compat providers).
        let mut no_tools = request("no-tools");
        no_tools.tools = Vec::new();
        no_tools.tool_choice = ToolChoice::Required;
        let body = build_body(&route(), &no_tools).expect("body");
        assert!(body.get("tools").is_none());
        assert!(body.get("tool_choice").is_none());

        // Tools with a named tool choice emit both.
        let mut named = request("named-tool");
        named.tool_choice = ToolChoice::Named("weather__lookup".to_owned());
        let body = build_body(&route(), &named).expect("body");
        assert!(body.get("tools").is_some());
        assert_eq!(
            body.get("tool_choice"),
            Some(&json!({"type": "function", "function": {"name": "weather__lookup"}}))
        );
    }

    #[test]
    fn non_follow_on_route_is_rejected() {
        let mut identity = route();
        identity.phase = PhaseKind::OpenrouterApi;
        let error = build_request(&identity, &request("x")).expect_err("inconsistent route");
        assert_eq!(
            error.kind,
            OpenAiCompatRequestErrorKind::InvalidRouteIdentity
        );

        let mut openrouter = route();
        openrouter.phase = PhaseKind::OpenrouterApi;
        openrouter.provider = ProviderService::Openrouter;
        openrouter.endpoint = EndpointIdentity {
            origin: "https://openrouter.ai".to_owned(),
            path: "/api/v1/chat/completions".to_owned(),
        };
        openrouter.catalog_source = CatalogSourceKind::OpenrouterCatalog;
        openrouter.billing_authority = BillingAuthorityClass::OpenrouterMeteredApi;
        openrouter.model_namespace = ModelNamespace::Openrouter;
        openrouter.model_id = "openrouter/auto".to_owned();
        assert_eq!(
            build_request(&openrouter, &request("x"))
                .expect_err("unsupported")
                .kind,
            OpenAiCompatRequestErrorKind::UnsupportedRoute
        );
    }

    fn deepseek_route() -> RouteIdentity {
        let mut identity = route();
        identity.phase = PhaseKind::DeepseekApi;
        identity.provider = ProviderService::Deepseek;
        identity.endpoint = EndpointIdentity {
            origin: router_protocol::DEEPSEEK_ORIGIN.to_owned(),
            path: router_protocol::DEEPSEEK_CHAT_PATH.to_owned(),
        };
        identity.catalog_source = CatalogSourceKind::DeepseekCatalog;
        identity.billing_authority = BillingAuthorityClass::DeepseekMeteredApi;
        identity.model_namespace = ModelNamespace::Deepseek;
        identity.credential_ref = Some("env:STEADYMUX_DEEPSEEK_KEY".to_owned());
        identity.credential_refs = vec!["env:STEADYMUX_DEEPSEEK_KEY".to_owned()];
        identity
    }

    #[test]
    fn deepseek_route_builds_verified_wire_shape() {
        // PRD-001y AC-02: the DeepSeek direct path is an OpenAI-compatible
        // chat route against `https://api.deepseek.com/chat/completions` with a
        // bearer secret reference, verified at build time against the published
        // API documentation. The model id comes from the route (catalog
        // driven), never from a literal in this adapter.
        let identity = deepseek_route();
        identity.validate().expect("valid deepseek route");
        let spec = build_request(&identity, &request("deepseek-canary")).expect("request");
        assert_eq!(spec.method, HttpMethod::Post);
        assert_eq!(spec.origin, router_protocol::DEEPSEEK_ORIGIN);
        assert_eq!(spec.path, router_protocol::DEEPSEEK_CHAT_PATH);
        assert_eq!(spec.accept, "text/event-stream");
        assert_eq!(spec.bearer_secret.as_str(), "env:STEADYMUX_DEEPSEEK_KEY");
        assert!(spec.network.require_tls);
        assert_eq!(
            spec.body().get("model"),
            Some(&json!("kimi-k2.6")),
            "the request model must come from the route identity, not a literal"
        );
    }

    #[test]
    fn deepseek_adapter_has_no_hardcoded_model_id_and_catalog_stays_reachable() {
        // PR-001y AC-03: model ids are catalog-driven, never hardcoded. The
        // adapter source must not contain any DeepSeek model literal, and the
        // OpenRouter catalog entries that keep DeepSeek reachable must remain
        // intact in the daemon capsule.
        // The non-test adapter body is the production code that must never
        // carry a hardcoded model id. The forbidden literals are assembled
        // from fragments so the test module itself never contains the exact
        // model strings it is checking for (the file is read via
        // `include_str!` above).
        let adapter_source = include_str!("openai_compat_chat.rs");
        let adapter_core = adapter_source
            .split("#[cfg(test)]")
            .next()
            .expect("adapter body precedes the test module");
        let forbidden = [
            concat!("deepseek", "-", "chat"),
            concat!("deepseek", "-", "reasoner"),
            concat!("deepseek", "-", "v4", "-", "flash"),
            concat!("deepseek", "-", "v4", "-", "pro"),
            concat!("deepseek", "-", "r1"),
        ];
        for model_literal in forbidden {
            assert!(
                !adapter_core.contains(model_literal),
                "adapter hardcodes the model literal {model_literal}"
            );
        }
        assert!(is_follow_on_phase(PhaseKind::DeepseekApi));

        // The OpenRouter catalog capsule still hosts the three DeepSeek
        // entries and the Ollama Cloud capsule the DeepSeek cloud alias, so
        // DeepSeek stays catalog-reachable without a native model literal.
        let openrouter_capsule =
            include_str!("../../router-daemon/src/openrouter_phase_capsule.rs");
        for entry in [
            "deepseek/deepseek-v4-flash",
            "deepseek/deepseek-v4-pro",
            "deepseek/deepseek-r1",
        ] {
            assert!(
                openrouter_capsule.contains(entry),
                "OpenRouter catalog lost the DeepSeek entry {entry}"
            );
        }
        let ollama_capsule = include_str!("../../router-daemon/src/ollama_cloud_phase_capsule.rs");
        assert!(
            ollama_capsule.contains("deepseek-v4-flash:cloud"),
            "Ollama Cloud catalog lost the DeepSeek cloud alias"
        );
    }

    #[test]
    fn sse_parses_fragmented_stream_usage_once_and_terminal_exactly_once() {
        let transcript = concat!(
            ": heartbeat\r\n\r\n",
            "data: {\"choices\":\r\n",
            "data: [{\"delta\":{\"content\":\"hé\"},\"finish_reason\":null}]}\r\n\r\n",
            "data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":\"call-provider\",\"function\":{\"name\":\"weather__lookup\",\"arguments\":\"{\"}}]}}]}\r\n\r\n",
            "data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"function\":{\"arguments\":\"\\\"city\\\":\\\"Paris\\\"}\"}}]},\"finish_reason\":\"tool_calls\"}]}\r\n\r\n",
            "data: {\"choices\":[],\"usage\":{\"prompt_tokens\":11,\"completion_tokens\":7,\"total_tokens\":18}}\r\n\r\n",
            "data: [DONE]"
        );
        let attempt = AttemptId::parse("attempt-kimi").expect("attempt");
        let mut parser = OpenAiCompatSseParser::new(&attempt);
        let mut fragmented = Vec::new();
        for byte in transcript.as_bytes() {
            fragmented.extend(parser.feed(std::slice::from_ref(byte)).expect("fragment"));
        }
        fragmented.extend(parser.finish().expect("complete"));
        assert!(matches!(
            &fragmented[0],
            OpenAiCompatStreamEvent::ContentDelta { text } if text == "hé"
        ));
        assert_eq!(
            fragmented
                .iter()
                .filter(|event| matches!(event, OpenAiCompatStreamEvent::Usage(_)))
                .count(),
            1
        );
        assert!(matches!(
            fragmented.last(),
            Some(OpenAiCompatStreamEvent::Done)
        ));
        let mut whole = OpenAiCompatSseParser::new(&attempt);
        let mut whole_events = whole.feed(transcript.as_bytes()).expect("whole stream");
        whole_events.extend(whole.finish().expect("complete"));
        assert_eq!(fragmented, whole_events);
    }

    #[test]
    fn usage_and_terminal_are_each_exactly_once() {
        let attempt = AttemptId::parse("attempt-terminal").expect("attempt");
        let usage = b"data: {\"choices\":[],\"usage\":{\"prompt_tokens\":1,\"completion_tokens\":2,\"total_tokens\":3}}\n\n";
        let mut duplicate_usage = OpenAiCompatSseParser::new(&attempt);
        duplicate_usage.feed(usage).expect("first usage");
        assert_eq!(
            duplicate_usage
                .feed(usage)
                .expect_err("duplicate usage")
                .kind,
            OpenAiCompatStreamErrorKind::DuplicateUsage
        );
        let mut missing_usage = OpenAiCompatSseParser::new(&attempt);
        assert_eq!(
            missing_usage
                .feed(b"data: [DONE]\n\n")
                .expect_err("usage required")
                .kind,
            OpenAiCompatStreamErrorKind::MissingUsage
        );
        let mut duplicate_terminal = OpenAiCompatSseParser::new(&attempt);
        duplicate_terminal.feed(usage).expect("usage");
        duplicate_terminal
            .feed(b"data: [DONE]\n\n")
            .expect("terminal");
        assert_eq!(
            duplicate_terminal
                .feed(b"data: [DONE]\n\n")
                .expect_err("duplicate terminal")
                .kind,
            OpenAiCompatStreamErrorKind::DuplicateTerminal
        );
    }

    #[test]
    fn status_and_transport_failures_include_429_and_ambiguity() {
        for (status, class) in [
            (400, OpenAiCompatFailureClass::InvalidRequest),
            (401, OpenAiCompatFailureClass::Authentication),
            (403, OpenAiCompatFailureClass::Authorization),
            (429, OpenAiCompatFailureClass::RateLimited),
            (503, OpenAiCompatFailureClass::UpstreamServer),
        ] {
            let failure =
                normalize_http_response(status, None, Some("999999"), false).expect_err("failure");
            assert_eq!(failure.class, class);
            assert!(failure.retry_after.is_some_and(|retry| retry.capped));
        }
        assert!(
            normalize_http_response(200, Some("text/event-stream; charset=utf-8"), None, false)
                .is_ok()
        );
        assert_eq!(
            normalize_http_response(200, Some("application/json"), None, false)
                .expect_err("content type")
                .class,
            OpenAiCompatFailureClass::MismatchedContentType
        );
        assert_eq!(
            normalize_transport_failure(TransportFailureSignal::Timeout, BodyWriteState::Written,)
                .class,
            OpenAiCompatFailureClass::AmbiguousDispatch
        );
        assert_eq!(
            normalize_transport_failure(
                TransportFailureSignal::Cancelled,
                BodyWriteState::NotWritten,
            )
            .class,
            OpenAiCompatFailureClass::Cancelled
        );
    }

    #[test]
    fn rate_limit_retry_requires_one_short_explicit_delay_before_visibility() {
        let rate_limited =
            normalize_http_response(429, None, Some("2"), false).expect_err("rate limited");
        assert_eq!(
            authorized_rate_limit_retry(rate_limited, ReplayDecision::Allowed, false),
            OpenAiCompatRetryDecision::RetryOnce {
                retry_after_seconds: 2
            }
        );
        let too_long =
            normalize_http_response(429, None, Some("3"), false).expect_err("rate limited");
        assert_eq!(
            authorized_rate_limit_retry(too_long, ReplayDecision::Allowed, false),
            OpenAiCompatRetryDecision::NoRetry
        );
        let post_visible =
            normalize_http_response(429, None, Some("2"), false).expect_err("rate limited");
        assert_eq!(
            authorized_rate_limit_retry(
                post_visible,
                ReplayDecision::Denied(ReplayBoundary::PublicContent),
                false,
            ),
            OpenAiCompatRetryDecision::NoRetry
        );
    }
}