fix(proxy): omit terminal [DONE] when stream terminates abnormally

crates/aisix-proxy/src/chat.rs

@@ -1153,6 +1153,15 @@ where
         // the bookkeeping robust to a provider that double-emits.
         // `chunk.id` / `chunk.model` / `chunk.finish_reason` use
         // last-seen-wins because those are stable per stream.
+        //
+        // Per docs/api-proxy.md §5: "If the upstream stream
+        // terminates abnormally, aisix sends a final error chunk
+        // and closes the response without `[DONE]`." Track whether
+        // an error has been yielded so we can skip the closing
+        // `[DONE]` — without this skip, a downstream SDK that
+        // treats `[DONE]` as clean-completion would mis-interpret
+        // the truncated response as a successful one.
+        let mut errored = false;
         while let Some(item) = upstream.next().await {
             let ev = match item {
                 Ok(chunk) => {
@@ -1193,14 +1202,21 @@ where
                     let rendered = render_chunk(created, chunk);
                     match serde_json::to_string(&rendered) {
                         Ok(json) => Event::default().data(json),
-                        Err(err) => Event::default()
-                            .event("error")
-                            .data(err.to_string()),
+                        Err(err) => {
+                            errored = true;
+                            Event::default()
+                                .event("error")
+                                .data(error_frame_payload("internal_error", &err.to_string()))
+                        }
                     }
                 }
-                Err(err) => Event::default()
-                    .event("error")
-                    .data(err.to_string()),
+                Err(err) => {
+                    errored = true;
+                    let etype = err.error_type();
+                    Event::default()
+                        .event("error")
+                        .data(error_frame_payload(etype, &err.to_string()))
+                }
             };
             yield Ok::<_, Infallible>(ev);
         }
@@ -1214,10 +1230,39 @@ where
         // accumulator's numeric fields stay 0; on_complete callers
         // must treat 0 as "no signal" (cp-api does — its pricing
         // catalog falls back to the standard rate when absent).
-        yield Ok::<_, Infallible>(Event::default().data("[DONE]"));
+        //
+        // Per docs §5: skip `[DONE]` on abnormal termination so SDK
+        // consumers can detect truncation. The `errored` flag is
+        // set by the loop above whenever an error event is yielded.
+        if !errored {
+            yield Ok::<_, Infallible>(Event::default().data("[DONE]"));
+        }
         // `guard` drops here. On client disconnect, the generator
         // drops at the suspension point inside the loop; Drop fires
         // there with whatever StreamCompletion has been captured up
         // to that point.
     }
 }
+
+/// Build the `data:` payload for an SSE `event: error` frame.
+/// The OpenAI Node SDK calls `JSON.parse(sse.data)` BEFORE checking
+/// `sse.event === "error"`, so a plain-string payload yields a
+/// `SyntaxError` ("Could not parse message into JSON: ...") on the
+/// SDK side rather than the typed `APIError` callers expect. Emit
+/// the OpenAI error envelope shape per
+/// <https://platform.openai.com/docs/guides/error-codes/api-errors>:
+/// `{"error": {"message": "...", "type": "..."}}`.
+fn error_frame_payload(error_type: &str, message: &str) -> String {
+    serde_json::to_string(&serde_json::json!({
+        "error": {
+            "message": message,
+            "type": error_type,
+        }
+    }))
+    // unreachable in practice — `serde_json::to_string` of a
+    // `Value` cannot fail. The fallback emits a minimal valid
+    // envelope so SDK consumers' `JSON.parse` still succeeds.
+    .unwrap_or_else(|_| {
+        r#"{"error":{"message":"error","type":"internal_error"}}"#.into()
+    })
+}

crates/aisix-proxy/src/lib.rs

@@ -616,6 +616,96 @@ data: [DONE]\n\n";
         );
     }
 
