feat: enqueue final error to ui on streaming error and timeout (#83)

<img width="1423" alt="image"
src="https://github.com/landing-ai/vision-agent-ui/assets/132986242/c2cb87c1-aed5-42a4-802d-1b75bab0968e">

app/api/vision-agent/route.ts

@@ -11,7 +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 TIMEOUT_MILI_SECONDS = 2 * 60 * 1000;
 
 const uploadBase64 = async (
   base64: string,
@@ -172,9 +172,21 @@ export const POST = withLogging(
                   request,
                   '__Agent_timeout__',
                 );
-                controller.error(
-                  `Haven't received any response in last ${TIMEOUT_MILI_SECONDS / 60000} minutes, agent timed out.`,
+                controller.enqueue(
+                  encoder.encode(
+                    JSON.stringify({
+                      type: 'final_error',
+                      status: 'failed',
+                      payload: {
+                        name: 'AgentTimeout',
+                        value: `Haven't received any response in last ${TIMEOUT_MILI_SECONDS / 60000} minutes, agent timed out.`,
+                        traceback_raw: [],
+                      },
+                    }) + '\n',
+                  ),
                 );
+                controller.close();
+                return;
               }
             } else {
               time = Date.now();
@@ -187,6 +199,9 @@ export const POST = withLogging(
             ) => {
               try {
                 const msg = JSON.parse(line);
+                if (msg.type === 'final_code' || msg.type === 'final_error') {
+                  done = true;
+                }
                 if (
                   msg.type !== 'final_code' &&
                   (msg.type !== 'code' ||
@@ -231,7 +246,6 @@ export const POST = withLogging(
                   if (mp4) result.results[index].mp4 = resp;
                 }
                 msg.payload.result = JSON.stringify(result);
-                done = true;
                 return JSON.stringify(msg);
               } catch (e) {
                 errorCallback?.(e as Error);

lib/utils/content.ts

@@ -72,55 +72,6 @@ type ReflectionBody =
       payload: { feedback: string; success: boolean };
     };
 
-type MessageBody =
-  | PlansBody
-  | ToolsBody
-  | CodeBody
-  | ReflectionBody
-  | PrismaJson.FinalChatResult;
-
-const getMessageTitle = (json: MessageBody) => {
-  switch (json.type) {
-    case 'plans':
-      if (json.status === 'started') {
-        return '🎬 Start generating plans...\n';
-      } else {
-        return '✅ Going to run the following plan(s) in sequence:\n';
-      }
-    case 'tools':
-      if (json.status === 'started') {
-        return '🎬 Start retrieving tools...\n';
-      } else {
-        return '✅ Tools retrieved:\n';
-      }
-    case 'code':
-      if (json.status === 'started') {
-        return '🎬 Start generating code...\n';
-      } else if (json.status === 'running') {
-        return '🎬 Code generated, start execution... ';
-      } else if (json.status === 'completed') {
-        return '✅ Code executed successfully. ';
-      } else {
-        return '❌ Code execution failed. ';
-      }
-    case 'self_reflection':
-      if (json.status === 'started') {
-        return '🎬 Start self reflection...\n';
-      } else if (json.status === 'completed') {
-        return '✅ Self reflection completed: \n';
-      } else {
-        return '❌ Self reflection failed: \n';
-      }
-    case 'final_code':
-      if (json.status === 'completed') {
-        return '✅ The vision agent has concluded the chat, the last execution is successful. \n';
-      } else {
-        return '❌ The vision agent has concluded the chat, the last execution is failed. \n';
-      }
-    default:
-      throw 'Not supported type';
-  }
-};
 
 export type CodeResult = {
   code: string;
@@ -130,7 +81,7 @@ export type CodeResult = {
 
 export type ChunkBody =
   | {
-      type: 'plans' | 'tools' | 'code' | 'final_code';
+      type: 'plans' | 'tools' | 'code' | 'final_code' | 'final_error';
       status: 'started' | 'completed' | 'failed' | 'running';
       payload:
         | Array<Record<string, string>> // PlansBody | ToolsBody