From 074f90c3915b477486210e429cfc4581815e034b Mon Sep 17 00:00:00 2001 From: Scott Werner Date: Thu, 11 Jun 2026 09:12:02 -0400 Subject: [PATCH] refactor(llm): consolidate the HTTP transport behind the codec seam (#491) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit PR 6 of the gateway refactor series (after #481, #488, #487, #489): collapse the four per-adapter transport copies into one `transport` module. Net −157 lines, and every cross-adapter duplication flagged in the #487/#488 simplify findings is resolved here. ## What moved where **`transport.rs` (new)** — how bytes travel, dialect-blind: - `HttpTransport` (promoted from `providers::http_api::HttpApi`): client, auth key, base URL, timeouts - `LineReader` + `parse_retry_after` + `parse_rate_limit_headers` (moved from `providers::common`, re-export shims kept there for the frozen fabro-cli imports; `LineReader::new` keeps its 2-arg signature) - `complete_via_http` / `send_for_body`: blocking send with the shared timeout/error/status warn logs, non-2xx mapped through `Codec::decode_error` - `stream_via_http` + one SSE decode loop, parameterized by `SseFraming::{EventBlocks, DataLines}` — replaces the four verbatim `StreamLoop` + unfold copies and the four divergent framers (anthropic's `parse_sse_block`, openai's `parse_sse_message`, the inline data-line handling in openai_compatible/gemini, and fabro_server's private block parser) **`codec/mod.rs`** — gains the dialect-neutral pure helpers `parse_error_body` and `extract_system_prompt` (moved from `providers::common`), so the codec layer no longer imports from the transport-side providers module. **Adapters** — shrink to auth + route config + codec composition. `send_and_read_response` and its `error_code_field` parameter are deleted: the dialect error-body key now lives only in the codecs, and any future `decode_error` override applies to blocking and streaming paths alike. ## Unified SSE framing semantics (deliberate decisions) The four framers disagreed on edge cases; the shared framer picks one behavior, stated here rather than chosen silently: - data payloads are trimmed; multi-line `data:` payloads join with `\n`; CRLF tolerated in both modes - comment (`:`), blank, and non-data lines are skipped - events with an **empty payload are dropped** rather than handed to the decoder — previously anthropic would error the whole stream on a bare `data:` line and openai_compatible would feed the decoder an empty string (also an error); openai/gemini already skipped All streaming wire snapshots pass unmodified through the shared loop, and the framer has direct unit tests for these cases. ## Behavior notes (beyond the framing edge cases) - **Error values are byte-identical**: `Codec::decode_error`'s default is exactly the `parse_error_body("type")` + `error_from_status_code` path the deleted call sites inlined; gemini's gRPC-aware override is what its paths already used. - **Logging only**: gemini's blocking paths gain the shared timeout/error/status warn logs (they had none); count-tokens requests are uniformly tagged `operation="input_token_count"` (previously only openai's was). The openai count-tokens logging pin passes unchanged. - gemini's timeout error message now uses the configured provider name instead of a hardcoded `gemini:` prefix (visible only on custom-named gemini routes). ## Verification - `cargo nextest run --workspace`: green except the 5 pre-existing fabro-workflow failures that fail identically on main (environment-dependent, unrelated) - fabro-llm: 545 passed — all PR 0 wire snapshots unmodified - clippy `-D warnings` + pinned-nightly fmt clean - fabro-cli compiles against the frozen `providers::common::{LineReader, parse_retry_after}` paths Next in the series: PR 7 (codec on the route in fabro-model) — route vocabulary + the route-equivalence table test. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Fable 5 --- .../src/codec/anthropic_messages/encode.rs | 6 +- .../src/codec/gemini_generate/encode.rs | 3 +- .../src/codec/gemini_generate/mod.rs | 3 +- lib/crates/fabro-llm/src/codec/mod.rs | 182 +++++- lib/crates/fabro-llm/src/lib.rs | 1 + .../fabro-llm/src/providers/anthropic.rs | 135 +--- lib/crates/fabro-llm/src/providers/common.rs | 455 +------------ .../fabro-llm/src/providers/fabro_server.rs | 60 +- lib/crates/fabro-llm/src/providers/gemini.rs | 141 +---- .../fabro-llm/src/providers/http_api.rs | 60 -- lib/crates/fabro-llm/src/providers/mod.rs | 1 - lib/crates/fabro-llm/src/providers/openai.rs | 145 +---- .../src/providers/openai_compatible.rs | 109 +--- lib/crates/fabro-llm/src/transport.rs | 596 ++++++++++++++++++ 14 files changed, 869 insertions(+), 1028 deletions(-) delete mode 100644 lib/crates/fabro-llm/src/providers/http_api.rs create mode 100644 lib/crates/fabro-llm/src/transport.rs diff --git a/lib/crates/fabro-llm/src/codec/anthropic_messages/encode.rs b/lib/crates/fabro-llm/src/codec/anthropic_messages/encode.rs index 3a2347ef3..360f4210a 100644 --- a/lib/crates/fabro-llm/src/codec/anthropic_messages/encode.rs +++ b/lib/crates/fabro-llm/src/codec/anthropic_messages/encode.rs @@ -9,8 +9,7 @@ use base64::engine::general_purpose::STANDARD as BASE64_STANDARD; use super::SYNTHETIC_TOOL_NAME; use super::wire::{ApiMessage, ApiRequest, ApiToolDef, CacheControl, CountTokensRequest}; -use crate::codec::{AnthropicVersion, CodecCtx, EncodedRequest}; -use crate::providers::common; +use crate::codec::{AnthropicVersion, CodecCtx, EncodedRequest, extract_system_prompt}; use crate::types::{ ContentPart, Message, ReasoningEffort, ReasoningEffortFeature, Request, ResponseFormatType, Role, Speed, ThinkingData, ToolChoice, ToolDefinition, @@ -71,7 +70,7 @@ fn build_headers(ctx: &CodecCtx<'_>) -> Vec<(String, String)> { fn build_request(ctx: &CodecCtx<'_>, stream: bool) -> ApiRequest { let request = ctx.request; - let (system, other_messages) = common::extract_system_prompt(&request.messages); + let (system, other_messages) = extract_system_prompt(&request.messages); let mut api_messages = translate_messages(&other_messages); // `ToolChoice::None` omits the tools entirely instead of sending a choice. @@ -544,6 +543,7 @@ mod tests { use super::*; use crate::codec::CodecParams; + use crate::providers::common; use crate::types::{AudioData, DocumentData, ResponseFormat}; // --- Test helpers -------------------------------------------------------- diff --git a/lib/crates/fabro-llm/src/codec/gemini_generate/encode.rs b/lib/crates/fabro-llm/src/codec/gemini_generate/encode.rs index cad715622..d85ab691e 100644 --- a/lib/crates/fabro-llm/src/codec/gemini_generate/encode.rs +++ b/lib/crates/fabro-llm/src/codec/gemini_generate/encode.rs @@ -11,8 +11,7 @@ use base64::engine::general_purpose::STANDARD as BASE64_STANDARD; use super::wire::{ ApiRequest, Content, GeminiFunctionDecl, GeminiToolGroup, GenerationOptions, SystemInstruction, }; -use crate::codec::{CodecCtx, EncodedRequest}; -use crate::providers::common::extract_system_prompt; +use crate::codec::{CodecCtx, EncodedRequest, extract_system_prompt}; use crate::types::{ ContentPart, Message, ResponseFormat, ResponseFormatType, Role, ToolChoice, ToolDefinition, }; diff --git a/lib/crates/fabro-llm/src/codec/gemini_generate/mod.rs b/lib/crates/fabro-llm/src/codec/gemini_generate/mod.rs index 112911879..dab98c4c9 100644 --- a/lib/crates/fabro-llm/src/codec/gemini_generate/mod.rs +++ b/lib/crates/fabro-llm/src/codec/gemini_generate/mod.rs @@ -13,9 +13,8 @@ mod encode; mod stream; mod wire; -use crate::codec::{Codec, CodecCtx, EncodedRequest, StreamDecoder}; +use crate::codec::{Codec, CodecCtx, EncodedRequest, StreamDecoder, parse_error_body}; use crate::error::Error; -use crate::providers::common::parse_error_body; use crate::types::{RateLimitInfo, Response}; /// Codec for the Gemini `generateContent` wire dialect. diff --git a/lib/crates/fabro-llm/src/codec/mod.rs b/lib/crates/fabro-llm/src/codec/mod.rs index d8e6c3b18..e8d8d6e2a 100644 --- a/lib/crates/fabro-llm/src/codec/mod.rs +++ b/lib/crates/fabro-llm/src/codec/mod.rs @@ -18,8 +18,7 @@ pub(crate) mod openai_responses; use fabro_model::Model; use crate::error::{Error, error_from_status_code}; -use crate::providers::common::parse_error_body; -use crate::types::{RateLimitInfo, Request, Response, StreamEvent}; +use crate::types::{Message, RateLimitInfo, Request, Response, Role, StreamEvent}; /// Per-request context. Borrowed — the codec reads what it needs and returns. pub(crate) struct CodecCtx<'a> { @@ -185,3 +184,182 @@ pub(crate) trait StreamDecoder: Send + 'static { /// gemini — synthesize `Finish` unconditionally if not yet finished. fn finish(&mut self) -> Vec; } + +// --- Dialect-neutral translation helpers +// --------------------------------------- + +/// Parse an error response body, extracting the message and error code. +/// +/// `error_code_field` is the JSON field name for the error code (e.g. "type" or +/// "status"). +#[must_use] +pub(crate) fn parse_error_body( + body: &str, + error_code_field: &str, +) -> (String, Option, Option) { + serde_json::from_str::(body).map_or_else( + |_| (body.to_string(), None, None), + |v| { + let message = v + .get("error") + .and_then(|e| e.get("message")) + .and_then(serde_json::Value::as_str) + // Codex endpoint returns {"detail": "..."} instead of {"error": {"message": "..."}} + .or_else(|| v.get("detail").and_then(serde_json::Value::as_str)) + .unwrap_or("Unknown error") + .to_string(); + let error_code = v + .get("error") + .and_then(|e| e.get(error_code_field)) + .and_then(serde_json::Value::as_str) + .map(String::from); + (message, error_code, Some(v)) + }, + ) +} + +/// Extract system and developer messages from a message list. +/// +/// Returns the joined system prompt and the remaining messages. +/// Per spec, Developer role messages are merged with system messages +/// for Anthropic and Gemini. +#[must_use] +pub(crate) fn extract_system_prompt(messages: &[Message]) -> (Option, Vec<&Message>) { + let mut system_parts = Vec::new(); + let mut other = Vec::new(); + for msg in messages { + if msg.role == Role::System || msg.role == Role::Developer { + let text = msg.text(); + if !text.trim().is_empty() { + system_parts.push(text); + } + } else { + other.push(msg); + } + } + let system = if system_parts.is_empty() { + None + } else { + Some(system_parts.join("\n")) + }; + (system, other) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::types::ContentPart; + + // --- parse_error_body --- + + #[test] + fn parse_error_body_valid_json() { + let body = r#"{"error":{"message":"rate limited","type":"rate_limit_error"}}"#; + let (msg, code, raw) = parse_error_body(body, "type"); + assert_eq!(msg, "rate limited"); + assert_eq!(code.as_deref(), Some("rate_limit_error")); + assert!(raw.is_some()); + } + + #[test] + fn parse_error_body_missing_error_field() { + let body = r#"{"status":"fail"}"#; + let (msg, code, raw) = parse_error_body(body, "type"); + assert_eq!(msg, "Unknown error"); + assert_eq!(code, None); + assert!(raw.is_some()); + } + + #[test] + fn parse_error_body_not_json() { + let body = "Internal Server Error"; + let (msg, code, raw) = parse_error_body(body, "type"); + assert_eq!(msg, "Internal Server Error"); + assert_eq!(code, None); + assert!(raw.is_none()); + } + + #[test] + fn parse_error_body_different_code_field() { + let body = r#"{"error":{"message":"bad","status":"INVALID_ARGUMENT"}}"#; + let (msg, code, _) = parse_error_body(body, "status"); + assert_eq!(msg, "bad"); + assert_eq!(code.as_deref(), Some("INVALID_ARGUMENT")); + } + + #[test] + fn parse_error_body_no_message() { + let body = r#"{"error":{"type":"server_error"}}"#; + let (msg, code, _) = parse_error_body(body, "type"); + assert_eq!(msg, "Unknown error"); + assert_eq!(code.as_deref(), Some("server_error")); + } + + // --- extract_system_prompt --- + + #[test] + fn extract_system_prompt_no_system() { + let msgs = vec![Message::user("hello")]; + let (sys, other) = extract_system_prompt(&msgs); + assert_eq!(sys, None); + assert_eq!(other.len(), 1); + } + + #[test] + fn extract_system_prompt_system_only() { + let msgs = vec![Message::system("Be helpful"), Message::user("hi")]; + let (sys, other) = extract_system_prompt(&msgs); + assert_eq!(sys.as_deref(), Some("Be helpful")); + assert_eq!(other.len(), 1); + assert_eq!(other[0].role, Role::User); + } + + #[test] + fn extract_system_prompt_multiple_system() { + let msgs = vec![ + Message::system("Rule 1"), + Message::system("Rule 2"), + Message::user("hi"), + ]; + let (sys, other) = extract_system_prompt(&msgs); + assert_eq!(sys.as_deref(), Some("Rule 1\nRule 2")); + assert_eq!(other.len(), 1); + } + + #[test] + fn extract_system_prompt_developer_role() { + let dev = Message { + role: Role::Developer, + content: vec![ContentPart::text("dev instructions")], + name: None, + tool_call_id: None, + }; + let msgs = vec![dev, Message::user("hi")]; + let (sys, other) = extract_system_prompt(&msgs); + assert_eq!(sys.as_deref(), Some("dev instructions")); + assert_eq!(other.len(), 1); + } + + #[test] + fn extract_system_prompt_ignores_whitespace_system_and_developer() { + let dev = Message { + role: Role::Developer, + content: vec![ContentPart::text(" \n\t ")], + name: None, + tool_call_id: None, + }; + let msgs = vec![Message::system(" "), dev, Message::user("hi")]; + let (sys, other) = extract_system_prompt(&msgs); + assert_eq!(sys, None); + assert_eq!(other.len(), 1); + assert_eq!(other[0].role, Role::User); + } + + #[test] + fn extract_system_prompt_empty() { + let msgs: Vec = vec![]; + let (sys, other) = extract_system_prompt(&msgs); + assert_eq!(sys, None); + assert!(other.is_empty()); + } +} diff --git a/lib/crates/fabro-llm/src/lib.rs b/lib/crates/fabro-llm/src/lib.rs index afe1bc306..7476d95be 100644 --- a/lib/crates/fabro-llm/src/lib.rs +++ b/lib/crates/fabro-llm/src/lib.rs @@ -11,6 +11,7 @@ pub mod providers; pub mod retry; pub mod token_count; pub mod tools; +pub(crate) mod transport; pub mod types; pub use error::{Error, ProviderErrorDetail, ProviderErrorKind, Result}; diff --git a/lib/crates/fabro-llm/src/providers/anthropic.rs b/lib/crates/fabro-llm/src/providers/anthropic.rs index beb2847ec..d65494b34 100644 --- a/lib/crates/fabro-llm/src/providers/anthropic.rs +++ b/lib/crates/fabro-llm/src/providers/anthropic.rs @@ -1,19 +1,15 @@ use std::sync::Arc; use fabro_model::{Catalog, ReasoningEffortFeature}; -use futures::stream; use crate::attachments::{self, AttachmentPolicy}; use crate::codec::anthropic_messages::{AnthropicMessages, anthropic_option}; -use crate::codec::{ - AnthropicVersion, Codec, CodecCtx, CodecParams, EncodedRequest, RawEvent, StreamDecoder, -}; +use crate::codec::{AnthropicVersion, Codec, CodecCtx, CodecParams, EncodedRequest}; use crate::error::Error; use crate::provider::{self, ProviderAdapter, StreamEventStream}; -use crate::providers::common::{ - self as common, parse_rate_limit_headers, parse_retry_after, send_and_read_response, -}; +use crate::providers::common::{self as common}; use crate::token_count::{InputTokenCount, InputTokenCountMethod}; +use crate::transport::{self, HttpTransport, SseFraming}; use crate::types::{AdapterTimeout, Request, Response, StreamEvent}; const DEFAULT_BASE_URL: &str = "https://api.anthropic.com/v1"; @@ -25,7 +21,7 @@ const DEFAULT_BASE_URL: &str = "https://api.anthropic.com/v1"; /// between the direct-Anthropic and Kimi-over-anthropic behaviors. All wire /// translation lives in the codec. pub struct Adapter { - pub(crate) http: super::http_api::HttpApi, + pub(crate) http: HttpTransport, provider_name: String, catalog: Option>, } @@ -39,7 +35,7 @@ impl Adapter { #[must_use] pub fn new_optional_auth(api_key: Option) -> Self { Self { - http: super::http_api::HttpApi::new_optional(api_key, DEFAULT_BASE_URL), + http: HttpTransport::new_optional(api_key, DEFAULT_BASE_URL), provider_name: "anthropic".to_string(), catalog: None, } @@ -207,16 +203,6 @@ enum AuthScheme { Bearer, } -/// State driving the streaming byte loop: the codec's decoder plus the line -/// reader, with a buffer that flattens batched events into individual items. -struct StreamLoop { - decoder: Box, - line_reader: super::common::LineReader, - pending: std::collections::VecDeque, - done: bool, - finished_emitted: bool, -} - /// The `provider_options.anthropic.thinking.type` value, if any. fn anthropic_thinking_type(provider_options: Option<&serde_json::Value>) -> Option<&str> { anthropic_option(provider_options, "thinking") @@ -224,34 +210,6 @@ fn anthropic_thinking_type(provider_options: Option<&serde_json::Value>) -> Opti .and_then(serde_json::Value::as_str) } -/// Parse an SSE event block (lines separated within a `\n\n`-delimited chunk) -/// into `(event_type, data)`. Returns `None` for blocks with no `data:` lines -/// (e.g. heartbeat comments). Borrows from the block — Anthropic events carry -/// a single `data:` line, so the hot path allocates nothing. -fn parse_sse_block(event_block: &str) -> Option<(&str, std::borrow::Cow<'_, str>)> { - let mut event_type = ""; - let mut data: Option> = None; - - for line in event_block.lines() { - if let Some(rest) = line.strip_prefix("event:") { - event_type = rest.trim(); - } else if let Some(rest) = line.strip_prefix("data:") { - let rest = rest.trim(); - data = Some(match data { - None => std::borrow::Cow::Borrowed(rest), - Some(prev) => { - let mut joined = prev.into_owned(); - joined.push('\n'); - joined.push_str(rest); - std::borrow::Cow::Owned(joined) - } - }); - } - } - - data.map(|data| (event_type, data)) -} - #[async_trait::async_trait] impl ProviderAdapter for Adapter { fn name(&self) -> &str { @@ -281,7 +239,8 @@ impl ProviderAdapter for Adapter { if let Some(t) = self.http.request_timeout { req = req.timeout(t); } - let (body, _headers) = send_and_read_response(req, &self.provider_name, "type").await?; + let (body, _headers) = + transport::send_for_body(req, "input_token_count", &codec, &ctx).await?; let input_tokens = codec.decode_count_tokens(&body)?; Ok(Some(InputTokenCount { @@ -313,9 +272,7 @@ impl ProviderAdapter for Adapter { if let Some(t) = self.http.request_timeout { req = req.timeout(t); } - let (body, headers) = send_and_read_response(req, &self.provider_name, "type").await?; - let rate_limit = parse_rate_limit_headers(&headers); - codec.decode_response(&body, &ctx, rate_limit) + transport::complete_via_http(req, &codec, &ctx).await } async fn stream(&self, request: &Request) -> Result { @@ -328,74 +285,14 @@ impl ProviderAdapter for Adapter { let ctx = self.codec_ctx(&resolved, &deployment_id, &route.codec_params); let encoded = codec.encode(&ctx, true)?; - let http_resp = self - .build_http_request(&encoded, &route) - .send() - .await - .map_err(|e| Error::network(e.to_string(), e))?; - - let status = http_resp.status(); - if !status.is_success() { - let retry_after = parse_retry_after(http_resp.headers()); - let body = http_resp - .text() - .await - .map_err(|e| Error::network(e.to_string(), e))?; - return Err(codec.decode_error(status.as_u16(), &body, &ctx, retry_after)); - } - - let rate_limit = parse_rate_limit_headers(http_resp.headers()); - let stream_read_timeout = self.http.stream_read_timeout; - let decoder = codec.stream_decoder(&ctx, rate_limit); - - let out = stream::unfold( - StreamLoop { - decoder, - line_reader: super::common::LineReader::new(http_resp, stream_read_timeout), - pending: std::collections::VecDeque::new(), - done: false, - finished_emitted: false, - }, - |mut state| async move { - loop { - if let Some(event) = state.pending.pop_front() { - return Some((Ok(event), state)); - } - - if state.done { - if state.finished_emitted { - return None; - } - state.finished_emitted = true; - let events = state.decoder.finish(); - if events.is_empty() { - return None; - } - state.pending.extend(events); - continue; - } - - match state.line_reader.read_next_chunk("\n\n").await { - Ok(Some(block)) => { - let Some((event_type, data)) = parse_sse_block(&block) else { - continue; - }; - match state.decoder.on_event(RawEvent { - event: Some(event_type), - data: &data, - }) { - Ok(events) => state.pending.extend(events), - Err(e) => return Some((Err(e), state)), - } - } - Ok(None) => state.done = true, - Err(e) => return Some((Err(e), state)), - } - } - }, - ); - - Ok(Box::pin(out)) + transport::stream_via_http( + self.build_http_request(&encoded, &route), + &codec, + &ctx, + SseFraming::EventBlocks, + self.http.stream_read_timeout, + ) + .await } fn supports_tool_choice(&self, mode: &str) -> bool { diff --git a/lib/crates/fabro-llm/src/providers/common.rs b/lib/crates/fabro-llm/src/providers/common.rs index 883a4634d..8871ff3f7 100644 --- a/lib/crates/fabro-llm/src/providers/common.rs +++ b/lib/crates/fabro-llm/src/providers/common.rs @@ -1,13 +1,8 @@ use base64::Engine; use base64::engine::general_purpose::STANDARD as BASE64_STANDARD; -use fabro_http::HeaderMap; use fabro_model::{Catalog, Model}; use fabro_static::EnvVars; -use tokio::{fs, time}; -use tracing::warn; - -use crate::error::{Error, error_from_status_code}; -use crate::types::{Message, RateLimitInfo, Role}; +use tokio::fs; #[must_use] pub fn catalog_model<'a>(catalog: Option<&'a Catalog>, model: &str) -> Option<&'a Model> { @@ -21,63 +16,6 @@ pub fn api_model_id(catalog: Option<&Catalog>, model: &str) -> String { .map_or_else(|| model.to_string(), |settings| settings.api_id.clone()) } -/// Parse an error response body, extracting the message and error code. -/// -/// `error_code_field` is the JSON field name for the error code (e.g. "type" or -/// "status"). -#[must_use] -pub fn parse_error_body( - body: &str, - error_code_field: &str, -) -> (String, Option, Option) { - serde_json::from_str::(body).map_or_else( - |_| (body.to_string(), None, None), - |v| { - let message = v - .get("error") - .and_then(|e| e.get("message")) - .and_then(serde_json::Value::as_str) - // Codex endpoint returns {"detail": "..."} instead of {"error": {"message": "..."}} - .or_else(|| v.get("detail").and_then(serde_json::Value::as_str)) - .unwrap_or("Unknown error") - .to_string(); - let error_code = v - .get("error") - .and_then(|e| e.get(error_code_field)) - .and_then(serde_json::Value::as_str) - .map(String::from); - (message, error_code, Some(v)) - }, - ) -} - -/// Extract system and developer messages from a message list. -/// -/// Returns the joined system prompt and the remaining messages. -/// Per spec, Developer role messages are merged with system messages -/// for Anthropic and Gemini. -#[must_use] -pub fn extract_system_prompt(messages: &[Message]) -> (Option, Vec<&Message>) { - let mut system_parts = Vec::new(); - let mut other = Vec::new(); - for msg in messages { - if msg.role == Role::System || msg.role == Role::Developer { - let text = msg.text(); - if !text.trim().is_empty() { - system_parts.push(text); - } - } else { - other.push(msg); - } - } - let system = if system_parts.is_empty() { - None - } else { - Some(system_parts.join("\n")) - }; - (system, other) -} - /// Check if a URL string looks like a local file path. #[must_use] pub fn is_file_path(url: &str) -> bool { @@ -135,189 +73,13 @@ pub async fn load_file_as_base64(path: &str) -> Result<(String, String), std::io Ok((BASE64_STANDARD.encode(&data), mime)) } -/// Extract the `Retry-After` header value from an HTTP response as seconds. -#[must_use] -pub fn parse_retry_after(headers: &HeaderMap) -> Option { - headers - .get("retry-after") - .and_then(|v| v.to_str().ok()) - .and_then(|s| s.parse::().ok()) -} - -/// Parse `x-ratelimit-*` headers into a `RateLimitInfo`. -/// -/// Returns `None` if no rate limit headers are present. -#[must_use] -pub fn parse_rate_limit_headers(headers: &HeaderMap) -> Option { - fn header_i64(headers: &HeaderMap, name: &str) -> Option { - headers - .get(name) - .and_then(|v| v.to_str().ok()) - .and_then(|s| s.parse::().ok()) - } - - fn header_str(headers: &HeaderMap, name: &str) -> Option { - headers - .get(name) - .and_then(|v| v.to_str().ok()) - .map(String::from) - } - - let requests_remaining = header_i64(headers, "x-ratelimit-remaining-requests"); - let requests_limit = header_i64(headers, "x-ratelimit-limit-requests"); - let tokens_remaining = header_i64(headers, "x-ratelimit-remaining-tokens"); - let tokens_limit = header_i64(headers, "x-ratelimit-limit-tokens"); - let reset_at = header_str(headers, "x-ratelimit-reset-requests") - .or_else(|| header_str(headers, "x-ratelimit-reset-tokens")); - - if requests_remaining.is_none() - && requests_limit.is_none() - && tokens_remaining.is_none() - && tokens_limit.is_none() - && reset_at.is_none() - { - return None; - } - - Some(RateLimitInfo { - requests_remaining, - requests_limit, - tokens_remaining, - tokens_limit, - reset_at, - }) -} - -/// Send an HTTP request, read the response body, and return it along with the -/// response headers. -/// -/// Returns an error on non-success status. -/// -/// # Errors -/// -/// Returns `Error::Network` on connection failure or `Error::Provider` on -/// non-success status. -pub async fn send_and_read_response( - request: fabro_http::RequestBuilder, - provider: &str, - error_code_field: &str, -) -> Result<(String, HeaderMap), Error> { - send_and_read_response_with_operation(request, provider, error_code_field, "provider_request") - .await -} - -pub(crate) async fn send_and_read_response_with_operation( - request: fabro_http::RequestBuilder, - provider: &str, - error_code_field: &str, - operation: &str, -) -> Result<(String, HeaderMap), Error> { - let http_resp = request.send().await.map_err(|e| { - if e.is_timeout() { - warn!(provider = %provider, operation = %operation, error = %e, "Provider request timed out"); - Error::request_timeout(format!("{provider}: {e}"), e) - } else { - warn!(provider = %provider, operation = %operation, error = %e, "Provider network error"); - Error::network(e.to_string(), e) - } - })?; - - let status = http_resp.status(); - let retry_after = parse_retry_after(http_resp.headers()); - let headers = http_resp.headers().clone(); - let body = http_resp - .text() - .await - .map_err(|e| Error::network(e.to_string(), e))?; - - if !status.is_success() { - warn!(provider = %provider, operation = %operation, status = status.as_u16(), "Provider returned error"); - let (msg, code, raw) = parse_error_body(&body, error_code_field); - return Err(error_from_status_code( - status.as_u16(), - msg, - provider.to_string(), - code, - raw, - retry_after, - )); - } - - Ok((body, headers)) -} - -/// Shared line reader for SSE streams. -/// -/// Buffers bytes from a `fabro_http::Response` and splits them by a -/// configurable delimiter (e.g. `"\n"` for Gemini/OpenAI-compatible, `"\n\n"` -/// for Anthropic/OpenAI SSE event blocks). -pub struct LineReader { - response: fabro_http::Response, - buffer: String, - stream_read_timeout: Option, -} - -impl LineReader { - pub fn new( - response: fabro_http::Response, - stream_read_timeout: Option, - ) -> Self { - Self { - response, - buffer: String::new(), - stream_read_timeout, - } - } - - /// Read the next complete segment delimited by `delimiter`. - /// - /// Returns `Ok(Some(segment))` for each complete segment, `Ok(None)` when - /// the stream is exhausted, or `Err` on I/O or timeout errors. When the - /// stream ends with data remaining in the buffer, the leftover is returned - /// as a final segment. - pub async fn read_next_chunk(&mut self, delimiter: &str) -> Result, Error> { - loop { - if let Some(pos) = self.buffer.find(delimiter) { - let segment = self.buffer[..pos].to_string(); - self.buffer = self.buffer[pos + delimiter.len()..].to_string(); - return Ok(Some(segment)); - } - - let chunk_result = match self.stream_read_timeout { - Some(timeout) => time::timeout(timeout, self.response.chunk()).await, - None => Ok(self.response.chunk().await), - }; - match chunk_result { - Ok(Ok(Some(bytes))) => { - let text = String::from_utf8_lossy(&bytes); - self.buffer.push_str(&text); - } - Ok(Ok(None)) => { - if self.buffer.is_empty() { - return Ok(None); - } - let remaining = std::mem::take(&mut self.buffer); - return Ok(Some(remaining)); - } - Ok(Err(e)) => { - return Err(Error::stream_error(e.to_string(), e)); - } - Err(_) => { - warn!("Stream read timed out waiting for next event"); - return Err(Error::Stream { - message: "stream read timed out waiting for next event".to_string(), - source: None, - }); - } - } - } - } -} +// Transport pieces moved to `crate::transport`; re-exported here because +// fabro-cli imports them from this path (frozen public surface). +pub use crate::transport::{LineReader, parse_rate_limit_headers, parse_retry_after}; #[cfg(test)] mod tests { use super::*; - use crate::types::ContentPart; #[test] fn is_file_path_absolute() { @@ -359,213 +121,4 @@ mod tests { assert_eq!(mime_from_extension("file.xyz"), "application/octet-stream"); assert_eq!(mime_from_extension("noext"), "application/octet-stream"); } - - #[test] - fn parse_rate_limit_headers_all_present() { - let mut headers = HeaderMap::new(); - headers.insert("x-ratelimit-remaining-requests", "99".parse().unwrap()); - headers.insert("x-ratelimit-limit-requests", "100".parse().unwrap()); - headers.insert("x-ratelimit-remaining-tokens", "9000".parse().unwrap()); - headers.insert("x-ratelimit-limit-tokens", "10000".parse().unwrap()); - headers.insert( - "x-ratelimit-reset-requests", - "2024-01-01T00:00:00Z".parse().unwrap(), - ); - - let info = parse_rate_limit_headers(&headers).unwrap(); - assert_eq!(info.requests_remaining, Some(99)); - assert_eq!(info.requests_limit, Some(100)); - assert_eq!(info.tokens_remaining, Some(9000)); - assert_eq!(info.tokens_limit, Some(10000)); - assert_eq!(info.reset_at.as_deref(), Some("2024-01-01T00:00:00Z")); - } - - #[test] - fn parse_rate_limit_headers_none_present() { - let headers = HeaderMap::new(); - assert!(parse_rate_limit_headers(&headers).is_none()); - } - - #[test] - fn parse_rate_limit_headers_partial() { - let mut headers = HeaderMap::new(); - headers.insert("x-ratelimit-remaining-requests", "50".parse().unwrap()); - - let info = parse_rate_limit_headers(&headers).unwrap(); - assert_eq!(info.requests_remaining, Some(50)); - assert_eq!(info.requests_limit, None); - assert_eq!(info.tokens_remaining, None); - assert_eq!(info.tokens_limit, None); - assert_eq!(info.reset_at, None); - } - - #[test] - fn parse_rate_limit_headers_reset_tokens_fallback() { - let mut headers = HeaderMap::new(); - headers.insert("x-ratelimit-limit-tokens", "5000".parse().unwrap()); - headers.insert( - "x-ratelimit-reset-tokens", - "2024-06-01T12:00:00Z".parse().unwrap(), - ); - - let info = parse_rate_limit_headers(&headers).unwrap(); - assert_eq!(info.tokens_limit, Some(5000)); - assert_eq!(info.reset_at.as_deref(), Some("2024-06-01T12:00:00Z")); - } - - #[test] - fn parse_rate_limit_headers_invalid_values_ignored() { - let mut headers = HeaderMap::new(); - headers.insert( - "x-ratelimit-remaining-requests", - "not-a-number".parse().unwrap(), - ); - headers.insert("x-ratelimit-limit-tokens", "10000".parse().unwrap()); - - let info = parse_rate_limit_headers(&headers).unwrap(); - assert_eq!(info.requests_remaining, None); - assert_eq!(info.tokens_limit, Some(10000)); - } - - // --- parse_error_body --- - - #[test] - fn parse_error_body_valid_json() { - let body = r#"{"error":{"message":"rate limited","type":"rate_limit_error"}}"#; - let (msg, code, raw) = parse_error_body(body, "type"); - assert_eq!(msg, "rate limited"); - assert_eq!(code.as_deref(), Some("rate_limit_error")); - assert!(raw.is_some()); - } - - #[test] - fn parse_error_body_missing_error_field() { - let body = r#"{"status":"fail"}"#; - let (msg, code, raw) = parse_error_body(body, "type"); - assert_eq!(msg, "Unknown error"); - assert_eq!(code, None); - assert!(raw.is_some()); - } - - #[test] - fn parse_error_body_not_json() { - let body = "Internal Server Error"; - let (msg, code, raw) = parse_error_body(body, "type"); - assert_eq!(msg, "Internal Server Error"); - assert_eq!(code, None); - assert!(raw.is_none()); - } - - #[test] - fn parse_error_body_different_code_field() { - let body = r#"{"error":{"message":"bad","status":"INVALID_ARGUMENT"}}"#; - let (msg, code, _) = parse_error_body(body, "status"); - assert_eq!(msg, "bad"); - assert_eq!(code.as_deref(), Some("INVALID_ARGUMENT")); - } - - #[test] - fn parse_error_body_no_message() { - let body = r#"{"error":{"type":"server_error"}}"#; - let (msg, code, _) = parse_error_body(body, "type"); - assert_eq!(msg, "Unknown error"); - assert_eq!(code.as_deref(), Some("server_error")); - } - - // --- extract_system_prompt --- - - #[test] - fn extract_system_prompt_no_system() { - let msgs = vec![Message::user("hello")]; - let (sys, other) = extract_system_prompt(&msgs); - assert_eq!(sys, None); - assert_eq!(other.len(), 1); - } - - #[test] - fn extract_system_prompt_system_only() { - let msgs = vec![Message::system("Be helpful"), Message::user("hi")]; - let (sys, other) = extract_system_prompt(&msgs); - assert_eq!(sys.as_deref(), Some("Be helpful")); - assert_eq!(other.len(), 1); - assert_eq!(other[0].role, Role::User); - } - - #[test] - fn extract_system_prompt_multiple_system() { - let msgs = vec![ - Message::system("Rule 1"), - Message::system("Rule 2"), - Message::user("hi"), - ]; - let (sys, other) = extract_system_prompt(&msgs); - assert_eq!(sys.as_deref(), Some("Rule 1\nRule 2")); - assert_eq!(other.len(), 1); - } - - #[test] - fn extract_system_prompt_developer_role() { - let dev = Message { - role: Role::Developer, - content: vec![ContentPart::text("dev instructions")], - name: None, - tool_call_id: None, - }; - let msgs = vec![dev, Message::user("hi")]; - let (sys, other) = extract_system_prompt(&msgs); - assert_eq!(sys.as_deref(), Some("dev instructions")); - assert_eq!(other.len(), 1); - } - - #[test] - fn extract_system_prompt_ignores_whitespace_system_and_developer() { - let dev = Message { - role: Role::Developer, - content: vec![ContentPart::text(" \n\t ")], - name: None, - tool_call_id: None, - }; - let msgs = vec![Message::system(" "), dev, Message::user("hi")]; - let (sys, other) = extract_system_prompt(&msgs); - assert_eq!(sys, None); - assert_eq!(other.len(), 1); - assert_eq!(other[0].role, Role::User); - } - - #[test] - fn extract_system_prompt_empty() { - let msgs: Vec = vec![]; - let (sys, other) = extract_system_prompt(&msgs); - assert_eq!(sys, None); - assert!(other.is_empty()); - } - - // --- parse_retry_after --- - - #[test] - fn parse_retry_after_valid() { - let mut headers = HeaderMap::new(); - headers.insert("retry-after", "2.5".parse().unwrap()); - assert_eq!(parse_retry_after(&headers), Some(2.5)); - } - - #[test] - fn parse_retry_after_missing() { - let headers = HeaderMap::new(); - assert_eq!(parse_retry_after(&headers), None); - } - - #[test] - fn parse_retry_after_invalid() { - let mut headers = HeaderMap::new(); - headers.insert("retry-after", "not-a-number".parse().unwrap()); - assert_eq!(parse_retry_after(&headers), None); - } - - #[test] - fn parse_retry_after_integer() { - let mut headers = HeaderMap::new(); - headers.insert("retry-after", "5".parse().unwrap()); - assert_eq!(parse_retry_after(&headers), Some(5.0)); - } } diff --git a/lib/crates/fabro-llm/src/providers/fabro_server.rs b/lib/crates/fabro-llm/src/providers/fabro_server.rs index a737ff8f1..2029b504c 100644 --- a/lib/crates/fabro-llm/src/providers/fabro_server.rs +++ b/lib/crates/fabro-llm/src/providers/fabro_server.rs @@ -4,7 +4,7 @@ use tracing::{debug, error}; use crate::error::{Error, error_from_status_code}; use crate::provider::{ProviderAdapter, StreamEventStream}; -use crate::providers::common::LineReader; +use crate::transport::{LineReader, parse_sse_block}; use crate::types::{FinishReason, Message, Request, Response, StreamEvent, TokenCounts}; /// Provider adapter that routes LLM requests through an fabro server's @@ -166,24 +166,22 @@ impl ProviderAdapter for Adapter { loop { match reader.read_next_chunk("\n\n").await { Ok(Some(block)) => { - if let Some((event_type, data)) = parse_sse_block(&block) { - if event_type == "stream_event" { - match serde_json::from_str::(&data) { - Ok(event) => return Some((Ok(event), reader)), - Err(e) => { - return Some(( - Err(Error::stream_error( - format!("failed to parse stream event: {e}"), - e, - )), - reader, - )); - } + if let Some((Some("stream_event"), data)) = parse_sse_block(&block) { + match serde_json::from_str::(&data) { + Ok(event) => return Some((Ok(event), reader)), + Err(e) => { + return Some(( + Err(Error::stream_error( + format!("failed to parse stream event: {e}"), + e, + )), + reader, + )); } } - // Skip non-stream_event SSE events } - // Empty or unparsable block — keep reading. + // Empty, unparsable, or non-stream_event block — keep + // reading. } Ok(None) => return None, Err(e) => return Some((Err(e), reader)), @@ -200,29 +198,6 @@ fn redacted_url_for_log(url: &str) -> String { .map_or_else(|_| "".to_string(), |url| url.redacted_string()) } -/// Parse a single SSE event block into `(event_type, data)`. -/// -/// Returns `None` if the block doesn't contain both an `event:` and `data:` -/// line. -fn parse_sse_block(block: &str) -> Option<(String, String)> { - let mut event_type = None; - let mut data_lines: Vec<&str> = Vec::new(); - - for line in block.lines() { - if let Some(value) = line.strip_prefix("event:") { - event_type = Some(value.trim().to_string()); - } else if let Some(value) = line.strip_prefix("data:") { - data_lines.push(value.trim()); - } - } - - let event_type = event_type?; - if data_lines.is_empty() { - return None; - } - Some((event_type, data_lines.join("\n"))) -} - // --------------------------------------------------------------------------- // Tests // --------------------------------------------------------------------------- @@ -464,7 +439,7 @@ data: {\"type\":\"stream_start\"}\n\ fn parse_sse_block_valid() { let block = "event: stream_event\ndata: {\"type\":\"stream_start\"}"; let (event_type, data) = parse_sse_block(block).unwrap(); - assert_eq!(event_type, "stream_event"); + assert_eq!(event_type, Some("stream_event")); assert_eq!(data, "{\"type\":\"stream_start\"}"); } @@ -474,10 +449,13 @@ data: {\"type\":\"stream_start\"}\n\ assert!(parse_sse_block(block).is_none()); } + /// A block without an `event:` line parses with `event = None`; the + /// stream loop's `Some("stream_event")` match is what filters it out. #[test] fn parse_sse_block_missing_event() { let block = "data: {\"type\":\"stream_start\"}"; - assert!(parse_sse_block(block).is_none()); + let (event_type, _) = parse_sse_block(block).unwrap(); + assert_eq!(event_type, None); } #[test] diff --git a/lib/crates/fabro-llm/src/providers/gemini.rs b/lib/crates/fabro-llm/src/providers/gemini.rs index 1f345227a..2b1f9ee4b 100644 --- a/lib/crates/fabro-llm/src/providers/gemini.rs +++ b/lib/crates/fabro-llm/src/providers/gemini.rs @@ -1,19 +1,18 @@ use std::sync::Arc; -use fabro_http::HeaderMap; use fabro_model::Catalog; -use futures::stream; use crate::attachments::{self, AttachmentPolicy}; use crate::codec::gemini_generate::GeminiGenerate; -use crate::codec::{Codec, CodecCtx, CodecParams, EncodedRequest, RawEvent, StreamDecoder}; +use crate::codec::{Codec, CodecCtx, CodecParams, EncodedRequest}; use crate::error::Error; use crate::provider::{ ProviderAdapter, StreamEventStream, validate_standard_speed, validate_tool_choice, }; -use crate::providers::common::{self as common, parse_rate_limit_headers, parse_retry_after}; +use crate::providers::common::{self as common}; use crate::token_count::{InputTokenCount, InputTokenCountMethod}; -use crate::types::{AdapterTimeout, Request, Response, StreamEvent}; +use crate::transport::{self, HttpTransport, SseFraming}; +use crate::types::{AdapterTimeout, Request, Response}; const DEFAULT_BASE_URL: &str = "https://generativelanguage.googleapis.com/v1beta"; @@ -25,7 +24,7 @@ const DEFAULT_BASE_URL: &str = "https://generativelanguage.googleapis.com/v1beta /// Gemini has no route variance (single auth scheme, count-tokens always /// available, no forced streaming), so there is no route config. pub struct Adapter { - pub(crate) http: super::http_api::HttpApi, + pub(crate) http: HttpTransport, provider_name: String, catalog: Option>, } @@ -39,7 +38,7 @@ impl Adapter { #[must_use] pub fn new_optional_auth(api_key: Option) -> Self { Self { - http: super::http_api::HttpApi::new_optional(api_key, DEFAULT_BASE_URL), + http: HttpTransport::new_optional(api_key, DEFAULT_BASE_URL), provider_name: "gemini".to_string(), catalog: None, } @@ -124,47 +123,6 @@ impl Adapter { } req.json(&encoded.body) } - - /// Send an HTTP request and read the Gemini response body, mapping non-2xx - /// responses through the codec's error decode (gRPC status aware). - async fn send_for_body( - &self, - request: fabro_http::RequestBuilder, - codec: &GeminiGenerate, - ctx: &CodecCtx<'_>, - ) -> Result<(String, HeaderMap), Error> { - let http_resp = request.send().await.map_err(|e| { - if e.is_timeout() { - Error::request_timeout(format!("gemini: {e}"), e) - } else { - Error::network(e.to_string(), e) - } - })?; - - let status = http_resp.status(); - let retry_after = parse_retry_after(http_resp.headers()); - let headers = http_resp.headers().clone(); - let body = http_resp - .text() - .await - .map_err(|e| Error::network(e.to_string(), e))?; - - if !status.is_success() { - return Err(codec.decode_error(status.as_u16(), &body, ctx, retry_after)); - } - - Ok((body, headers)) - } -} - -/// State driving the streaming byte loop: the codec's decoder plus the line -/// reader, with a buffer that flattens batched events into individual items. -struct StreamLoop { - decoder: Box, - line_reader: super::common::LineReader, - pending: std::collections::VecDeque, - done: bool, - finished_emitted: bool, } #[async_trait::async_trait] @@ -201,7 +159,8 @@ impl ProviderAdapter for Adapter { if let Some(t) = self.http.request_timeout { req = req.timeout(t); } - let (body, _headers) = self.send_for_body(req, &codec, &ctx).await?; + let (body, _headers) = + transport::send_for_body(req, "input_token_count", &codec, &ctx).await?; let input_tokens = codec.decode_count_tokens(&body)?; Ok(Some(InputTokenCount { @@ -227,9 +186,7 @@ impl ProviderAdapter for Adapter { if let Some(t) = self.http.request_timeout { req = req.timeout(t); } - let (body, headers) = self.send_for_body(req, &codec, &ctx).await?; - let rate_limit = parse_rate_limit_headers(&headers); - codec.decode_response(&body, &ctx, rate_limit) + transport::complete_via_http(req, &codec, &ctx).await } async fn stream(&self, request: &Request) -> Result { @@ -239,81 +196,17 @@ impl ProviderAdapter for Adapter { let codec = GeminiGenerate; let deployment_id = common::api_model_id(self.catalog.as_deref(), &resolved.model); let params = CodecParams::default(); - let stream_read_timeout = self.http.stream_read_timeout; let ctx = self.codec_ctx(&resolved, &deployment_id, ¶ms); let encoded = codec.encode(&ctx, true)?; - let http_resp = self - .build_http_request(&encoded) - .send() - .await - .map_err(|e| Error::network(e.to_string(), e))?; - - let status = http_resp.status(); - if !status.is_success() { - let retry_after = parse_retry_after(http_resp.headers()); - let body = http_resp - .text() - .await - .map_err(|e| Error::network(e.to_string(), e))?; - return Err(codec.decode_error(status.as_u16(), &body, &ctx, retry_after)); - } - let rate_limit = parse_rate_limit_headers(http_resp.headers()); - let decoder = codec.stream_decoder(&ctx, rate_limit); - - let out = stream::unfold( - StreamLoop { - decoder, - line_reader: super::common::LineReader::new(http_resp, stream_read_timeout), - pending: std::collections::VecDeque::new(), - done: false, - finished_emitted: false, - }, - |mut state| async move { - loop { - if let Some(event) = state.pending.pop_front() { - return Some((Ok(event), state)); - } - - if state.done { - if state.finished_emitted { - return None; - } - state.finished_emitted = true; - let events = state.decoder.finish(); - if events.is_empty() { - return None; - } - state.pending.extend(events); - continue; - } - - match state.line_reader.read_next_chunk("\n").await { - Ok(Some(line)) => { - // Gemini uses data-only SSE: lines starting with - // "data:" carry the payload; everything else - // (empty lines, comments, event: lines) is - // ignored. - let line = line.trim_end_matches('\r'); - let Some(data) = line.strip_prefix("data:").map(str::trim) else { - continue; - }; - if data.is_empty() { - continue; - } - match state.decoder.on_event(RawEvent { event: None, data }) { - Ok(events) => state.pending.extend(events), - Err(e) => return Some((Err(e), state)), - } - } - Ok(None) => state.done = true, - Err(e) => return Some((Err(e), state)), - } - } - }, - ); - - Ok(Box::pin(out)) + transport::stream_via_http( + self.build_http_request(&encoded), + &codec, + &ctx, + SseFraming::DataLines, + self.http.stream_read_timeout, + ) + .await } } diff --git a/lib/crates/fabro-llm/src/providers/http_api.rs b/lib/crates/fabro-llm/src/providers/http_api.rs deleted file mode 100644 index 6939acbc6..000000000 --- a/lib/crates/fabro-llm/src/providers/http_api.rs +++ /dev/null @@ -1,60 +0,0 @@ -use std::collections::HashMap; -use std::time::Duration; - -use crate::types::AdapterTimeout; - -/// Shared HTTP infrastructure for provider adapters. -/// -/// Holds the API key, base URL, reqwest client, default headers, and timeout -/// configuration that every provider needs. Provider-specific fields live on -/// the adapter struct itself. -pub struct HttpApi { - pub(crate) api_key: Option, - pub(crate) base_url: String, - pub(crate) default_headers: HashMap, - pub(crate) client: fabro_http::HttpClient, - pub(crate) request_timeout: Option, - pub(crate) stream_read_timeout: Option, -} - -impl HttpApi { - fn build_client(timeout: AdapterTimeout) -> fabro_http::HttpClient { - fabro_http::HttpClientBuilder::new() - .connect_timeout(Duration::from_secs_f64(timeout.connect)) - .build() - .expect("LLM HTTP client should build") - } - - #[must_use] - pub fn new(api_key: impl Into, base_url: impl Into) -> Self { - Self::new_optional(Some(api_key.into()), base_url) - } - - #[must_use] - pub fn new_optional(api_key: Option, base_url: impl Into) -> Self { - let timeout = AdapterTimeout::default(); - let client = Self::build_client(timeout); - Self { - api_key, - base_url: base_url.into(), - default_headers: HashMap::new(), - client, - request_timeout: timeout.request.map(Duration::from_secs_f64), - stream_read_timeout: timeout.stream_read.map(Duration::from_secs_f64), - } - } - - #[must_use] - pub fn with_timeout(mut self, timeout: AdapterTimeout) -> Self { - self.client = Self::build_client(timeout); - self.request_timeout = timeout.request.map(Duration::from_secs_f64); - self.stream_read_timeout = timeout.stream_read.map(Duration::from_secs_f64); - self - } - - #[must_use] - pub fn with_default_headers(mut self, headers: HashMap) -> Self { - self.default_headers = headers; - self - } -} diff --git a/lib/crates/fabro-llm/src/providers/mod.rs b/lib/crates/fabro-llm/src/providers/mod.rs index 13d105b2d..ecc2ee96c 100644 --- a/lib/crates/fabro-llm/src/providers/mod.rs +++ b/lib/crates/fabro-llm/src/providers/mod.rs @@ -2,7 +2,6 @@ pub mod anthropic; pub mod common; pub mod fabro_server; pub mod gemini; -pub mod http_api; pub mod openai; pub mod openai_compatible; diff --git a/lib/crates/fabro-llm/src/providers/openai.rs b/lib/crates/fabro-llm/src/providers/openai.rs index 24e02a6ac..ad25d01a9 100644 --- a/lib/crates/fabro-llm/src/providers/openai.rs +++ b/lib/crates/fabro-llm/src/providers/openai.rs @@ -1,20 +1,17 @@ use std::sync::Arc; use fabro_model::Catalog; -use futures::stream; use crate::attachments::{self, AttachmentPolicy}; use crate::codec::openai_responses::OpenAiResponses; -use crate::codec::{Codec, CodecCtx, CodecParams, EncodedRequest, RawEvent, StreamDecoder}; +use crate::codec::{Codec, CodecCtx, CodecParams, EncodedRequest}; use crate::error::Error; use crate::provider::{ ProviderAdapter, StreamEventStream, validate_standard_speed, validate_tool_choice, }; -use crate::providers::common::{ - self as common, parse_rate_limit_headers, parse_retry_after, send_and_read_response, - send_and_read_response_with_operation, -}; +use crate::providers::common::{self as common}; use crate::token_count::{InputTokenCount, InputTokenCountMethod}; +use crate::transport::{self, HttpTransport, SseFraming}; use crate::types::{AdapterTimeout, Request, Response, StreamEvent}; const DEFAULT_BASE_URL: &str = "https://api.openai.com/v1"; @@ -30,7 +27,7 @@ const DEFAULT_BASE_URL: &str = "https://api.openai.com/v1"; /// Completions) to properly surface reasoning tokens, built-in tools, and /// server-side state. pub struct Adapter { - pub(crate) http: super::http_api::HttpApi, + pub(crate) http: HttpTransport, org_id: Option, project_id: Option, provider_name: String, @@ -48,7 +45,7 @@ impl Adapter { #[must_use] pub fn new_optional_auth(api_key: Option) -> Self { Self { - http: super::http_api::HttpApi::new_optional(api_key, DEFAULT_BASE_URL), + http: HttpTransport::new_optional(api_key, DEFAULT_BASE_URL), org_id: None, project_id: None, provider_name: "openai".to_string(), @@ -195,49 +192,6 @@ impl Adapter { } } -/// State driving the streaming byte loop: the codec's decoder plus the line -/// reader, with a buffer that flattens batched events into individual items. -struct StreamLoop { - decoder: Box, - line_reader: super::common::LineReader, - pending: std::collections::VecDeque, - done: bool, - finished_emitted: bool, -} - -/// Parse a single SSE message block into an (`event_type`, `data`) pair. -/// -/// Each SSE message consists of one or more lines (`event:` and `data:` -/// prefixed). Returns `None` if the block has no `data:` lines. -fn parse_sse_message(message_block: &str) -> Option<(Option, String)> { - let mut current_event: Option = None; - let mut current_data = String::new(); - - for line in message_block.lines() { - if let Some(stripped) = line.strip_prefix("event: ") { - current_event = Some(stripped.to_string()); - } else if let Some(stripped) = line.strip_prefix("event:") { - current_event = Some(stripped.trim().to_string()); - } else if let Some(stripped) = line.strip_prefix("data: ") { - if !current_data.is_empty() { - current_data.push('\n'); - } - current_data.push_str(stripped); - } else if let Some(stripped) = line.strip_prefix("data:") { - if !current_data.is_empty() { - current_data.push('\n'); - } - current_data.push_str(stripped.trim()); - } - } - - if current_data.is_empty() { - None - } else { - Some((current_event, current_data)) - } -} - #[async_trait::async_trait] impl ProviderAdapter for Adapter { fn name(&self) -> &str { @@ -272,13 +226,8 @@ impl ProviderAdapter for Adapter { if let Some(t) = self.http.request_timeout { req = req.timeout(t); } - let (body, _headers) = send_and_read_response_with_operation( - req, - &self.provider_name, - "type", - "input_token_count", - ) - .await?; + let (body, _headers) = + transport::send_for_body(req, "input_token_count", &codec, &ctx).await?; let input_tokens = codec.decode_count_tokens(&body)?; Ok(Some(InputTokenCount { @@ -310,9 +259,7 @@ impl ProviderAdapter for Adapter { if let Some(t) = self.http.request_timeout { req = req.timeout(t); } - let (body, headers) = send_and_read_response(req, &self.provider_name, "type").await?; - let rate_limit = parse_rate_limit_headers(&headers); - codec.decode_response(&body, &ctx, rate_limit) + transport::complete_via_http(req, &codec, &ctx).await } async fn stream(&self, request: &Request) -> Result { @@ -325,74 +272,14 @@ impl ProviderAdapter for Adapter { let ctx = self.codec_ctx(&resolved, &deployment_id, ¶ms); let encoded = codec.encode(&ctx, true)?; - let http_resp = self - .build_http_request(&encoded) - .send() - .await - .map_err(|e| Error::network(e.to_string(), e))?; - - let status = http_resp.status(); - if !status.is_success() { - let retry_after = parse_retry_after(http_resp.headers()); - let body = http_resp - .text() - .await - .map_err(|e| Error::network(e.to_string(), e))?; - return Err(codec.decode_error(status.as_u16(), &body, &ctx, retry_after)); - } - - let rate_limit = parse_rate_limit_headers(http_resp.headers()); - let stream_read_timeout = self.http.stream_read_timeout; - let decoder = codec.stream_decoder(&ctx, rate_limit); - - let out = stream::unfold( - StreamLoop { - decoder, - line_reader: super::common::LineReader::new(http_resp, stream_read_timeout), - pending: std::collections::VecDeque::new(), - done: false, - finished_emitted: false, - }, - |mut state| async move { - loop { - if let Some(event) = state.pending.pop_front() { - return Some((Ok(event), state)); - } - - if state.done { - if state.finished_emitted { - return None; - } - state.finished_emitted = true; - let events = state.decoder.finish(); - if events.is_empty() { - return None; - } - state.pending.extend(events); - continue; - } - - match state.line_reader.read_next_chunk("\n\n").await { - Ok(Some(block)) => { - let Some((event_type, data)) = parse_sse_message(&block) else { - continue; - }; - match state.decoder.on_event(RawEvent { - event: event_type.as_deref(), - data: &data, - }) { - Ok(events) => state.pending.extend(events), - Err(e) => return Some((Err(e), state)), - } - } - Ok(None) => state.done = true, - Err(e) => return Some((Err(e), state)), - } - } - }, - ); - - Ok(Box::pin(out)) + transport::stream_via_http( + self.build_http_request(&encoded), + &codec, + &ctx, + SseFraming::EventBlocks, + self.http.stream_read_timeout, + ) + .await } } diff --git a/lib/crates/fabro-llm/src/providers/openai_compatible.rs b/lib/crates/fabro-llm/src/providers/openai_compatible.rs index 801ad4669..1e17dd825 100644 --- a/lib/crates/fabro-llm/src/providers/openai_compatible.rs +++ b/lib/crates/fabro-llm/src/providers/openai_compatible.rs @@ -1,19 +1,16 @@ -use std::collections::VecDeque; use std::sync::Arc; use fabro_model::Catalog; -use futures::stream; use crate::codec::openai_compatible::OpenAiCompatible; -use crate::codec::{Codec, CodecCtx, CodecParams, RawEvent, StreamDecoder}; +use crate::codec::{Codec, CodecCtx, CodecParams}; use crate::error::Error; use crate::provider::{ ProviderAdapter, StreamEventStream, validate_standard_speed, validate_tool_choice, }; -use crate::providers::common::{ - api_model_id, parse_rate_limit_headers, parse_retry_after, send_and_read_response, -}; -use crate::types::{AdapterTimeout, Request, Response, StreamEvent}; +use crate::providers::common::api_model_id; +use crate::transport::{self, HttpTransport, SseFraming}; +use crate::types::{AdapterTimeout, Request, Response}; /// `OpenAI`-compatible Chat Completions adapter (Section 7.10). /// @@ -27,7 +24,7 @@ use crate::types::{AdapterTimeout, Request, Response, StreamEvent}; /// auth, base URL, and the streaming byte loop, and delegates all wire /// translation to the codec. pub struct Adapter { - pub(crate) http: super::http_api::HttpApi, + pub(crate) http: HttpTransport, provider_name: String, catalog: Option>, } @@ -41,7 +38,7 @@ impl Adapter { #[must_use] pub fn new_optional_auth(api_key: Option, base_url: impl Into) -> Self { Self { - http: super::http_api::HttpApi::new_optional(api_key, base_url), + http: HttpTransport::new_optional(api_key, base_url), provider_name: "openai-compatible".to_string(), catalog: None, } @@ -130,20 +127,6 @@ impl Adapter { } } -/// State driving the streaming byte loop: the codec's decoder plus the line -/// reader, with a small buffer that flattens batched events into individual -/// stream items. -struct StreamLoop { - decoder: Box, - line_reader: super::common::LineReader, - /// Events decoded but not yet yielded. - pending: VecDeque, - /// Byte stream exhausted. - done: bool, - /// `finish()` already drained. - finished_emitted: bool, -} - #[async_trait::async_trait] impl ProviderAdapter for Adapter { fn name(&self) -> &str { @@ -171,9 +154,7 @@ impl ProviderAdapter for Adapter { req = req.timeout(t); } - let (body, headers) = send_and_read_response(req, &self.provider_name, "type").await?; - let rate_limit = parse_rate_limit_headers(&headers); - codec.decode_response(&body, &ctx, rate_limit) + transport::complete_via_http(req, &codec, &ctx).await } async fn stream(&self, request: &Request) -> Result { @@ -185,73 +166,13 @@ impl ProviderAdapter for Adapter { let ctx = self.codec_ctx(request, &deployment_id, ¶ms); let req = self.encoded_request(&codec, &ctx, true)?; - let http_resp = req - .send() - .await - .map_err(|e| Error::network(e.to_string(), e))?; - - let status = http_resp.status(); - if !status.is_success() { - let retry_after = parse_retry_after(http_resp.headers()); - let body = http_resp - .text() - .await - .map_err(|e| Error::network(e.to_string(), e))?; - return Err(codec.decode_error(status.as_u16(), &body, &ctx, retry_after)); - } - - let rate_limit = parse_rate_limit_headers(http_resp.headers()); - let stream_read_timeout = self.http.stream_read_timeout; - let decoder = codec.stream_decoder(&ctx, rate_limit); - let line_reader = super::common::LineReader::new(http_resp, stream_read_timeout); - - let out = stream::unfold( - StreamLoop { - decoder, - line_reader, - pending: VecDeque::new(), - done: false, - finished_emitted: false, - }, - |mut state| async move { - loop { - if let Some(event) = state.pending.pop_front() { - return Some((Ok(event), state)); - } - - if state.done { - if state.finished_emitted { - return None; - } - state.finished_emitted = true; - state.pending = state.decoder.finish().into(); - if state.pending.is_empty() { - return None; - } - continue; - } - - match state.line_reader.read_next_chunk("\n").await { - Ok(Some(line)) => { - let line = line.trim(); - if line.is_empty() || line.starts_with(':') { - continue; - } - let Some(data) = line.strip_prefix("data:").map(str::trim) else { - continue; - }; - match state.decoder.on_event(RawEvent { event: None, data }) { - Ok(events) => state.pending = events.into(), - Err(e) => return Some((Err(e), state)), - } - } - Ok(None) => state.done = true, - Err(e) => return Some((Err(e), state)), - } - } - }, - ); - - Ok(Box::pin(out)) + transport::stream_via_http( + req, + &codec, + &ctx, + SseFraming::DataLines, + self.http.stream_read_timeout, + ) + .await } } diff --git a/lib/crates/fabro-llm/src/transport.rs b/lib/crates/fabro-llm/src/transport.rs new file mode 100644 index 000000000..1ed1ee558 --- /dev/null +++ b/lib/crates/fabro-llm/src/transport.rs @@ -0,0 +1,596 @@ +//! The HTTP transport shared by every provider adapter: how request bytes +//! travel, not what they say. +//! +//! A transport owns the HTTP client, timeouts, the streaming byte loop, and +//! SSE framing. It knows nothing about wire dialects — bodies, endpoints, and +//! error shapes arrive from (and return to) a [`Codec`]. Adapters shrink to +//! auth + route config composed over these helpers. +//! +//! The split mirrors `codec/mod.rs`: a codec knows *what the bytes say*; this +//! module knows *how they travel*. + +use std::borrow::Cow; +use std::collections::{HashMap, VecDeque}; +use std::time::Duration; + +use fabro_http::HeaderMap; +use futures::stream; +use tokio::time; +use tracing::warn; + +use crate::codec::{Codec, CodecCtx, RawEvent, StreamDecoder}; +use crate::error::Error; +use crate::provider::StreamEventStream; +use crate::types::{AdapterTimeout, RateLimitInfo, Response, StreamEvent}; + +// --- HTTP client + configuration +// ---------------------------------------------- + +/// Shared HTTP infrastructure for provider adapters. +/// +/// Holds the API key, base URL, reqwest client, default headers, and timeout +/// configuration that every provider needs. Provider-specific fields live on +/// the adapter struct itself. +pub(crate) struct HttpTransport { + pub(crate) api_key: Option, + pub(crate) base_url: String, + pub(crate) default_headers: HashMap, + pub(crate) client: fabro_http::HttpClient, + pub(crate) request_timeout: Option, + pub(crate) stream_read_timeout: Option, +} + +impl HttpTransport { + fn build_client(timeout: AdapterTimeout) -> fabro_http::HttpClient { + fabro_http::HttpClientBuilder::new() + .connect_timeout(Duration::from_secs_f64(timeout.connect)) + .build() + .expect("LLM HTTP client should build") + } + + #[must_use] + pub(crate) fn new_optional(api_key: Option, base_url: impl Into) -> Self { + let timeout = AdapterTimeout::default(); + let client = Self::build_client(timeout); + Self { + api_key, + base_url: base_url.into(), + default_headers: HashMap::new(), + client, + request_timeout: timeout.request.map(Duration::from_secs_f64), + stream_read_timeout: timeout.stream_read.map(Duration::from_secs_f64), + } + } + + #[must_use] + pub(crate) fn with_timeout(mut self, timeout: AdapterTimeout) -> Self { + self.client = Self::build_client(timeout); + self.request_timeout = timeout.request.map(Duration::from_secs_f64); + self.stream_read_timeout = timeout.stream_read.map(Duration::from_secs_f64); + self + } + + #[must_use] + pub(crate) fn with_default_headers(mut self, headers: HashMap) -> Self { + self.default_headers = headers; + self + } +} + +// --- Response header parsing +// --------------------------------------------------- + +/// Extract the `Retry-After` header value from an HTTP response as seconds. +#[must_use] +pub fn parse_retry_after(headers: &HeaderMap) -> Option { + headers + .get("retry-after") + .and_then(|v| v.to_str().ok()) + .and_then(|s| s.parse::().ok()) +} + +/// Parse `x-ratelimit-*` headers into a `RateLimitInfo`. +/// +/// Returns `None` if no rate limit headers are present. +#[must_use] +pub fn parse_rate_limit_headers(headers: &HeaderMap) -> Option { + fn header_i64(headers: &HeaderMap, name: &str) -> Option { + headers + .get(name) + .and_then(|v| v.to_str().ok()) + .and_then(|s| s.parse::().ok()) + } + + fn header_str(headers: &HeaderMap, name: &str) -> Option { + headers + .get(name) + .and_then(|v| v.to_str().ok()) + .map(String::from) + } + + let requests_remaining = header_i64(headers, "x-ratelimit-remaining-requests"); + let requests_limit = header_i64(headers, "x-ratelimit-limit-requests"); + let tokens_remaining = header_i64(headers, "x-ratelimit-remaining-tokens"); + let tokens_limit = header_i64(headers, "x-ratelimit-limit-tokens"); + let reset_at = header_str(headers, "x-ratelimit-reset-requests") + .or_else(|| header_str(headers, "x-ratelimit-reset-tokens")); + + if requests_remaining.is_none() + && requests_limit.is_none() + && tokens_remaining.is_none() + && tokens_limit.is_none() + && reset_at.is_none() + { + return None; + } + + Some(RateLimitInfo { + requests_remaining, + requests_limit, + tokens_remaining, + tokens_limit, + reset_at, + }) +} + +// --- Blocking requests +// ----------------------------------------------------------- + +/// Send a blocking request and decode the response through the codec: +/// `send_for_body` + rate-limit headers + [`Codec::decode_response`]. +pub(crate) async fn complete_via_http( + request: fabro_http::RequestBuilder, + codec: &dyn Codec, + ctx: &CodecCtx<'_>, +) -> Result { + let (body, headers) = send_for_body(request, "provider_request", codec, ctx).await?; + let rate_limit = parse_rate_limit_headers(&headers); + codec.decode_response(&body, ctx, rate_limit) +} + +/// Send an HTTP request and read the response body plus headers, mapping +/// non-2xx responses through [`Codec::decode_error`]. `operation` tags the +/// warning logs (`provider_request`, `input_token_count`). +pub(crate) async fn send_for_body( + request: fabro_http::RequestBuilder, + operation: &str, + codec: &dyn Codec, + ctx: &CodecCtx<'_>, +) -> Result<(String, HeaderMap), Error> { + let provider = ctx.provider_name; + let http_resp = request.send().await.map_err(|e| { + if e.is_timeout() { + warn!(provider = %provider, operation = %operation, error = %e, "Provider request timed out"); + Error::request_timeout(format!("{provider}: {e}"), e) + } else { + warn!(provider = %provider, operation = %operation, error = %e, "Provider network error"); + Error::network(e.to_string(), e) + } + })?; + + let status = http_resp.status(); + let retry_after = parse_retry_after(http_resp.headers()); + let headers = http_resp.headers().clone(); + let body = http_resp + .text() + .await + .map_err(|e| Error::network(e.to_string(), e))?; + + if !status.is_success() { + warn!(provider = %provider, operation = %operation, status = status.as_u16(), "Provider returned error"); + return Err(codec.decode_error(status.as_u16(), &body, ctx, retry_after)); + } + + Ok((body, headers)) +} + +// --- Streaming +// ------------------------------------------------------------------- + +/// How a route frames its SSE byte stream into decoder events. +#[derive(Clone, Copy, Debug)] +pub(crate) enum SseFraming { + /// `\n\n`-delimited blocks carrying `event:` + `data:` lines (anthropic, + /// openai responses). + EventBlocks, + /// Newline-delimited `data:` lines; comments, blank lines, and non-data + /// fields are skipped (openai_compatible, gemini). + DataLines, +} + +impl SseFraming { + fn delimiter(self) -> &'static str { + match self { + Self::EventBlocks => "\n\n", + Self::DataLines => "\n", + } + } +} + +/// Send a streaming request and decode its SSE byte stream through the +/// codec's [`StreamDecoder`]. A non-2xx response is mapped through +/// [`Codec::decode_error`] before any bytes flow. +pub(crate) async fn stream_via_http( + request: fabro_http::RequestBuilder, + codec: &dyn Codec, + ctx: &CodecCtx<'_>, + framing: SseFraming, + stream_read_timeout: Option, +) -> Result { + let http_resp = request + .send() + .await + .map_err(|e| Error::network(e.to_string(), e))?; + + let status = http_resp.status(); + if !status.is_success() { + let retry_after = parse_retry_after(http_resp.headers()); + let body = http_resp + .text() + .await + .map_err(|e| Error::network(e.to_string(), e))?; + return Err(codec.decode_error(status.as_u16(), &body, ctx, retry_after)); + } + + let rate_limit = parse_rate_limit_headers(http_resp.headers()); + let decoder = codec.stream_decoder(ctx, rate_limit); + Ok(decode_sse_stream( + http_resp, + decoder, + framing, + stream_read_timeout, + )) +} + +/// State driving the streaming byte loop: the codec's decoder plus the line +/// reader, with a buffer that flattens batched events into individual items. +struct StreamLoop { + decoder: Box, + line_reader: LineReader, + /// Events decoded but not yet yielded. + pending: VecDeque, + /// Byte stream exhausted. + done: bool, + /// `finish()` already drained. + finished_emitted: bool, +} + +/// Drive `decoder` over the SSE byte stream of `response`: frame each chunk, +/// feed it to the decoder, flatten batched events, and drain +/// [`StreamDecoder::finish`] at byte-stream end. +fn decode_sse_stream( + response: fabro_http::Response, + decoder: Box, + framing: SseFraming, + stream_read_timeout: Option, +) -> StreamEventStream { + let out = stream::unfold( + StreamLoop { + decoder, + line_reader: LineReader::new(response, stream_read_timeout), + pending: VecDeque::new(), + done: false, + finished_emitted: false, + }, + move |mut state| async move { + loop { + if let Some(event) = state.pending.pop_front() { + return Some((Ok(event), state)); + } + + if state.done { + if state.finished_emitted { + return None; + } + state.finished_emitted = true; + state.pending.extend(state.decoder.finish()); + if state.pending.is_empty() { + return None; + } + continue; + } + + match state.line_reader.read_next_chunk(framing.delimiter()).await { + Ok(Some(chunk)) => { + let Some((event, data)) = frame_sse_chunk(framing, &chunk) else { + continue; + }; + match state.decoder.on_event(RawEvent { event, data: &data }) { + Ok(events) => state.pending.extend(events), + Err(e) => return Some((Err(e), state)), + } + } + Ok(None) => state.done = true, + Err(e) => return Some((Err(e), state)), + } + } + }, + ); + Box::pin(out) +} + +/// Frame one delimiter-separated chunk into an SSE `(event, data)` pair. +/// Returns `None` for chunks with no payload to decode: heartbeat comments, +/// blank lines, non-data fields, and empty `data:` payloads. +fn frame_sse_chunk(framing: SseFraming, chunk: &str) -> Option<(Option<&str>, Cow<'_, str>)> { + match framing { + SseFraming::EventBlocks => parse_sse_block(chunk), + SseFraming::DataLines => { + let data = chunk.trim().strip_prefix("data:")?.trim(); + if data.is_empty() { + return None; + } + Some((None, Cow::Borrowed(data))) + } + } +} + +/// Parse an SSE event block (lines within a `\n\n`-delimited chunk) into +/// `(event_type, data)`. Multi-line `data:` payloads are joined with `\n`; +/// the common single-line case borrows from the block. Returns `None` for +/// blocks with no non-empty payload (e.g. heartbeat comments). +pub(crate) fn parse_sse_block(block: &str) -> Option<(Option<&str>, Cow<'_, str>)> { + let mut event: Option<&str> = None; + let mut data: Option> = None; + + for line in block.lines() { + if let Some(rest) = line.strip_prefix("event:") { + event = Some(rest.trim()); + } else if let Some(rest) = line.strip_prefix("data:") { + let rest = rest.trim(); + data = Some(match data { + None => Cow::Borrowed(rest), + Some(prev) => { + let mut joined = prev.into_owned(); + joined.push('\n'); + joined.push_str(rest); + Cow::Owned(joined) + } + }); + } + } + + let data = data?; + if data.is_empty() { + return None; + } + Some((event, data)) +} + +// --- Byte-stream reading ----------------------------------------------------- + +/// Shared line reader for SSE streams. +/// +/// Buffers bytes from a `fabro_http::Response` and splits them by a +/// configurable delimiter (e.g. `"\n"` for Gemini/OpenAI-compatible, `"\n\n"` +/// for Anthropic/OpenAI SSE event blocks). +pub struct LineReader { + response: fabro_http::Response, + buffer: String, + stream_read_timeout: Option, +} + +impl LineReader { + pub fn new(response: fabro_http::Response, stream_read_timeout: Option) -> Self { + Self { + response, + buffer: String::new(), + stream_read_timeout, + } + } + + /// Read the next complete segment delimited by `delimiter`. + /// + /// Returns `Ok(Some(segment))` for each complete segment, `Ok(None)` when + /// the stream is exhausted, or `Err` on I/O or timeout errors. When the + /// stream ends with data remaining in the buffer, the leftover is returned + /// as a final segment. + pub async fn read_next_chunk(&mut self, delimiter: &str) -> Result, Error> { + loop { + if let Some(pos) = self.buffer.find(delimiter) { + let segment = self.buffer[..pos].to_string(); + self.buffer = self.buffer[pos + delimiter.len()..].to_string(); + return Ok(Some(segment)); + } + + let chunk_result = match self.stream_read_timeout { + Some(timeout) => time::timeout(timeout, self.response.chunk()).await, + None => Ok(self.response.chunk().await), + }; + match chunk_result { + Ok(Ok(Some(bytes))) => { + let text = String::from_utf8_lossy(&bytes); + self.buffer.push_str(&text); + } + Ok(Ok(None)) => { + if self.buffer.is_empty() { + return Ok(None); + } + let remaining = std::mem::take(&mut self.buffer); + return Ok(Some(remaining)); + } + Ok(Err(e)) => { + return Err(Error::stream_error(e.to_string(), e)); + } + Err(_) => { + warn!("Stream read timed out waiting for next event"); + return Err(Error::Stream { + message: "stream read timed out waiting for next event".to_string(), + source: None, + }); + } + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parse_rate_limit_headers_all_present() { + let mut headers = HeaderMap::new(); + headers.insert("x-ratelimit-remaining-requests", "99".parse().unwrap()); + headers.insert("x-ratelimit-limit-requests", "100".parse().unwrap()); + headers.insert("x-ratelimit-remaining-tokens", "9000".parse().unwrap()); + headers.insert("x-ratelimit-limit-tokens", "10000".parse().unwrap()); + headers.insert( + "x-ratelimit-reset-requests", + "2024-01-01T00:00:00Z".parse().unwrap(), + ); + + let info = parse_rate_limit_headers(&headers).unwrap(); + assert_eq!(info.requests_remaining, Some(99)); + assert_eq!(info.requests_limit, Some(100)); + assert_eq!(info.tokens_remaining, Some(9000)); + assert_eq!(info.tokens_limit, Some(10000)); + assert_eq!(info.reset_at.as_deref(), Some("2024-01-01T00:00:00Z")); + } + + #[test] + fn parse_rate_limit_headers_none_present() { + let headers = HeaderMap::new(); + assert!(parse_rate_limit_headers(&headers).is_none()); + } + + #[test] + fn parse_rate_limit_headers_partial() { + let mut headers = HeaderMap::new(); + headers.insert("x-ratelimit-remaining-requests", "50".parse().unwrap()); + + let info = parse_rate_limit_headers(&headers).unwrap(); + assert_eq!(info.requests_remaining, Some(50)); + assert_eq!(info.requests_limit, None); + assert_eq!(info.tokens_remaining, None); + assert_eq!(info.tokens_limit, None); + assert_eq!(info.reset_at, None); + } + + #[test] + fn parse_rate_limit_headers_reset_tokens_fallback() { + let mut headers = HeaderMap::new(); + headers.insert("x-ratelimit-limit-tokens", "5000".parse().unwrap()); + headers.insert( + "x-ratelimit-reset-tokens", + "2024-06-01T12:00:00Z".parse().unwrap(), + ); + + let info = parse_rate_limit_headers(&headers).unwrap(); + assert_eq!(info.tokens_limit, Some(5000)); + assert_eq!(info.reset_at.as_deref(), Some("2024-06-01T12:00:00Z")); + } + + #[test] + fn parse_rate_limit_headers_invalid_values_ignored() { + let mut headers = HeaderMap::new(); + headers.insert( + "x-ratelimit-remaining-requests", + "not-a-number".parse().unwrap(), + ); + headers.insert("x-ratelimit-limit-tokens", "10000".parse().unwrap()); + + let info = parse_rate_limit_headers(&headers).unwrap(); + assert_eq!(info.requests_remaining, None); + assert_eq!(info.tokens_limit, Some(10000)); + } + + // --- parse_retry_after --- + + #[test] + fn parse_retry_after_valid() { + let mut headers = HeaderMap::new(); + headers.insert("retry-after", "2.5".parse().unwrap()); + assert_eq!(parse_retry_after(&headers), Some(2.5)); + } + + #[test] + fn parse_retry_after_missing() { + let headers = HeaderMap::new(); + assert_eq!(parse_retry_after(&headers), None); + } + + #[test] + fn parse_retry_after_invalid() { + let mut headers = HeaderMap::new(); + headers.insert("retry-after", "not-a-number".parse().unwrap()); + assert_eq!(parse_retry_after(&headers), None); + } + + #[test] + fn parse_retry_after_integer() { + let mut headers = HeaderMap::new(); + headers.insert("retry-after", "5".parse().unwrap()); + assert_eq!(parse_retry_after(&headers), Some(5.0)); + } + + // --- frame_sse_chunk: event blocks --- + + #[test] + fn parse_sse_block_event_and_data() { + let block = "event: message_start\ndata: {\"a\":1}"; + let (event, data) = parse_sse_block(block).unwrap(); + assert_eq!(event, Some("message_start")); + assert_eq!(data, "{\"a\":1}"); + } + + #[test] + fn parse_sse_block_data_without_event() { + let block = "data: {\"a\":1}"; + let (event, data) = parse_sse_block(block).unwrap(); + assert_eq!(event, None); + assert_eq!(data, "{\"a\":1}"); + } + + #[test] + fn parse_sse_block_joins_multiple_data_lines() { + let block = "event: e\ndata: line1\ndata: line2"; + let (event, data) = parse_sse_block(block).unwrap(); + assert_eq!(event, Some("e")); + assert_eq!(data, "line1\nline2"); + } + + #[test] + fn parse_sse_block_skips_comment_only_block() { + assert!(parse_sse_block(": heartbeat").is_none()); + assert!(parse_sse_block("event: ping").is_none()); + assert!(parse_sse_block("").is_none()); + } + + #[test] + fn parse_sse_block_skips_empty_data_payload() { + assert!(parse_sse_block("data:").is_none()); + assert!(parse_sse_block("event: e\ndata: ").is_none()); + } + + #[test] + fn parse_sse_block_trims_crlf() { + let block = "event: e\r\ndata: {\"a\":1}\r"; + let (event, data) = parse_sse_block(block).unwrap(); + assert_eq!(event, Some("e")); + assert_eq!(data, "{\"a\":1}"); + } + + // --- frame_sse_chunk: data lines --- + + #[test] + fn data_lines_strips_prefix_and_trims() { + let (event, data) = frame_sse_chunk(SseFraming::DataLines, "data: {\"a\":1}\r").unwrap(); + assert_eq!(event, None); + assert_eq!(data, "{\"a\":1}"); + } + + #[test] + fn data_lines_passes_done_sentinel() { + let (_, data) = frame_sse_chunk(SseFraming::DataLines, "data: [DONE]").unwrap(); + assert_eq!(data, "[DONE]"); + } + + #[test] + fn data_lines_skips_comments_blanks_and_other_fields() { + assert!(frame_sse_chunk(SseFraming::DataLines, ": keep-alive").is_none()); + assert!(frame_sse_chunk(SseFraming::DataLines, "").is_none()); + assert!(frame_sse_chunk(SseFraming::DataLines, "event: x").is_none()); + assert!(frame_sse_chunk(SseFraming::DataLines, "data:").is_none()); + } +}