fix: online stream message issue (#71)

.gitignore

@@ -38,4 +38,5 @@ yarn-error.log*
 .env*.local
 
 # ts
-tsconfig.tsbuildinfo
+tsconfig.tsbuildinfo
+certificates

app/api/vision-agent/route.ts

@@ -1,7 +1,8 @@
-import { StreamingTextResponse } from 'ai';
+import { StreamingTextResponse, experimental_StreamData } from 'ai';
 
 // import { auth } from '@/auth';
 import { MessageUI, SignedPayload } from '@/lib/types';
+
 import { logger, withLogging } from '@/lib/logger';
 import { CLEANED_SEPARATOR } from '@/lib/constants';
 import { cleanAnswerMessage, cleanInputMessage } from '@/lib/utils/content';
@@ -145,29 +146,21 @@ export const POST = withLogging(
         );
       }
     }
+    // const streamData = new experimental_StreamData();
 
     if (fetchResponse.body) {
       const encoder = new TextEncoder();
       const decoder = new TextDecoder('utf-8');
       let maxChunkSize = 0;
-      const stream = fetchResponse.body.pipeThrough(
-        new TransformStream({
-          flush: () => {
-            logger.info(
-              session,
-              {
-                message: 'Streaming finished',
-                maxChunkSize,
-              },
-              request,
-              '__AGENT_DONE',
-            );
-          },
-          transform: async (chunk, controller) => {
-            const data = decoder.decode(chunk, { stream: true });
+      const stream = new ReadableStream({
+        async start(controller) {
+          // const parser = createParser(streamParser);
+          for await (const chunk of fetchResponse.body as any) {
+            const data = decoder.decode(chunk);
             maxChunkSize = Math.max(data.length, maxChunkSize);
             const lines = data.split('\n');
             const results = [];
+            let done = false;
             for (let line of lines) {
               if (!line.trim()) {
                 continue;
@@ -194,6 +187,7 @@ export const POST = withLogging(
                 }
                 msg.payload.result = JSON.stringify(result);
                 results.push(JSON.stringify(msg));
+                done = true;
               } catch (e) {
                 console.error(e);
                 logger.error(
@@ -203,6 +197,7 @@ export const POST = withLogging(
                   },
                   request,
                 );
+                controller.error(e);
               }
             }
             controller.enqueue(
@@ -210,10 +205,26 @@ export const POST = withLogging(
                 results.length === 0 ? '' : results.join('\n') + '\n',
               ),
             );
-          },
-        }),
-      );
-      return new StreamingTextResponse(stream);
+            if (done) {
+              logger.info(
+                session,
+                {
+                  message: 'Streaming finished',
+                  maxChunkSize,
+                },
+                request,
+                '__AGENT_DONE',
+              );
+              controller.close();
+            }
+          }
+        },
+      });
+      return new Response(stream, {
+        headers: {
+          'Content-Type': 'application/x-ndjson',
+        },
+      });
     } else {
       return fetchResponse;
     }

components/chat/ChatMessage.tsx

@@ -1,5 +1,4 @@
 import { useMemo, useState } from 'react';
-import { cn } from '@/lib/utils';
 import { CodeBlock } from '@/components/ui/CodeBlock';
 import {
   IconCheckCircle,
@@ -114,7 +113,7 @@ const ChunkPayloadAction: React.FC<{
             <IconListUnordered />
           </Button>
         </DialogTrigger>
-        <DialogContent className="max-w-5xl">
+        <DialogContent className="max-w-5xl" onOpenAutoFocus={e => e.preventDefault()}>
           <Table>
             <TableHeader>
               <TableRow className="border-primary/50">