Refine AI chat streaming API

Author: IzumiSy

docs/api/create-ai-gateway-client.md

@@ -34,11 +34,11 @@ function createAIGatewayClient(config: {
 
 ```typescript
 interface AIGatewayClient {
-  chatCompletionStream(request: AIGatewayChatRequest): AsyncGenerator<string, void, unknown>;
+  chatCompletionStream(request: AIGatewayChatRequest): AsyncIterable<string>;
 }
 ```
 
-The generator yields text deltas. Concatenate them to build the assistant response.
+The iterable yields text deltas. Concatenate them to build the assistant response.
 
 ## Related Types
 
@@ -51,6 +51,7 @@ interface AIGatewayChatMessage {
 interface AIGatewayChatRequest {
   model: string;
   messages: AIGatewayChatMessage[];
+  stream?: boolean;
   signal?: AbortSignal;
 }
 ```
@@ -83,8 +84,9 @@ console.log(deltas.join(""));
 
 ## Notes
 
-- Non-Gemini models are consumed as streaming SSE responses
-- `gemini-*` models are consumed as JSON responses and yielded once
+- `stream` defaults to `true`
+- Pass `stream: false` when the endpoint returns a single JSON response instead of SSE
+- Gemini-backed routes often need `stream: false`
 - The public API is intentionally text-only
 
 ## Related

docs/api/use-ai-chat.md

@@ -10,11 +10,11 @@ React hook for simple text-only chat on top of `createAIGatewayClient`.
 ## Signature
 
 ```typescript
-const useAIChat: (config: { client: AIGatewayClient; model: string }) => {
+const useAIChat: (config: { client: AIGatewayClient; model: string; stream?: boolean }) => {
   messages: AIChatMessage[];
   status: "ready" | "submitted" | "streaming" | "error";
   error?: Error;
-  sendMessage: (message: { text: string } | string) => Promise<void>;
+  sendMessage: (message: string) => Promise<boolean>;
   stop: () => void;
 };
 ```
@@ -43,8 +43,9 @@ interface AIChatMessage {
 
 ### `sendMessage()`
 
-- **Type:** `(message: { text: string } | string) => Promise<void>`
+- **Type:** `(message: string) => Promise<boolean>`
 - **Description:** Appends a user message and streams the assistant response
+- **Returns:** `true` when the request completes successfully, `false` when the call is ignored, stopped, or fails
 
 ### `stop()`
 
@@ -81,7 +82,7 @@ export function ChatScreen() {
       ))}
 
       <button
-        onClick={() => sendMessage("Hello")}
+        onClick={() => void sendMessage("Hello")}
         disabled={status === "submitted" || status === "streaming"}
       >
         Send
@@ -97,9 +98,11 @@ export function ChatScreen() {
 
 ## Notes
 
+- `stream` defaults to `true`
+- Pass `stream: false` when the endpoint returns a single JSON response instead of SSE
 - The hook is intentionally text-only
 - System prompts and custom history shaping should use the low-level client directly
-- `stop()` keeps any already-streamed assistant text
+- `stop()` keeps any already-streamed assistant text and ignores late chunks from the stopped request
 
 ## Related
 

docs/concepts/authentication.md

@@ -185,7 +185,7 @@ function ChatScreen() {
         </div>
       ))}
 
-      <button onClick={() => sendMessage("Hello")}>Send</button>
+      <button onClick={() => void sendMessage("Hello")}>Send</button>
       <button onClick={stop} disabled={status !== "submitted" && status !== "streaming"}>
         Stop
       </button>

packages/core/src/ai/client.test.ts

@@ -82,7 +82,7 @@ describe("createAIGatewayClient", () => {
     });
   });
 
-  it("yields a single text chunk for gemini models", async () => {
+  it("yields a single text chunk for json responses", async () => {
     const authClient = createMockAuthClient(
       new Response(
         JSON.stringify({
@@ -105,7 +105,10 @@ describe("createAIGatewayClient", () => {
       authClient,
     });
 
-    const deltas = await collectStream(client, createRequest({ model: "gemini-2.5-flash" }));
+    const deltas = await collectStream(
+      client,
+      createRequest({ model: "gemini-2.5-flash", stream: false }),
+    );
 
     expect(deltas).toEqual(["Grounded answer"]);
     expect(authClient.fetch).toHaveBeenCalledWith(
@@ -120,6 +123,7 @@ describe("createAIGatewayClient", () => {
     ).toEqual({
       model: "gemini-2.5-flash",
       messages: [{ role: "user", content: "Hello" }],
+      stream: false,
     });
   });
 
@@ -149,7 +153,7 @@ describe("createAIGatewayClient", () => {
     );
   });
 
-  it("throws when gemini responses do not include assistant text", async () => {
+  it("throws when json responses do not include assistant text", async () => {
     const authClient = createMockAuthClient(
       new Response(JSON.stringify({ choices: [{ message: { content: [] } }] }), { status: 200 }),
     );
@@ -159,7 +163,7 @@ describe("createAIGatewayClient", () => {
     });
 
     await expect(
-      collectStream(client, createRequest({ model: "gemini-2.5-flash" })),
-    ).rejects.toThrow("AI Gateway Gemini response did not include assistant text.");
+      collectStream(client, createRequest({ model: "gemini-2.5-flash", stream: false })),
+    ).rejects.toThrow("AI Gateway JSON response did not include assistant text.");
   });
 });

packages/core/src/ai/client.ts

@@ -8,11 +8,15 @@ export interface AIGatewayChatMessage {
 export interface AIGatewayChatRequest {
   model: string;
   messages: AIGatewayChatMessage[];
+  stream?: boolean;
   signal?: AbortSignal;
 }
 
 export interface AIGatewayClient {
-  chatCompletionStream(request: AIGatewayChatRequest): AsyncGenerator<string, void, unknown>;
+  /**
+   * Honor request.signal when possible so callers can stop work early.
+   */
+  chatCompletionStream(request: AIGatewayChatRequest): AsyncIterable<string>;
 }
 
 interface OpenAIStreamChunk {
@@ -39,8 +43,8 @@ export function createAIGatewayClient(config: {
 
   return {
     async *chatCompletionStream(request) {
-      if (isGeminiModel(request.model)) {
-        yield* streamGeminiResponse({
+      if (request.stream === false) {
+        yield* streamJSONResponse({
           endpoint,
           authClient: config.authClient,
           request,
@@ -96,7 +100,7 @@ async function* streamOpenAICompatibleResponse(input: {
   }
 }
 
-async function* streamGeminiResponse(input: {
+async function* streamJSONResponse(input: {
   endpoint: string;
   authClient: AuthClient;
   request: AIGatewayChatRequest;
@@ -110,17 +114,18 @@ async function* streamGeminiResponse(input: {
     body: JSON.stringify({
       model: input.request.model,
       messages: input.request.messages,
+      stream: false,
     }),
     signal: input.request.signal,
   });
 
-  await assertOK(response, "AI Gateway Gemini request");
+  await assertOK(response, "AI Gateway JSON request");
 
   const payload = parseJSON<OpenAIFinalResponse>(await response.text(), "AI Gateway JSON response");
   const text = extractText(payload.choices?.[0]?.message?.content);
 
   if (!text) {
-    throw new Error("AI Gateway Gemini response did not include assistant text.");
+    throw new Error("AI Gateway JSON response did not include assistant text.");
   }
 
   yield text;
@@ -229,10 +234,6 @@ async function* iterateSSEDataEvents(
   }
 }
 
-function isGeminiModel(model: string): boolean {
-  return model.startsWith("gemini-");
-}
-
 function withTrailingSlash(value: string): string {
   return value.endsWith("/") ? value : `${value}/`;
 }

packages/core/src/ai/use-ai-chat.test.tsx

@@ -66,6 +66,7 @@ describe("useAIChat", () => {
     expect(client.chatCompletionStream).toHaveBeenCalledWith({
       model: "gpt-5-mini",
       messages: [{ role: "user", content: "Hi" }],
+      stream: true,
       signal: expect.any(AbortSignal),
     });
 
@@ -131,6 +132,78 @@ describe("useAIChat", () => {
     });
   });
 
+  it("ignores stale chunks after stop and allows the next send", async () => {
+    let releaseStoppedRequest!: () => void;
+    const stoppedRequestGate = new Promise<void>((resolve) => {
+      releaseStoppedRequest = resolve;
+    });
+
+    const client = {
+      chatCompletionStream: vi.fn(async function* ({ messages }) {
+        const text = messages.at(-1)?.content;
+
+        if (text === "First") {
+          yield "Partial";
+          await stoppedRequestGate;
+          yield " stale";
+          return;
+        }
+
+        yield "Fresh";
+      }),
+    } satisfies AIGatewayClient;
+
+    const { result } = renderHook(() => useAIChat({ client, model: "gpt-5-mini" }));
+
+    let firstSendPromise: Promise<boolean> | undefined;
+    await act(async () => {
+      firstSendPromise = result.current.sendMessage("First");
+    });
+
+    await waitFor(() => {
+      expect(result.current.status).toBe("streaming");
+      expect(result.current.messages).toEqual([
+        { id: "id-1", role: "user", content: "First" },
+        { id: "id-2", role: "assistant", content: "Partial" },
+      ]);
+    });
+
+    act(() => {
+      result.current.stop();
+    });
+
+    await waitFor(() => {
+      expect(result.current.status).toBe("ready");
+    });
+
+    await act(async () => {
+      await expect(result.current.sendMessage("Second")).resolves.toBe(true);
+    });
+
+    await waitFor(() => {
+      expect(result.current.messages).toEqual([
+        { id: "id-1", role: "user", content: "First" },
+        { id: "id-2", role: "assistant", content: "Partial" },
+        { id: "id-3", role: "user", content: "Second" },
+        { id: "id-4", role: "assistant", content: "Fresh" },
+      ]);
+    });
+
+    releaseStoppedRequest();
+    await act(async () => {
+      await expect(firstSendPromise).resolves.toBe(false);
+    });
+
+    await waitFor(() => {
+      expect(result.current.messages).toEqual([
+        { id: "id-1", role: "user", content: "First" },
+        { id: "id-2", role: "assistant", content: "Partial" },
+        { id: "id-3", role: "user", content: "Second" },
+        { id: "id-4", role: "assistant", content: "Fresh" },
+      ]);
+    });
+  });
+
   it("sets error state and recovers on the next successful send", async () => {
     let shouldFail = true;
     const client = {