fix(#71): extract SSE streaming passthrough into shared helexa-stream

The true-streaming SSE passthrough (Body::from_stream, no full-response
buffering, with chunk-observation hooks) was cortex-only. helexa-router
(#69) needs the same mechanism to proxy a chat-completions/messages
stream verbatim to a selected cortex. Extract it once.

New `crates/helexa-stream` owns the *mechanism* (kept HTTP-free
cortex-core untouched — it would have forced axum/reqwest/futures onto
every cortex-core consumer):

- `forward_streaming(client, url, headers, body, observer)` — POST and
  stream the response back chunk-for-chunk; status-agnostic, so a
  non-2xx (e.g. cortex 429) is passed through with status+headers
  intact (the #69 backpressure-passthrough requirement).
- `ChunkObserver` trait + `ObservedStream` wrapper — feeds each chunk to
  the observer, calls `finish` exactly once on clean end or on drop
  (client disconnect).
- `BodyTail` (bounded tail accumulator) + `last_count_for` (trailing
  OpenAI `usage` extraction) — the reusable pieces an observer uses.

cortex keeps its *policy*: `proxy.rs` now supplies a `CortexMetrics`
observer (per-request token metrics + per-principal reservation settle),
its logging contract, and the error envelope, driving the shared
mechanism. `proxy::last_count_for` is re-exported so `handlers`/
`anthropic_sse` call sites are unchanged. No behaviour change — the
existing cortex `streaming.rs` tests pass as-is.

helexa-stream tests prove chunk-for-chunk incremental delivery, observer
finish-once, usage extraction, and non-2xx passthrough.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01F6o3ddqmYNh9kzdwq6eowh

Cargo.lock

@@ -800,6 +800,7 @@ dependencies = [
  "cortex-core",
  "eventsource-stream",
  "futures",
+ "helexa-stream",
  "metrics",
  "metrics-exporter-prometheus",
  "reqwest",
@@ -1922,6 +1923,19 @@ dependencies = [
  "tracing-subscriber",
 ]
 
+[[package]]
+name = "helexa-stream"
+version = "0.1.16"
+dependencies = [
+ "async-stream",
+ "axum",
+ "futures",
+ "reqwest",
+ "thiserror 2.0.18",
+ "tokio",
+ "tokio-stream",
+]
+
 [[package]]
 name = "hermit-abi"
 version = "0.5.2"

Cargo.toml

@@ -7,6 +7,7 @@ members = [
     "crates/neuron",
     "crates/helexa-acp",
     "crates/helexa-bench",
+    "crates/helexa-stream",
 ]
 
 [workspace.package]

crates/cortex-gateway/Cargo.toml

@@ -6,6 +6,7 @@ license.workspace = true
 
 [dependencies]
 cortex-core.workspace = true
+helexa-stream = { path = "../helexa-stream" }
 async-trait.workspace = true
 tokio.workspace = true
 axum.workspace = true

crates/cortex-gateway/src/proxy.rs

@@ -1,21 +1,27 @@
 //! Streaming HTTP reverse proxy to neuron backends.
 //!
-//! For streaming requests, SSE chunks are forwarded as they arrive.
-//! The proxy captures timing information for metrics but does not
-//! buffer the full response.
+//! The streaming *mechanism* — forward an SSE body chunk-for-chunk without
+//! buffering, observing the bytes for metrics — lives in the shared
+//! [`helexa_stream`] crate (#71), so cortex and helexa-router use one
+//! implementation. This module supplies cortex's *policy*: the
+//! [`CortexMetrics`] observer (per-request token metrics + per-principal
+//! reservation settle), cortex's logging contract, and the cortex error
+//! envelope. The usage-extraction helper is re-exported from the shared
+//! crate so existing call sites keep working.
 
 use crate::router::RouteDecision;
-use anyhow::Result;
-use axum::body::Body;
-use axum::http::{HeaderMap, StatusCode};
+use axum::http::HeaderMap;
+use axum::http::StatusCode;
 use axum::response::{IntoResponse, Response};
-use futures::Stream;
-use futures::stream::BoxStream;
+use helexa_stream::{BodyTail, ChunkObserver, StreamError};
 use reqwest::Client;
-use std::pin::Pin;
-use std::task::{Context, Poll};
 use std::time::Instant;
 
+/// Re-export the shared usage-extraction helper. Several cortex modules
+/// (`handlers`, `anthropic_sse`) pull token counts out of a buffered body
+/// tail via this function; it lives in `helexa-stream` now.
+pub use helexa_stream::last_count_for;
+
 /// Proxy a request body to the resolved backend node and stream the response.
 ///
 /// Logging contract: every call emits exactly one structured event at
@@ -42,66 +48,41 @@ pub async fn forward_request(
         "proxying request"
     );
 
-    let mut req_builder = client.post(&url).body(body);
+    let observer = CortexMetrics::new(model_id, &route.node_name, request_start, usage_sink);
 
-    // Forward relevant headers.
-    for (key, value) in headers.iter() {
-        if key == "host" || key == "content-length" {
-            continue; // reqwest sets these
-        }
-        req_builder = req_builder.header(key, value);
-    }
+    let response = helexa_stream::forward_streaming(client, &url, headers, body, observer)
+        .await
+        .map_err(|e| {
+            match &e {
+                StreamError::Upstream(err) => tracing::warn!(
+                    node = %route.node_name,
+                    url = %url,
+                    error = %err,
+                    "proxy: upstream request failed (network)"
+                ),
+                StreamError::ResponseBuild(err) => tracing::warn!(
+                    node = %route.node_name,
+                    url = %url,
+                    error = %err,
+                    "proxy: failed to build response"
+                ),
+            }
+            ProxyError::from(e)
+        })?;
 
-    let upstream_resp = match req_builder.send().await {
-        Ok(r) => r,
-        Err(e) => {
-            tracing::warn!(
-                node = %route.node_name,
-                url = %url,
-                error = %e,
-                "proxy: upstream request failed (network)"
-            );
-            return Err(ProxyError::Upstream(e));
-        }
-    };
-
-    let upstream_status = upstream_resp.status();
-    if !upstream_status.is_success() {
+    if !response.status().is_success() {
         // Streaming body — can't snippet without breaking the stream
         // pass-through. Log status + URL; the client still gets the
         // upstream status, just without the leaked body.
         tracing::warn!(
             node = %route.node_name,
             url = %url,
-            status = upstream_status.as_u16(),
+            status = response.status().as_u16(),
             "proxy: upstream returned non-2xx"
         );
     }
 
-    let status = StatusCode::from_u16(upstream_status.as_u16()).unwrap_or(StatusCode::BAD_GATEWAY);
-
-    let resp_headers = upstream_resp.headers().clone();
-    let stream = TokenMetricsStream::new(
-        Box::pin(upstream_resp.bytes_stream()),
-        TokenMetrics::new(model_id, &route.node_name, request_start, usage_sink),
-    );
-
-    let body = Body::from_stream(stream);
-
-    let mut response = Response::builder().status(status);
-    for (key, value) in resp_headers.iter() {
-        response = response.header(key, value);
-    }
-
-    response.body(body).map_err(|e| {
-        tracing::warn!(
-            node = %route.node_name,
-            url = %url,
-            error = %e,
-            "proxy: failed to build response"
-        );
-        ProxyError::ResponseBuild(e.to_string())
-    })
+    Ok(response)
 }
 
 #[derive(Debug, thiserror::Error)]
@@ -112,6 +93,15 @@ pub enum ProxyError {
     ResponseBuild(String),
 }
 
+impl From<StreamError> for ProxyError {
+    fn from(e: StreamError) -> Self {
+        match e {
+            StreamError::Upstream(err) => ProxyError::Upstream(err),
+            StreamError::ResponseBuild(msg) => ProxyError::ResponseBuild(msg),
+        }
+    }
+}
+
 impl IntoResponse for ProxyError {
     fn into_response(self) -> Response {
         let (status, code, message) = match &self {
@@ -139,9 +129,10 @@ impl IntoResponse for ProxyError {
 //
 // The proxy never buffers or re-serialises the upstream body — chunks
 // are forwarded verbatim. For metrics it observes each chunk's arrival
-// time and keeps a bounded tail of the body text, from which the final
-// OpenAI `usage` object (present on the last SSE chunk and on
-// non-streaming JSON bodies alike) yields engine-truth token counts.
+// time and keeps a bounded tail of the body text (via the shared
+// `helexa_stream::BodyTail`), from which the final OpenAI `usage` object
+// (present on the last SSE chunk and on non-streaming JSON bodies alike)
+// yields engine-truth token counts.
 //
 // Emitted per request, labelled {model, node}:
 //   cortex_time_to_first_token_seconds  (histogram) — first body chunk
@@ -155,37 +146,15 @@ impl IntoResponse for ProxyError {
 /// non-streaming bodies.
 const TAIL_CAP_BYTES: usize = 64 * 1024;
 
-/// Find the value of the LAST `"key": <integer>` occurrence in `tail`.
-/// Pure and chunk-boundary-safe (the tail is contiguous appended text).
-/// The quoted-needle form means `completion_tokens` never matches
-/// `completion_tokens_details`.
-pub(crate) fn last_count_for(tail: &str, key: &str) -> Option<u64> {
-    let needle = format!("\"{key}\"");
-    let mut result = None;
-    for (idx, _) in tail.match_indices(&needle) {
-        let rest = tail[idx + needle.len()..].trim_start();
-        let Some(rest) = rest.strip_prefix(':') else {
-            continue;
-        };
-        let rest = rest.trim_start();
-        let digits: &str = &rest[..rest
-            .char_indices()
-            .find(|(_, c)| !c.is_ascii_digit())
-            .map(|(i, _)| i)
-            .unwrap_or(rest.len())];
-        if let Ok(v) = digits.parse::<u64>() {
-            result = Some(v);
-        }
-    }
-    result
-}
-
-struct TokenMetrics {
+/// cortex's [`ChunkObserver`]: per-request token metrics plus the
+/// per-principal reservation settle. Drives cortex policy over the shared
+/// streaming mechanism.
+struct CortexMetrics {
     labels: [(&'static str, String); 2],
     request_start: Instant,
     first_chunk: Option<Instant>,
     last_chunk: Option<Instant>,
-    tail: String,
+    tail: BodyTail,
     finished: bool,
     /// Per-principal metering hook (#51). Invoked exactly once in `finish`
     /// with the observed `(prompt, completion)` so the reservation can be
@@ -193,7 +162,7 @@ struct TokenMetrics {
     usage_sink: Option<crate::metering::UsageSink>,
 }
 
-impl TokenMetrics {
+impl CortexMetrics {
     fn new(
         model_id: &str,
         node_name: &str,
@@ -208,26 +177,19 @@ impl TokenMetrics {
             request_start,
             first_chunk: None,
             last_chunk: None,
-            tail: String::new(),
+            tail: BodyTail::new(TAIL_CAP_BYTES),
             finished: false,
             usage_sink,
         }
     }
+}
 
+impl ChunkObserver for CortexMetrics {
     fn observe(&mut self, chunk: &[u8]) {
         let now = Instant::now();
         self.first_chunk.get_or_insert(now);
         self.last_chunk = Some(now);
-        self.tail.push_str(&String::from_utf8_lossy(chunk));
-        if self.tail.len() > TAIL_CAP_BYTES {
-            // Keep the newest half; the usage object is always at the
-            // very end of the body. Split at a char boundary.
-            let mut cut = self.tail.len() - TAIL_CAP_BYTES / 2;
-            while !self.tail.is_char_boundary(cut) {
-                cut += 1;
-            }
-            self.tail.drain(..cut);
-        }
+        self.tail.push(chunk);
     }
 
     /// Emit the metrics exactly once — called on clean stream end and
@@ -239,8 +201,8 @@ impl TokenMetrics {
         }
         self.finished = true;
 
-        let prompt = last_count_for(&self.tail, "prompt_tokens");
-        let completion = last_count_for(&self.tail, "completion_tokens");
+        let prompt = last_count_for(self.tail.as_str(), "prompt_tokens");
+        let completion = last_count_for(self.tail.as_str(), "completion_tokens");
 
         // Per-model metrics — only when body chunks actually arrived.
         if let Some(first) = self.first_chunk {
@@ -280,97 +242,3 @@ impl TokenMetrics {
         }
     }
 }
-
-/// Pass-through stream wrapper that feeds [`TokenMetrics`]. Emits on
-/// clean end-of-stream; the Drop impl covers client disconnects.
-struct TokenMetricsStream {
-    inner: BoxStream<'static, Result<bytes::Bytes, reqwest::Error>>,
-    metrics: TokenMetrics,
-}
-
-impl TokenMetricsStream {
-    fn new(
-        inner: BoxStream<'static, Result<bytes::Bytes, reqwest::Error>>,
-        metrics: TokenMetrics,
-    ) -> Self {
-        Self { inner, metrics }
-    }
-}
-
-impl Stream for TokenMetricsStream {
-    type Item = Result<bytes::Bytes, reqwest::Error>;
-
-    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
-        let this = self.get_mut();
-        match this.inner.as_mut().poll_next(cx) {
-            Poll::Ready(Some(Ok(chunk))) => {
-                this.metrics.observe(&chunk);
-                Poll::Ready(Some(Ok(chunk)))
-            }
-            Poll::Ready(Some(Err(e))) => Poll::Ready(Some(Err(e))),
-            Poll::Ready(None) => {
-                this.metrics.finish();
-                Poll::Ready(None)
-            }
-            Poll::Pending => Poll::Pending,
-        }
-    }
-}
-
-impl Drop for TokenMetricsStream {
-    fn drop(&mut self) {
-        self.metrics.finish();
-    }
-}
-
-#[cfg(test)]
-mod tests {
-    use super::last_count_for;
-
-    #[test]
-    fn extracts_counts_from_final_sse_usage_chunk() {
-        let tail = concat!(
-            "data: {\"choices\":[{\"delta\":{\"content\":\"hi\"}}]}\n\n",
-            "data: {\"choices\":[],\"usage\":{\"prompt_tokens\":225,",
-            "\"completion_tokens\":42,\"total_tokens\":267}}\n\n",
-            "data: [DONE]\n\n"
-        );
-        assert_eq!(last_count_for(tail, "prompt_tokens"), Some(225));
-        assert_eq!(last_count_for(tail, "completion_tokens"), Some(42));
-    }
-
-    #[test]
-    fn extracts_counts_from_non_streaming_body() {
-        let tail = "{\"choices\":[{\"message\":{\"content\":\"hi\"}}],\
-                    \"usage\":{\"prompt_tokens\": 12, \"completion_tokens\": 7}}";
-        assert_eq!(last_count_for(tail, "prompt_tokens"), Some(12));
-        assert_eq!(last_count_for(tail, "completion_tokens"), Some(7));
-    }
-
-    #[test]
-    fn ignores_details_variants_and_takes_last_occurrence() {
-        // completion_tokens_details must not shadow completion_tokens,
-        // and the LAST usage object wins (matters when content echoes
-        // a usage-shaped string earlier in the stream).
-        let tail = concat!(
-            "data: {\"usage\":{\"completion_tokens\":1}}\n\n",
-            "data: {\"usage\":{\"completion_tokens\":99,",
-            "\"completion_tokens_details\":{\"reasoning_tokens\":3}}}\n\n"
-        );
-        assert_eq!(last_count_for(tail, "completion_tokens"), Some(99));
-    }
-
-    #[test]
-    fn absent_keys_yield_none() {
-        assert_eq!(
-            last_count_for("data: [DONE]\n\n", "completion_tokens"),
-            None
-        );
-        assert_eq!(last_count_for("", "prompt_tokens"), None);
-        // key present but non-numeric value
-        assert_eq!(
-            last_count_for("\"completion_tokens\": null", "completion_tokens"),
-            None
-        );
-    }
-}

crates/helexa-stream/Cargo.toml

@@ -0,0 +1,21 @@
+[package]
+name = "helexa-stream"
+version.workspace = true
+edition.workspace = true
+license.workspace = true
+repository.workspace = true
+
+[lib]
+name = "helexa_stream"
+path = "src/lib.rs"
+
+[dependencies]
+axum = { workspace = true }
+reqwest = { workspace = true }
+futures = { workspace = true }
+thiserror = { workspace = true }
+
+[dev-dependencies]
+tokio = { workspace = true }
+tokio-stream = { workspace = true }
+async-stream = "0.3"

crates/helexa-stream/src/lib.rs

@@ -0,0 +1,290 @@
+//! Shared streaming reverse-proxy mechanism (#71).
+//!
+//! cortex and helexa-router both need to proxy an OpenAI/Anthropic SSE
+//! response from a downstream backend **verbatim** — chunks forwarded as
+//! they arrive, never buffering the full body — while observing the bytes
+//! for metrics/metering. This crate owns that mechanism so there is one
+//! implementation, not one per tier.
+//!
+//! The split is mechanism vs policy:
+//!
+//! - **Mechanism (here):** [`forward_streaming`] POSTs to a backend and
+//!   streams the response body back through an [`ObservedStream`], which
+//!   feeds every chunk to a caller-supplied [`ChunkObserver`] and calls
+//!   [`ChunkObserver::finish`] exactly once on clean end-of-stream or on
+//!   drop (client disconnect mid-stream). [`BodyTail`] and
+//!   [`last_count_for`] are the reusable pieces an observer uses to pull
+//!   the trailing OpenAI `usage` object out of the streamed bytes.
+//! - **Policy (caller):** what to *do* with the observed bytes — which
+//!   metric names to emit, which labels, whether to settle a per-principal
+//!   reservation — lives in the consumer's `ChunkObserver` impl, not here.
+//!
+//! The proxy is status-agnostic: a non-2xx upstream response (e.g. a
+//! cortex `429 rate_limit_exceeded`) is streamed back with its status and
+//! headers intact, so honest backpressure reaches the client unchanged.
+//! Only a network failure or a malformed response build is an error.
+
+use axum::body::{Body, Bytes};
+use axum::http::{HeaderMap, StatusCode};
+use axum::response::Response;
+use futures::Stream;
+use futures::stream::BoxStream;
+use reqwest::Client;
+use std::pin::Pin;
+use std::task::{Context, Poll};
+
+/// Observes the bytes of a streamed proxy response without altering them.
+///
+/// `observe` is called for each forwarded chunk; `finish` is called
+/// exactly once — on clean end-of-stream or on drop — and implementations
+/// must be idempotent (the [`ObservedStream`] guards against a double call,
+/// but a `finish` that runs side effects should still self-guard).
+pub trait ChunkObserver: Send + Unpin + 'static {
+    /// A body chunk has been forwarded downstream. The slice is the exact
+    /// bytes the client receives.
+    fn observe(&mut self, chunk: &[u8]);
+
+    /// The stream has ended (cleanly or via client disconnect). Called once.
+    fn finish(&mut self);
+}
+
+/// A bounded accumulator for the tail of a streamed body.
+///
+/// The OpenAI `usage` object rides on the final SSE chunk (and sits at the
+/// end of a non-streaming JSON body), so retaining a generous tail is
+/// enough to recover token counts via [`last_count_for`]; the cap bounds
+/// memory on huge bodies. Appends are char-boundary-safe.
+#[derive(Debug)]
+pub struct BodyTail {
+    tail: String,
+    cap: usize,
+}
+
+impl BodyTail {
+    /// Create a tail retaining at most `cap` bytes.
+    pub fn new(cap: usize) -> Self {
+        Self {
+            tail: String::new(),
+            cap,
+        }
+    }
+
+    /// Append a chunk, trimming from the front past the cap. When trimming,
+    /// the newest half is kept (the usage object is always at the very end).
+    pub fn push(&mut self, chunk: &[u8]) {
+        self.tail.push_str(&String::from_utf8_lossy(chunk));
+        if self.tail.len() > self.cap {
+            let mut cut = self.tail.len() - self.cap / 2;
+            while !self.tail.is_char_boundary(cut) {
+                cut += 1;
+            }
+            self.tail.drain(..cut);
+        }
+    }
+
+    /// The retained tail text.
+    pub fn as_str(&self) -> &str {
+        &self.tail
+    }
+}
+
+/// Find the value of the LAST `"key": <integer>` occurrence in `tail`.
+///
+/// Pure and chunk-boundary-safe (the tail is contiguous appended text).
+/// The quoted-needle form means `completion_tokens` never matches
+/// `completion_tokens_details`, and taking the last occurrence means the
+/// final `usage` object wins even if content earlier in the stream echoed
+/// a usage-shaped string.
+pub fn last_count_for(tail: &str, key: &str) -> Option<u64> {
+    let needle = format!("\"{key}\"");
+    let mut result = None;
+    for (idx, _) in tail.match_indices(&needle) {
+        let rest = tail[idx + needle.len()..].trim_start();
+        let Some(rest) = rest.strip_prefix(':') else {
+            continue;
+        };
+        let rest = rest.trim_start();
+        let digits: &str = &rest[..rest
+            .char_indices()
+            .find(|(_, c)| !c.is_ascii_digit())
+            .map(|(i, _)| i)
+            .unwrap_or(rest.len())];
+        if let Ok(v) = digits.parse::<u64>() {
+            result = Some(v);
+        }
+    }
+    result
+}
+
+/// Error from [`forward_streaming`]. Distinguishes a network/transport
+/// failure reaching the backend from a failure assembling the downstream
+/// response. A non-2xx upstream *status* is not an error — it is streamed
+/// through verbatim.
+#[derive(Debug, thiserror::Error)]
+pub enum StreamError {
+    #[error("upstream request failed")]
+    Upstream(reqwest::Error),
+    #[error("failed to build response")]
+    ResponseBuild(String),
+}
+
+/// POST `body` to `url` and stream the response back verbatim through
+/// `observer`.
+///
+/// Request headers are forwarded except `host` / `content-length` (reqwest
+/// sets these). The returned [`Response`] carries the upstream status and
+/// headers unchanged — including non-2xx — with a body that streams the
+/// upstream bytes chunk-for-chunk, feeding each chunk to `observer`.
+pub async fn forward_streaming<O: ChunkObserver>(
+    client: &Client,
+    url: &str,
+    headers: HeaderMap,
+    body: Bytes,
+    observer: O,
+) -> Result<Response, StreamError> {
+    let mut req_builder = client.post(url).body(body);
+    for (key, value) in headers.iter() {
+        if key == "host" || key == "content-length" {
+            continue; // reqwest sets these
+        }
+        req_builder = req_builder.header(key, value);
+    }
+
+    let upstream = req_builder.send().await.map_err(StreamError::Upstream)?;
+
+    let status =
+        StatusCode::from_u16(upstream.status().as_u16()).unwrap_or(StatusCode::BAD_GATEWAY);
+    let resp_headers = upstream.headers().clone();
+
+    let stream = ObservedStream::new(Box::pin(upstream.bytes_stream()), observer);
+    let body = Body::from_stream(stream);
+
+    let mut response = Response::builder().status(status);
+    for (key, value) in resp_headers.iter() {
+        response = response.header(key, value);
+    }
+    response
+        .body(body)
+        .map_err(|e| StreamError::ResponseBuild(e.to_string()))
+}
+
+/// Pass-through stream wrapper that feeds a [`ChunkObserver`]. Forwards
+/// each chunk verbatim, calls `observe` per chunk, and `finish` once on
+/// clean end-of-stream; the `Drop` impl covers client disconnects.
+pub struct ObservedStream<O: ChunkObserver> {
+    inner: BoxStream<'static, Result<Bytes, reqwest::Error>>,
+    observer: O,
+    finished: bool,
+}
+
+impl<O: ChunkObserver> ObservedStream<O> {
+    /// Wrap a byte stream with an observer.
+    pub fn new(inner: BoxStream<'static, Result<Bytes, reqwest::Error>>, observer: O) -> Self {
+        Self {
+            inner,
+            observer,
+            finished: false,
+        }
+    }
+
+    fn finish(&mut self) {
+        if self.finished {
+            return;
+        }
+        self.finished = true;
+        self.observer.finish();
+    }
+}
+
+impl<O: ChunkObserver> Stream for ObservedStream<O> {
+    type Item = Result<Bytes, reqwest::Error>;
+
+    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        let this = self.get_mut();
+        match this.inner.as_mut().poll_next(cx) {
+            Poll::Ready(Some(Ok(chunk))) => {
+                this.observer.observe(&chunk);
+                Poll::Ready(Some(Ok(chunk)))
+            }
+            Poll::Ready(Some(Err(e))) => Poll::Ready(Some(Err(e))),
+            Poll::Ready(None) => {
+                this.finish();
+                Poll::Ready(None)
+            }
+            Poll::Pending => Poll::Pending,
+        }
+    }
+}
+
+impl<O: ChunkObserver> Drop for ObservedStream<O> {
+    fn drop(&mut self) {
+        self.finish();
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn extracts_counts_from_final_sse_usage_chunk() {
+        let tail = concat!(
+            "data: {\"choices\":[{\"delta\":{\"content\":\"hi\"}}]}\n\n",
+            "data: {\"choices\":[],\"usage\":{\"prompt_tokens\":225,",
+            "\"completion_tokens\":42,\"total_tokens\":267}}\n\n",
+            "data: [DONE]\n\n"
+        );
+        assert_eq!(last_count_for(tail, "prompt_tokens"), Some(225));
+        assert_eq!(last_count_for(tail, "completion_tokens"), Some(42));
+    }
+
+    #[test]
+    fn extracts_counts_from_non_streaming_body() {
+        let tail = "{\"choices\":[{\"message\":{\"content\":\"hi\"}}],\
+                    \"usage\":{\"prompt_tokens\": 12, \"completion_tokens\": 7}}";
+        assert_eq!(last_count_for(tail, "prompt_tokens"), Some(12));
+        assert_eq!(last_count_for(tail, "completion_tokens"), Some(7));
+    }
+
+    #[test]
+    fn ignores_details_variants_and_takes_last_occurrence() {
+        // completion_tokens_details must not shadow completion_tokens,
+        // and the LAST usage object wins (matters when content echoes
+        // a usage-shaped string earlier in the stream).
+        let tail = concat!(
+            "data: {\"usage\":{\"completion_tokens\":1}}\n\n",
+            "data: {\"usage\":{\"completion_tokens\":99,",
+            "\"completion_tokens_details\":{\"reasoning_tokens\":3}}}\n\n"
+        );
+        assert_eq!(last_count_for(tail, "completion_tokens"), Some(99));
+    }
+
+    #[test]
+    fn absent_keys_yield_none() {
+        assert_eq!(
+            last_count_for("data: [DONE]\n\n", "completion_tokens"),
+            None
+        );
+        assert_eq!(last_count_for("", "prompt_tokens"), None);
+        // key present but non-numeric value
+        assert_eq!(
+            last_count_for("\"completion_tokens\": null", "completion_tokens"),
+            None
+        );
+    }
+
+    #[test]
+    fn body_tail_retains_usage_after_cap_trim() {
+        // Cap small enough that the filler forces several front-trims, but
+        // (as in production, where cap ≫ the usage object) large enough that
+        // the trailing usage object survives the newest-half retention.
+        let mut tail = BodyTail::new(512);
+        for _ in 0..100 {
+            tail.push(b"data: {\"choices\":[{\"delta\":{\"content\":\"x\"}}]}\n\n");
+        }
+        assert!(tail.as_str().len() <= 512, "cap must bound the tail");
+        tail.push(b"data: {\"usage\":{\"prompt_tokens\":5,\"completion_tokens\":9}}\n\n");
+        assert_eq!(last_count_for(tail.as_str(), "prompt_tokens"), Some(5));
+        assert_eq!(last_count_for(tail.as_str(), "completion_tokens"), Some(9));
+    }
+}

crates/helexa-stream/tests/streaming.rs

@@ -0,0 +1,162 @@
+//! Integration tests for the shared streaming proxy (#71): proves a backend
+//! SSE response is forwarded chunk-for-chunk (no buffering), the observer
+//! sees every byte and finishes once, and non-2xx is streamed through with
+//! its status intact — the behaviours both cortex and helexa-router rely on.
+
+use axum::Router;
+use axum::body::Body;
+use axum::http::{HeaderMap, StatusCode};
+use axum::response::Response;
+use axum::routing::post;
+use helexa_stream::{BodyTail, ChunkObserver, forward_streaming, last_count_for};
+use std::sync::{Arc, Mutex};
+use std::time::{Duration, Instant};
+use tokio::net::TcpListener;
+
+/// Observer that records what it saw, for assertions.
+#[derive(Clone, Default)]
+struct RecordingObserver {
+    inner: Arc<Mutex<Recorded>>,
+}
+
+#[derive(Default)]
+struct Recorded {
+    chunks: usize,
+    finished: usize,
+    tail: String,
+}
+
+impl ChunkObserver for RecordingObserver {
+    fn observe(&mut self, chunk: &[u8]) {
+        let mut r = self.inner.lock().unwrap();
+        r.chunks += 1;
+        r.tail.push_str(&String::from_utf8_lossy(chunk));
+    }
+    fn finish(&mut self) {
+        self.inner.lock().unwrap().finished += 1;
+    }
+}
+
+/// Mock backend that streams 5 SSE chunks with 30ms gaps, then a usage
+/// chunk and `[DONE]`.
+async fn sse_handler() -> Response {
+    let chunks: Vec<&'static str> = vec![
+        "data: {\"choices\":[{\"delta\":{\"content\":\"a\"}}]}\n\n",
+        "data: {\"choices\":[{\"delta\":{\"content\":\"b\"}}]}\n\n",
+        "data: {\"choices\":[{\"delta\":{\"content\":\"c\"}}]}\n\n",
+        "data: {\"choices\":[{\"delta\":{\"content\":\"d\"}}]}\n\n",
+        "data: {\"choices\":[{\"delta\":{\"content\":\"e\"}}]}\n\n",
+        "data: {\"choices\":[],\"usage\":{\"prompt_tokens\":11,\"completion_tokens\":5}}\n\n",
+        "data: [DONE]\n\n",
+    ];
+    let stream = async_stream::stream! {
+        for c in chunks {
+            tokio::time::sleep(Duration::from_millis(30)).await;
+            yield Ok::<_, std::io::Error>(axum::body::Bytes::from_static(c.as_bytes()));
+        }
+    };
+    Response::new(Body::from_stream(stream))
+}
+
+async fn rate_limited_handler() -> Response {
+    Response::builder()
+        .status(StatusCode::TOO_MANY_REQUESTS)
+        .body(Body::from("{\"error\":{\"type\":\"rate_limit_exceeded\"}}"))
+        .unwrap()
+}
+
+async fn spawn_backend(router: Router) -> String {
+    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+    let addr = listener.local_addr().unwrap();
+    tokio::spawn(async move {
+        axum::serve(listener, router).await.unwrap();
+    });
+    format!("http://{addr}")
+}
+
+#[tokio::test]
+async fn streams_chunks_incrementally_and_observes_usage() {
+    let base = spawn_backend(Router::new().route("/v1/chat/completions", post(sse_handler))).await;
+    let observer = RecordingObserver::default();
+    let probe = observer.clone();
+
+    let client = reqwest::Client::new();
+    let resp = forward_streaming(
+        &client,
+        &format!("{base}/v1/chat/completions"),
+        HeaderMap::new(),
+        axum::body::Bytes::from_static(b"{\"model\":\"x\",\"stream\":true}"),
+        observer,
+    )
+    .await
+    .expect("forward ok");
+
+    assert_eq!(resp.status(), StatusCode::OK);
+
+    // Read the proxied body as a stream, timestamping arrivals.
+    let mut body = resp.into_body().into_data_stream();
+    let mut arrivals: Vec<Instant> = Vec::new();
+    let mut collected = String::new();
+    use futures::StreamExt;
+    while let Some(item) = body.next().await {
+        let bytes = item.unwrap();
+        arrivals.push(Instant::now());
+        collected.push_str(&String::from_utf8_lossy(&bytes));
+    }
+
+    // Incremental delivery: first and last chunk are meaningfully apart
+    // (5×30ms gaps), proving no full-response buffering.
+    let spread = *arrivals.last().unwrap() - arrivals[0];
+    assert!(
+        spread >= Duration::from_millis(100),
+        "expected incremental delivery, spread was {spread:?}"
+    );
+
+    // The client received the terminator and the usage object verbatim.
+    assert!(collected.contains("data: [DONE]"));
+
+    // The observer saw the bytes and finished exactly once.
+    let r = probe.inner.lock().unwrap();
+    assert!(r.chunks >= 5, "observer saw {} chunks", r.chunks);
+    assert_eq!(r.finished, 1, "finish must run exactly once");
+    assert_eq!(last_count_for(&r.tail, "prompt_tokens"), Some(11));
+    assert_eq!(last_count_for(&r.tail, "completion_tokens"), Some(5));
+}
+
+#[tokio::test]
+async fn non_2xx_is_streamed_through_verbatim() {
+    let base =
+        spawn_backend(Router::new().route("/v1/chat/completions", post(rate_limited_handler)))
+            .await;
+    let observer = RecordingObserver::default();
+    let probe = observer.clone();
+
+    let client = reqwest::Client::new();
+    let resp = forward_streaming(
+        &client,
+        &format!("{base}/v1/chat/completions"),
+        HeaderMap::new(),
+        axum::body::Bytes::new(),
+        observer,
+    )
+    .await
+    .expect("forward ok");
+
+    // Backpressure status reaches the client unchanged.
+    assert_eq!(resp.status(), StatusCode::TOO_MANY_REQUESTS);
+    let body = axum::body::to_bytes(resp.into_body(), usize::MAX)
+        .await
+        .unwrap();
+    assert!(String::from_utf8_lossy(&body).contains("rate_limit_exceeded"));
+
+    // finish still runs once even with a tiny non-streaming body.
+    assert_eq!(probe.inner.lock().unwrap().finished, 1);
+}
+
+#[test]
+fn body_tail_smoke() {
+    let mut tail = BodyTail::new(128);
+    tail.push(b"hello ");
+    tail.push(b"world");
+    assert_eq!(tail.as_str(), "hello world");
+}