fix(azure): apply stream_timeout per chunk, not as a whole-response ceiling

crates/aisix-provider-azure-openai/Cargo.toml

@@ -30,5 +30,9 @@ bytes.workspace = true
 http.workspace = true
 
 [dev-dependencies]
-tokio = { workspace = true, features = ["macros", "rt", "time"] }
+# `net` + `io-util` back the hand-rolled SSE upstream in
+# `chat_stream_delivers_a_long_stream_whose_gaps_stay_within_budget`, which
+# needs to emit frames on a schedule — wiremock can only delay a response as
+# a whole, so it cannot distinguish per-chunk from total-duration timeouts.
+tokio = { workspace = true, features = ["macros", "rt", "time", "net", "io-util"] }
 wiremock.workspace = true

crates/aisix-provider-azure-openai/src/bridge.rs

@@ -782,11 +782,13 @@ impl Bridge for AzureOpenAiBridge {
         let bridge_name = self.name;
         let request_id_for_log = ctx.request_id.clone();
 
-        // Audit H2: thread the deadline into the streaming loop so a
-        // slow / hanging upstream body can't wedge the connection
-        // after headers arrive. The post-headers POST already covered
-        // the initial wait; this covers the per-chunk wait too.
-        let stream_deadline = ctx.deadline.map(|d| started + d);
+        // Audit H2: thread the budget into the streaming loop so a slow /
+        // hanging upstream body can't wedge the connection after headers
+        // arrive. It bounds the wait for EACH chunk and resets after every
+        // successful read — `stream_timeout` is defined as the maximum gap
+        // *between* chunks, so a long-but-healthy response must not be cut
+        // off just because its total duration exceeds one gap's budget.
+        let per_chunk_timeout = ctx.deadline;
 
         let byte_stream = resp.bytes_stream();
         let stream = build_chunk_stream(
@@ -795,8 +797,7 @@ impl Bridge for AzureOpenAiBridge {
             done_marker_policy,
             bridge_name,
             request_id_for_log,
-            stream_deadline,
-            started,
+            per_chunk_timeout,
         );
         Ok(Box::pin(stream))
     }
@@ -808,8 +809,7 @@ fn build_chunk_stream<S>(
     done_marker_policy: Option<StreamDoneMarker>,
     bridge_name: &'static str,
     request_id: String,
-    deadline: Option<Instant>,
-    started: Instant,
+    per_chunk_timeout: Option<Duration>,
 ) -> impl futures::Stream<Item = Result<ChatChunk, BridgeError>> + Send
 where
     S: futures::Stream<Item = reqwest::Result<bytes::Bytes>> + Send + 'static,
@@ -819,15 +819,16 @@ where
         let mut stream = Box::pin(byte_stream);
         let mut done_marker_seen = false;
         'outer: loop {
-            // Audit H2: enforce deadline on each per-chunk wait. The
-            // first POST is already covered upstream; this covers the
-            // body-streaming phase. `None` deadline disables timeout.
-            let next = match deadline {
-                Some(d) => match tokio::time::timeout_at(d.into(), stream.next()).await {
+            // Audit H2: bound each individual chunk wait. `tokio::time::timeout`
+            // (not `timeout_at`) so the budget restarts on every chunk — the
+            // configured value is the maximum inter-chunk gap, not a ceiling on
+            // the response as a whole. `None` disables the timeout.
+            let next = match per_chunk_timeout {
+                Some(d) => match tokio::time::timeout(d, stream.next()).await {
                     Ok(item) => item,
                     Err(_) => {
                         Err(BridgeError::Timeout {
-                            elapsed_ms: started.elapsed().as_millis() as u64,
+                            elapsed_ms: d.as_millis() as u64,
                             cause: String::new(),
                         })?;
                         unreachable!()
@@ -2008,6 +2009,184 @@ mod tests {
         }
     }
 
+    /// AISIX-Cloud#1122: `stream_timeout` is defined as the maximum gap
+    /// *between* chunks (see `Model::stream_timeout`), but the budget was
+    /// applied as `timeout_at(started + d)` — one absolute instant reused
+    /// for every chunk, i.e. a ceiling on the whole response. A long but
+    /// perfectly healthy stream was cut off with a Timeout once its total
+    /// duration passed a single gap's budget.
+    ///
+    /// Here every gap (80ms) is comfortably inside the budget (250ms)
+    /// while the total (~320ms) is past it. Per-chunk semantics deliver
+    /// the whole stream; the absolute deadline truncated it.
+    #[tokio::test]
+    async fn stream_budget_bounds_each_chunk_gap_not_the_whole_response() {
+        use futures::stream::StreamExt as _;
+
+        let gap = Duration::from_millis(80);
+        let budget = Duration::from_millis(250);
+        let frames: Vec<bytes::Bytes> = vec![
+            bytes::Bytes::from(
+                "data: {\"id\":\"x\",\"model\":\"m\",\"choices\":[{\"delta\":{\"content\":\"a\"},\"finish_reason\":null}]}\n\n",
+            ),
+            bytes::Bytes::from(
+                "data: {\"id\":\"x\",\"model\":\"m\",\"choices\":[{\"delta\":{\"content\":\"b\"},\"finish_reason\":null}]}\n\n",
+            ),
+            bytes::Bytes::from(
+                "data: {\"id\":\"x\",\"model\":\"m\",\"choices\":[{\"delta\":{\"content\":\"c\"},\"finish_reason\":\"stop\"}]}\n\n",
+            ),
+            bytes::Bytes::from("data: [DONE]\n\n"),
+        ];
+        let byte_stream = futures::stream::iter(frames).then(move |b| async move {
+            tokio::time::sleep(gap).await;
+            Ok::<bytes::Bytes, reqwest::Error>(b)
+        });
+
+        let started = Instant::now();
+        let mut stream = Box::pin(build_chunk_stream(
+            byte_stream,
+            None,
+            None,
+            "azure-openai",
+            "req-per-chunk".to_string(),
+            Some(budget),
+        ));
+
+        let mut content = String::new();
+        while let Some(item) = stream.next().await {
+            let chunk = item.expect("no chunk may time out — every gap is within budget");
+            if let Some(text) = chunk.delta.content.as_deref() {
+                content.push_str(text);
+            }
+        }
+
+        assert_eq!(content, "abc", "the whole stream must be delivered");
+        assert!(
+            started.elapsed() > budget,
+            "the test is only meaningful if the total run outlasts one gap's \
+             budget (elapsed {:?}, budget {budget:?})",
+            started.elapsed()
+        );
+    }
+
+    /// The same contract as the test above, but driven through the public
+    /// `chat_stream` entry point instead of `build_chunk_stream`, so it
+    /// also covers `chat_stream` actually forwarding `ctx.deadline` — the
+    /// helper-level test alone would still pass if that wiring were
+    /// dropped.
+    ///
+    /// Needs an upstream that emits frames on a schedule, which wiremock
+    /// cannot do (`set_delay` delays the response as a whole, which is
+    /// exactly the distinction under test), so this hand-rolls a one-shot
+    /// SSE server: no `content-length`, `connection: close`, body framed
+    /// by EOF.
+    #[tokio::test]
+    async fn chat_stream_delivers_a_long_stream_whose_gaps_stay_within_budget() {
+        use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _};
+
+        let gap = Duration::from_millis(80);
+        let budget = Duration::from_millis(250);
+
+        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
+        let addr = listener.local_addr().unwrap();
+        tokio::spawn(async move {
+            let (mut sock, _) = listener.accept().await.unwrap();
+            // Drain the request headers so the client's write completes.
+            let mut buf = [0u8; 4096];
+            let _ = sock.read(&mut buf).await;
+            sock.write_all(
+                b"HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\nconnection: close\r\n\r\n",
+            )
+            .await
+            .unwrap();
+            sock.flush().await.unwrap();
+            for text in ["a", "b", "c"] {
+                tokio::time::sleep(gap).await;
+                let frame = format!(
+                    "data: {{\"id\":\"x\",\"model\":\"m\",\"choices\":[{{\"delta\":{{\"content\":\"{text}\"}},\"finish_reason\":null}}]}}\n\n"
+                );
+                sock.write_all(frame.as_bytes()).await.unwrap();
+                sock.flush().await.unwrap();
+            }
+            tokio::time::sleep(gap).await;
+            sock.write_all(b"data: [DONE]\n\n").await.unwrap();
+            sock.flush().await.unwrap();
+            sock.shutdown().await.unwrap();
+        });
+
+        let bridge = AzureOpenAiBridge::new()
+            .with_url_override(mock_chat_url(&format!("http://{addr}"), "gpt4o-prod"));
+        let mut ctx = canonical_test_ctx();
+        ctx.deadline = Some(budget);
+        let req = ChatFormat::new("my-azure-gpt4", vec![ChatMessage::user("hi")]);
+
+        let started = Instant::now();
+        let mut stream = bridge.chat_stream(&req, &ctx).await.expect("stream opens");
+        let mut content = String::new();
+        while let Some(item) = stream.next().await {
+            let chunk = item.expect("no chunk may time out — every gap is within budget");
+            if let Some(text) = chunk.delta.content.as_deref() {
+                content.push_str(text);
+            }
+        }
+
+        assert_eq!(content, "abc", "the whole stream must be delivered");
+        assert!(
+            started.elapsed() > budget,
+            "the test is only meaningful if the total run outlasts one gap's \
+             budget (elapsed {:?}, budget {budget:?})",
+            started.elapsed()
+        );
+    }
+
+    /// The counterpart that actually pins the wiring: response headers
+    /// arrive immediately, so the connect-phase `with_deadline` is already
+    /// satisfied, and only then does a single gap exceed the budget. The
+    /// per-chunk timeout is therefore the only thing that can fire — so
+    /// this test fails if `chat_stream` ever stops forwarding
+    /// `ctx.deadline` into the stream, which the delivery test above
+    /// (and a helper-level test) would not catch.
+    #[tokio::test]
+    async fn chat_stream_times_out_when_one_gap_exceeds_the_budget() {
+        use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _};
+
+        let budget = Duration::from_millis(150);
+        let stall = Duration::from_millis(700);
+
+        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
+        let addr = listener.local_addr().unwrap();
+        tokio::spawn(async move {
+            let (mut sock, _) = listener.accept().await.unwrap();
+            let mut buf = [0u8; 4096];
+            let _ = sock.read(&mut buf).await;
+            // Headers first and flushed, so the bridge commits the stream.
+            sock.write_all(
+                b"HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\nconnection: close\r\n\r\n",
+            )
+            .await
+            .unwrap();
+            sock.flush().await.unwrap();
+            // ...then stall well past the budget before the first frame.
+            tokio::time::sleep(stall).await;
+            let _ = sock.write_all(b"data: [DONE]\n\n").await;
+        });
+
+        let bridge = AzureOpenAiBridge::new()
+            .with_url_override(mock_chat_url(&format!("http://{addr}"), "gpt4o-prod"));
+        let mut ctx = canonical_test_ctx();
+        ctx.deadline = Some(budget);
+        let req = ChatFormat::new("my-azure-gpt4", vec![ChatMessage::user("hi")]);
+
+        let mut stream = bridge
+            .chat_stream(&req, &ctx)
+            .await
+            .expect("headers arrive within the budget, so the stream opens");
+        match stream.next().await {
+            Some(Err(BridgeError::Timeout { .. })) => {}
+            other => panic!("expected a per-chunk Timeout, got {other:?}"),
+        }
+    }
+
     // ─── AAD (Entra ID) auth scheme (D6.6) ─────────────────────────
 
     /// JSON-shaped Azure secret triggering the AAD branch.