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;
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}));
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>,
}
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,
}
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() {
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());
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() {
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() {
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));
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
);
}
}