+    /// Issue #177: per docs/api-proxy.md §5, abnormal upstream
+    /// stream termination must close the response WITHOUT `[DONE]`
+    /// — SDK consumers that key off `[DONE]` for clean-completion
+    /// signal need to detect truncation. The previous behavior
+    /// emitted `[DONE]` after the SSE error event, masking
+    /// truncated responses as complete.
+    #[tokio::test]
+    async fn streaming_response_omits_done_when_upstream_returns_invalid_json_mid_stream() {
+        let upstream = MockServer::start().await;
+        // Upstream emits two valid SSE chunks then a malformed JSON
+        // payload. The malformed payload triggers `serde_json::from_str`
+        // to fail in the bridge's `build_chunk_stream`, surfacing as
+        // `BridgeError::UpstreamDecode` to `build_sse_stream`, which
+        // emits an SSE `event: error` frame. After that frame the
+        // proxy MUST NOT emit `[DONE]`.
+        let sse = "\
+data: {\"id\":\"up-1\",\"model\":\"gpt-4o\",\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\"},\"finish_reason\":null}]}\n\n\
+data: {\"id\":\"up-1\",\"model\":\"gpt-4o\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"partial \"},\"finish_reason\":null}]}\n\n\
+data: <not valid json>\n\n";
+        Mock::given(method("POST"))
+            .and(path("/chat/completions"))
+            .respond_with(
+                ResponseTemplate::new(200)
+                    .insert_header("content-type", "text/event-stream")
+                    .set_body_string(sse),
+            )
+            .mount(&upstream)
+            .await;
+
+        let hub = Arc::new(Hub::new());
+        hub.register(Provider::Openai, Arc::new(OpenAiBridge::new()));
+        let snap = seed_snapshot("my-gpt4", &["my-gpt4"], &upstream.uri());
+        let app = build_router(build_state(snap, hub));
+
+        let body = serde_json::json!({
+            "model": "my-gpt4",
+            "messages": [{"role": "user", "content": "hi"}],
+            "stream": true
+        });
+        let req = Request::builder()
+            .method("POST")
+            .uri("/v1/chat/completions")
+            .header("authorization", "Bearer sk-caller")
+            .header("content-type", "application/json")
+            .body(Body::from(body.to_string()))
+            .unwrap();
+        let resp = run(app, req).await;
+        assert_eq!(resp.status(), StatusCode::OK);
+
+        // Drain the wire bytes — at this layer we want byte-level
+        // assertions, not just decoded SSE events, because the
+        // contract being verified is "did `[DONE]` appear at all
+        // on the wire".
+        let mut body_stream = resp.into_body().into_data_stream();
+        let mut wire = Vec::new();
+        while let Some(chunk) = body_stream.next().await {
+            wire.extend_from_slice(chunk.unwrap().as_ref());
+        }
+        let wire_str = String::from_utf8(wire).expect("SSE bytes are utf8");
+
+        // Per docs §5: NO `[DONE]` after abnormal termination.
+        assert!(
+            !wire_str.contains("data: [DONE]"),
+            "abnormal termination MUST close without [DONE]; got wire:\n{wire_str}"
+        );
+        // The error event MUST be emitted so SDK consumers see a
+        // failure signal.
+        assert!(
+            wire_str.contains("event: error"),
+            "abnormal termination MUST emit `event: error`; got wire:\n{wire_str}"
+        );
+        // The error payload MUST be valid OpenAI-envelope JSON
+        // (the SDK does `JSON.parse(sse.data)` BEFORE checking
+        // event type, so plain-string payloads yield a SyntaxError
+        // instead of the typed APIError callers expect).
+        let err_event_idx = wire_str.find("event: error\n").unwrap();
+        let after_err = &wire_str[err_event_idx + "event: error\n".len()..];
+        let data_line = after_err
+            .lines()
+            .find(|l| l.starts_with("data: "))
+            .expect("error event followed by a data line");
+        let json_payload = &data_line["data: ".len()..];
+        let parsed: serde_json::Value = serde_json::from_str(json_payload)
+            .expect("error frame data must be valid OpenAI-envelope JSON");
+        assert!(
+            parsed.get("error").is_some(),
+            "error frame data must be `{{\"error\": {{...}}}}` shape; got {json_payload}"
+        );
+    }
+
     // ---- regression coverage for issue #107 -------------------------
     // Pre-fix only /v1/chat/completions enforced rate-limit / budget;
     // every other LLM endpoint silently bypassed both. The test below

tests/e2e/src/cases/streaming-edges-e2e.test.ts

