diff --git a/Cargo.lock b/Cargo.lock index 8c340a44..8ff4e67d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -127,6 +127,7 @@ dependencies = [ "anyhow", "async-trait", "bytes", + "dashmap", "eventsource-stream", "futures", "futures-util", diff --git a/crates/aisix-core/src/models/model.rs b/crates/aisix-core/src/models/model.rs index 15b2d2a7..9eb3c4f9 100644 --- a/crates/aisix-core/src/models/model.rs +++ b/crates/aisix-core/src/models/model.rs @@ -19,7 +19,7 @@ static MODEL_ID_RE: Lazy = /// Supported provider prefixes. The `model` field must start with one of these /// followed by `/`. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] #[serde(rename_all = "lowercase")] pub enum Provider { Openai, diff --git a/crates/aisix-gateway/Cargo.toml b/crates/aisix-gateway/Cargo.toml index aea5c295..ebfa4d17 100644 --- a/crates/aisix-gateway/Cargo.toml +++ b/crates/aisix-gateway/Cargo.toml @@ -24,3 +24,7 @@ thiserror.workspace = true anyhow.workspace = true tracing.workspace = true http.workspace = true +dashmap.workspace = true + +[dev-dependencies] +tokio = { workspace = true, features = ["macros", "rt"] } diff --git a/crates/aisix-gateway/src/bridge.rs b/crates/aisix-gateway/src/bridge.rs new file mode 100644 index 00000000..554b1ec1 --- /dev/null +++ b/crates/aisix-gateway/src/bridge.rs @@ -0,0 +1,220 @@ +//! The [`Bridge`] trait — what every provider crate implements. +//! +//! A Bridge is the provider-specific adapter between the gateway's +//! normalised [`ChatFormat`] and whichever upstream API shape the vendor +//! requires. Bridges are held in [`crate::hub::Hub`] and selected by the +//! Model's [`aisix_core::Provider`] enum. +//! +//! Responsibilities of a Bridge: +//! - Translate `ChatFormat` → upstream request body +//! - Perform the HTTP call (authorisation, timeouts, retries at transport) +//! - For streaming requests, produce a `Stream` +//! - For non-streaming, produce a full [`ChatResponse`] +//! - Surface errors as typed [`BridgeError`] variants so the proxy layer +//! can map them to consistent OpenAI-style error envelopes +//! +//! The trait is deliberately `async_trait` rather than GATs — ergonomic +//! wins outweigh the boxing cost on the provider path. + +use aisix_core::Model; +use async_trait::async_trait; +use futures::stream::BoxStream; +use std::time::Duration; + +use crate::chat::{ChatChunk, ChatFormat, ChatResponse}; + +/// Context carried through the whole request lifecycle. +/// +/// The proxy layer fills this in after it has authenticated the request +/// and resolved the target model. Bridges read from it but do not mutate +/// it — the fields relevant to the transport (auth, timeout) are owned +/// references to structures in the current [`aisix_core::AisixSnapshot`]. +#[derive(Debug, Clone)] +pub struct BridgeContext { + /// Correlation id propagated into traces and error envelopes. + pub request_id: String, + /// The resolved upstream model — the Bridge reads provider_config + /// (api_key, api_base) and the upstream model name from here. + pub model: std::sync::Arc, + /// Deadline for the entire upstream call. Bridges are expected to + /// honour this by cancelling any in-flight HTTP request. + pub deadline: Option, +} + +impl BridgeContext { + pub fn new(request_id: impl Into, model: std::sync::Arc) -> Self { + Self { + request_id: request_id.into(), + model, + deadline: None, + } + } + + pub fn with_deadline(mut self, deadline: Duration) -> Self { + self.deadline = Some(deadline); + self + } +} + +/// Error surfaced by any Bridge. Each variant maps to a stable +/// client-visible HTTP status and OpenAI-style error code so the proxy +/// layer can translate without further inspection. +#[derive(Debug, thiserror::Error)] +pub enum BridgeError { + #[error("upstream request timed out after {elapsed_ms}ms")] + Timeout { elapsed_ms: u64 }, + #[error("upstream returned HTTP {status}: {message}")] + UpstreamStatus { status: u16, message: String }, + #[error("upstream returned an unparseable body: {0}")] + UpstreamDecode(String), + #[error("bridge is misconfigured: {0}")] + Config(String), + #[error("transport error: {0}")] + Transport(String), + #[error("upstream cancelled the response mid-stream")] + StreamAborted, +} + +impl BridgeError { + /// Stable HTTP status mapping. The proxy layer uses this to build + /// its OpenAI-compatible `{error:{message,type,...}}` envelope. + pub fn http_status(&self) -> u16 { + match self { + BridgeError::Timeout { .. } => 504, + BridgeError::UpstreamStatus { status, .. } => { + // We only forward 4xx directly; everything else collapses + // to 502 so clients don't see upstream 5xx bleed through. + if (400..500).contains(status) { + *status + } else { + 502 + } + } + BridgeError::UpstreamDecode(_) => 502, + BridgeError::Config(_) => 500, + BridgeError::Transport(_) => 502, + BridgeError::StreamAborted => 502, + } + } + + /// Stable error-type token, mirroring OpenAI's error.type field. + pub fn error_type(&self) -> &'static str { + match self { + BridgeError::Timeout { .. } => "timeout", + BridgeError::UpstreamStatus { .. } => "upstream_error", + BridgeError::UpstreamDecode(_) => "upstream_decode_error", + BridgeError::Config(_) => "config_error", + BridgeError::Transport(_) => "transport_error", + BridgeError::StreamAborted => "stream_aborted", + } + } +} + +/// A live stream of chunks. Boxed so the Bridge trait stays object-safe +/// (the Hub holds `Arc` values). +pub type ChatChunkStream = BoxStream<'static, Result>; + +/// The provider-agnostic chat operation. Implementors live in the +/// individual `aisix-provider-*` crates. +#[async_trait] +pub trait Bridge: Send + Sync + 'static { + /// Human-readable name used in logs and metrics labels. Stable across + /// upgrades so dashboards don't break. + fn name(&self) -> &'static str; + + /// Non-streaming call: one request, one response. + async fn chat( + &self, + req: &ChatFormat, + ctx: &BridgeContext, + ) -> Result; + + /// Streaming call: one request, a stream of deltas. + async fn chat_stream( + &self, + req: &ChatFormat, + ctx: &BridgeContext, + ) -> Result; +} + +#[cfg(test)] +mod tests { + use super::*; + use aisix_core::models::Provider; + + #[test] + fn timeout_maps_to_504() { + let e = BridgeError::Timeout { elapsed_ms: 30_000 }; + assert_eq!(e.http_status(), 504); + assert_eq!(e.error_type(), "timeout"); + } + + #[test] + fn upstream_4xx_passes_through_5xx_collapses_to_502() { + let e400 = BridgeError::UpstreamStatus { + status: 429, + message: "rate limit".into(), + }; + assert_eq!(e400.http_status(), 429); + + let e500 = BridgeError::UpstreamStatus { + status: 503, + message: "busy".into(), + }; + assert_eq!(e500.http_status(), 502); + + let e3xx = BridgeError::UpstreamStatus { + status: 301, + message: "redirect".into(), + }; + // Non-4xx collapses too — redirects we don't follow are 502-worthy. + assert_eq!(e3xx.http_status(), 502); + } + + #[test] + fn transport_and_decode_errors_collapse_to_502() { + assert_eq!( + BridgeError::Transport("connection refused".into()).http_status(), + 502, + ); + assert_eq!( + BridgeError::UpstreamDecode("bad json".into()).http_status(), + 502, + ); + } + + #[test] + fn config_error_maps_to_500() { + assert_eq!( + BridgeError::Config("missing api_key".into()).http_status(), + 500 + ); + } + + #[test] + fn context_defaults_no_deadline_with_helper_setter() { + let m = std::sync::Arc::new(sample_model()); + let ctx = BridgeContext::new("req-1", m.clone()); + assert_eq!(ctx.request_id, "req-1"); + assert!(ctx.deadline.is_none()); + let ctx = ctx.with_deadline(Duration::from_secs(30)); + assert_eq!(ctx.deadline, Some(Duration::from_secs(30))); + } + + fn sample_model() -> Model { + serde_json::from_str( + r#"{ + "name": "test", + "model": "openai/gpt-4o", + "provider_config": {"api_key": "sk-x"} + }"#, + ) + .unwrap() + } + + #[test] + fn sample_model_parses_and_routes_to_openai() { + let m = sample_model(); + assert_eq!(m.provider(), Some(Provider::Openai)); + } +} diff --git a/crates/aisix-gateway/src/chat.rs b/crates/aisix-gateway/src/chat.rs new file mode 100644 index 00000000..98402176 --- /dev/null +++ b/crates/aisix-gateway/src/chat.rs @@ -0,0 +1,264 @@ +//! Provider-agnostic chat request / response types. +//! +//! The gateway normalises every client request into a [`ChatFormat`] and +//! hands it to whichever [`crate::bridge::Bridge`] implementation matches +//! the target provider. The response shape (either a full [`ChatResponse`] +//! or a stream of [`ChatChunk`]s) is symmetric: providers emit the normalised +//! form and the proxy layer re-encodes into whatever the client expects +//! (defaulting to OpenAI-compatible JSON). +//! +//! These types are deliberately a superset of OpenAI's shape because that +//! is the most permissive of the four providers we're targeting; fields +//! that don't map cleanly to a specific upstream become the provider's +//! responsibility to drop or translate. + +use serde::{Deserialize, Serialize}; + +/// Role of a chat message. Matches OpenAI's taxonomy; providers that only +/// support system/user/assistant are expected to reject `Tool` at their +/// own boundary rather than silently collapsing roles. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "lowercase")] +pub enum Role { + System, + User, + Assistant, + Tool, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct ChatMessage { + pub role: Role, + pub content: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub name: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub tool_call_id: Option, +} + +impl ChatMessage { + pub fn system(content: impl Into) -> Self { + Self { + role: Role::System, + content: content.into(), + name: None, + tool_call_id: None, + } + } + + pub fn user(content: impl Into) -> Self { + Self { + role: Role::User, + content: content.into(), + name: None, + tool_call_id: None, + } + } + + pub fn assistant(content: impl Into) -> Self { + Self { + role: Role::Assistant, + content: content.into(), + name: None, + tool_call_id: None, + } + } +} + +/// Normalised chat completion request. +/// +/// `model` is the **public-facing** name from the Admin API (e.g. +/// `"my-gpt4"`), not the upstream model id. The gateway resolves this to +/// an `aisix_core::Model` before calling a Bridge; the Bridge receives +/// only the resolved [`crate::bridge::BridgeContext`] and translates the +/// `ChatFormat` to the provider's own request shape. +/// +/// Unknown top-level fields are **not** rejected — OpenAI's API adds +/// params regularly (e.g. `top_k`, `seed`, `presence_penalty`), and each +/// Bridge is responsible for forwarding or ignoring them. Extras land in +/// the `extra` map via `#[serde(flatten)]`. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ChatFormat { + pub model: String, + pub messages: Vec, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub temperature: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub top_p: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub max_tokens: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub stream: Option, + /// Free-form extra fields the client sent. We don't strip unknown + /// params at the gateway — each Bridge decides what to forward. + #[serde(default, skip_serializing_if = "serde_json::Map::is_empty", flatten)] + pub extra: serde_json::Map, +} + +impl ChatFormat { + pub fn new(model: impl Into, messages: Vec) -> Self { + Self { + model: model.into(), + messages, + temperature: None, + top_p: None, + max_tokens: None, + stream: None, + extra: serde_json::Map::new(), + } + } + + pub fn is_streaming(&self) -> bool { + self.stream.unwrap_or(false) + } +} + +/// Why a completion finished. Unknown upstream reasons collapse to +/// [`FinishReason::Other`] carrying the original string. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum FinishReason { + Stop, + Length, + ContentFilter, + ToolCalls, + Other(String), +} + +#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct UsageStats { + pub prompt_tokens: u32, + pub completion_tokens: u32, + pub total_tokens: u32, +} + +impl UsageStats { + pub fn new(prompt: u32, completion: u32) -> Self { + Self { + prompt_tokens: prompt, + completion_tokens: completion, + total_tokens: prompt.saturating_add(completion), + } + } +} + +/// Full (non-streaming) chat response. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct ChatResponse { + pub id: String, + pub model: String, + pub message: ChatMessage, + pub finish_reason: FinishReason, + pub usage: UsageStats, +} + +/// One streamed delta event. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct ChatChunk { + pub id: String, + pub model: String, + pub delta: ChatDelta, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub finish_reason: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub usage: Option, +} + +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct ChatDelta { + #[serde(default, skip_serializing_if = "Option::is_none")] + pub role: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub content: Option, +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn chat_format_round_trips_through_json() { + let f = ChatFormat { + model: "my-gpt4".into(), + messages: vec![ + ChatMessage::system("you are helpful"), + ChatMessage::user("hi"), + ], + temperature: Some(0.2), + top_p: None, + max_tokens: Some(100), + stream: Some(true), + extra: serde_json::Map::new(), + }; + + let json = serde_json::to_string(&f).unwrap(); + let back: ChatFormat = serde_json::from_str(&json).unwrap(); + assert_eq!(back.model, "my-gpt4"); + assert_eq!(back.messages.len(), 2); + assert_eq!(back.temperature, Some(0.2)); + assert!(back.is_streaming()); + } + + #[test] + fn extras_capture_unknown_top_level_fields() { + // `top_k` isn't a known field — it lands in `extra` so the Bridge + // can decide whether to forward it to the upstream provider. + let json = r#"{ + "model": "my-gpt4", + "messages": [], + "top_k": 40 + }"#; + let f: ChatFormat = serde_json::from_str(json).unwrap(); + assert_eq!(f.extra.get("top_k").and_then(|v| v.as_u64()), Some(40)); + } + + #[test] + fn is_streaming_defaults_to_false_when_unset() { + let f = ChatFormat::new("m", vec![]); + assert!(!f.is_streaming()); + } + + #[test] + fn finish_reason_known_variants_are_snake_case() { + let stop: FinishReason = serde_json::from_str(r#""stop""#).unwrap(); + let content_filter: FinishReason = serde_json::from_str(r#""content_filter""#).unwrap(); + assert_eq!(stop, FinishReason::Stop); + assert_eq!(content_filter, FinishReason::ContentFilter); + } + + #[test] + fn usage_stats_saturates_total() { + let u = UsageStats::new(u32::MAX, 10); + assert_eq!(u.total_tokens, u32::MAX); + } + + #[test] + fn message_constructors_set_role() { + assert_eq!(ChatMessage::system("x").role, Role::System); + assert_eq!(ChatMessage::user("x").role, Role::User); + assert_eq!(ChatMessage::assistant("x").role, Role::Assistant); + } + + #[test] + fn chat_chunk_omits_optional_fields_on_wire() { + let chunk = ChatChunk { + id: "cmpl-1".into(), + model: "m".into(), + delta: ChatDelta { + role: None, + content: Some("hello".into()), + }, + finish_reason: None, + usage: None, + }; + let json = serde_json::to_string(&chunk).unwrap(); + assert!(!json.contains("\"finish_reason\"")); + assert!(!json.contains("\"usage\"")); + assert!(json.contains("\"content\":\"hello\"")); + } +} diff --git a/crates/aisix-gateway/src/hub.rs b/crates/aisix-gateway/src/hub.rs new file mode 100644 index 00000000..64caf9fc --- /dev/null +++ b/crates/aisix-gateway/src/hub.rs @@ -0,0 +1,166 @@ +//! The [`Hub`] dispatches `ChatFormat` requests to the matching +//! [`Bridge`] based on the target Model's `Provider` enum. +//! +//! Hubs are constructed once at startup (spec §1 step 7 — before the +//! proxy router is built) and hold an `Arc` per provider. +//! Lookups are `O(1)` — a 4-entry hashmap keyed on the Provider enum. +//! +//! There is no fallback logic here — that is the proxy layer's job and +//! lands in its own PR. The Hub exists purely to resolve Provider → +//! Bridge cheaply and consistently. + +use aisix_core::models::Provider; +use dashmap::DashMap; +use std::sync::Arc; + +use crate::bridge::Bridge; + +/// Registry of providers → bridges. +/// +/// `DashMap` lets us register bridges after construction (useful for tests +/// and for future dynamic-reload scenarios) without taking out a lock on +/// the lookup path. +#[derive(Default)] +pub struct Hub { + bridges: DashMap>, +} + +impl Hub { + pub fn new() -> Self { + Self::default() + } + + /// Register a bridge for a provider. Overwrites any previous entry, + /// which is what we want during live reconfigure — the etcd watcher + /// can swap a broken bridge without tearing down the Hub. + pub fn register(&self, provider: Provider, bridge: Arc) { + self.bridges.insert(provider, bridge); + } + + pub fn get(&self, provider: Provider) -> Option> { + self.bridges.get(&provider).map(|r| r.clone()) + } + + pub fn providers(&self) -> Vec { + self.bridges.iter().map(|r| *r.key()).collect() + } + + pub fn len(&self) -> usize { + self.bridges.len() + } + + pub fn is_empty(&self) -> bool { + self.bridges.is_empty() + } +} + +impl std::fmt::Debug for Hub { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("Hub") + .field("providers", &self.providers()) + .finish() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::bridge::{BridgeContext, BridgeError, ChatChunkStream}; + use crate::chat::{ChatFormat, ChatMessage, ChatResponse, FinishReason, UsageStats}; + use async_trait::async_trait; + use futures::stream; + + /// Minimal Bridge that short-circuits to a canned response. Used to + /// verify Hub wiring without dragging in reqwest or a real provider. + struct StubBridge { + name: &'static str, + } + + #[async_trait] + impl Bridge for StubBridge { + fn name(&self) -> &'static str { + self.name + } + + async fn chat( + &self, + req: &ChatFormat, + _ctx: &BridgeContext, + ) -> Result { + Ok(ChatResponse { + id: "stub-1".into(), + model: req.model.clone(), + message: ChatMessage::assistant("stubbed"), + finish_reason: FinishReason::Stop, + usage: UsageStats::new(0, 0), + }) + } + + async fn chat_stream( + &self, + _req: &ChatFormat, + _ctx: &BridgeContext, + ) -> Result { + Ok(Box::pin(stream::iter(Vec::new()))) + } + } + + #[test] + fn empty_hub_returns_none_for_any_provider() { + let hub = Hub::new(); + assert!(hub.is_empty()); + assert!(hub.get(Provider::Openai).is_none()); + } + + #[test] + fn register_and_get_round_trip() { + let hub = Hub::new(); + hub.register( + Provider::Openai, + Arc::new(StubBridge { + name: "stub-openai", + }), + ); + let b = hub.get(Provider::Openai).unwrap(); + assert_eq!(b.name(), "stub-openai"); + } + + #[test] + fn register_overwrites_previous_bridge_for_same_provider() { + let hub = Hub::new(); + hub.register(Provider::Openai, Arc::new(StubBridge { name: "v1" })); + hub.register(Provider::Openai, Arc::new(StubBridge { name: "v2" })); + assert_eq!(hub.len(), 1); + assert_eq!(hub.get(Provider::Openai).unwrap().name(), "v2"); + } + + #[test] + fn providers_returns_all_registered_keys() { + let hub = Hub::new(); + hub.register(Provider::Openai, Arc::new(StubBridge { name: "a" })); + hub.register(Provider::Anthropic, Arc::new(StubBridge { name: "b" })); + let mut ps = hub.providers(); + ps.sort_by_key(|p| format!("{p:?}")); + assert_eq!(ps.len(), 2); + } + + #[tokio::test] + async fn registered_bridge_is_callable() { + let hub = Hub::new(); + hub.register(Provider::Openai, Arc::new(StubBridge { name: "stub" })); + let bridge = hub.get(Provider::Openai).unwrap(); + + let m = std::sync::Arc::new( + serde_json::from_str::( + r#"{"name":"t","model":"openai/gpt-4o","provider_config":{"api_key":"k"}}"#, + ) + .unwrap(), + ); + let ctx = BridgeContext::new("req-1", m); + let req = ChatFormat::new("t", vec![ChatMessage::user("hi")]); + + let resp = bridge.chat(&req, &ctx).await.unwrap(); + assert_eq!(resp.message.content, "stubbed"); + assert_eq!(resp.finish_reason, FinishReason::Stop); + } +} diff --git a/crates/aisix-gateway/src/lib.rs b/crates/aisix-gateway/src/lib.rs index 5ef3a55c..44717ec3 100644 --- a/crates/aisix-gateway/src/lib.rs +++ b/crates/aisix-gateway/src/lib.rs @@ -1,4 +1,33 @@ -//! aisix-gateway — Hub-and-Bridge request/response pipeline. +//! aisix-gateway — the Hub-and-Bridge core. +//! +//! This crate holds the provider-agnostic primitives shared by every +//! `aisix-provider-*` crate and by the proxy router: +//! +//! - [`chat`] — normalised `ChatFormat`, `ChatMessage`, `ChatResponse`, +//! streaming `ChatChunk`, and the usage/finish-reason taxonomy. +//! - [`bridge`] — the `Bridge` trait every provider implements, plus +//! `BridgeContext` and typed `BridgeError` with stable HTTP status +//! mapping. +//! - [`hub`] — a small registry keyed on [`aisix_core::models::Provider`] +//! that dispatches `ChatFormat` to the right `Bridge`. +//! - [`sse`] — a provider-agnostic SSE line decoder. Bridges that stream +//! over SSE feed it raw bytes and pull typed events back out. +//! +//! The concrete HTTP transport lives in the provider crates — keeping +//! this crate free of `reqwest` at the public-API level makes it testable +//! without wiremock. #![forbid(unsafe_code)] #![deny(rust_2018_idioms)] + +pub mod bridge; +pub mod chat; +pub mod hub; +pub mod sse; + +pub use bridge::{Bridge, BridgeContext, BridgeError, ChatChunkStream}; +pub use chat::{ + ChatChunk, ChatDelta, ChatFormat, ChatMessage, ChatResponse, FinishReason, Role, UsageStats, +}; +pub use hub::Hub; +pub use sse::{SseDecoder, SseEvent}; diff --git a/crates/aisix-gateway/src/sse.rs b/crates/aisix-gateway/src/sse.rs new file mode 100644 index 00000000..afa018f2 --- /dev/null +++ b/crates/aisix-gateway/src/sse.rs @@ -0,0 +1,210 @@ +//! Server-Sent Events (SSE) line decoder used by streaming Bridges. +//! +//! All four target providers (OpenAI, Anthropic, Gemini, DeepSeek) emit +//! completions as SSE with the OpenAI-style shape: +//! +//! ```text +//! data: {"choices":[…]} +//! data: {"choices":[…]} +//! data: [DONE] +//! ``` +//! +//! The decoder is feed-driven: callers push raw chunks from the HTTP body +//! stream and pull typed [`SseEvent`]s back out. State survives partial +//! messages that straddle chunk boundaries. We intentionally don't use +//! `eventsource-stream` directly here — its interface is reqwest-flavoured +//! and forces a particular body shape; this decoder works against any +//! `&[u8]` and keeps the Bridge trait HTTP-client-agnostic. +//! +//! The decoder handles the subset of the SSE spec that all four providers +//! actually emit: +//! - `data:` lines (everything else is ignored) +//! - `\n\n` separator marking a complete event +//! - UTF-8 only (all four providers are JSON-over-SSE) + +use std::borrow::Cow; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum SseEvent { + /// A `data:` payload (everything between `data: ` and the message + /// terminator, concatenated if the provider split over multiple + /// lines — per the SSE spec). + Data(String), + /// The OpenAI-style sentinel `[DONE]`. Called out separately so + /// Bridges don't have to string-match. + Done, +} + +#[derive(Debug, Default)] +pub struct SseDecoder { + buffer: String, + current_data: String, +} + +impl SseDecoder { + pub fn new() -> Self { + Self::default() + } + + /// Feed a chunk of bytes. Returns every complete event that was + /// unlocked by this feed. Partial messages remain buffered until a + /// subsequent call supplies the rest. + pub fn feed<'a>(&mut self, bytes: impl Into>) -> Vec { + let bytes = bytes.into(); + // Non-UTF-8 bytes are replaced rather than erroring — upstreams + // that break encoding still surface a best-effort event so a + // single bad byte doesn't kill the whole stream. + let chunk = String::from_utf8_lossy(&bytes); + self.buffer.push_str(&chunk); + + let mut events = Vec::new(); + // Event terminator is \n\n; process one message at a time. + while let Some(idx) = self.buffer.find("\n\n") { + let message: String = self.buffer.drain(..idx + 2).collect(); + self.decode_message(&message, &mut events); + } + events + } + + /// Flush any buffered trailing bytes as the final event. Call once + /// the HTTP body has ended; returns `None` if nothing is buffered. + pub fn finish(&mut self) -> Option { + if self.buffer.trim().is_empty() && self.current_data.is_empty() { + return None; + } + // Treat the tail as a terminated message so the decoder emits + // whatever it had collected. + let tail = std::mem::take(&mut self.buffer); + let mut events = Vec::new(); + self.decode_message(&format!("{tail}\n\n"), &mut events); + events.into_iter().next() + } + + fn decode_message(&mut self, message: &str, out: &mut Vec) { + for line in message.lines() { + let line = line.trim_end_matches('\r'); + if line.is_empty() { + continue; + } + // Only `data:` lines are relevant for our providers. + if let Some(rest) = line.strip_prefix("data:") { + let data = rest.strip_prefix(' ').unwrap_or(rest); + if !self.current_data.is_empty() { + self.current_data.push('\n'); + } + self.current_data.push_str(data); + } + // Silently skip other field lines (event:, id:, retry:, comments). + } + + if self.current_data.is_empty() { + return; + } + + let finished = std::mem::take(&mut self.current_data); + if finished.trim() == "[DONE]" { + out.push(SseEvent::Done); + } else { + out.push(SseEvent::Data(finished)); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn single_event_is_emitted_on_terminator() { + let mut d = SseDecoder::new(); + let ev = d.feed(b"data: {\"x\":1}\n\n".as_slice()); + assert_eq!(ev, vec![SseEvent::Data(r#"{"x":1}"#.into())]); + } + + #[test] + fn done_sentinel_is_decoded_separately() { + let mut d = SseDecoder::new(); + let ev = d.feed(b"data: [DONE]\n\n".as_slice()); + assert_eq!(ev, vec![SseEvent::Done]); + } + + #[test] + fn events_split_across_feeds_are_reassembled() { + let mut d = SseDecoder::new(); + let first = d.feed(b"data: {\"x".as_slice()); + let second = d.feed(b"\":1}\n\n".as_slice()); + assert!(first.is_empty()); + assert_eq!(second, vec![SseEvent::Data(r#"{"x":1}"#.into())]); + } + + #[test] + fn multiple_data_lines_concatenate_with_newline() { + let mut d = SseDecoder::new(); + let ev = d.feed(b"data: line1\ndata: line2\n\n".as_slice()); + assert_eq!(ev, vec![SseEvent::Data("line1\nline2".into())]); + } + + #[test] + fn non_data_fields_are_skipped() { + let mut d = SseDecoder::new(); + let ev = d.feed(b"event: ping\ndata: payload\nid: 42\n\n".as_slice()); + assert_eq!(ev, vec![SseEvent::Data("payload".into())]); + } + + #[test] + fn multiple_events_in_one_feed() { + let mut d = SseDecoder::new(); + let ev = d.feed(b"data: a\n\ndata: b\n\ndata: [DONE]\n\n".as_slice()); + assert_eq!( + ev, + vec![ + SseEvent::Data("a".into()), + SseEvent::Data("b".into()), + SseEvent::Done, + ] + ); + } + + #[test] + fn crlf_line_endings_are_tolerated() { + let mut d = SseDecoder::new(); + // Note: SSE event terminator is \n\n per spec; we don't promise + // \r\n\r\n support because no target provider emits it. Here we + // verify stray \r at end-of-line doesn't leak into the payload. + let ev = d.feed(b"data: hello\r\n\n".as_slice()); + assert_eq!(ev, vec![SseEvent::Data("hello".into())]); + } + + #[test] + fn finish_emits_trailing_unterminated_event() { + let mut d = SseDecoder::new(); + let mid = d.feed(b"data: tail-only".as_slice()); + assert!(mid.is_empty()); + let finale = d.finish(); + assert_eq!(finale, Some(SseEvent::Data("tail-only".into()))); + } + + #[test] + fn finish_on_empty_buffer_returns_none() { + let mut d = SseDecoder::new(); + assert!(d.finish().is_none()); + } + + #[test] + fn invalid_utf8_is_lossily_decoded() { + let mut d = SseDecoder::new(); + let bytes: Vec = b"data: " + .iter() + .copied() + .chain([0xff_u8, b'\n', b'\n']) + .collect(); + let ev = d.feed(bytes.as_slice()); + // 0xff becomes U+FFFD. + assert_eq!(ev.len(), 1); + if let SseEvent::Data(payload) = &ev[0] { + assert!(payload.contains('\u{FFFD}')); + } else { + panic!("expected Data event"); + } + } +}