improve stream completion handling and database sync; add error handling for SQL; updated some tests

whoiskatrin · threepointone · commit a66c0c8a1aa3 · 2025-11-07T12:20:09.000Z
diff --git a/packages/agents/src/resumable-stream-manager.ts b/packages/agents/src/resumable-stream-manager.ts
@@ -380,25 +380,77 @@ export class ResumableStreamManager<Message extends ChatMessage = ChatMessage> {
     const assistantMessageId = `assistant_${nanoid()}`;
     let buffer = "";
 
+    let completedNaturally = false;
     try {
       while (true) {
+        // Check if stream was completed by onFinish callback
+        const currentState = this._activeStreams.get(streamId);
+        if (currentState?.completed) {
+          // Ensure database state is synchronized
+          try {
+            this.sql`
+              update cf_ai_http_chat_streams
+              set fetching = 0, completed = 1, updated_at = current_timestamp
+              where stream_id = ${streamId}
+            `;
+          } catch (sqlError) {
+            console.error(
+              `[ResumableStreamManager] Error syncing completion state for ${streamId}:`,
+              sqlError
+            );
+          }
+
+          completedNaturally = true;
+          break;
+        }
+
         const { done, value } = await reader.read();
-        if (done) break;
+        if (done) {
+          completedNaturally = true;
+          break;
+        }
 
-        // Store raw chunk with sequence number
-        const seq = streamState.seq++;
+        // Store raw chunk with sequence number atomically
         const chunkBase64 = btoa(String.fromCharCode(...value));
-        this.sql`
-          insert into cf_ai_http_chat_chunks (stream_id, seq, chunk)
-          values (${streamId}, ${seq}, ${chunkBase64})
-        `;
 
-        // Update sequence in stream state
-        this.sql`
-          update cf_ai_http_chat_streams
-          set seq = ${streamState.seq}, updated_at = current_timestamp
-          where stream_id = ${streamId}
-        `;
+        try {
+          // Atomically get next sequence number and insert chunk
+          const seqResult = this.sql`
+            update cf_ai_http_chat_streams
+            set seq = seq + 1, updated_at = current_timestamp
+            where stream_id = ${streamId}
+            returning seq
+          `;
+
+          const seq = Number(seqResult[0]?.seq) || streamState.seq++;
+
+          this.sql`
+            insert into cf_ai_http_chat_chunks (stream_id, seq, chunk)
+            values (${streamId}, ${seq}, ${chunkBase64})
+          `;
+
+          // Update in-memory state to match database
+          streamState.seq = seq + 1;
+        } catch (sqlError) {
+          console.error(
+            `[ResumableStreamManager] SQL error for stream ${streamId}:`,
+            sqlError
+          );
+          // Fall back to in-memory sequence if SQL fails
+          const seq = streamState.seq++;
+          try {
+            this.sql`
+              insert into cf_ai_http_chat_chunks (stream_id, seq, chunk)
+              values (${streamId}, ${seq}, ${chunkBase64})
+            `;
+          } catch (fallbackError) {
+            console.error(
+              `[ResumableStreamManager] Fallback SQL error for stream ${streamId}:`,
+              fallbackError
+            );
+            // Continue processing even if storage fails
+          }
+        }
 
         // Parse for assistant message content
         const chunk = decoder.decode(value, { stream: true });
@@ -456,33 +508,28 @@ export class ResumableStreamManager<Message extends ChatMessage = ChatMessage> {
         await persistMessages([...messages, assistantMessage]);
       }
     } finally {
-      // Mark stream as completed
-      streamState.fetching = false;
-      streamState.completed = true;
-
-      this.sql`
-        update cf_ai_http_chat_streams
-        set fetching = 0, completed = 1, updated_at = current_timestamp
-        where stream_id = ${streamId}
-      `;
-
-      // Close all readers/writers
-      for (const readerOrWriter of streamState.readers) {
+      // Only mark as completed if stream finished naturally, not if interrupted
+      if (completedNaturally) {
+        this._markStreamCompleted(streamId);
+      } else {
+        // Stream was interrupted - update fetching state but don't mark as completed
+        const currentState = this._activeStreams.get(streamId);
+        if (currentState) {
+          currentState.fetching = false;
+        }
         try {
-          if (readerOrWriter instanceof WritableStreamDefaultWriter) {
-            readerOrWriter.close();
-          } else {
-            // Handle ReadableStreamDefaultController
-            if (
-              "close" in readerOrWriter &&
-              typeof readerOrWriter.close === "function"
-            ) {
-              readerOrWriter.close();
-            }
-          }
-        } catch {}
+          this.sql`
+            update cf_ai_http_chat_streams
+            set fetching = 0, updated_at = current_timestamp
+            where stream_id = ${streamId}
+          `;
+        } catch (sqlError) {
+          console.error(
+            `[ResumableStreamManager] Error updating fetching state for ${streamId}:`,
+            sqlError
+          );
+        }
       }
-      streamState.readers.clear();
     }
   }
 
@@ -566,13 +613,29 @@ export class ResumableStreamManager<Message extends ChatMessage = ChatMessage> {
           }
         });
 
