fix: when api sending multiple chunks together with incomplete json (#52)

app/api/vision-agent/route.ts

@@ -10,6 +10,46 @@ import { cleanAnswerMessage, cleanInputMessage } from '@/lib/messageUtils';
 export const dynamic = 'force-dynamic';
 export const maxDuration = 300; // This function can run for a maximum of 5 minutes
 
+const parseLine = (line: string) => {
+  try {
+    const json = JSON.parse(line);
+    let message = (json.log ?? '') + '\n';
+    if (json.task || json.plan || json.reflection) {
+      const arr = json.plan
+        ? json.plan
+        : json.task
+          ? [json.task]
+          : [json.reflection];
+      const keys = Object.keys(arr[0]);
+      message += '\n';
+      message += '| ' + keys.join(' | ') + ' |' + '\n';
+      message += new Array(keys.length + 1).fill('|').join(' :- ') + '\n';
+      arr.forEach((obj: any) => {
+        message += '| ' + keys.map(key => obj[key]).join(' | ') + ' |' + '\n';
+      });
+      message += '\n';
+    }
+    if (json.tools) {
+      message += '\n';
+      message += '| ' + 'Descriptions' + ' |' + '\n';
+      message += '| ' + ':-' + ' |' + '\n';
+      json.tools.forEach((tool: string) => {
+        message += '| ' + tool + ' |' + '\n';
+      });
+      message += '\n';
+    }
+    if (json.code) {
+      message += `\`\`\`python\n${json.code}\n\`\`\`\n`;
+    }
+    if (json.result) {
+      message += `\`\`\`\n${json.result}\n\`\`\`\n`;
+    }
+    return { message };
+  } catch (e) {
+    return { error: e };
+  }
+};
+
 export const POST = withLogging(
   async (
     session,
@@ -110,74 +150,41 @@ export const POST = withLogging(
 
     if (fetchResponse.body) {
       const encoder = new TextEncoder();
-      const decoder = new TextDecoder();
+      const decoder = new TextDecoder('utf-8');
+      let buffer = '';
       const stream = fetchResponse.body.pipeThrough(
         new TransformStream({
           transform: async (chunk, controller) => {
-            const data = decoder.decode(chunk);
-            data.split('\n').forEach(line => {
+            const data = decoder.decode(chunk, { stream: true });
+            buffer += data;
+            let lines = buffer.split('\n');
+            buffer = lines.pop() ?? ''; // Save the last incomplete line back to the buffer
+            for (let line of lines) {
               if (!line.trim()) {
-                return;
+                continue;
               }
-              try {
-                const json = JSON.parse(line);
-                let message = (json.log ?? '') + '\n';
-                if (json.task || json.plan || json.reflection) {
-                  const arr = json.plan
-                    ? json.plan
-                    : json.task
-                      ? [json.task]
-                      : [json.reflection];
-                  const keys = Object.keys(arr[0]);
-                  message += '\n';
-                  message += '| ' + keys.join(' | ') + ' |' + '\n';
-                  message +=
-                    new Array(keys.length + 1).fill('|').join(' :- ') + '\n';
-                  arr.forEach((obj: any) => {
-                    message +=
-                      '| ' +
-                      keys.map(key => obj[key]).join(' | ') +
-                      ' |' +
-                      '\n';
-                  });
-                  message += '\n';
+              if (line.trim()) {
+                const { message, error } = parseLine(line.trim());
+                if (message) {
+                  controller.enqueue(encoder.encode(message));
+                } else if (error) {
+                  logger.error(
+                    session,
+                    { message: (error as Error).message, data },
+                    request,
+                  );
+                  controller.error(error);
+                  controller.terminate();
                 }
-                if (json.tools) {
-                  message += '\n';
-                  message += '| ' + 'Descriptions' + ' |' + '\n';
-                  message += '| ' + ':-' + ' |' + '\n';
-                  json.tools.forEach((tool: string) => {
-                    message += '| ' + tool + ' |' + '\n';
-                  });
-                  message += '\n';
-                }
-                if (json.code) {
-                  message += `\`\`\`python\n${json.code}\n\`\`\`\n`;
-                }
-                if (json.result) {
-                  message += `\`\`\`\n${json.result}\n\`\`\`\n`;
-                }
-
-                logger.info(
-                  session,
-                  {
-                    message,
-                  },
-                  request,
-                  '__AGENT_RESPONSE',
-                );
+              }
+            }
+            if (buffer.trim()) {
+              const { message } = parseLine(buffer.trim());
+              if (message) {
+                buffer = '';
                 controller.enqueue(encoder.encode(message));
-              } catch (e) {
-                console.log(data);
-                logger.error(
-                  session,
-                  { message: (e as Error).message, data },
-                  request,
-                );
-                controller.error(e);
-                controller.terminate();
               }
-            });
+            }
           },
         }),
       );