mirror of
https://github.com/fabro-sh/fabro.git
synced 2026-10-07 03:00:29 +00:00
refactor(llm): consolidate the HTTP transport behind the codec seam (#491)
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 <noreply@anthropic.com>
This commit is contained in:
Scott Werner
GitHub
parent
548c1574d2
commit
074f90c391
14 changed files with 869 additions and 1028 deletions
|
|
@ -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 --------------------------------------------------------
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
};
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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<StreamEvent>;
|
||||
}
|
||||
|
||||
// --- 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<String>, Option<serde_json::Value>) {
|
||||
serde_json::from_str::<serde_json::Value>(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<String>, 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<Message> = vec![];
|
||||
let (sys, other) = extract_system_prompt(&msgs);
|
||||
assert_eq!(sys, None);
|
||||
assert!(other.is_empty());
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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};
|
||||
|
|
|
|||
|
|
@ -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<Arc<Catalog>>,
|
||||
}
|
||||
|
|
@ -39,7 +35,7 @@ impl Adapter {
|
|||
#[must_use]
|
||||
pub fn new_optional_auth(api_key: Option<String>) -> 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<dyn StreamDecoder>,
|
||||
line_reader: super::common::LineReader,
|
||||
pending: std::collections::VecDeque<StreamEvent>,
|
||||
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<std::borrow::Cow<'_, str>> = 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<StreamEventStream, Error> {
|
||||
|
|
@ -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 {
|
||||
|
|
|
|||
|
|
@ -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<String>, Option<serde_json::Value>) {
|
||||
serde_json::from_str::<serde_json::Value>(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<String>, 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<f64> {
|
||||
headers
|
||||
.get("retry-after")
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.and_then(|s| s.parse::<f64>().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<RateLimitInfo> {
|
||||
fn header_i64(headers: &HeaderMap, name: &str) -> Option<i64> {
|
||||
headers
|
||||
.get(name)
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.and_then(|s| s.parse::<i64>().ok())
|
||||
}
|
||||
|
||||
fn header_str(headers: &HeaderMap, name: &str) -> Option<String> {
|
||||
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<std::time::Duration>,
|
||||
}
|
||||
|
||||
impl LineReader {
|
||||
pub fn new(
|
||||
response: fabro_http::Response,
|
||||
stream_read_timeout: Option<std::time::Duration>,
|
||||
) -> 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<Option<String>, 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<Message> = 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));
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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::<StreamEvent>(&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::<StreamEvent>(&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(|_| "<invalid url>".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]
|
||||
|
|
|
|||
|
|
@ -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<Arc<Catalog>>,
|
||||
}
|
||||
|
|
@ -39,7 +38,7 @@ impl Adapter {
|
|||
#[must_use]
|
||||
pub fn new_optional_auth(api_key: Option<String>) -> 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<dyn StreamDecoder>,
|
||||
line_reader: super::common::LineReader,
|
||||
pending: std::collections::VecDeque<StreamEvent>,
|
||||
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<StreamEventStream, Error> {
|
||||
|
|
@ -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
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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<String>,
|
||||
pub(crate) base_url: String,
|
||||
pub(crate) default_headers: HashMap<String, String>,
|
||||
pub(crate) client: fabro_http::HttpClient,
|
||||
pub(crate) request_timeout: Option<Duration>,
|
||||
pub(crate) stream_read_timeout: Option<Duration>,
|
||||
}
|
||||
|
||||
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<String>, base_url: impl Into<String>) -> Self {
|
||||
Self::new_optional(Some(api_key.into()), base_url)
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn new_optional(api_key: Option<String>, base_url: impl Into<String>) -> 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<String, String>) -> Self {
|
||||
self.default_headers = headers;
|
||||
self
|
||||
}
|
||||
}
|
||||
|
|
@ -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;
|
||||
|
||||
|
|
|
|||
|
|
@ -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<String>,
|
||||
project_id: Option<String>,
|
||||
provider_name: String,
|
||||
|
|
@ -48,7 +45,7 @@ impl Adapter {
|
|||
#[must_use]
|
||||
pub fn new_optional_auth(api_key: Option<String>) -> 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<dyn StreamDecoder>,
|
||||
line_reader: super::common::LineReader,
|
||||
pending: std::collections::VecDeque<StreamEvent>,
|
||||
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>, String)> {
|
||||
let mut current_event: Option<String> = 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<StreamEventStream, Error> {
|
||||
|
|
@ -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
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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<Arc<Catalog>>,
|
||||
}
|
||||
|
|
@ -41,7 +38,7 @@ impl Adapter {
|
|||
#[must_use]
|
||||
pub fn new_optional_auth(api_key: Option<String>, base_url: impl Into<String>) -> 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<dyn StreamDecoder>,
|
||||
line_reader: super::common::LineReader,
|
||||
/// Events decoded but not yet yielded.
|
||||
pending: VecDeque<StreamEvent>,
|
||||
/// 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<StreamEventStream, Error> {
|
||||
|
|
@ -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
|
||||
}
|
||||
}
|
||||
|
|
|
|||
596
lib/crates/fabro-llm/src/transport.rs
Normal file
596
lib/crates/fabro-llm/src/transport.rs
Normal file
|
|
@ -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<String>,
|
||||
pub(crate) base_url: String,
|
||||
pub(crate) default_headers: HashMap<String, String>,
|
||||
pub(crate) client: fabro_http::HttpClient,
|
||||
pub(crate) request_timeout: Option<Duration>,
|
||||
pub(crate) stream_read_timeout: Option<Duration>,
|
||||
}
|
||||
|
||||
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<String>, base_url: impl Into<String>) -> 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<String, String>) -> 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<f64> {
|
||||
headers
|
||||
.get("retry-after")
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.and_then(|s| s.parse::<f64>().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<RateLimitInfo> {
|
||||
fn header_i64(headers: &HeaderMap, name: &str) -> Option<i64> {
|
||||
headers
|
||||
.get(name)
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.and_then(|s| s.parse::<i64>().ok())
|
||||
}
|
||||
|
||||
fn header_str(headers: &HeaderMap, name: &str) -> Option<String> {
|
||||
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<Response, Error> {
|
||||
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<Duration>,
|
||||
) -> Result<StreamEventStream, Error> {
|
||||
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<dyn StreamDecoder>,
|
||||
line_reader: LineReader,
|
||||
/// Events decoded but not yet yielded.
|
||||
pending: VecDeque<StreamEvent>,
|
||||
/// 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<dyn StreamDecoder>,
|
||||
framing: SseFraming,
|
||||
stream_read_timeout: Option<Duration>,
|
||||
) -> 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<Cow<'_, str>> = 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<Duration>,
|
||||
}
|
||||
|
||||
impl LineReader {
|
||||
pub fn new(response: fabro_http::Response, stream_read_timeout: Option<Duration>) -> 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<Option<String>, 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());
|
||||
}
|
||||
}
|
||||
Loading…
Add table
Reference in a new issue