@@ -12,31 +12,30 @@ import {
   type SpawnedApp,
 } from "../harness/index.js";
 
-// E2E: streaming-edge case. One production failure mode pinned
-// here, derived from `docs/api-proxy.md` §5:
+// E2E: streaming-edge cases derived from `docs/api-proxy.md` §5:
 //
 //   - Client abort mid-stream — caller starts a streaming chat
 //     completion, then aborts the request via AbortSignal. After
 //     the abort the gateway must remain HEALTHY: subsequent
 //     streaming calls from the same caller run to completion.
-//     This pins the "no resource leak / no parser corruption"
-//     property that's externally observable from the client.
+//     Pins the "no resource leak / no parser corruption" property
+//     that's externally observable from the client.
 //
-// (The "upstream disconnect mid-stream" case — partial chunks
-// reach the caller, then iterator surfaces error per docs §5
-// "aisix sends a final error chunk and closes without [DONE]" —
-// is held back. The gateway today short-circuits to a request-
-// time 502 instead of forwarding partial chunks. See follow-up
-// issue.)
+//   - Upstream disconnects mid-stream — mock upstream emits N
+//     SSE chunks then closes the TCP connection without `[DONE]`.
+//     Per docs §5 ("If the upstream stream terminates abnormally,
+//     aisix sends a final error chunk and closes the response
+//     without `[DONE]`"), the caller MUST see partial chunks +
+//     iterator-time error + no synthetic `finish_reason:"stop"`.
 //
-// Note on scope: this case verifies gateway LIVENESS post-abort,
-// not upstream-side disconnect propagation. The harness has no
-// signal for "upstream observed the client disconnect", so a
-// regression where the gateway holds the upstream connection open
-// (silently consuming chunks after the client aborted) would
-// pass this test. Filing that as a separate harness-extension
-// task; today's coverage is "gateway stays alive", which is the
-// load-bearing user-visible contract.
+// Note on scope (client-abort case): verifies gateway LIVENESS
+// post-abort, not upstream-side disconnect propagation. The
+// harness has no signal for "upstream observed the client
+// disconnect", so a regression where the gateway holds the
+// upstream connection open (silently consuming chunks after the
+// client aborted) would pass this test. Filing that as a separate
+// harness-extension task; today's coverage is "gateway stays
+// alive", which is the load-bearing user-visible contract.
 //
 // References:
 // - Gateway's own streaming contract: `docs/api-proxy.md` §5
@@ -184,4 +183,123 @@ describe("streaming edges e2e: client abort mid-stream", () => {
     expect(followupFinishReason).toBe("stop");
   });
 
