feat: add timeout and remove png from code result (#80)

the error message could not be surfaced somehow and instead got network
error, but it timeout as expected

app/api/vision-agent/route.ts

@@ -11,6 +11,7 @@ import { getPresignedUrl } from '@/lib/aws';
 // export const runtime = 'edge';
 export const dynamic = 'force-dynamic';
 export const maxDuration = 300; // This function can run for a maximum of 5 minutes
+const TIMEOUT_MILI_SECONDS = 5 * 60 * 1000;
 
 const uploadBase64 = async (
   base64: string,
@@ -91,7 +92,7 @@ export const POST = withLogging(
 
     const fetchResponse = await fetch(
       `https://api.dev.landing.ai/v1/agent/chat?agent_class=vision_agent&self_reflection=false`,
-      // `http://localhost:5001/v1/agent/chat?agent_class=vision_agent&self_reflection=${enableSelfReflection}`,
+      // `http://localhost:5001/v1/agent/chat?agent_class=vision_agent&self_reflection=false`,
       {
         method: 'POST',
         headers: {
@@ -148,6 +149,7 @@ export const POST = withLogging(
       const decoder = new TextDecoder('utf-8');
       let maxChunkSize = 0;
       let buffer = '';
+      let time = Date.now();
       const stream = new ReadableStream({
         async start(controller) {
           // const parser = createParser(streamParser);
@@ -158,6 +160,23 @@ export const POST = withLogging(
             const lines = buffer
               .split('\n')
               .filter(line => line.trim().length > 0);
+            if (lines.length === 0) {
+              if (Date.now() - time > TIMEOUT_MILI_SECONDS) {
+                logger.info(
+                  session,
+                  {
+                    message: 'Agent timed out',
+                  },
+                  request,
+                  '__Agent_timeout__',
+                );
+                controller.error(
+                  `Haven't received any response in last ${TIMEOUT_MILI_SECONDS / 60000} minutes, agent timed out.`,
+                );
+              }
+            } else {
+              time = Date.now();
+            }
             buffer = lines.pop() ?? ''; // Save the last incomplete line back to the buffer
             let done = false;
             const parseLine = async (
@@ -166,9 +185,27 @@ export const POST = withLogging(
             ) => {
               try {
                 const msg = JSON.parse(line);
-                if (msg.type !== 'final_code') {
+                if (
+                  msg.type !== 'final_code' &&
+                  (msg.type !== 'code' ||
+                    msg.status === 'started' ||
+                    msg.status === 'running')
+                ) {
                   return line;
                 }
+                if (msg.type === 'code') {
+                  msg.payload.result = JSON.stringify({
+                    ...msg.payload.result,
+                    results: msg.payload.result.results.map((_result: any) => {
+                      return {
+                        ..._result,
+                        png: undefined,
+                        mp4: undefined,
+                      };
+                    }),
+                  });
+                  return JSON.stringify(msg);
+                }
                 const result = JSON.parse(msg.payload.result) as ResultPayload;
                 for (let index = 0; index < result.results.length; index++) {
                   const png = result.results[index].png ?? '';