-        // 2. If still fetching, add to live readers
+        // 2. Check if stream is truly complete by verifying both in-memory and database state
         const currentState = this._activeStreams.get(streamId);
-        if (currentState?.fetching) {
-          currentState.readers.add(writer);
+
+        // Get the latest database state to ensure consistency
+        const dbState = this.sql`
+          select fetching, completed from cf_ai_http_chat_streams
+          where stream_id = ${streamId}
+        `[0] as unknown as
+          | Pick<StreamStateRow, "fetching" | "completed">
+          | undefined;
+
+        const isStillFetching =
+          currentState?.fetching || dbState?.fetching === 1;
+        const isCompleted = currentState?.completed && dbState?.completed === 1;
+
+        if (isStillFetching && !isCompleted) {
+          // Stream is still active, join as live reader
+          if (currentState) {
+            currentState.readers.add(writer);
+          }
           await this._backfillGaps(streamId, writer, lastSeenSeq + 1);
         } else {
-          // Stream is complete
+          // Stream is complete, close writer
           await writer.close();
         }
       } catch (error) {
@@ -687,11 +750,20 @@ export class ResumableStreamManager<Message extends ChatMessage = ChatMessage> {
       streamState.readers.clear();
     }
 
-    this.sql`
-      update cf_ai_http_chat_streams
-      set fetching = 0, completed = 1, updated_at = current_timestamp
-      where stream_id = ${streamId}
-    `;
+    // Update database state with error handling
+    try {
+      this.sql`
+        update cf_ai_http_chat_streams
+        set fetching = 0, completed = 1, updated_at = current_timestamp
+        where stream_id = ${streamId}
+      `;
+    } catch (sqlError) {
+      console.error(
+        `[ResumableStreamManager] Error marking stream ${streamId} completed:`,
+        sqlError
+      );
+      // Stream is still marked as completed in memory even if SQL fails
+    }
 
     // Clean up from memory after some time
     setTimeout(() => {
diff --git a/packages/agents/src/tests/resumable-streaming.test.ts b/packages/agents/src/tests/resumable-streaming.test.ts
@@ -1,8 +1,7 @@
 import { createExecutionContext, env } from "cloudflare:test";
 import { describe, it, expect, beforeEach } from "vitest";
-import worker, { type Env, TestResumableStreamAgent } from "./worker";
+import worker, { type Env } from "./worker";
 import { nanoid } from "nanoid";
-import type { ExecutionContext } from "@cloudflare/workers-types";
 
 declare module "cloudflare:test" {
   interface ProvidedEnv extends Env {}
@@ -12,7 +11,7 @@ async function makeRequest(
   agentId: string,
   method: string,
   path: string,
-  body?: any
+  body?: unknown
 ) {
   const ctx = createExecutionContext();
   const url = `http://example.com/agents/resumable-stream-agent/${agentId}${path}`;
@@ -31,7 +30,7 @@ async function makeRequest(
 
 async function readPartialStreamChunks(
   response: Response,
-  maxChunks: number = 3
+  maxChunks = 3
 ): Promise<{
   chunks: string[];
   reader: ReadableStreamDefaultReader<Uint8Array>;
@@ -375,18 +374,17 @@ describe("Resumable Streaming - Multiple Clients", () => {
     const text1 = extractTextFromSSE(chunk1);
     await reader1.cancel();
 
-    // Resume and interrupt again
+    // Wait a bit to let the stream complete in background
+    await new Promise((resolve) => setTimeout(resolve, 300));
+
+    // Resume and read from completed stream (should have all chunks now)
     const response2 = await makeRequest(
       agentId,
       "GET",
       `/stream/${customStreamId}`
     );
-    const { chunks: chunk2, reader: reader2 } = await readPartialStreamChunks(
-      response2,
-      6
-    );
+    const chunk2 = await readStreamChunks(response2);
     const text2 = extractTextFromSSE(chunk2);
-    await reader2.cancel();
 
     // Final resume; should get complete stream
     const response3 = await makeRequest(
@@ -816,7 +814,7 @@ describe("Resumable Streaming - Error Handling and Edge Cases", () => {
           if (rawChunks.length > 10) break;
         }
         await reader.cancel();
-      } catch (error) {
+      } catch (_error) {
         // Expected when canceling
       }
     }
diff --git a/packages/agents/src/tests/worker.ts b/packages/agents/src/tests/worker.ts
@@ -11,7 +11,11 @@ import {
   type WSMessage
 } from "../index.ts";
 import { AIChatAgent } from "../ai-chat-agent.ts";
-import type { UIMessage as ChatMessage } from "ai";
+import type {
+  UIMessage as ChatMessage,
+  StreamTextOnFinishCallback,
+  ToolSet
+} from "ai";
 import type { MCPClientConnection } from "../mcp/client-connection";
 import { AIHttpChatAgent } from "../ai-chat-agent-http.ts";
 
@@ -404,7 +408,7 @@ export class TestResumableStreamAgent extends AIHttpChatAgent<
   // Track requests for testing
   requestHistory: Array<{ method: string; url: string; body?: unknown }> = [];
 
-  constructor(ctx: any, env: Env) {
+  constructor(ctx: DurableObjectState, env: Env) {
     super(ctx, env);
 
     // Set up some mock responses
@@ -420,7 +424,7 @@ export class TestResumableStreamAgent extends AIHttpChatAgent<
   }
 
   async onChatMessage(
-    onFinish: any,
+    onFinish: StreamTextOnFinishCallback<ToolSet>,
     options?: { streamId?: string }
   ): Promise<Response | undefined> {
     // Track the request
@@ -443,7 +447,7 @@ export class TestResumableStreamAgent extends AIHttpChatAgent<
       .join(" ");
 
     // Find mock response or use default
-    let responseText =
+    const responseText =
       this.mockResponses.get(content.toLowerCase()) || `Echo: ${content}`;
 
     // Create a streaming response
@@ -468,8 +472,22 @@ export class TestResumableStreamAgent extends AIHttpChatAgent<
         controller.enqueue(new TextEncoder().encode("data: [DONE]\n\n"));
         controller.close();
 
-        // Call onFinish
-        await onFinish();
+        // Call onFinish with mock result for testing
+        await onFinish({
+          text: responseText,
+          toolCalls: [],
+          toolResults: [],
+          finishReason: "stop",
+          usage: {
+            promptTokens: 0,
+            completionTokens: 0,
+            totalTokens: 0
+          },
+          rawResponse: {
+            headers: {}
+          },
+          warnings: []
+        } as unknown as Parameters<typeof onFinish>[0]);
       }
     });
 
diff --git a/packages/agents/src/tests/wrangler.jsonc b/packages/agents/src/tests/wrangler.jsonc
@@ -8,6 +8,11 @@
     "enable_nodejs_http_modules",
     "enable_nodejs_perf_hooks_module"
   ],
+
+  "observability": {
+    "enabled": true,
+    "head_sampling_rate": 1
+  },
   "durable_objects": {
     "bindings": [
       {
