diff --git a/Cargo.lock b/Cargo.lock index da0db2f..78db81a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1231,7 +1231,7 @@ dependencies = [ "tokio", "tokio-rustls", "tower-service", - "webpki-roots 1.0.8", + "webpki-roots 1.0.10", ] [[package]] @@ -2114,7 +2114,7 @@ dependencies = [ "wasm-bindgen-futures", "wasm-streams", "web-sys", - "webpki-roots 1.0.8", + "webpki-roots 1.0.10", ] [[package]] @@ -2207,9 +2207,9 @@ dependencies = [ [[package]] name = "rustls-webpki" -version = "0.103.15" +version = "0.103.16" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f3c3cf1d8b1e7d4927e2d154c3fcb02979afb9939629c62cd9048d4f07b60ac2" +checksum = "4443b1b6d39d7c663139942cc4558262562d2c84123bb6282ef177009ca8b281" dependencies = [ "ring", "rustls-pki-types", @@ -3244,14 +3244,14 @@ version = "0.26.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "521bc38abb08001b01866da9f51eb7c5d647a19260e00054a8c7fd5f9e57f7a9" dependencies = [ - "webpki-roots 1.0.8", + "webpki-roots 1.0.10", ] [[package]] name = "webpki-roots" -version = "1.0.8" +version = "1.0.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bf85cb06032201fa7c6f829d7db5a7e5aa45bcc0655327713065f6f0576731bf" +checksum = "1c1c8e20c8a38fc7d65abfe74eb7280e4aa667e8eec8bc84996d2874810748d9" dependencies = [ "rustls-pki-types", ] diff --git a/crates/engines/src/http_transport.rs b/crates/engines/src/http_transport.rs index 7b2c659..864f5dd 100644 --- a/crates/engines/src/http_transport.rs +++ b/crates/engines/src/http_transport.rs @@ -195,43 +195,63 @@ impl Transport for HttpTransport { || ct.starts_with("application/vnd.amazon.eventstream") }); let headers = std::mem::take(resp.headers_mut()); - if is_sse { - use futures::TryStreamExt; - let stream = Box::pin(resp.bytes_stream().map_err(|e| StreamFault { - timeout: e.is_timeout(), - message: e.without_url().to_string(), - })); + if !is_sse && !req.stream { + let bytes = resp.bytes().await.map_err(|e| { + let e = e.without_url(); + // both stay 502 (failover-eligible); the code split keeps the external 408-vs-424 class + let what = if e.is_timeout() { + "read upstream body timed out" + } else { + "read upstream body failed" + }; + GatewayError::new(upstream_fault_code(e.is_timeout()), 502, what).with_source(e) + })?; return Ok(UpstreamResponse { status, - body: idle_capped(stream, timeout), + body: UpstreamBody::Json(bytes), headers, }); } - let read = resp.bytes(); - let bytes = if req.stream { - tokio::time::timeout(timeout, read).await.map_err(|_| { + use futures::{StreamExt, TryStreamExt}; + let mut stream: futures::stream::BoxStream<'static, _> = + Box::pin(resp.bytes_stream().map_err(|e| StreamFault { + timeout: e.is_timeout(), + message: e.without_url().to_string(), + })); + if !is_sse { + let deadline = tokio::time::Instant::now() + timeout; + let late = |_| { GatewayError::new( gw_consts::ErrCode::FED_RESP_TIMEOUT, 502, - format!("upstream body not read within {:?}", timeout), + format!("upstream body not read within {timeout:?}"), ) - })? - } else { - read.await - }; - let bytes = bytes.map_err(|e| { - let e = e.without_url(); - // both stay 502 (failover-eligible); the code split keeps the external 408-vs-424 class - let what = if e.is_timeout() { - "read upstream body timed out" - } else { - "read upstream body failed" }; - GatewayError::new(upstream_fault_code(e.is_timeout()), 502, what).with_source(e) - })?; + let head = tokio::time::timeout_at(deadline, stream.try_next()) + .await + .map_err(late)? + .map_err(StreamFault::into_error)? + .unwrap_or_default(); + if status >= 400 || !head.trim_ascii_start().starts_with(b"data:") { + let body = stream.try_fold(Vec::from(head), |mut body, b| async move { + body.extend_from_slice(&b); + Ok(body) + }); + let body = tokio::time::timeout_at(deadline, body) + .await + .map_err(late)? + .map_err(StreamFault::into_error)?; + return Ok(UpstreamResponse { + status, + body: UpstreamBody::Json(body.into()), + headers, + }); + } + stream = Box::pin(futures::stream::once(async { Ok(head) }).chain(stream)); + } Ok(UpstreamResponse { status, - body: UpstreamBody::Json(bytes), + body: idle_capped(stream, timeout), headers, }) } @@ -394,6 +414,25 @@ mod tests { (format!("http://{addr}/v1/chat/completions"), hits) } + async fn serve_once(status: u16, content_type: &'static str, body: &'static str) -> String { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + tokio::spawn(async move { + let (mut sock, _) = listener.accept().await.unwrap(); + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + let mut buf = [0u8; 1024]; + let _ = sock.read(&mut buf).await; + let resp = format!( + "HTTP/1.1 {status} X\r\nContent-Type: {content_type}\r\n\ + Content-Length: {}\r\nConnection: close\r\n\r\n{body}", + body.len() + ); + let _ = sock.write_all(resp.as_bytes()).await; + let _ = sock.shutdown().await; + }); + format!("http://{addr}/v1/chat/completions") + } + fn request(url: String) -> UpstreamRequest { UpstreamRequest { protocol: gw_consts::Protocol::OpenaiChat, @@ -490,6 +529,32 @@ mod tests { assert_eq!(hits.load(Ordering::SeqCst), 2); } + #[tokio::test] + async fn a_stream_request_body_is_read_by_its_first_bytes() { + let sse = "data: {\"choices\":[]}\n\ndata: [DONE]\n\n"; + for (status, body, want_sse) in [ + (200, sse, true), + (200, "{\"ok\":true}", false), + (502, sse, false), + ] { + let mut req = request(serve_once(status, "application/json", body).await); + req.stream = true; + let t = HttpTransport::new(Duration::from_secs(5)).unwrap(); + let resp = t.send(req).await.unwrap(); + assert_eq!( + matches!(resp.body, UpstreamBody::SseStream(_)), + want_sse, + "{status} {body}" + ); + let got = match resp.buffered().await.unwrap().body { + UpstreamBody::Json(b) => b.to_vec(), + UpstreamBody::Sse(b) => b, + UpstreamBody::SseStream(_) => unreachable!("buffered drains the stream"), + }; + assert_eq!(got, body.as_bytes(), "{status}"); + } + } + #[test] fn retry_after_is_honoured_and_capped() { let mut h = reqwest::header::HeaderMap::new(); diff --git a/docs/providers.md b/docs/providers.md index 63e06bd..0252666 100644 --- a/docs/providers.md +++ b/docs/providers.md @@ -31,8 +31,10 @@ models: | `xai` | `https://api.x.ai` | openai-chat, responses, image, video, realtime | `Bearer` (Grok: `reasoning_effort` per model — grok-4.6 `low`–`xhigh`, grok-4.3 also `none`, others reject values they don't list; usage carries `prompt_tokens_details.cached_tokens` and `completion_tokens_details.reasoning_tokens`, and the chat wire counts those reasoning tokens **outside** `completion_tokens` — it adds them into `total_tokens` instead, so the gateway bills them as output from that arithmetic (its Responses wire, OpenRouter and Bedrock all count them inside); every chat, Responses and image reply also carries `usage.cost_in_usd_ticks`, recorded as the vendor cost; xAI's own Anthropic-compatible surface is deprecated, so Anthropic clients reach Grok through this gateway's `/v1/messages` cross-protocol path; verified live: grok-4.3/4.6/4.20 chat + stream + effort tiers, `/v1/messages` both ways, grok-4.5 through `protocol: responses` natively and from the chat/messages surfaces incl. tool loops, grok-imagine-image-2.0 per-image units, and grok-voice-latest through `/v1/realtime` (`protocol: realtime`; xAI's `response.done` carries an empty `usage`, so the turn bills the delivered-output estimate — transcript tokens plus audio bytes/4 — as `estimated`, input unmetered: price the model per output unit accordingly); grok-imagine-video-1.5 through `/v1/videos/generations` + `GET /v1/videos/{id}` (`protocol: video`, unit price per generated second, vendor cost from `usage.cost_in_usd_ticks`); files/collections and vendor batches are not wired) | | `siliconflow` | `https://api.siliconflow.cn` | openai-chat, embeddings, rerank, tts, stt, image, video | `Bearer` (Qwen3 `enable_thinking`, DeepSeek/GLM/Kimi/MiniMax hosted models, bge/Qwen3 embeddings and rerankers, CosyVoice TTS, SenseVoice STT, Kolors images — all verified live) | -Any other OpenAI-compatible vendor (Qwen, Ollama, vLLM, a relay) uses -`kind: openai` with an `endpoint:` override: +Any other OpenAI-compatible vendor (Qwen, Ollama, vLLM, Parasail, a relay) uses +`kind: openai` with an `endpoint:` override. A streamed reply is read as SSE +when its body opens with `data:`, whatever its `content-type` says (Parasail +labels its streams `application/json`): ```yaml providers: