1

prisma/services/deepThink/expert.ts

@@ -1,5 +1,6 @@
 import { ModelOption, ExpertResult } from '../../types';
 import { getExpertSystemInstruction } from './prompts';
+import { withRetry } from '../utils/retry';
 
 export const streamExpertResponse = async (
   ai: any,
@@ -10,7 +11,10 @@ export const streamExpertResponse = async (
   signal: AbortSignal,
   onChunk: (text: string, thought: string) => void
 ): Promise<void> => {
-  const streamResult = await ai.models.generateContentStream({
+  // We wrap the stream initiation in retry. 
+  // If the stream is successfully established but fails during iteration, 
+  // we catch that separately.
+  const streamResult = await withRetry(() => ai.models.generateContentStream({
     model: model,
     contents: expert.prompt,
     config: {
@@ -21,23 +25,30 @@ export const streamExpertResponse = async (
           includeThoughts: true 
       }
     }
-  });
+  }));
 
-  for await (const chunk of streamResult) {
-     if (signal.aborted) break;
+  try {
+    for await (const chunk of streamResult) {
+       if (signal.aborted) break;
 
-     let chunkText = "";
-     let chunkThought = "";
+       let chunkText = "";
+       let chunkThought = "";
 
-     if (chunk.candidates?.[0]?.content?.parts) {
-         for (const part of chunk.candidates[0].content.parts) {
-             if (part.thought) {
-                 chunkThought += (part.text || "");
-             } else if (part.text) {
-                 chunkText += part.text;
-             }
-         }
-         onChunk(chunkText, chunkThought);
-     }
+       if (chunk.candidates?.[0]?.content?.parts) {
+           for (const part of chunk.candidates[0].content.parts) {
+               if (part.thought) {
+                   chunkThought += (part.text || "");
+               } else if (part.text) {
+                   chunkText += part.text;
+               }
+           }
+           onChunk(chunkText, chunkThought);
+       }
+    }
+  } catch (streamError) {
+    console.error(`Stream interrupted for expert ${expert.role}:`, streamError);
+    // We don't retry mid-stream automatically here to avoid complex state management,
+    // but the initial connection is protected by withRetry.
+    throw streamError;
   }
 };