+  test("upstream disconnects mid-stream: partial chunks reach caller, iterator throws, no [DONE]", async (ctx) => {
+    if (!etcdReachable || !app || !admin) {
+      ctx.skip();
+      return;
+    }
+
+    // Mock upstream emits 2 SSE chunks then drops the connection
+    // (`disconnectAfterEvents: 2`). No `finish_reason: "stop"`,
+    // no `[DONE]` from the upstream — premature close per docs §5.
+    //
+    // `eventDelayMs: 200` between writes ensures both chunks flush
+    // to the gateway BEFORE `res.destroy()` aborts the connection.
+    // Without the delay, write buffering on the mock side can
+    // cause the destroy to land before the second chunk reaches
+    // the wire, making the test flaky on which chunks the gateway
+    // actually receives. Matches the peer client-abort case for
+    // consistency.
+    const upstream = await startOpenAiUpstream({
+      streamEvents: [
+        '{"id":"disc","object":"chat.completion.chunk","model":"gpt-4o-mini","choices":[{"index":0,"delta":{"role":"assistant"},"finish_reason":null}]}',
+        '{"id":"disc","object":"chat.completion.chunk","model":"gpt-4o-mini","choices":[{"index":0,"delta":{"content":"partial "},"finish_reason":null}]}',
+        // Never emitted (disconnect happens at i >= 2):
+        '{"id":"disc","object":"chat.completion.chunk","model":"gpt-4o-mini","choices":[{"index":0,"delta":{"content":"never "},"finish_reason":null}]}',
+        '{"id":"disc","object":"chat.completion.chunk","model":"gpt-4o-mini","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}',
+        "[DONE]",
+      ],
+      eventDelayMs: 200,
+      disconnectAfterEvents: 2,
+    });
+    upstreams.push(upstream);
+
+    const pk = await admin.createProviderKey({
+      display_name: "stream-disc-pk",
+      secret: "sk-mock",
+      api_base: `${upstream.baseUrl}/v1`,
+    });
+    await admin.createModel({
+      display_name: "stream-disc",
+      provider: "openai",
+      model_name: "gpt-4o-mini",
+      provider_key_id: pk.id,
+    });
+
+    const client = new OpenAI({
+      apiKey: CALLER_PLAINTEXT,
+      baseURL: `${app.proxyUrl}/v1`,
+      maxRetries: 0,
+    });
+
+    // Use ProxyClient.listModels for snapshot-readiness gating —
+    // chat-completions probes don't work cleanly here because the
+    // mock upstream is configured for streaming with disconnects,
+    // so any chat-completions probe (with or without `stream:true`)
+    // would match the streaming-cutoff failure mode the test is
+    // meant to verify.
+    const probe = new ProxyClient(app.proxyUrl, CALLER_PLAINTEXT);
+    await waitConfigPropagation(async () => {
+      const r = await probe.listModels();
+      if (r.status !== 200) return false;
+      const data = (r.body as { data?: Array<{ id?: string }> }).data ?? [];
+      return data.some((m) => m.id === "stream-disc");
+    });
+
+    // Per docs §5 strict contract:
+    //   - 2 partial chunks reach the caller (the bytes the upstream
+    //     emitted before the close)
+    //   - Iterator THROWS during streaming (gateway emits SSE
+    //     `event: error` chunk, OpenAI SDK surfaces it as iteration
+    //     error)
+    //   - NO synthetic `finish_reason: "stop"` injection (would
+    //     turn truncated responses into apparently-complete ones)
+    //   - NO synthetic `[DONE]` after error (would make the SDK
+    //     treat the truncated stream as a clean completion)
+    const collected: string[] = [];
+    let sawFinish = false;
+    let surfacedError = false;
+    let errCtor = "";
+    let errMessage = "";
+
+    const stream = await client.chat.completions.create({
+      model: "stream-disc",
+      messages: [{ role: "user", content: "give me content" }],
+      stream: true,
+    });
+    try {
+      for await (const chunk of stream) {
+        const delta = chunk.choices[0]?.delta;
+        if (delta?.content) collected.push(delta.content);
+        if (chunk.choices[0]?.finish_reason) sawFinish = true;
+      }
+    } catch (e) {
+      surfacedError = true;
+      errCtor =
+        (e as { constructor?: { name?: string } })?.constructor?.name ?? "";
+      errMessage = (e as { message?: string })?.message ?? "";
+    }
+
+    // Partial chunks reached the caller. Chunk 2 carried the
+    // string "partial "; chunk 1 was role-only with no content.
+    expect(collected.join("")).toBe("partial ");
+    // No fake completion signal. The upstream's mid-stream close
+    // happened BEFORE the chunk that carries finish_reason:"stop".
+    expect(sawFinish).toBe(false);
+    // Iterator surfaced an error. The OpenAI Node SDK throws on
+    // an SSE `event: error` chunk during stream iteration.
+    expect(surfacedError).toBe(true);
+
+    // Tighten: distinguish typed `APIError` from `SyntaxError`.
+    // Without this, a regression that re-introduces a plain-string
+    // `data:` payload on the SSE error frame would still pass
+    // `surfacedError === true` because the SDK's `JSON.parse(sse.data)`
+    // (in `streaming.ts`, called BEFORE the `sse.event === "error"`
+    // check) would throw `SyntaxError("Could not parse message into
+    // JSON: ...")` instead of the typed `APIError` callers expect.
+    // See OpenAI Node SDK <https://github.com/openai/openai-node/blob/main/src/streaming.ts>
+    // for the parse-then-classify ordering.
+    expect(errCtor).not.toBe("SyntaxError");
+    expect(errMessage).not.toContain("Could not parse message into JSON");
+  });
 });