From 54b8757fbd2485f2c257205251c910d90ec10925 Mon Sep 17 00:00:00 2001 From: Mac Anderson Date: Mon, 10 Aug 2026 14:08:46 -0700 Subject: [PATCH 1/3] =?UTF-8?q?feat(stella-model):=20streaming=E2=86=92non?= =?UTF-8?q?-streaming=20fallback=20on=20hung=20or=20empty=20streams,=20wit?= =?UTF-8?q?h=20a=20first-byte=20deadline?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A provider stream that hangs before its first byte (a proxy buffering the SSE body) or comes back as a 200 with an empty stream used to burn the whole retry budget re-issuing the identical streaming request into the same broken pipe — the idle bound only fired after 120s per attempt, and nothing ever tried the transport that would have worked. Now the shared chat-completions adapter bounds the FIRST body read by a distinct 90s first-byte deadline (http::FIRST_BYTE_TIMEOUT), classifies the two nothing-arrived shapes as fallback-eligible, and arms a bounded per-session latch (stream_recovery::StreamRecovery) so the caller's ordinary retry of the faulted attempt re-issues the byte-identical payload with stream: false through the unary read bound. The recovery is deliberately split across two attempts rather than hidden inside one provider call: the faulted attempt fails retryably, so the engine's retry ladder stays the single owner of retries and bills the discarded attempt through its existing UsageIncomplete observer. The latch confirms only on unary evidence (Streaming → Probing → Confirmed), reverts on a failed probe, and never arms for a mid-stream death with salvage — partially streamed content keeps its retry-as-a-stream path and partial-usage accounting, so there is nothing to tombstone. Provider parity (invariant 8) gains its third axis, StreamFallbackPosture: the zai-family ids declare UnaryFallback with wiremock witnesses, the other streaming dialects declare the gap, Bedrock is AlwaysUnary. zai.rs's aggregation half moves to zai/stream.rs (shared assembly rules with zai/unary.rs), retiring its god-file baseline entry. Closes #2686 --- AGENTS.md | 14 +- crates/stella-cli/src/config/tests.rs | 18 + crates/stella-model/README.md | 8 +- crates/stella-model/src/http.rs | 63 +- crates/stella-model/src/lib.rs | 4 + crates/stella-model/src/provider_parity.rs | 173 +++++- crates/stella-model/src/stream_recovery.rs | 173 ++++++ crates/stella-model/src/zai.rs | 537 ++++-------------- crates/stella-model/src/zai/stream.rs | 506 +++++++++++++++++ crates/stella-model/src/zai/tests.rs | 36 +- .../src/zai/tests/stream_fallback.rs | 292 ++++++++++ .../stella-model/src/zai/tests/zai_effort.rs | 35 ++ crates/stella-model/src/zai/unary.rs | 192 +++++++ scripts/file-size-baseline.txt | 1 - 14 files changed, 1579 insertions(+), 473 deletions(-) create mode 100644 crates/stella-model/src/stream_recovery.rs create mode 100644 crates/stella-model/src/zai/stream.rs create mode 100644 crates/stella-model/src/zai/tests/stream_fallback.rs create mode 100644 crates/stella-model/src/zai/unary.rs diff --git a/AGENTS.md b/AGENTS.md index 4ffff483d..c7f681c2f 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -294,7 +294,7 @@ Append; do not renumber. `scripts/check-invariants.sh` enforces both halves. *after* the stable prefix (see `crates/stella-cli/src/agent.rs::build_system_prompt` and `crates/stella-cli/src/memory.rs` for the L-E8 discipline). 8. **Provider feature parity is declared, not assumed.** Providers diverge - in sneaky ways, and this is guarded on **two axes** today in + in sneaky ways, and this is guarded on **three axes** today in `crates/stella-model/src/provider_parity.rs`: - **`CachePosture`** — how the prompt cache is engaged/observed (Anthropic's cache is explicit opt-in; DeepSeek spells its cache-hit @@ -307,10 +307,16 @@ Append; do not renumber. `scripts/check-invariants.sh` enforces both halves. (`reasoning[_]effort`) honor a pinned effort; the shared adapter drops it for `Unsupported` providers (bedrock/deepseek/local) — and a pinned effort against one surfaces a one-line boot notice, never a silent drop. + - **`StreamFallbackPosture`** — how a provider recovers when its + streaming path is broken (a stream hung before its first byte, or a + 200 with an empty stream). The shared chat-completions adapter arms a + bounded per-session latch and re-issues the retried attempt as a unary + request (#2686); the other streaming dialects declare the gap; Bedrock + is already unary. Each provider id declares a posture on **every** axis and, for a - controllable/opt-in/implicit posture, names the **witness test** proving - it on the wire. Tests enforce each matrix from both sides: `stella-cli`'s + controllable/opt-in/implicit/fallback posture, names the **witness test** + proving it on the wire. Tests enforce each matrix from both sides: `stella-cli`'s config tests fail if a seeded provider lacks a row on either axis, and `stella-model`'s parity tests fail if a row's witness test no longer exists. Adding a provider — or a new divergent feature axis — means @@ -547,7 +553,7 @@ a plan needs and the part that rarely changes: |---|---| | `stella-cli` | `src/command_deck.rs`, `src/agent.rs`, `src/agent/tests.rs`, `src/fleet_cmd.rs` | | `stella-core` | `src/driver/tests.rs`, `src/driver.rs`, `src/bus.rs` | -| `stella-model` | `src/openai.rs`, `src/zai/tests.rs`, `src/anthropic/tests.rs`, `src/zai.rs` | +| `stella-model` | `src/openai.rs`, `src/zai/tests.rs`, `src/anthropic/tests.rs` | | `stella-pipeline` | `src/pipeline.rs`, `src/pipeline/tests.rs` | | `stella-store` | `src/tests.rs`, `src/lib.rs`, `src/usage.rs` | | `stella-tools` | `src/registry.rs`, `src/scripts.rs` | diff --git a/crates/stella-cli/src/config/tests.rs b/crates/stella-cli/src/config/tests.rs index bc9c7014c..5e5a47e79 100644 --- a/crates/stella-cli/src/config/tests.rs +++ b/crates/stella-cli/src/config/tests.rs @@ -75,6 +75,24 @@ fn every_seeded_provider_declares_a_reasoning_posture() { } } +/// The stream-fallback sibling (#2686): every seeded provider must declare +/// how it recovers when its streaming path is broken — a unary fallback +/// (with the witness proving the retried attempt goes out non-streaming), a +/// declared streaming-only gap, or an adapter that is already unary. Same +/// completeness contract as the other two axes. +#[test] +fn every_seeded_provider_declares_a_stream_fallback_posture() { + for provider in PROVIDERS.iter().chain(std::iter::once(&LOCAL_PROVIDER)) { + assert!( + stella_model::provider_parity::stream_fallback_posture(provider.id).is_some(), + "provider `{}` has no StreamFallbackPosture row in \ + stella-model/src/provider_parity.rs — add it (with a witness test for a \ + UnaryFallback row, or a note for a no-fallback posture) in this PR", + provider.id + ); + } +} + #[test] fn alias_env_var_resolves_when_the_primary_is_unset() { // Synthetic provider with unique var names so parallel tests can't diff --git a/crates/stella-model/README.md b/crates/stella-model/README.md index 7d3414541..f0a648470 100644 --- a/crates/stella-model/README.md +++ b/crates/stella-model/README.md @@ -85,7 +85,6 @@ the file. - [`src/anthropic/tests.rs`](src/anthropic/tests.rs) - [`src/openai.rs`](src/openai.rs) -- [`src/zai.rs`](src/zai.rs) - [`src/zai/tests.rs`](src/zai/tests.rs) A ceiling can move only via `make file-size-update`, which lands as a @@ -99,14 +98,15 @@ escape hatch for an irreducible line (a module declaration in an oversized |---|---| | [`src/lib.rs`](src/lib.rs) | The crate map (read this first) and the public re-exports: `Provider`, `Catalog`, `ApiKey`, the cache-economics helpers. | | [`src/provider.rs`](src/provider.rs) | Two-line re-export of the port. Open it to be reminded where the trait actually lives. | -| [`src/zai.rs`](src/zai.rs) (+ [`src/zai/tests.rs`](src/zai/tests.rs), [`src/zai/tests/error_classify.rs`](src/zai/tests/error_classify.rs)) | The shared OpenAI Chat Completions adapter. One adapter serving Z.ai/GLM, xAI, DeepSeek, OpenRouter, `local`, and settings-defined gateways; per-identity behavior (OpenRouter's root `cache_control` + sticky `session_id`, GLM's `thinking`, xAI's `reasoning_effort`) is gated on `self.id` inside it. | +| [`src/zai.rs`](src/zai.rs) (+ [`src/zai/stream.rs`](src/zai/stream.rs), [`src/zai/unary.rs`](src/zai/unary.rs), [`src/zai/tests.rs`](src/zai/tests.rs), [`src/zai/tests/error_classify.rs`](src/zai/tests/error_classify.rs)) | The shared OpenAI Chat Completions adapter. One adapter serving Z.ai/GLM, xAI, DeepSeek, OpenRouter, `local`, and settings-defined gateways; per-identity behavior (OpenRouter's root `cache_control` + sticky `session_id`, GLM's `thinking`, xAI's `reasoning_effort`) is gated on `self.id` inside it. `stream.rs` is the SSE aggregation half, `unary.rs` the non-streaming fallback a broken streaming path latches onto (`src/stream_recovery.rs`, #2686). | | [`src/anthropic.rs`](src/anthropic.rs) (+ [`src/anthropic/tests.rs`](src/anthropic/tests.rs)), [`src/openai.rs`](src/openai.rs), [`src/gemini.rs`](src/gemini.rs), [`src/vertex.rs`](src/vertex.rs), [`src/bedrock.rs`](src/bedrock.rs) | One adapter per structurally distinct wire dialect: Messages API, Responses API, `generateContent` (direct and Vertex's project-scoped enterprise path, sharing wire types and the stream aggregator), Bedrock Converse. All follow the same shape — `new(ApiKey, model)` capturing catalog pricing, `with_base_url`, `http::client()`, `SseDecoder`, `http::classify_http_status`, `impl Provider` (plus `complete_observed`, the mid-stream tool-call announcement the engine speculates on — every streaming adapter implements it; Bedrock, which is unary, implements it as one terminal `text_delta` and announces no tool calls). Open the one whose vendor you are debugging. | | [`src/catalog.rs`](src/catalog.rs) | `Catalog`, `CatalogEntry`, `Pricing`, `ToolDialect`, and the compile-time seed. The only sanctioned slug → model resolution. | | [`src/credential.rs`](src/credential.rs) (+ [`src/credential/aux.rs`](src/credential/aux.rs)) | `ApiKey` (the non-`Display` secret wrapper), the flag → env → file → prompt chain, `CredentialsFile` (keys *and* the `[credential_fields.]` companions), the multi-variable resolvers `VertexAddressing` / `BedrockCredentials`, and `AuxCredentials` — the redacting, zeroizing set a host uses to carry the values a provider needs beyond one key (Bedrock's secret access key, session token, and region). | -| [`src/provider_parity.rs`](src/provider_parity.rs) | `CachePosture` / `ReasoningPosture` — the per-provider matrix. A new provider id lands here or tests fail. | +| [`src/provider_parity.rs`](src/provider_parity.rs) | `CachePosture` / `ReasoningPosture` / `StreamFallbackPosture` — the per-provider matrix. A new provider id lands here or tests fail. | | [`src/cache_economics.rs`](src/cache_economics.rs) | Cache savings arithmetic (`Pricing::cache_savings_usd`) and `diagnose_cache`, which reads the parity matrix to tell an opt-in bug from prefix instability. | | [`src/sse.rs`](src/sse.rs) | Dependency-free SSE line parser + incremental UTF-8 decoder every streaming adapter feeds. | -| [`src/http.rs`](src/http.rs) | Crate-private plumbing: the timeout-bounded `reqwest` clients, `classify_http_status`, and the two shared stream-failure errors. Change error retryability here, not in an adapter. | +| [`src/http.rs`](src/http.rs) | Crate-private plumbing: the timeout-bounded `reqwest` clients, the first-byte/idle stream-read bounds, `classify_http_status`, and the two shared stream-failure errors. Change error retryability here, not in an adapter. | +| [`src/stream_recovery.rs`](src/stream_recovery.rs) | Crate-private: the streaming→non-streaming fallback latch (`StreamRecovery`) and the fault classification (`StreamFault`) any adapter with a unary fallback path shares (#2686). | | [`src/attachment.rs`](src/attachment.rs) | Crate-private. Turns `Attachment`s into dialect-neutral `WirePart`s once, so each adapter only maps parts onto its own JSON. | | [`src/modelsdev.rs`](src/modelsdev.rs), [`src/provider_listing.rs`](src/provider_listing.rs) | Fetch-and-parse only: the models.dev master list and provider-native `/models` discovery. Deciding what to store belongs to `stella-cli`. Best-effort by contract — a dead or shape-drifted endpoint returns an `Err(String)` the caller reports and moves past, so one provider can never fail a refresh of the others. | diff --git a/crates/stella-model/src/http.rs b/crates/stella-model/src/http.rs index d3bc95561..4d7c58763 100644 --- a/crates/stella-model/src/http.rs +++ b/crates/stella-model/src/http.rs @@ -24,6 +24,21 @@ pub(crate) const CONNECT_TIMEOUT: Duration = Duration::from_secs(10); /// the turn forever. pub(crate) const STREAM_IDLE_TIMEOUT: Duration = Duration::from_secs(120); +/// How long to wait for the **first** chunk of a stream's body before +/// treating the stream itself as broken. Distinct from — and shorter than — +/// [`STREAM_IDLE_TIMEOUT`], because the two waits mean different things: a +/// gap *between* fragments is a model thinking, while a response that has +/// sent its headers and then not one body byte is the signature of a proxy +/// buffering the SSE body — a path fault the same request completes fine +/// over a non-streaming call (issue #2686). 90 seconds matches the zappy +/// comparator's stream watchdog and sits well above any observed +/// time-to-first-token (which grows with prompt-processing time on large +/// cache writes, not with generation length), while discovering a buffering +/// proxy 30s sooner than the idle bound and ~13x sooner than the engine's +/// 816s model deadline. A trip is fallback-eligible: see +/// `crate::stream_recovery`. +pub(crate) const FIRST_BYTE_TIMEOUT: Duration = Duration::from_secs(90); + /// A `reqwest::Client` with [`CONNECT_TIMEOUT`] applied, plus a per-read /// stall bound of [`STREAM_IDLE_TIMEOUT`]. The read timeout closes the gap /// [`next_with_timeout`] cannot see: the wait between a successful connect @@ -414,13 +429,43 @@ pub(crate) fn classify_http_status( } } +/// One bounded stream read, with the four outcomes kept distinct. The +/// distinction [`next_with_timeout`] collapses — a stall versus a transport +/// fault — is exactly the bit the streaming→non-streaming fallback needs: a +/// stall before the first byte is a broken *streaming path* (fallback +/// material), while a reset says nothing about streaming specifically. +pub(crate) enum StreamRead { + /// The next item arrived within the bound. + Item(T), + /// Clean end of stream. + End, + /// Nothing arrived within the bound — the stream is stalled. + Idle, + /// The transport failed mid-read. + Failed(String), +} + +/// Await the next stream item, bounded by `idle`, reporting the outcome as a +/// [`StreamRead`]. `idle` is a parameter (rather than reading +/// [`STREAM_IDLE_TIMEOUT`] directly) purely so the timeout path is +/// unit-testable in milliseconds. +pub(crate) async fn next_stream_read(stream: &mut S, idle: Duration) -> StreamRead +where + S: Stream> + Unpin, +{ + match tokio::time::timeout(idle, stream.next()).await { + Ok(Some(Ok(item))) => StreamRead::Item(item), + Ok(Some(Err(e))) => StreamRead::Failed(e.to_string()), + Ok(None) => StreamRead::End, + Err(_elapsed) => StreamRead::Idle, + } +} + /// Await the next stream item, bounded by `idle`. Maps a stalled stream (no /// item within `idle`) and any transport error to a **retryable** /// `ProviderError::Transport`, and a clean end-of-stream to `Ok(None)`. -/// -/// `idle` is a parameter (rather than reading [`STREAM_IDLE_TIMEOUT`] -/// directly) purely so the timeout path is unit-testable in milliseconds; -/// adapters always pass [`STREAM_IDLE_TIMEOUT`]. +/// A thin classification over [`next_stream_read`] for the adapters that +/// don't need to tell the two failure shapes apart. pub(crate) async fn next_with_timeout( stream: &mut S, idle: Duration, @@ -428,11 +473,11 @@ pub(crate) async fn next_with_timeout( where S: Stream> + Unpin, { - match tokio::time::timeout(idle, stream.next()).await { - Ok(Some(Ok(item))) => Ok(Some(item)), - Ok(Some(Err(e))) => Err(ProviderError::transport(e.to_string())), - Ok(None) => Ok(None), - Err(_elapsed) => Err(ProviderError::transport(format!( + match next_stream_read(stream, idle).await { + StreamRead::Item(item) => Ok(Some(item)), + StreamRead::End => Ok(None), + StreamRead::Failed(message) => Err(ProviderError::transport(message)), + StreamRead::Idle => Err(ProviderError::transport(format!( "stream idle timeout: no data for {}s", idle.as_secs() ))), diff --git a/crates/stella-model/src/lib.rs b/crates/stella-model/src/lib.rs index c05682b13..6a3fd0d10 100644 --- a/crates/stella-model/src/lib.rs +++ b/crates/stella-model/src/lib.rs @@ -34,6 +34,9 @@ //! provider-native `/models` discovery, both fetch-and-parse only. //! - [`sse`] — the shared, dependency-free SSE + incremental UTF-8 decoder //! every streaming adapter feeds. +//! - `stream_recovery` (crate-private) — the streaming→non-streaming +//! fallback latch armed when a stream hangs before its first byte or +//! comes back empty (#2686). pub mod anthropic; pub(crate) mod attachment; pub mod bedrock; @@ -49,6 +52,7 @@ pub mod provider; pub mod provider_listing; pub mod provider_parity; pub mod sse; +pub(crate) mod stream_recovery; pub mod vertex; pub mod zai; diff --git a/crates/stella-model/src/provider_parity.rs b/crates/stella-model/src/provider_parity.rs index 14223ba1e..08b9c9d77 100644 --- a/crates/stella-model/src/provider_parity.rs +++ b/crates/stella-model/src/provider_parity.rs @@ -12,7 +12,7 @@ //! row is missing (`stella-cli` config tests), duplicated, or names a witness //! test that no longer exists in the adapter sources. //! -//! Two axes are guarded today, both born from the same shape of silent +//! Three axes are guarded today, all born from the same shape of silent //! per-provider divergence: //! - [`CachePosture`] — how a provider's prompt cache is engaged and observed. //! - [`ReasoningPosture`] — how a provider's reasoning/thinking budget is @@ -21,6 +21,10 @@ //! honored a reasoning preference on the shared chat-completions adapter, //! so a pinned `effort` was *silently dropped* for xAI, DeepSeek, and local //! — the exact "nothing enforces the omission stays deliberate" gap. +//! - [`StreamFallbackPosture`] — how a provider recovers when its streaming +//! path is broken (hung before the first byte, or an empty stream): a +//! unary fallback (#2686), a declared streaming-only gap, or an adapter +//! that is already unary. //! //! **The law for new providers:** adding a provider id means adding a row on //! every axis here in the same PR, and a `Controllable`/`OptIn`/`Implicit` @@ -331,6 +335,122 @@ pub fn reasoning_posture(provider_id: &str) -> Option<&'static ReasoningPosture> .map(|(_, posture)| posture) } +/// How a provider recovers when its *streaming path* is broken — the stream +/// hangs before its first byte (a proxy buffering the SSE body) or comes +/// back empty (a gateway answering 200 with no data). The third axis of the +/// matrix (#2686), with the same law as the other two: behavior that +/// diverges per provider is declared, never assumed. +#[derive(Debug)] +pub enum StreamFallbackPosture { + /// The adapter arms a bounded per-session latch on a fallback-eligible + /// stream fault and re-issues the retried attempt as a unary + /// (non-streaming) request for the same payload — see + /// `crate::stream_recovery` for the state machine. + UnaryFallback { + /// The wire mechanism, for humans reading the matrix. + mechanism: &'static str, + /// Name of the test function that proves the fallback: the faulted + /// streaming attempt fails retryably and the retry completes over + /// `stream: false`. Checked for existence by this module's tests. + witness: &'static str, + }, + /// The adapter streams and has no unary fallback path (yet): a broken + /// streaming path fails the attempt with its ordinary classification. + /// Allowed only with a note a reviewer can check. + StreamingOnly { note: &'static str }, + /// The adapter is already unary — there is no stream to fall back from. + AlwaysUnary { note: &'static str }, +} + +/// One stream-fallback row per provider id constructible by the CLI — same +/// completeness contract as the other two axes. Settings-defined custom +/// providers inherit the shared OpenAI-compatible adapter and its fallback, +/// so they need no row of their own. +pub static STREAM_FALLBACK_POSTURE: &[(&str, StreamFallbackPosture)] = &[ + ( + "anthropic", + StreamFallbackPosture::StreamingOnly { + note: "the Messages adapter has no unary parse path yet; extending the shared \ + fallback latch to this dialect is tracked in #2746", + }, + ), + ( + "bedrock", + StreamFallbackPosture::AlwaysUnary { + note: "the adapter calls Converse, not ConverseStream — every completion is \ + already unary, so there is no stream to fall back from", + }, + ), + ( + "openrouter", + StreamFallbackPosture::UnaryFallback { + mechanism: "shared chat-completions adapter: retried attempt re-issues the \ + byte-identical body with stream: false through the unary read bound", + witness: "an_empty_stream_falls_back_to_a_non_streaming_request", + }, + ), + ( + "openai", + StreamFallbackPosture::StreamingOnly { + note: "the Responses adapter has no unary parse path yet; extending the shared \ + fallback latch to this dialect is tracked in #2746", + }, + ), + ( + "gemini", + StreamFallbackPosture::StreamingOnly { + note: "streamGenerateContent has no unary parse path yet (generateContent would \ + be the fallback); tracked in #2746", + }, + ), + ( + "vertex", + StreamFallbackPosture::StreamingOnly { + note: "shares gemini's streaming aggregator and its gap; tracked in #2746", + }, + ), + ( + "zai", + StreamFallbackPosture::UnaryFallback { + mechanism: "retried attempt re-issues the byte-identical body with stream: false \ + through the unary read bound (http::unary_client)", + witness: "a_stream_hung_before_its_first_byte_falls_back_to_a_non_streaming_request", + }, + ), + ( + "xai", + StreamFallbackPosture::UnaryFallback { + mechanism: "shared chat-completions adapter fallback (see the zai row)", + witness: "a_stream_hung_before_its_first_byte_falls_back_to_a_non_streaming_request", + }, + ), + ( + "deepseek", + StreamFallbackPosture::UnaryFallback { + mechanism: "shared chat-completions adapter fallback (see the zai row)", + witness: "a_stream_hung_before_its_first_byte_falls_back_to_a_non_streaming_request", + }, + ), + ( + "local", + StreamFallbackPosture::UnaryFallback { + mechanism: "shared chat-completions adapter fallback (see the zai row) — local \ + gateways and proxies are where SSE buffering is most likely", + witness: "an_empty_stream_falls_back_to_a_non_streaming_request", + }, + ), +]; + +/// The declared stream-fallback posture for `provider_id`, or `None` for an +/// id the matrix doesn't know — which the `stella-cli` completeness test +/// turns into a hard failure for any seeded provider. +pub fn stream_fallback_posture(provider_id: &str) -> Option<&'static StreamFallbackPosture> { + STREAM_FALLBACK_POSTURE + .iter() + .find(|(id, _)| *id == provider_id) + .map(|(_, posture)| posture) +} + /// The tier `effort` is really served as by `provider_id`, when the adapter /// cannot put that tier on the wire distinctly — `None` when it reaches the /// wire as itself. @@ -370,7 +490,7 @@ mod tests { /// is a false alarm rather than the rotted proof it exists to catch. The /// parent `tests.rs` is over the file-size ratchet, so those splits keep /// happening — the list has to follow them. - fn adapter_sources() -> [&'static str; 14] { + fn adapter_sources() -> [&'static str; 15] { [ include_str!("anthropic/tests.rs"), include_str!("anthropic/tests/cache_breakpoints.rs"), @@ -383,6 +503,7 @@ mod tests { include_str!("zai/tests/error_classify.rs"), include_str!("zai/tests/openrouter_effort.rs"), include_str!("zai/tests/openrouter_stream.rs"), + include_str!("zai/tests/stream_fallback.rs"), include_str!("zai/tests/stream_frame.rs"), include_str!("zai/tests/vision.rs"), include_str!("zai/tests/zai_effort.rs"), @@ -504,6 +625,38 @@ mod tests { } } + /// The stream-fallback sibling: every `UnaryFallback` row must name a + /// test that exists in the adapter sources, proving the retried attempt + /// really goes out non-streaming. The no-fallback variants carry a note, + /// not a witness. + #[test] + fn every_stream_fallback_witness_test_exists_in_the_adapter_sources() { + let sources = adapter_sources(); + for (id, posture) in STREAM_FALLBACK_POSTURE { + let witness = match posture { + StreamFallbackPosture::UnaryFallback { witness, .. } => witness, + StreamFallbackPosture::StreamingOnly { .. } + | StreamFallbackPosture::AlwaysUnary { .. } => continue, + }; + let needle = format!("fn {witness}("); + assert!( + sources.iter().any(|source| source.contains(&needle)), + "stream-fallback witness for `{id}` not found in adapter sources: {witness}" + ); + } + } + + #[test] + fn stream_fallback_provider_ids_are_unique() { + let mut seen = std::collections::BTreeSet::new(); + for (id, _) in STREAM_FALLBACK_POSTURE { + assert!( + seen.insert(id), + "duplicate stream-fallback-posture row for `{id}`" + ); + } + } + #[test] fn provider_ids_are_unique() { let mut seen = std::collections::BTreeSet::new(); @@ -523,18 +676,24 @@ mod tests { } } - /// Both axes must cover exactly the same set of provider ids — a provider - /// present on one axis but not the other is a matrix hole. + /// Every axis must cover exactly the same set of provider ids — a + /// provider present on one axis but not another is a matrix hole. #[test] - fn both_axes_cover_the_same_provider_ids() { + fn all_axes_cover_the_same_provider_ids() { let cache: std::collections::BTreeSet<_> = CACHE_POSTURE.iter().map(|(id, _)| *id).collect(); let reasoning: std::collections::BTreeSet<_> = REASONING_POSTURE.iter().map(|(id, _)| *id).collect(); + let fallback: std::collections::BTreeSet<_> = + STREAM_FALLBACK_POSTURE.iter().map(|(id, _)| *id).collect(); assert_eq!( cache, reasoning, "cache and reasoning matrices cover different provider ids" ); + assert_eq!( + cache, fallback, + "cache and stream-fallback matrices cover different provider ids" + ); } #[test] @@ -547,5 +706,9 @@ mod tests { assert!(reasoning_posture(id).is_some()); } assert!(reasoning_posture("no-such-provider").is_none()); + for (id, _) in STREAM_FALLBACK_POSTURE { + assert!(stream_fallback_posture(id).is_some()); + } + assert!(stream_fallback_posture("no-such-provider").is_none()); } } diff --git a/crates/stella-model/src/stream_recovery.rs b/crates/stella-model/src/stream_recovery.rs new file mode 100644 index 000000000..f8303f271 --- /dev/null +++ b/crates/stella-model/src/stream_recovery.rs @@ -0,0 +1,173 @@ +//! Streaming→non-streaming fallback state, shared by any adapter that can +//! re-issue a completion as a unary request when its stream is broken. +//! +//! Two broken-stream shapes motivate this (issue #2686): a proxy that +//! buffers the SSE body, so the response hangs before its first byte; and a +//! gateway that answers 200 with an **empty stream** (EOF before any data). +//! Both are properties of the *path*, not of the request — the same payload +//! completes fine over a non-streaming call — and both used to burn the +//! whole retry budget re-issuing the identical streaming request into the +//! same broken pipe. +//! +//! The recovery is deliberately split across two attempts rather than hidden +//! inside one provider call: the faulted streaming attempt **fails with a +//! retryable error**, so the caller's retry machinery (which already bills +//! every discarded attempt through its `UsageIncomplete` observer and bounds +//! the attempt count) stays the single owner of retries and accounting. This +//! latch only changes what the *next* attempt sends. An adapter must consult +//! [`StreamRecovery::use_unary`] per attempt, report a qualifying stream +//! fault via [`StreamRecovery::note_stream_fault`], and report every unary +//! attempt's outcome via [`StreamRecovery::note_unary_outcome`]. +//! +//! The state machine is a bounded latch, after zappy's "reactive resilience" +//! watchdog-and-fallback (adapted to this workspace's accounting rules): +//! +//! ```text +//! Streaming ──stream fault (hung/empty, nothing arrived)──▶ Probing +//! Probing ──unary attempt succeeds──▶ Confirmed (sticky for this session) +//! Probing ──unary attempt fails──▶ Streaming +//! ``` +//! +//! `Confirmed` is sticky because the trigger faults are deterministic: a +//! proxy that buffers SSE buffers it on every request, so re-probing the +//! stream would pay the full first-byte deadline once per step for the rest +//! of the session. A failed probe reverts instead of latching — one +//! unproven fault must not condemn streaming (and with it mid-stream +//! observation/speculation) when the provider was simply down, and the +//! ping-pong this allows is bounded by the caller's retry policy like every +//! other attempt. + +use std::sync::atomic::{AtomicU8, Ordering}; + +use stella_protocol::ProviderError; + +/// One streaming attempt's failure, classified for the fallback decision. +/// +/// `fallback_eligible` is true only for the two shapes where the retry +/// loses nothing by going out unary: the stream hung or died having +/// delivered **no completion signal whatsoever** (no content, no usage +/// frame, no terminal marker). A mid-stream death with salvage keeps its +/// existing retry-as-a-stream classification, so partially-streamed content +/// is never the fallback's problem — there is nothing to tombstone. +#[derive(Debug)] +pub(crate) struct StreamFault { + pub(crate) error: ProviderError, + pub(crate) fallback_eligible: bool, +} + +impl StreamFault { + pub(crate) fn ineligible(error: ProviderError) -> Self { + Self { + error, + fallback_eligible: false, + } + } +} + +/// Every plain `ProviderError` raised inside a stream aggregator (malformed +/// frames, in-band error frames, truncated tool input) is ineligible: only +/// the aggregator's explicit hung/empty exits construct an eligible fault. +impl From for StreamFault { + fn from(error: ProviderError) -> Self { + Self::ineligible(error) + } +} + +const STREAMING: u8 = 0; +const PROBING: u8 = 1; +const CONFIRMED: u8 = 2; + +/// Per-provider-instance fallback latch. One instance lives one session, so +/// the latch's lifetime is the session's — exactly the scope on which a +/// broken streaming path is a stable fact. +/// +/// All transitions are CAS-guarded: concurrent calls through one provider +/// instance may race, and the worst a lost race can do is leave the latch in +/// the state the other caller proved, never skip a state. +#[derive(Debug, Default)] +pub(crate) struct StreamRecovery { + state: AtomicU8, +} + +impl StreamRecovery { + /// Whether the next attempt should be a unary (non-streaming) request. + pub(crate) fn use_unary(&self) -> bool { + self.state.load(Ordering::Acquire) != STREAMING + } + + /// A streaming attempt faulted in a fallback-eligible way (hung before + /// its first byte, or ended as an empty stream): arm the probe so the + /// caller's retry of this attempt goes out unary. + /// + /// Returns whether this call armed it — `false` when a probe was already + /// armed or confirmed, so the caller can word its error without promising + /// a switch some other attempt already made. + pub(crate) fn note_stream_fault(&self) -> bool { + self.state + .compare_exchange(STREAMING, PROBING, Ordering::AcqRel, Ordering::Acquire) + .is_ok() + } + + /// A unary attempt resolved. Success while probing confirms the latch + /// for the rest of the session; failure while probing reverts to + /// streaming (the fault evidently wasn't the stream's). A confirmed + /// latch is sticky — a later transient unary failure must not send the + /// session back to a stream already proven broken. + pub(crate) fn note_unary_outcome(&self, success: bool) { + let next = if success { CONFIRMED } else { STREAMING }; + let _ = self + .state + .compare_exchange(PROBING, next, Ordering::AcqRel, Ordering::Acquire); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn starts_streaming_and_arms_a_probe_on_the_first_fault() { + let recovery = StreamRecovery::default(); + assert!(!recovery.use_unary(), "healthy sessions stream"); + assert!(recovery.note_stream_fault(), "first fault arms the probe"); + assert!(recovery.use_unary(), "the retry must go out unary"); + assert!( + !recovery.note_stream_fault(), + "an already-armed probe reports no new transition" + ); + } + + #[test] + fn a_successful_probe_confirms_the_latch_for_good() { + let recovery = StreamRecovery::default(); + recovery.note_stream_fault(); + recovery.note_unary_outcome(true); + assert!(recovery.use_unary(), "confirmed sessions stay unary"); + // Sticky: a later transient unary failure must not send the session + // back to a stream already proven broken. + recovery.note_unary_outcome(false); + assert!(recovery.use_unary(), "confirmation is sticky"); + } + + #[test] + fn a_failed_probe_reverts_to_streaming() { + let recovery = StreamRecovery::default(); + recovery.note_stream_fault(); + recovery.note_unary_outcome(false); + assert!( + !recovery.use_unary(), + "one unproven fault must not condemn streaming when the provider \ + was simply down" + ); + } + + #[test] + fn a_unary_outcome_without_a_probe_is_a_no_op() { + let recovery = StreamRecovery::default(); + recovery.note_unary_outcome(true); + assert!( + !recovery.use_unary(), + "success reported outside a probe must not invent a latch" + ); + } +} diff --git a/crates/stella-model/src/zai.rs b/crates/stella-model/src/zai.rs index 46c8eb0e2..23953e31e 100644 --- a/crates/stella-model/src/zai.rs +++ b/crates/stella-model/src/zai.rs @@ -2,24 +2,30 @@ //! 5.2's tool-call dialect (`openai-json`: an accumulating `tool_calls` //! array keyed by index, arguments streamed as string fragments). -use std::collections::BTreeMap; use std::sync::atomic::{AtomicU64, Ordering}; +use std::time::Duration; use async_trait::async_trait; use serde::{Deserialize, Serialize}; use serde_json::Value; use stella_protocol::{ - CompletionMessage, CompletionRequestRef, CompletionResult, CompletionUsage, FinishReason, - MessageRole, ProviderError, ReasoningEffort, ToolCall, + CompletionMessage, CompletionRequestRef, CompletionResult, MessageRole, ProviderError, + ReasoningEffort, }; +// Used by the assembled result the delivery-path modules now own, but still +// named unqualified throughout the test tree, which imports `super::*`. +#[cfg(test)] +use stella_protocol::{CompletionUsage, FinishReason, ToolCall}; use crate::catalog::{Catalog, Pricing}; use crate::credential::ApiKey; use crate::http; use crate::provider::{Provider, ToolCallObserver}; -use crate::sse::SseDecoder; +use crate::stream_recovery::{StreamFault, StreamRecovery}; pub(crate) mod effort; +mod stream; +mod unary; use effort::{xai_reasoning_effort, xai_supports_reasoning_effort, zai_reasoning_effort}; // `map_zai_effort` is asserted on directly by the wire tests, which reach it // through this module's namespace like every other helper here. Its xAI sibling @@ -69,6 +75,23 @@ pub struct ZaiProvider { /// the OpenAI adapter's `prompt_cache_key` lifecycle. Volatile by design: /// it rides as a request parameter and never enters the cached bytes. session_id: String, + /// The streaming→non-streaming fallback latch (#2686): armed when a + /// stream hangs before its first byte or comes back empty, consulted per + /// attempt so the retry of a faulted attempt goes out unary. See + /// [`crate::stream_recovery`] for the state machine and its bounds. + recovery: StreamRecovery, + /// The client the unary fallback dispatches through. Separate from + /// [`Self::client`] because a non-streaming call has no first token to + /// reset the per-read clock: the whole generation must fit inside one + /// read, so it needs [`http::unary_client`]'s 600s bound where the + /// streaming client's 120s per-chunk bound would fail every completion + /// slower than two minutes as retryable Transport (#547's lesson, + /// learned on Bedrock). + unary_client: reqwest::Client, + /// [`http::FIRST_BYTE_TIMEOUT`] in production; a field so the + /// hung-stream path is testable in milliseconds (the same reason + /// `next_with_timeout` takes `idle` as a parameter). + first_byte_deadline: Duration, } /// Process-wide monotonic suffix guaranteeing two [`ZaiProvider`] @@ -136,9 +159,20 @@ impl ZaiProvider { extra_headers: Vec::new(), usage_accounting: false, session_id: new_session_id(), + recovery: StreamRecovery::default(), + unary_client: http::unary_client(), + first_byte_deadline: http::FIRST_BYTE_TIMEOUT, } } + /// Shrink the first-byte deadline so the hung-stream fallback is + /// testable in milliseconds instead of 90 seconds of wall clock. + #[cfg(test)] + pub(crate) fn with_first_byte_deadline(mut self, deadline: Duration) -> Self { + self.first_byte_deadline = deadline; + self + } + /// The session-stable sticky-routing id minted for this construction. /// Test-only: the fleet-distinctness witness inspects it on the builder /// (no live gateway), and the wire-gating tests assert it appears only on @@ -1030,20 +1064,88 @@ impl ZaiProvider { /// One request/stream/aggregate cycle. `force_default_reasoning` drops the /// `reasoning` object from the body so an endpoint that mandates reasoning /// applies its own default instead of rejecting the request. + /// + /// Streams by default, but consults the fallback latch first: the retry + /// of an attempt whose stream hung before its first byte or came back + /// empty goes out as a **unary** request for the same payload instead + /// (#2686) — see [`crate::stream_recovery`] for the latch's states and + /// bounds, and [`unary`] for that path. async fn complete_attempt( &self, req: CompletionRequestRef<'_>, observer: Option<&dyn ToolCallObserver>, force_default_reasoning: bool, ) -> Result { + if self.recovery.use_unary() { + return self + .complete_unary_attempt(req, force_default_reasoning) + .await; + } + let body = self.build_body(req, force_default_reasoning, true); + let response = self.dispatch(&self.client, &body).await?; + let (text, tool_calls, usage, finish_reason, reported_cost_usd) = + stream::aggregate_zai_stream( + response, + &self.label, + observer, + self.pricing.as_ref(), + self.first_byte_deadline, + ) + .await + .map_err(|fault| self.absorb_stream_fault(fault))?; + // A gateway-reported cost (OpenRouter usage accounting) is + // authoritative; catalog list pricing is the estimate for providers + // that don't report one. + let cost_usd = reported_cost_usd + .unwrap_or_else(|| self.pricing.map(|p| p.cost_usd(&usage)).unwrap_or(0.0)); + Ok(CompletionResult { + text, + tool_calls, + usage, + model: self.model.clone(), + cost_usd, + finish_reason, + }) + } + + /// Route a [`StreamFault`] back into the retry ladder. An eligible fault + /// (hung before its first byte / empty stream) arms the fallback latch — + /// the caller's retry of this attempt will go out unary — and its error + /// is retryable Transport, so the ordinary retry machinery both drives + /// that switch and bills this discarded attempt through its + /// `UsageIncomplete` observer, exactly like every other doomed attempt. + /// The message names the switch only when THIS fault armed the latch, so + /// an error never promises a fallback some other attempt already made. + fn absorb_stream_fault(&self, fault: StreamFault) -> ProviderError { + if fault.fallback_eligible && self.recovery.note_stream_fault() { + return match fault.error { + ProviderError::Transport { message, partial } => ProviderError::Transport { + message: format!("{message}; retrying this attempt as a non-streaming request"), + partial, + }, + other => other, + }; + } + fault.error + } + + /// The one request body both delivery paths serialize — `stream` is the + /// only field on which they differ, so the unary fallback re-issues the + /// byte-identical payload minus the stream flag. + fn build_body( + &self, + req: CompletionRequestRef<'_>, + force_default_reasoning: bool, + stream: bool, + ) -> ZaiRequest<'_> { // "Include" semantics: every override is `None` unless the caller // set it, and `None` never reaches the wire — a request without // params serializes byte-identical to the pre-params body. let params = req.params.unwrap_or_default(); - let body = ZaiRequest { + ZaiRequest { model: &self.model, messages: to_zai_messages(req.messages, &self.model), - stream: true, + stream, max_tokens: reasoning_aware_max_tokens(req.max_output_tokens, req.reasoning), temperature: req.temperature, top_p: params.top_p, @@ -1107,17 +1209,26 @@ impl ZaiProvider { .serves_openrouter() .then_some(ZaiCacheControl { kind: "ephemeral" }), session_id: self.serves_openrouter().then_some(self.session_id.as_str()), - }; + } + } - let mut request = self - .client + /// POST `body` and run the shared non-success ladder (vendor 429 + /// pre-check first). Returns the successful response for the caller — + /// streaming or unary — to consume. The two delivery paths differ only + /// in their client's read bound, so `client` is a parameter. + async fn dispatch( + &self, + client: &reqwest::Client, + body: &ZaiRequest<'_>, + ) -> Result { + let mut request = client .post(format!("{}/chat/completions", self.base_url)) .bearer_auth(self.api_key.reveal()); for (name, value) in &self.extra_headers { request = request.header(*name, value); } let response = request - .json(&body) + .json(body) .send() .await .map_err(|e| ProviderError::transport(e.to_string()))?; @@ -1153,412 +1264,8 @@ impl ZaiProvider { &self.model, )); } - - let (text, tool_calls, usage, finish_reason, reported_cost_usd) = - aggregate_zai_stream(response, &self.label, observer, self.pricing.as_ref()).await?; - // A gateway-reported cost (OpenRouter usage accounting) is - // authoritative; catalog list pricing is the estimate for providers - // that don't report one. - let cost_usd = reported_cost_usd - .unwrap_or_else(|| self.pricing.map(|p| p.cost_usd(&usage)).unwrap_or(0.0)); - Ok(CompletionResult { - text, - tool_calls, - usage, - model: self.model.clone(), - cost_usd, - finish_reason, - }) - } -} - -/// Accumulator for one in-progress streamed tool call, keyed by the -/// provider's `index` field until it's complete. -#[derive(Default)] -struct ToolCallAccumulator { - id: String, - name: String, - arguments: String, - /// Whether this call was already announced to a [`ToolCallObserver`]. - /// OpenAI-style tool calls stream sequentially by index, so a call is - /// complete the moment a HIGHER index appears — that boundary announces - /// it exactly once. - /// - /// The stream's LAST call has no higher index behind it, so its boundary - /// is the chunk that carries `finish_reason: "tool_calls"` instead — - /// which arrives before the final usage frame and `[DONE]`, with an - /// end-of-stream fallback for a server that jumps straight to `[DONE]`. - /// Without that boundary a turn that makes exactly one tool call — the - /// common shape for this agent — was never announced at all, and - /// `stella-core`'s speculative execution never fired for any single-call - /// turn on this dialect. The window is smaller than what the dialects - /// with a per-call terminator get (`response.function_call_arguments.done` - /// on `openai.rs`, `content_block_stop` on `anthropic.rs`, whole - /// `functionCall` parts on `gemini.rs`), but not empty. This flag is what - /// keeps every boundary exactly-once: a call announced at one boundary is - /// skipped at every later one. - announced: bool, -} - -/// Announce every un-announced accumulator below `next_index` to the -/// observer. Only calls whose arguments already parse are announced — a -/// call the end-of-stream assembly would hand the `Null` repair sentinel -/// must never reach speculative execution. Announced calls re-parse the -/// same bytes at final assembly, so an announced call and its committed -/// twin are structurally identical. -fn announce_completed_below( - observer: &dyn ToolCallObserver, - tool_calls: &mut BTreeMap, - next_index: usize, -) { - for (_, acc) in tool_calls.range_mut(..next_index) { - if acc.announced { - continue; - } - acc.announced = true; - if acc.id.is_empty() { - continue; - } - let trimmed = acc.arguments.trim(); - let input = if trimmed.is_empty() { - Some(Value::Object(serde_json::Map::new())) - } else { - serde_json::from_str(trimmed).ok() - }; - if let Some(input) = input { - observer.tool_call_streamed(&ToolCall { - call_id: acc.id.clone(), - name: acc.name.clone(), - input, - }); - } - } -} - -async fn aggregate_zai_stream( - response: reqwest::Response, - label: &str, - observer: Option<&dyn ToolCallObserver>, - pricing: Option<&Pricing>, -) -> Result< - ( - String, - Vec, - CompletionUsage, - Option, - Option, - ), - ProviderError, -> { - let mut decoder = SseDecoder::new(); - let mut text = String::new(); - // Chain-of-thought streamed under `reasoning_content`, kept separate from - // the answer. Used only as a fallback when `content` never arrives, so a - // reasoning-only turn is visible instead of blank. - let mut reasoning = String::new(); - let mut usage = CompletionUsage::default(); - let mut tool_calls: BTreeMap = BTreeMap::new(); - // Set once any choice reports `finish_reason: "length"` — the output was - // cut off at the token limit, so a tool call whose argument JSON didn't - // finish streaming is truncated, not merely malformed. - let mut truncated_at_token_limit = false; - // Gateway-reported per-call cost (OpenRouter usage accounting), from the - // final usage frame. `None` when the endpoint doesn't report one. - let mut reported_cost_usd: Option = None; - let mut usage_seen = false; - let mut terminal_seen = false; - let mut stream = response.bytes_stream(); - - while let Some(chunk) = http::next_with_timeout(&mut stream, http::STREAM_IDLE_TIMEOUT) - .await - .map_err(|e| http::attach_partial(e, &usage, &text, pricing))? - { - decoder - .push_bytes(&chunk) - .map_err(|e| ProviderError::Malformed(e.to_string()))?; - for event in decoder.poll() { - let data = event.data.trim(); - if data.is_empty() { - continue; - } - if data == "[DONE]" { - terminal_seen = true; - continue; - } - // Two failure modes hide behind one `from_str`, and they must not - // share a branch. A frame that is not JSON at all is a keep-alive - // or a gateway comment: ignoring it is correct. A frame that IS - // valid JSON but does not fit the chunk shape is a dialect - // deviation, and swallowing it loses whatever it carried. Since - // every field in this tree defaults, an unknown or empty object - // still parses — so reaching the error arm means a real type - // mismatch on a field that matters, and the likeliest casualty is - // `tool_calls`. Dropping that produces a turn with no calls, which - // the driver reads as a clean completion: the run reports success - // having done nothing. The `error` field above documents this exact - // failure happening once already. Fail loudly instead. - let Ok(value) = serde_json::from_str::(data) else { - continue; // not JSON: keep-alive / ping / gateway comment - }; - let parsed: ZaiStreamChunk = match serde_json::from_value(value) { - Ok(v) => v, - Err(e) => { - return Err(ProviderError::Malformed(format!( - "{label}: unparseable stream frame ({e}); refusing to \ - treat a dropped frame as an empty turn" - ))); - } - }; - // A mid-stream error frame aborts the turn with a typed error — - // never a truncated Ok with the partial text seen so far. - if let Some(err) = &parsed.error { - return Err(classify_zai_stream_error(err, label)); - } - if let Some(u) = parsed.usage { - usage_seen = true; - usage.input_tokens = u.prompt_tokens; - usage.output_tokens = u.completion_tokens; - let details = u.prompt_tokens_details.unwrap_or_default(); - // Two wire spellings for the same fact: the OpenAI-compatible - // details object, or DeepSeek's native top-level field — take - // whichever the server spoke (never both on one endpoint). - usage.cached_input_tokens = if details.cached_tokens > 0 { - details.cached_tokens - } else { - u.prompt_cache_hit_tokens - }; - usage.cache_write_tokens = details.cache_write_tokens; - // Only overwrite when this frame carried the breakdown: a - // later usage frame without the detail object must not erase - // a count an earlier one reported. - if let Some(reasoning) = - u.completion_tokens_details.and_then(|d| d.reasoning_tokens) - { - usage.reasoning_tokens = Some(reasoning); - } - if u.cost.is_some() { - reported_cost_usd = u.cost; - } - } - for choice in parsed.choices { - if choice.finish_reason.as_deref() == Some("length") { - truncated_at_token_limit = true; - } - if let Some(content) = choice.delta.content { - // `content` is the user-visible answer; thinking rides - // `reasoning_delta` below. The two channels stay separate - // all the way to the transcript. - if let Some(observer) = observer { - observer.text_delta(&content); - } - text.push_str(&content); - } - // GLM's name for chain-of-thought and OpenRouter's normalized - // one. Both announce as thinking, so a reasoning model's - // deliberation is visible live (collapsed, dimmed) instead of - // the turn looking idle until the answer lands. - if let Some(rc) = choice.delta.reasoning_content { - if let Some(observer) = observer { - observer.reasoning_delta(&rc); - } - reasoning.push_str(&rc); - } - if let Some(r) = choice.delta.reasoning { - if let Some(observer) = observer { - observer.reasoning_delta(&r); - } - reasoning.push_str(&r); - } - for (position, tc_delta) in choice - .delta - .tool_calls - .unwrap_or_default() - .into_iter() - .enumerate() - { - // Absent `index` falls back to the fragment's position in - // this chunk — see `ZaiStreamToolCallDelta`. - let index = tc_delta.index.unwrap_or(position); - // A delta for index N proves every lower index finished - // streaming (the dialect emits calls sequentially) — - // the moment those calls can be announced for - // speculative execution. - if let Some(observer) = observer { - announce_completed_below(observer, &mut tool_calls, index); - } - let acc = tool_calls.entry(index).or_default(); - if let Some(id) = tc_delta.id { - acc.id = id; - } - if let Some(function) = tc_delta.function { - if let Some(name) = function.name { - // Accumulated, not assigned: the dialect contract - // (pinned by - // `complete_reassembles_a_streamed_tool_call_split_across_many_chunks`) - // allows the name to arrive in fragments like the - // arguments do, so append is the correct - // reassembly. A gateway that instead repeats the - // WHOLE name per fragment would garble this — no - // such gateway has been observed; if one appears, - // dedupe the exact-repeat case rather than - // switching to last-wins. - acc.name.push_str(&name); - } - if let Some(args) = function.arguments { - // Liveness only (see - // `ToolCallObserver::tool_input_delta`): a - // call-only generation must still register as - // producing against the idle deadline. - if let Some(observer) = observer { - observer.tool_input_delta(); - } - acc.arguments.push_str(&args); - } - } - } - // The chunk carrying `finish_reason: "tool_calls"` is this - // dialect's end-of-calls terminator: the LAST call has no - // higher index behind it, so this is its completion boundary. - // Announcing here — before the final usage frame and `[DONE]` - // — is what opens the speculative-execution window for a - // single-call turn; `announced` keeps calls already announced - // at an index boundary from repeating. - if choice.finish_reason.as_deref() == Some("tool_calls") - && let Some(observer) = observer - { - announce_completed_below(observer, &mut tool_calls, usize::MAX); - } - } - } - } - - // EOF without the `[DONE]` sentinel is a disconnect, not a completion — - // whatever accumulated is a half-answer, and even a `finish_reason` seen - // earlier can't prove the stream wasn't cut after it. Retryable - // Transport, upholding the same "never a truncated Ok" promise as the - // mid-stream error-frame path above. - if !terminal_seen { - return Err(http::attach_partial( - http::stream_ended_before_terminal(label, "[DONE]"), - &usage, - &text, - pricing, - )); - } - - // Fallback terminator: a server that ends the stream without ever sending - // a `finish_reason: "tool_calls"` chunk still completed its last call — - // `[DONE]` proves the stream is whole, so announce whatever is still open - // before final assembly (`announced` makes this a no-op when the finish - // chunk already fired). Skipped when the token limit cut the stream: a - // truncated call must never reach speculative execution, and final - // assembly below turns it into a terminal error instead. - if !truncated_at_token_limit && let Some(observer) = observer { - announce_completed_below(observer, &mut tool_calls, usize::MAX); + Ok(response) } - - usage.reported = usage_seen; - - // OpenAI-style tool calls stream sequentially by index, so when the - // stream reports `finish_reason: "length"` only the highest-index call - // can be the one the token limit cut. Pinning truncation there keeps the - // blame on the right call — an earlier call whose JSON is broken is the - // model's own malformed output and still gets the repair sentinel below. - let truncated_index = if truncated_at_token_limit { - tool_calls.keys().next_back().copied() - } else { - None - }; - - let mut calls = Vec::with_capacity(tool_calls.len()); - for (index, acc) in tool_calls { - let truncated = Some(index) == truncated_index; - let trimmed = acc.arguments.trim(); - let input = if trimmed.is_empty() { - if truncated { - // The limit landed after the call's id/name but before any - // argument fragment: executing it with `{}` would fail on - // missing parameters and re-enter the same unwinnable - // retry-retruncate loop as a mid-payload cut. - return Err(http::truncated_tool_input_error( - label, - &acc.name, - "", - "finish_reason=length", - )); - } - // A no-argument tool call arrives as `arguments: ""`; that is an - // empty object, not null — a downstream tool deserializing its - // input as an object must not be handed `null`. - Value::Object(serde_json::Map::new()) - } else { - match serde_json::from_str(trimmed) { - Ok(value) => value, - // The stream stopped at the token limit MID-arguments: the - // JSON is truncated, not the model's own broken syntax. - // Terminal and turn-aborting — mirroring openai.rs's - // `response.incomplete` handling — because retrying the - // identical request re-truncates identically (the reported - // "stuck-loop" defect). - Err(_) if truncated => { - return Err(http::truncated_tool_input_error( - label, - &acc.name, - trimmed, - "finish_reason=length", - )); - } - // A *non-empty* body that fails to parse without being the - // truncated call is the model's own broken JSON (GLM emits - // these): fall back to the `Value::Null` sentinel - // `driver.rs::execute_with_repair` checks for, so the repair - // loop — documented as tuned to GLM's failure shapes — can - // ask the model to retry. Mirrors anthropic.rs. - Err(_) => Value::Null, - } - }; - calls.push(ToolCall { - call_id: acc.id, - name: acc.name, - input, - }); - } - - // Reasoning-only fallback: if the model emitted no answer `content` but did - // stream chain-of-thought, surface the reasoning as the text so the turn is - // never blank. Normal turns keep `content` as the answer and ignore it. - // - // `calls.is_empty()` is load-bearing, not defensive. A tool-calling turn - // legitimately has empty `content` — that is the *normal* shape for every - // Anthropic model routed through OpenRouter, which streams `content: ""` - // alongside `reasoning` and the tool call. Without this guard the fallback - // fired on essentially every reasoning-model tool turn and published the - // model's private chain-of-thought as its user-visible answer (users saw - // "The user is asking… I should clarify…" third-person deliberation in - // place of a reply). A turn that called a tool is not blank, so it never - // needs the fallback. - // `!truncated_at_token_limit` is the second load-bearing guard. "Otherwise - // blank" holds for a model that FINISHED thinking and wrote no answer, not one - // CUT OFF mid-thought — `stella_core::driver::truncation` handles that case and - // can only recognize it while `text` is still empty. Promoting anyway published - // an abandoned scratchpad as the answer, which the driver then kept for the turn. - if text.trim().is_empty() - && calls.is_empty() - && !reasoning.trim().is_empty() - && !truncated_at_token_limit - { - text = reasoning; - } - - let finish_reason = if truncated_at_token_limit { - Some(FinishReason::Length) - } else if !calls.is_empty() { - Some(FinishReason::ToolCalls) - } else { - Some(FinishReason::Stop) - }; - - Ok((text, calls, usage, finish_reason, reported_cost_usd)) } #[cfg(test)] diff --git a/crates/stella-model/src/zai/stream.rs b/crates/stella-model/src/zai/stream.rs new file mode 100644 index 000000000..5d4a1e6e9 --- /dev/null +++ b/crates/stella-model/src/zai/stream.rs @@ -0,0 +1,506 @@ +//! The streaming half of the shared OpenAI-compatible adapter: SSE +//! aggregation for one completion, split out of `zai.rs` (a grandfathered +//! god file closed to growth) when the streaming→non-streaming fallback +//! (#2686) landed. +//! +//! [`aggregate_zai_stream`] consumes one `chat/completions` SSE body into +//! the adapter's result tuple. Its error type is [`StreamFault`] rather than +//! a bare `ProviderError`, because the caller needs one extra bit: whether +//! the fault is **fallback-eligible** — the stream hung before its first +//! byte, or ended with nothing accumulated at all — which is what arms the +//! [`crate::stream_recovery::StreamRecovery`] latch so the retry of this +//! attempt goes out as a unary request. Every other failure (a mid-stream +//! death with content already salvaged, a malformed frame, an in-band error +//! frame) keeps its existing classification and never arms the latch. +//! +//! The per-call assembly rules the streaming and unary +//! ([`super::unary`]) paths must agree on — argument-JSON repair, the +//! reasoning-only text promotion, usage folding, the final finish reason — +//! live here as shared helpers so the two paths cannot drift. + +use std::collections::BTreeMap; +use std::time::Duration; + +use serde_json::Value; +use stella_protocol::{CompletionUsage, FinishReason, ProviderError, ToolCall}; + +use super::{ZaiUsage, classify_zai_stream_error}; +use crate::catalog::Pricing; +use crate::http; +use crate::provider::ToolCallObserver; +use crate::sse::SseDecoder; +use crate::stream_recovery::StreamFault; + +/// The stream chunk shape, deserialized per SSE `data:` frame — see the +/// field docs in `zai.rs` for the dialect quirks each one absorbs. +use super::ZaiStreamChunk; + +/// Fold one usage frame into the normalized envelope. Shared verbatim by +/// the streaming aggregator (final usage frame) and the unary fallback +/// (response-level `usage` object), so the two paths cannot disagree about +/// how this dialect's cache/reasoning telemetry is read. +pub(super) fn fold_usage( + frame: ZaiUsage, + usage: &mut CompletionUsage, + reported_cost_usd: &mut Option, +) { + usage.input_tokens = frame.prompt_tokens; + usage.output_tokens = frame.completion_tokens; + let details = frame.prompt_tokens_details.unwrap_or_default(); + // Two wire spellings for the same fact: the OpenAI-compatible + // details object, or DeepSeek's native top-level field — take + // whichever the server spoke (never both on one endpoint). + usage.cached_input_tokens = if details.cached_tokens > 0 { + details.cached_tokens + } else { + frame.prompt_cache_hit_tokens + }; + usage.cache_write_tokens = details.cache_write_tokens; + // Only overwrite when this frame carried the breakdown: a + // later usage frame without the detail object must not erase + // a count an earlier one reported. + if let Some(reasoning) = frame + .completion_tokens_details + .and_then(|d| d.reasoning_tokens) + { + usage.reasoning_tokens = Some(reasoning); + } + if frame.cost.is_some() { + *reported_cost_usd = frame.cost; + } +} + +/// Parse one completed tool call's accumulated argument string under this +/// dialect's repair rules — shared by both delivery paths. +/// +/// An empty body is an empty *object*, not null — a downstream tool +/// deserializing its input as an object must not be handed `null`. When the +/// token limit cut the call (`truncated`), executing it anyway would fail on +/// missing parameters and re-enter the same unwinnable retry-retruncate +/// loop as a mid-payload cut, so it is a terminal error either way the cut +/// landed. A *non-truncated* body that fails to parse is the model's own +/// broken JSON (GLM emits these): fall back to the `Value::Null` sentinel +/// `driver.rs::execute_with_repair` checks for, so the repair loop can ask +/// the model to retry. Mirrors anthropic.rs. +pub(super) fn tool_call_input( + label: &str, + name: &str, + raw: &str, + truncated: bool, +) -> Result { + let trimmed = raw.trim(); + if trimmed.is_empty() { + if truncated { + return Err(http::truncated_tool_input_error( + label, + name, + "", + "finish_reason=length", + )); + } + return Ok(Value::Object(serde_json::Map::new())); + } + match serde_json::from_str(trimmed) { + Ok(value) => Ok(value), + Err(_) if truncated => Err(http::truncated_tool_input_error( + label, + name, + trimmed, + "finish_reason=length", + )), + Err(_) => Ok(Value::Null), + } +} + +/// Reasoning-only fallback: whether to surface the chain-of-thought as the +/// answer text so the turn is never blank. Shared by both delivery paths. +/// +/// `calls.is_empty()` is load-bearing, not defensive. A tool-calling turn +/// legitimately has empty `content` — that is the *normal* shape for every +/// Anthropic model routed through OpenRouter, which streams `content: ""` +/// alongside `reasoning` and the tool call. Without this guard the fallback +/// fired on essentially every reasoning-model tool turn and published the +/// model's private chain-of-thought as its user-visible answer. A turn that +/// called a tool is not blank, so it never needs the fallback. +/// `!truncated` is the second load-bearing guard. "Otherwise blank" holds +/// for a model that FINISHED thinking and wrote no answer, not one CUT OFF +/// mid-thought — `stella_core::driver::truncation` handles that case and can +/// only recognize it while `text` is still empty. Promoting anyway published +/// an abandoned scratchpad as the answer. +pub(super) fn promote_reasoning_as_text( + text: &str, + calls: &[ToolCall], + reasoning: &str, + truncated: bool, +) -> bool { + text.trim().is_empty() && calls.is_empty() && !reasoning.trim().is_empty() && !truncated +} + +/// The finish reason both delivery paths report from the same two facts. +pub(super) fn final_finish_reason(truncated: bool, has_calls: bool) -> Option { + if truncated { + Some(FinishReason::Length) + } else if has_calls { + Some(FinishReason::ToolCalls) + } else { + Some(FinishReason::Stop) + } +} + +/// Accumulator for one in-progress streamed tool call, keyed by the +/// provider's `index` field until it's complete. +#[derive(Default)] +struct ToolCallAccumulator { + id: String, + name: String, + arguments: String, + /// Whether this call was already announced to a [`ToolCallObserver`]. + /// OpenAI-style tool calls stream sequentially by index, so a call is + /// complete the moment a HIGHER index appears — that boundary announces + /// it exactly once. + /// + /// The stream's LAST call has no higher index behind it, so its boundary + /// is the chunk that carries `finish_reason: "tool_calls"` instead — + /// which arrives before the final usage frame and `[DONE]`, with an + /// end-of-stream fallback for a server that jumps straight to `[DONE]`. + /// Without that boundary a turn that makes exactly one tool call — the + /// common shape for this agent — was never announced at all, and + /// `stella-core`'s speculative execution never fired for any single-call + /// turn on this dialect. The window is smaller than what the dialects + /// with a per-call terminator get (`response.function_call_arguments.done` + /// on `openai.rs`, `content_block_stop` on `anthropic.rs`, whole + /// `functionCall` parts on `gemini.rs`), but not empty. This flag is what + /// keeps every boundary exactly-once: a call announced at one boundary is + /// skipped at every later one. + announced: bool, +} + +/// Announce every un-announced accumulator below `next_index` to the +/// observer. Only calls whose arguments already parse are announced — a +/// call the end-of-stream assembly would hand the `Null` repair sentinel +/// must never reach speculative execution. Announced calls re-parse the +/// same bytes at final assembly, so an announced call and its committed +/// twin are structurally identical. +fn announce_completed_below( + observer: &dyn ToolCallObserver, + tool_calls: &mut BTreeMap, + next_index: usize, +) { + for (_, acc) in tool_calls.range_mut(..next_index) { + if acc.announced { + continue; + } + acc.announced = true; + if acc.id.is_empty() { + continue; + } + let trimmed = acc.arguments.trim(); + let input = if trimmed.is_empty() { + Some(Value::Object(serde_json::Map::new())) + } else { + serde_json::from_str(trimmed).ok() + }; + if let Some(input) = input { + observer.tool_call_streamed(&ToolCall { + call_id: acc.id.clone(), + name: acc.name.clone(), + input, + }); + } + } +} + +pub(super) async fn aggregate_zai_stream( + response: reqwest::Response, + label: &str, + observer: Option<&dyn ToolCallObserver>, + pricing: Option<&Pricing>, + first_byte: Duration, +) -> Result< + ( + String, + Vec, + CompletionUsage, + Option, + Option, + ), + StreamFault, +> { + let mut decoder = SseDecoder::new(); + let mut text = String::new(); + // Chain-of-thought streamed under `reasoning_content`, kept separate from + // the answer. Used only as a fallback when `content` never arrives, so a + // reasoning-only turn is visible instead of blank. + let mut reasoning = String::new(); + let mut usage = CompletionUsage::default(); + let mut tool_calls: BTreeMap = BTreeMap::new(); + // Set once any choice reports `finish_reason: "length"` — the output was + // cut off at the token limit, so a tool call whose argument JSON didn't + // finish streaming is truncated, not merely malformed. + let mut truncated_at_token_limit = false; + // Gateway-reported per-call cost (OpenRouter usage accounting), from the + // final usage frame. `None` when the endpoint doesn't report one. + let mut reported_cost_usd: Option = None; + let mut usage_seen = false; + let mut terminal_seen = false; + // The first body read runs against the (shorter) first-byte deadline + // rather than the inter-fragment idle bound: a response that has sent + // its headers and then not one body byte is a buffering proxy, not a + // thinking model (#2686). Any chunk at all — even a keep-alive — moves + // the stream onto the ordinary idle clock. + let mut awaiting_first_chunk = true; + let mut stream = response.bytes_stream(); + + loop { + let idle = if awaiting_first_chunk { + first_byte + } else { + http::STREAM_IDLE_TIMEOUT + }; + let chunk = match http::next_stream_read(&mut stream, idle).await { + http::StreamRead::Item(chunk) => chunk, + http::StreamRead::End => break, + // A transport fault (reset, TLS error) is NOT fallback-eligible: + // it says nothing about the streaming path specifically, and the + // ordinary retry may well succeed over the same stream. + http::StreamRead::Failed(message) => { + return Err(StreamFault::ineligible(http::attach_partial( + ProviderError::transport(message), + &usage, + &text, + pricing, + ))); + } + http::StreamRead::Idle => { + let error = if awaiting_first_chunk { + ProviderError::transport(format!( + "{label} stream hung before its first byte: no data within \ + the {}s first-byte deadline", + idle.as_secs() + )) + } else { + ProviderError::transport(format!( + "stream idle timeout: no data for {}s", + idle.as_secs() + )) + }; + // A hang is fallback-eligible exactly when there is nothing + // to lose: no content, no usage frame, no terminal marker. + // A stream that hung after real content is the existing + // mid-stream-death case — retried as a stream, its salvage + // attached. + let fallback_eligible = !terminal_seen + && !usage_seen + && text.is_empty() + && reasoning.is_empty() + && tool_calls.is_empty(); + return Err(StreamFault { + fallback_eligible, + error: http::attach_partial(error, &usage, &text, pricing), + }); + } + }; + awaiting_first_chunk = false; + decoder + .push_bytes(&chunk) + .map_err(|e| ProviderError::Malformed(e.to_string()))?; + for event in decoder.poll() { + let data = event.data.trim(); + if data.is_empty() { + continue; + } + if data == "[DONE]" { + terminal_seen = true; + continue; + } + // Two failure modes hide behind one `from_str`, and they must not + // share a branch. A frame that is not JSON at all is a keep-alive + // or a gateway comment: ignoring it is correct. A frame that IS + // valid JSON but does not fit the chunk shape is a dialect + // deviation, and swallowing it loses whatever it carried. Since + // every field in this tree defaults, an unknown or empty object + // still parses — so reaching the error arm means a real type + // mismatch on a field that matters, and the likeliest casualty is + // `tool_calls`. Dropping that produces a turn with no calls, which + // the driver reads as a clean completion: the run reports success + // having done nothing. The `error` field above documents this exact + // failure happening once already. Fail loudly instead. + let Ok(value) = serde_json::from_str::(data) else { + continue; // not JSON: keep-alive / ping / gateway comment + }; + let parsed: ZaiStreamChunk = match serde_json::from_value(value) { + Ok(v) => v, + Err(e) => { + return Err(ProviderError::Malformed(format!( + "{label}: unparseable stream frame ({e}); refusing to \ + treat a dropped frame as an empty turn" + )) + .into()); + } + }; + // A mid-stream error frame aborts the turn with a typed error — + // never a truncated Ok with the partial text seen so far. + if let Some(err) = &parsed.error { + return Err(classify_zai_stream_error(err, label).into()); + } + if let Some(u) = parsed.usage { + usage_seen = true; + fold_usage(u, &mut usage, &mut reported_cost_usd); + } + for choice in parsed.choices { + if choice.finish_reason.as_deref() == Some("length") { + truncated_at_token_limit = true; + } + if let Some(content) = choice.delta.content { + // `content` is the user-visible answer; thinking rides + // `reasoning_delta` below. The two channels stay separate + // all the way to the transcript. + if let Some(observer) = observer { + observer.text_delta(&content); + } + text.push_str(&content); + } + // GLM's name for chain-of-thought and OpenRouter's normalized + // one. Both announce as thinking, so a reasoning model's + // deliberation is visible live (collapsed, dimmed) instead of + // the turn looking idle until the answer lands. + if let Some(rc) = choice.delta.reasoning_content { + if let Some(observer) = observer { + observer.reasoning_delta(&rc); + } + reasoning.push_str(&rc); + } + if let Some(r) = choice.delta.reasoning { + if let Some(observer) = observer { + observer.reasoning_delta(&r); + } + reasoning.push_str(&r); + } + for (position, tc_delta) in choice + .delta + .tool_calls + .unwrap_or_default() + .into_iter() + .enumerate() + { + // Absent `index` falls back to the fragment's position in + // this chunk — see `ZaiStreamToolCallDelta`. + let index = tc_delta.index.unwrap_or(position); + // A delta for index N proves every lower index finished + // streaming (the dialect emits calls sequentially) — + // the moment those calls can be announced for + // speculative execution. + if let Some(observer) = observer { + announce_completed_below(observer, &mut tool_calls, index); + } + let acc = tool_calls.entry(index).or_default(); + if let Some(id) = tc_delta.id { + acc.id = id; + } + if let Some(function) = tc_delta.function { + if let Some(name) = function.name { + // Accumulated, not assigned: the dialect contract + // (pinned by + // `complete_reassembles_a_streamed_tool_call_split_across_many_chunks`) + // allows the name to arrive in fragments like the + // arguments do, so append is the correct + // reassembly. A gateway that instead repeats the + // WHOLE name per fragment would garble this — no + // such gateway has been observed; if one appears, + // dedupe the exact-repeat case rather than + // switching to last-wins. + acc.name.push_str(&name); + } + if let Some(args) = function.arguments { + // Liveness only (see + // `ToolCallObserver::tool_input_delta`): a + // call-only generation must still register as + // producing against the idle deadline. + if let Some(observer) = observer { + observer.tool_input_delta(); + } + acc.arguments.push_str(&args); + } + } + } + // The chunk carrying `finish_reason: "tool_calls"` is this + // dialect's end-of-calls terminator: the LAST call has no + // higher index behind it, so this is its completion boundary. + // Announcing here — before the final usage frame and `[DONE]` + // — is what opens the speculative-execution window for a + // single-call turn; `announced` keeps calls already announced + // at an index boundary from repeating. + if choice.finish_reason.as_deref() == Some("tool_calls") + && let Some(observer) = observer + { + announce_completed_below(observer, &mut tool_calls, usize::MAX); + } + } + } + } + + // EOF without the `[DONE]` sentinel is a disconnect, not a completion — + // whatever accumulated is a half-answer, and even a `finish_reason` seen + // earlier can't prove the stream wasn't cut after it. Retryable + // Transport, upholding the same "never a truncated Ok" promise as the + // mid-stream error-frame path above. When NOTHING accumulated it is the + // other broken-stream shape #2686 names — a gateway answering 200 with + // an empty stream — and is fallback-eligible: the retry loses nothing by + // going out unary. + if !terminal_seen { + let fallback_eligible = + !usage_seen && text.is_empty() && reasoning.is_empty() && tool_calls.is_empty(); + return Err(StreamFault { + fallback_eligible, + error: http::attach_partial( + http::stream_ended_before_terminal(label, "[DONE]"), + &usage, + &text, + pricing, + ), + }); + } + + // Fallback terminator: a server that ends the stream without ever sending + // a `finish_reason: "tool_calls"` chunk still completed its last call — + // `[DONE]` proves the stream is whole, so announce whatever is still open + // before final assembly (`announced` makes this a no-op when the finish + // chunk already fired). Skipped when the token limit cut the stream: a + // truncated call must never reach speculative execution, and final + // assembly below turns it into a terminal error instead. + if !truncated_at_token_limit && let Some(observer) = observer { + announce_completed_below(observer, &mut tool_calls, usize::MAX); + } + + usage.reported = usage_seen; + + // OpenAI-style tool calls stream sequentially by index, so when the + // stream reports `finish_reason: "length"` only the highest-index call + // can be the one the token limit cut. Pinning truncation there keeps the + // blame on the right call — an earlier call whose JSON is broken is the + // model's own malformed output and still gets the repair sentinel. + let truncated_index = if truncated_at_token_limit { + tool_calls.keys().next_back().copied() + } else { + None + }; + + let mut calls = Vec::with_capacity(tool_calls.len()); + for (index, acc) in tool_calls { + let truncated = Some(index) == truncated_index; + let input = tool_call_input(label, &acc.name, &acc.arguments, truncated)?; + calls.push(ToolCall { + call_id: acc.id, + name: acc.name, + input, + }); + } + + if promote_reasoning_as_text(&text, &calls, &reasoning, truncated_at_token_limit) { + text = reasoning; + } + + let finish_reason = final_finish_reason(truncated_at_token_limit, !calls.is_empty()); + + Ok((text, calls, usage, finish_reason, reported_cost_usd)) +} diff --git a/crates/stella-model/src/zai/tests.rs b/crates/stella-model/src/zai/tests.rs index 4818115de..d30ac66f3 100644 --- a/crates/stella-model/src/zai/tests.rs +++ b/crates/stella-model/src/zai/tests.rs @@ -1737,6 +1737,7 @@ fn xai_reasoning_effort_support_denies_only_the_original_grok4() { mod error_classify; mod openrouter_effort; mod openrouter_stream; +mod stream_fallback; mod stream_frame; mod vision; mod zai_effort; @@ -1858,38 +1859,3 @@ fn a_caps_flip_degrades_instead_of_aborting_the_turn() { ); } } - -/// The output allowance an un-capped reasoning turn gets, and the parity that -/// makes it correct: `anthropic.rs` picks 32k for a thinking model and this -/// adapter must not disagree. Reasoning spends from the same allowance as the -/// answer, so a provider's own non-thinking-sized default truncates the turn -/// mid-thought — and on this dialect that lands as a `finish_reason: length` -/// step carrying chain-of-thought and no tool call, the exact shape that ended -/// 30 of 30 cap-hit steps in the 2026-07-31 Terminal-Bench bundle. -#[test] -fn an_uncapped_reasoning_turn_gets_thinking_headroom() { - assert_eq!(reasoning_aware_max_tokens(None, Some(true)), Some(32_000)); -} - -/// A caller who pinned a cap keeps it verbatim — the default only fills the -/// hole where no preference was expressed. Overriding an explicit ceiling -/// would silently spend a caller's money on a budget they declined. -#[test] -fn a_pinned_cap_is_honored_whatever_the_reasoning_setting() { - for reasoning in [None, Some(false), Some(true)] { - assert_eq!( - reasoning_aware_max_tokens(Some(8_192), reasoning), - Some(8_192), - "an explicit cap survives reasoning={reasoning:?}" - ); - } -} - -/// Reasoning off (or unstated) leaves the field absent, so a request without -/// a cap serializes byte-identical to what this adapter has always sent — -/// the prompt-cache stability contract the params work established. -#[test] -fn a_non_reasoning_turn_still_sends_no_cap() { - assert_eq!(reasoning_aware_max_tokens(None, Some(false)), None); - assert_eq!(reasoning_aware_max_tokens(None, None), None); -} diff --git a/crates/stella-model/src/zai/tests/stream_fallback.rs b/crates/stella-model/src/zai/tests/stream_fallback.rs new file mode 100644 index 000000000..bdf98685e --- /dev/null +++ b/crates/stella-model/src/zai/tests/stream_fallback.rs @@ -0,0 +1,292 @@ +// SPDX-License-Identifier: AGPL-3.0-only +//! The streaming→non-streaming fallback (#2686): a stream that hangs before +//! its first byte or comes back empty fails retryably, arms the per-session +//! latch, and the retried attempt re-issues the byte-identical payload with +//! `stream: false`. Split out of the parent `tests.rs` (file-size gate), +//! mirroring the other delivery-path suites. + +use super::*; + +fn plain_request() -> CompletionRequest { + CompletionRequest { + messages: vec![CompletionMessage::user("say hello")], + max_output_tokens: None, + temperature: None, + effort: None, + tools: vec![], + reasoning: None, + params: None, + } +} + +/// Read one HTTP request (headers + `Content-Length` body) off a blocking +/// socket — just enough parser for the hand-rolled server below. +fn read_request(socket: &mut std::net::TcpStream) -> String { + use std::io::Read; + let mut buf = Vec::new(); + let mut chunk = [0u8; 4096]; + loop { + let n = socket.read(&mut chunk).unwrap_or(0); + if n == 0 { + break; + } + buf.extend_from_slice(&chunk[..n]); + let text = String::from_utf8_lossy(&buf); + if let Some(header_end) = text.find("\r\n\r\n") { + let content_length = text + .lines() + .find_map(|line| { + let (name, value) = line.split_once(':')?; + name.eq_ignore_ascii_case("content-length") + .then(|| value.trim().parse::().ok()) + .flatten() + }) + .unwrap_or(0); + if buf.len() >= header_end + 4 + content_length { + break; + } + } + } + String::from_utf8_lossy(&buf).into_owned() +} + +/// The one server shape wiremock cannot express, hand-rolled on a std +/// thread: a streaming request is answered with its HTTP headers and then +/// **not one body byte** — the socket is held open, exactly what a proxy +/// buffering the SSE body looks like from the client side. A non-streaming +/// request for the same path completes normally with `unary_body`. Returns +/// the base URL. +fn hang_streams_answer_unary(unary_body: &'static str) -> String { + use std::io::Write; + let listener = std::net::TcpListener::bind("127.0.0.1:0").expect("bind loopback"); + let addr = listener.local_addr().expect("local addr"); + std::thread::spawn(move || { + // Held open so the client sees a live-but-silent stream, never an + // EOF (which would be the *empty stream* shape instead). + let mut held = Vec::new(); + for conn in listener.incoming() { + let Ok(mut socket) = conn else { break }; + let request = read_request(&mut socket); + if request.contains("\"stream\":true") { + let _ = + socket.write_all(b"HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\n\r\n"); + let _ = socket.flush(); + held.push(socket); + } else { + let _ = write!( + socket, + "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\n\ + content-length: {}\r\nconnection: close\r\n\r\n{}", + unary_body.len(), + unary_body + ); + } + } + }); + format!("http://{addr}") +} + +/// The #2686 witness: a stream that hangs before its first byte fails the +/// attempt at the first-byte deadline — retryably, so the ordinary retry +/// machinery re-drives it — and that retry completes as a non-streaming +/// request for the same payload. On main the first call instead waited the +/// full 120s idle bound and every retry re-issued the identical streaming +/// request into the same buffering proxy. +#[tokio::test] +async fn a_stream_hung_before_its_first_byte_falls_back_to_a_non_streaming_request() { + let base_url = hang_streams_answer_unary( + r#"{"id":"cmpl-1","choices":[{"message":{"content":"recovered without streaming"},"finish_reason":"stop"}],"usage":{"prompt_tokens":8,"completion_tokens":5}}"#, + ); + let provider = ZaiProvider::new(ApiKey::new("sk-test-zai"), "glm-5.2") + .with_base_url(base_url) + .with_first_byte_deadline(Duration::from_millis(120)); + + let error = provider + .complete(plain_request()) + .await + .expect_err("a hung stream must fault at the first-byte deadline, not hang"); + assert!( + error.is_retryable(), + "the fault must be retryable so the retry ladder drives the fallback: {error:?}" + ); + let message = error.to_string(); + assert!(message.contains("first byte"), "{message}"); + assert!( + message.contains("non-streaming"), + "the error names the switch it armed: {message}" + ); + + // The engine retries a retryable fault through the same provider + // instance; the armed latch makes that retry unary. + let result = provider + .complete(plain_request()) + .await + .expect("the fallback attempt completes without streaming"); + assert_eq!(result.text, "recovered without streaming"); + assert_eq!(result.usage.input_tokens, 8); + assert_eq!(result.usage.output_tokens, 5); + assert!(result.usage.reported); + assert_eq!(result.finish_reason, Some(FinishReason::Stop)); +} + +/// The other broken-stream shape: a gateway answering 200 with an empty +/// stream (EOF before any data). Same latch, and the unary response here +/// carries a tool call, proving the fallback path parses the whole dialect +/// — not just text. +#[tokio::test] +async fn an_empty_stream_falls_back_to_a_non_streaming_request() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/chat/completions")) + .and(body_string_contains("\"stream\":true")) + .respond_with(ResponseTemplate::new(200).set_body_raw("", "text/event-stream")) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/chat/completions")) + .and(body_string_contains("\"stream\":false")) + .respond_with(ResponseTemplate::new(200).set_body_raw( + r#"{"choices":[{"message":{"content":"","tool_calls":[{"id":"call_1","type":"function","function":{"name":"read_file","arguments":"{\"path\":\"src/lib.rs\"}"}}]},"finish_reason":"tool_calls"}],"usage":{"prompt_tokens":12,"completion_tokens":9}}"#, + "application/json", + )) + .mount(&server) + .await; + + let provider = + ZaiProvider::new(ApiKey::new("sk-test-zai"), "glm-5.2").with_base_url(server.uri()); + + let error = provider + .complete(plain_request()) + .await + .expect_err("an empty stream is a fault, never an empty Ok"); + assert!(error.is_retryable(), "{error:?}"); + assert!(error.to_string().contains("non-streaming"), "{error}"); + + let result = provider + .complete(plain_request()) + .await + .expect("the fallback attempt completes"); + assert_eq!(result.tool_calls.len(), 1); + assert_eq!(result.tool_calls[0].name, "read_file"); + assert_eq!( + result.tool_calls[0].input, + serde_json::json!({"path": "src/lib.rs"}) + ); + assert_eq!(result.finish_reason, Some(FinishReason::ToolCalls)); + assert!(result.usage.reported); +} + +/// Control: a healthy stream never arms the latch — every request of the +/// session keeps `stream: true` on the wire, so the fallback costs nothing +/// when nothing is broken. +#[tokio::test] +async fn a_healthy_stream_never_arms_the_fallback() { + let server = MockServer::start().await; + let sse_body = concat!( + "data: {\"choices\":[{\"delta\":{\"content\":\"Hello!\"}}]}\n\n", + "data: {\"choices\":[{\"delta\":{}}],\"usage\":{\"prompt_tokens\":8,\"completion_tokens\":3}}\n\n", + "data: [DONE]\n\n", + ); + Mock::given(method("POST")) + .and(path("/chat/completions")) + .respond_with(ResponseTemplate::new(200).set_body_raw(sse_body, "text/event-stream")) + .mount(&server) + .await; + + let provider = + ZaiProvider::new(ApiKey::new("sk-test-zai"), "glm-5.2").with_base_url(server.uri()); + for _ in 0..2 { + let result = provider + .complete(plain_request()) + .await + .expect("healthy streams complete"); + assert_eq!(result.text, "Hello!"); + } + + let sent = server.received_requests().await.expect("recorded requests"); + assert_eq!(sent.len(), 2); + for request in &sent { + let body = String::from_utf8_lossy(&request.body); + assert!( + body.contains("\"stream\":true"), + "no request may lose streaming on a healthy session: {body}" + ); + } +} + +/// A stream that died AFTER delivering content is the existing mid-stream +/// death, not fallback material: there is salvage to bill and the retry +/// stays a stream. Guards the eligibility boundary against widening into +/// "any stream error goes unary". +#[tokio::test] +async fn a_stream_that_died_after_content_is_retried_as_a_stream_not_unary() { + let server = MockServer::start().await; + // A real delta arrives, then the connection closes without `[DONE]`. + let sse_body = "data: {\"choices\":[{\"delta\":{\"content\":\"Hel\"}}]}\n\n"; + Mock::given(method("POST")) + .and(path("/chat/completions")) + .respond_with(ResponseTemplate::new(200).set_body_raw(sse_body, "text/event-stream")) + .mount(&server) + .await; + + let provider = + ZaiProvider::new(ApiKey::new("sk-test-zai"), "glm-5.2").with_base_url(server.uri()); + for _ in 0..2 { + let error = provider + .complete(plain_request()) + .await + .expect_err("a half-answer must never commit as Ok"); + assert!( + !error.to_string().contains("non-streaming"), + "salvageable deaths must not arm the fallback: {error}" + ); + } + + let sent = server.received_requests().await.expect("recorded requests"); + assert_eq!(sent.len(), 2); + for request in &sent { + let body = String::from_utf8_lossy(&request.body); + assert!(body.contains("\"stream\":true"), "{body}"); + } +} + +/// A probe that fails reverts the latch: when the unary retry ALSO fails, +/// the fault evidently wasn't the streaming path's, and the session must +/// not stay pinned to a transport it has no evidence for. The wire +/// sequence proves the full cycle: stream → unary probe → stream again. +#[tokio::test] +async fn a_failed_unary_probe_reverts_the_session_to_streaming() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/chat/completions")) + .and(body_string_contains("\"stream\":true")) + .respond_with(ResponseTemplate::new(200).set_body_raw("", "text/event-stream")) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/chat/completions")) + .and(body_string_contains("\"stream\":false")) + .respond_with(ResponseTemplate::new(502).set_body_string("upstream down")) + .mount(&server) + .await; + + let provider = + ZaiProvider::new(ApiKey::new("sk-test-zai"), "glm-5.2").with_base_url(server.uri()); + for _ in 0..3 { + let _ = provider + .complete(plain_request()) + .await + .expect_err("every arm of this server fails"); + } + + let sent = server.received_requests().await.expect("recorded requests"); + let flags: Vec = sent + .iter() + .map(|request| String::from_utf8_lossy(&request.body).contains("\"stream\":true")) + .collect(); + assert_eq!( + flags, + vec![true, false, true], + "stream fault → unary probe → probe failed → back to streaming" + ); +} diff --git a/crates/stella-model/src/zai/tests/zai_effort.rs b/crates/stella-model/src/zai/tests/zai_effort.rs index f5d3f106e..c2dfb8a36 100644 --- a/crates/stella-model/src/zai/tests/zai_effort.rs +++ b/crates/stella-model/src/zai/tests/zai_effort.rs @@ -124,3 +124,38 @@ fn map_zai_effort_never_yields_minimal() { assert_eq!(map_zai_effort(ReasoningEffort::Xhigh), "high"); assert_eq!(map_zai_effort(ReasoningEffort::Max), "high"); } + +/// The output allowance an un-capped reasoning turn gets, and the parity that +/// makes it correct: `anthropic.rs` picks 32k for a thinking model and this +/// adapter must not disagree. Reasoning spends from the same allowance as the +/// answer, so a provider's own non-thinking-sized default truncates the turn +/// mid-thought — and on this dialect that lands as a `finish_reason: length` +/// step carrying chain-of-thought and no tool call, the exact shape that ended +/// 30 of 30 cap-hit steps in the 2026-07-31 Terminal-Bench bundle. +#[test] +fn an_uncapped_reasoning_turn_gets_thinking_headroom() { + assert_eq!(reasoning_aware_max_tokens(None, Some(true)), Some(32_000)); +} + +/// A caller who pinned a cap keeps it verbatim — the default only fills the +/// hole where no preference was expressed. Overriding an explicit ceiling +/// would silently spend a caller's money on a budget they declined. +#[test] +fn a_pinned_cap_is_honored_whatever_the_reasoning_setting() { + for reasoning in [None, Some(false), Some(true)] { + assert_eq!( + reasoning_aware_max_tokens(Some(8_192), reasoning), + Some(8_192), + "an explicit cap survives reasoning={reasoning:?}" + ); + } +} + +/// Reasoning off (or unstated) leaves the field absent, so a request without +/// a cap serializes byte-identical to what this adapter has always sent — +/// the prompt-cache stability contract the params work established. +#[test] +fn a_non_reasoning_turn_still_sends_no_cap() { + assert_eq!(reasoning_aware_max_tokens(None, Some(false)), None); + assert_eq!(reasoning_aware_max_tokens(None, None), None); +} diff --git a/crates/stella-model/src/zai/unary.rs b/crates/stella-model/src/zai/unary.rs new file mode 100644 index 000000000..ae83f8597 --- /dev/null +++ b/crates/stella-model/src/zai/unary.rs @@ -0,0 +1,192 @@ +//! The non-streaming delivery path of the shared OpenAI-compatible adapter — +//! the fallback a completion takes once the session's streaming path has +//! proven broken (hung before its first byte, or a 200 with an empty +//! stream; #2686, `crate::stream_recovery`). +//! +//! Same endpoint, same payload, `stream: false`: the response is one JSON +//! `chat.completion` object instead of an SSE body, dispatched through the +//! [`crate::http::unary_client`] read bound (the whole generation must fit +//! inside a single read — #547). Assembly reuses [`super::stream`]'s shared +//! rules verbatim, so the two paths cannot disagree on argument-JSON repair, +//! the reasoning-only promotion, usage folding, or the finish reason. The +//! price of this path is the loss of mid-stream observation: no text/ +//! reasoning previews and no speculative tool execution — which is why the +//! latch only confirms on evidence and never arms on an ordinary transport +//! fault. + +use serde::Deserialize; +use serde_json::Value; +use stella_protocol::{ + CompletionRequestRef, CompletionResult, CompletionUsage, ProviderError, ToolCall, +}; + +use super::{ZaiProvider, ZaiUsage, classify_zai_stream_error, stream}; + +/// One non-streamed `chat.completion` response. Every field defaults for +/// the same reason the stream-chunk tree's do: an unknown or partial object +/// must degrade to an explicit error below, never fail deserialization into +/// a silently empty turn. +#[derive(Deserialize, Debug, Default)] +struct ZaiUnaryResponse { + #[serde(default)] + choices: Vec, + #[serde(default)] + usage: Option, + /// Defensive: a gateway that reports an error inside a 200 body (the + /// unary sibling of the in-band SSE error frame). Classified by the same + /// shared classifier so the two paths agree on retryability. + #[serde(default)] + error: Option, +} + +#[derive(Deserialize, Debug, Default)] +struct ZaiUnaryChoice { + #[serde(default)] + message: ZaiUnaryMessage, + #[serde(default)] + finish_reason: Option, +} + +/// The assembled assistant message: the unary spelling of everything the +/// stream delivers as deltas, including both chain-of-thought field names +/// (GLM's `reasoning_content`, OpenRouter's normalized `reasoning`). +#[derive(Deserialize, Debug, Default)] +struct ZaiUnaryMessage { + #[serde(default)] + content: Option, + #[serde(default)] + reasoning_content: Option, + #[serde(default)] + reasoning: Option, + #[serde(default)] + tool_calls: Option>, +} + +#[derive(Deserialize, Debug, Default)] +struct ZaiUnaryToolCall { + #[serde(default)] + id: String, + #[serde(default)] + function: ZaiUnaryFunction, +} + +#[derive(Deserialize, Debug, Default)] +struct ZaiUnaryFunction { + #[serde(default)] + name: String, + #[serde(default)] + arguments: String, +} + +impl ZaiProvider { + /// One unary request/parse cycle — [`super::ZaiProvider::complete_attempt`]'s + /// other arm. Reports its outcome to the recovery latch: success while + /// probing confirms the fallback for the session, failure reverts to + /// streaming (see `crate::stream_recovery` for why neither is skipped). + pub(super) async fn complete_unary_attempt( + &self, + req: CompletionRequestRef<'_>, + force_default_reasoning: bool, + ) -> Result { + let body = self.build_body(req, force_default_reasoning, false); + let outcome = async { + let response = self.dispatch(&self.unary_client, &body).await?; + let payload = response + .text() + .await + .map_err(|e| ProviderError::transport(e.to_string()))?; + self.assemble_unary(&payload) + } + .await; + self.recovery.note_unary_outcome(outcome.is_ok()); + outcome + } + + /// Parse one `chat.completion` body into the adapter's result, under + /// exactly the shared assembly rules the streaming path uses. + fn assemble_unary(&self, payload: &str) -> Result { + let label = &self.label; + // Same fail-loudly contract as the stream-frame parse: every field + // defaults, so an error here is a real type mismatch on a field that + // matters — swallowing it would report a turn that did nothing as a + // clean completion. + let parsed: ZaiUnaryResponse = serde_json::from_str(payload).map_err(|e| { + ProviderError::Malformed(format!( + "{label}: unparseable non-streaming completion body ({e}); refusing \ + to treat a dropped response as an empty turn" + )) + })?; + if let Some(err) = &parsed.error { + return Err(classify_zai_stream_error(err, label)); + } + let Some(choice) = parsed.choices.into_iter().next() else { + // A 200 with no choices is the unary spelling of the empty + // stream. There is no third transport to fall back to, so it is + // an ordinary retryable fault. + return Err(ProviderError::transport(format!( + "{label} returned a completion with no choices" + ))); + }; + + let truncated_at_token_limit = choice.finish_reason.as_deref() == Some("length"); + let text = choice.message.content.unwrap_or_default(); + // Both chain-of-thought spellings, folded exactly as the stream does. + let mut reasoning = choice.message.reasoning_content.unwrap_or_default(); + reasoning.push_str(choice.message.reasoning.as_deref().unwrap_or_default()); + + let unary_calls = choice.message.tool_calls.unwrap_or_default(); + // Unary tool calls arrive whole, so only the LAST one can be the one + // a token limit cut — the same blame-pinning rule as the stream's + // highest-index accumulator. + let truncated_index = truncated_at_token_limit + .then(|| unary_calls.len().checked_sub(1)) + .flatten(); + let mut calls = Vec::with_capacity(unary_calls.len()); + for (index, call) in unary_calls.into_iter().enumerate() { + let truncated = Some(index) == truncated_index; + let input: Value = stream::tool_call_input( + label, + &call.function.name, + &call.function.arguments, + truncated, + )?; + calls.push(ToolCall { + call_id: call.id, + name: call.function.name, + input, + }); + } + + let text = if stream::promote_reasoning_as_text( + &text, + &calls, + &reasoning, + truncated_at_token_limit, + ) { + reasoning + } else { + text + }; + + let mut usage = CompletionUsage::default(); + let mut reported_cost_usd = None; + let usage_seen = parsed.usage.is_some(); + if let Some(frame) = parsed.usage { + stream::fold_usage(frame, &mut usage, &mut reported_cost_usd); + } + usage.reported = usage_seen; + let cost_usd = reported_cost_usd + .unwrap_or_else(|| self.pricing.map(|p| p.cost_usd(&usage)).unwrap_or(0.0)); + + let finish_reason = + stream::final_finish_reason(truncated_at_token_limit, !calls.is_empty()); + Ok(CompletionResult { + text, + tool_calls: calls, + usage, + model: self.model.clone(), + cost_usd, + finish_reason, + }) + } +} diff --git a/scripts/file-size-baseline.txt b/scripts/file-size-baseline.txt index 2cf056578..59d16429a 100644 --- a/scripts/file-size-baseline.txt +++ b/scripts/file-size-baseline.txt @@ -24,7 +24,6 @@ 3384 crates/stella-core/src/driver/tests.rs 1781 crates/stella-model/src/anthropic/tests.rs 2093 crates/stella-model/src/openai.rs -1565 crates/stella-model/src/zai.rs 1895 crates/stella-model/src/zai/tests.rs 3126 crates/stella-pipeline/src/pipeline.rs 2475 crates/stella-pipeline/src/pipeline/tests.rs From 0ee24e6357c02f00e13acdc2f384a52a16c570e0 Mon Sep 17 00:00:00 2001 From: Mac Anderson Date: Mon, 10 Aug 2026 14:53:12 -0700 Subject: [PATCH 2/3] fix(stella-model): classify a unary-fallback read-timeout as Terminal, not retryable Transport MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review catch on #2748: the fallback's unary dispatch mapped every send error to retryable Transport, so once the latch was Confirmed a generation longer than UNARY_READ_TIMEOUT would be re-issued identically until the retry budget died — the exact #547 storm, one adapter over. The unary path now routes send errors through http::classify_unary_dispatch_error, exactly as bedrock.rs does: a read-timeout consumed the whole generation bound and is Terminal; a connect failure generated nothing and stays retryable; the streaming dispatch's classification is untouched (there the same expiry is a header stall the next attempt may clear). Witnessed by a_unary_read_timeout_is_terminal_never_a_retry_storm, which drives dispatch with a millisecond read bound against a stalled mock and pins both directions. Refs #2686 --- crates/stella-model/src/zai.rs | 28 +++++++--- .../src/zai/tests/stream_fallback.rs | 55 +++++++++++++++++++ crates/stella-model/src/zai/unary.rs | 8 ++- 3 files changed, 81 insertions(+), 10 deletions(-) diff --git a/crates/stella-model/src/zai.rs b/crates/stella-model/src/zai.rs index 23953e31e..45aaf91cd 100644 --- a/crates/stella-model/src/zai.rs +++ b/crates/stella-model/src/zai.rs @@ -1082,7 +1082,7 @@ impl ZaiProvider { .await; } let body = self.build_body(req, force_default_reasoning, true); - let response = self.dispatch(&self.client, &body).await?; + let response = self.dispatch(&self.client, &body, false).await?; let (text, tool_calls, usage, finish_reason, reported_cost_usd) = stream::aggregate_zai_stream( response, @@ -1214,12 +1214,22 @@ impl ZaiProvider { /// POST `body` and run the shared non-success ladder (vendor 429 /// pre-check first). Returns the successful response for the caller — - /// streaming or unary — to consume. The two delivery paths differ only - /// in their client's read bound, so `client` is a parameter. + /// streaming or unary — to consume. The two delivery paths differ in + /// their client's read bound AND in what that bound's expiry means, so + /// both ride as parameters. + /// + /// `unary` selects the send-error classification (#547's other half): on + /// the unary client the read bound covers the ENTIRE generation, so its + /// expiry means the request was too long to serve — Terminal, because + /// re-issuing the identical request just waits out the full bound again + /// once per retry. On the streaming client the same expiry is only a + /// header stall (the first token would have reset the clock), which the + /// next attempt may well clear — retryable, as it always was. async fn dispatch( &self, client: &reqwest::Client, body: &ZaiRequest<'_>, + unary: bool, ) -> Result { let mut request = client .post(format!("{}/chat/completions", self.base_url)) @@ -1227,11 +1237,13 @@ impl ZaiProvider { for (name, value) in &self.extra_headers { request = request.header(*name, value); } - let response = request - .json(body) - .send() - .await - .map_err(|e| ProviderError::transport(e.to_string()))?; + let response = request.json(body).send().await.map_err(|e| { + if unary { + http::classify_unary_dispatch_error(&self.label, &e) + } else { + ProviderError::transport(e.to_string()) + } + })?; if response.status() == reqwest::StatusCode::TOO_MANY_REQUESTS { // Vendor pre-check ahead of the shared ladder — Z.ai overloads diff --git a/crates/stella-model/src/zai/tests/stream_fallback.rs b/crates/stella-model/src/zai/tests/stream_fallback.rs index bdf98685e..68aaa5a8e 100644 --- a/crates/stella-model/src/zai/tests/stream_fallback.rs +++ b/crates/stella-model/src/zai/tests/stream_fallback.rs @@ -250,6 +250,61 @@ async fn a_stream_that_died_after_content_is_retried_as_a_stream_not_unary() { } } +/// #547's other half, applied to the fallback path: on the unary client the +/// read bound covers the ENTIRE generation, so its expiry means the request +/// was too long to serve — re-issuing it identically just waits out the full +/// bound again once per retry (the exact storm #547 documented on Bedrock). +/// The unary dispatch must classify a read-timeout as non-retryable +/// `Terminal`, while the SAME expiry on the streaming dispatch stays +/// retryable (there it is only a header stall the next attempt may clear). +/// Exercised in milliseconds by handing `dispatch` a client with a tiny read +/// bound against a server that accepts and then stalls. +#[tokio::test] +async fn a_unary_read_timeout_is_terminal_never_a_retry_storm() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/chat/completions")) + // Accept, then hold the response — the stalled-generation shape, + // compressed so the test costs milliseconds instead of 600s. + .respond_with(ResponseTemplate::new(200).set_delay(Duration::from_secs(30))) + .mount(&server) + .await; + + let provider = + ZaiProvider::new(ApiKey::new("sk-test-zai"), "glm-5.2").with_base_url(server.uri()); + let stalled_client = reqwest::Client::builder() + .read_timeout(Duration::from_millis(50)) + .build() + .expect("client builds"); + let request = plain_request(); + let body = provider.build_body(request.as_borrowed(), false, false); + + let error = provider + .dispatch(&stalled_client, &body, true) + .await + .expect_err("the stalled unary dispatch must time out"); + assert!( + matches!(error, ProviderError::Terminal(_)), + "a unary read timeout must be Terminal, got {error:?}" + ); + assert!( + !error.is_retryable(), + "a retryable unary read timeout re-issues the identical too-long \ + request until the budget dies (#547): {error:?}" + ); + + // The converse, so the fix cannot over-reach: the same expiry on the + // STREAMING dispatch is a header stall and stays retryable. + let error = provider + .dispatch(&stalled_client, &body, false) + .await + .expect_err("the stalled streaming dispatch must time out"); + assert!( + error.is_retryable(), + "a streaming header stall generated nothing and must stay retryable: {error:?}" + ); +} + /// A probe that fails reverts the latch: when the unary retry ALSO fails, /// the fault evidently wasn't the streaming path's, and the session must /// not stay pinned to a transport it has no evidence for. The wire diff --git a/crates/stella-model/src/zai/unary.rs b/crates/stella-model/src/zai/unary.rs index ae83f8597..ca9cd4259 100644 --- a/crates/stella-model/src/zai/unary.rs +++ b/crates/stella-model/src/zai/unary.rs @@ -6,7 +6,11 @@ //! Same endpoint, same payload, `stream: false`: the response is one JSON //! `chat.completion` object instead of an SSE body, dispatched through the //! [`crate::http::unary_client`] read bound (the whole generation must fit -//! inside a single read — #547). Assembly reuses [`super::stream`]'s shared +//! inside a single read — #547) with #547's classification to match: a +//! unary read-timeout consumed the whole bound and re-issuing the identical +//! request just waits it out again, so it surfaces as non-retryable +//! `Terminal`, never as the retryable `Transport` that turned one wedged +//! Bedrock call into four full 600s attempts. Assembly reuses [`super::stream`]'s shared //! rules verbatim, so the two paths cannot disagree on argument-JSON repair, //! the reasoning-only promotion, usage folding, or the finish reason. The //! price of this path is the loss of mid-stream observation: no text/ @@ -90,7 +94,7 @@ impl ZaiProvider { ) -> Result { let body = self.build_body(req, force_default_reasoning, false); let outcome = async { - let response = self.dispatch(&self.unary_client, &body).await?; + let response = self.dispatch(&self.unary_client, &body, true).await?; let payload = response .text() .await From 53d6c35c9281ac9feedeb1c4a78a78d0ff1cda39 Mon Sep 17 00:00:00 2001 From: "vercel[bot]" <35613825+vercel[bot]@users.noreply.github.com> Date: Mon, 10 Aug 2026 22:00:06 +0000 Subject: [PATCH 3/3] Fix: The unary-fallback body read (`response.text()`) still maps a 600s read-timeout to a retryable `Transport` error, reintroducing the #547 retry-storm even after the `.send()` path was fixed. MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This commit fixes the issue reported at crates/stella-model/src/zai/unary.rs:100 ## Status after commit `0ee24e6` The commit landed a partial fix that matches most of the original suggestion: - `ZaiProvider::dispatch` now takes a `unary: bool` flag (`zai.rs:1232`) and classifies the `.send()` failure with `http::classify_unary_dispatch_error` when `unary == true` (`zai.rs:1242`), keeping the streaming path's send timeout retryable (`false`). - Callers were updated: streaming passes `false` (`zai.rs:1085`), unary passes `true` (`unary.rs:97`). ## Remaining bug The **body read** in `complete_unary_attempt` was left unchanged and still uses the retryable mapping: ```rust let payload = response .text() .await .map_err(|e| ProviderError::transport(e.to_string()))?; ``` The unary client is built with `http::unary_client()`, whose `.read_timeout(UNARY_READ_TIMEOUT)` = **600s** bounds the *entire* response — head *and* body arrive within that single read window. If the connection succeeds but the body streams too slowly (a generation slower than 600s, or an LB that accepts then black-holes mid-body), `response.text()` yields an `is_timeout()` error. **Failure mode / concrete trigger:** a Z.ai unary-fallback completion whose body does not fully arrive within 600s. The timeout surfaces as retryable `Transport`, the driver re-issues the *identical* request, and it times out again — repeating until the retry budget is drained. This is exactly the #547 failure mode the `dispatch` send-path fix now guards against; the body read is the last unguarded seam of the same 600s read bound. ## Fix Classify the `response.text()` error with `http::classify_unary_dispatch_error` (timeout → `Terminal`, connect/ordinary transport → retryable `Transport`), matching the already-fixed `dispatch` send path and `bedrock.rs:768`. Co-authored-by: Vercel Co-authored-by: macanderson --- crates/stella-model/src/zai/unary.rs | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/crates/stella-model/src/zai/unary.rs b/crates/stella-model/src/zai/unary.rs index ca9cd4259..1963c3d2b 100644 --- a/crates/stella-model/src/zai/unary.rs +++ b/crates/stella-model/src/zai/unary.rs @@ -95,10 +95,13 @@ impl ZaiProvider { let body = self.build_body(req, force_default_reasoning, false); let outcome = async { let response = self.dispatch(&self.unary_client, &body, true).await?; + // The body arrives inside the same 600s read bound as the head, + // so a timeout here is #547 too — classify it terminal rather than + // retryable, matching the dispatch send path. let payload = response .text() .await - .map_err(|e| ProviderError::transport(e.to_string()))?; + .map_err(|e| crate::http::classify_unary_dispatch_error(&self.label, &e))?; self.assemble_unary(&payload) } .await;