feat: update final_error to result and add cancel callback (#103)

<img width="1288" alt="image"
src="https://github.com/landing-ai/vision-agent-ui/assets/132986242/743a66c9-65f9-4698-9d67-3ef6eb209ef3">

<img width="1271" alt="image"
src="https://github.com/landing-ai/vision-agent-ui/assets/132986242/4ac27c6f-85b9-496b-b436-1921c2ba052c">

.github/workflows/playwright.yml

@@ -5,7 +5,7 @@ jobs:
   run-e2es:
     timeout-minutes: 5
     runs-on: ubuntu-latest
-    if: github.event_name == 'deployment_status' && github.event.deployment_status.state == 'success'
+    if: github.event_name == 'deployment_status' && github.event.deployment_status.state == 'success' && github.event.deployment_status.environment_url
     steps:
     - uses: actions/checkout@v4
     - name: Setup pnpm

app/api/vision-agent/route.ts

@@ -1,8 +1,5 @@
 import { StreamingTextResponse } from 'ai';
 
-// import { auth } from '@/auth';
-import { MessageUI } from '@/lib/types';
-
 import { logger, withLogging } from '@/lib/logger';
 import { getPresignedUrl } from '@/lib/aws';
 import { dbPostUpdateMessageResponse } from '@/lib/db/functions';
@@ -179,6 +176,43 @@ export const POST = withLogging(
     let time = Date.now();
     const results: PrismaJson.MessageBody[] = [];
     const stream = new ReadableStream({
+      async cancel(reason) {
+        logger.info(
+          session,
+          {
+            message: 'Streaming cancelled',
+            maxChunkSize,
+            reason,
+          },
+          request,
+          '__AGENT_STREAM_CANCELLED',
+        );
+        const processMsgs = results.filter(
+          res => res.type !== 'final_code',
+        ) as PrismaJson.AgentResponseBodies;
+        await dbPostUpdateMessageResponse(
+          messageId,
+          {
+            response: processMsgs.map(res => JSON.stringify(res)).join('\n'),
+            result: {
+              type: 'final_error',
+              status: 'failed',
+              payload: {
+                name:
+                  reason instanceof Error ? reason.name : 'Stream cancelled',
+                value:
+                  reason instanceof Error
+                    ? reason.stack ?? ''
+                    : 'Unknown error',
+                traceback_raw: [],
+              },
+            },
+            responseBody: processMsgs,
+            streamDuration: (Date.now() - time) / 1000,
+          },
+          false, // shouldRevalidatePath
+        );
+      },
       async start(controller) {
         const parseLine = async (
           line: string,
@@ -289,9 +323,10 @@ export const POST = withLogging(
           await dbPostUpdateMessageResponse(messageId, {
             response: processMsgs.map(res => JSON.stringify(res)).join('\n'),
             result: results.find(
-              res => res.type === 'final_code',
-            ) as PrismaJson.FinalCodeBody,
+              res => res.type === 'final_code' || res.type === 'final_error',
+            ) as PrismaJson.FinalResultBody,
             responseBody: processMsgs,
+            streamDuration: (Date.now() - time) / 1000,
           });
           logger.info(
             session,

components/ChatInterface.tsx

@@ -24,11 +24,12 @@ const ChatInterface: React.FC<ChatInterfaceProps> = ({ chat, userId }) => {
         data-state={messageCodeResult?.payload ? 'open' : 'closed'}
         className="pl-4 peer absolute right-0 inset-y-0 hidden translate-x-full data-[state=open]:translate-x-0 z-30 duration-300 ease-in-out xl:flex flex-col items-start xl:w-1/2 h-full dark:bg-zinc-950 overflow-auto"
       >
-        {messageCodeResult?.payload && (
-          <Card className="size-full overflow-auto">
-            <CodeResultDisplay codeResult={messageCodeResult.payload} />
-          </Card>
-        )}
+        {messageCodeResult?.type === 'final_code' &&
+          messageCodeResult.payload && (
+            <Card className="size-full overflow-auto">
+              <CodeResultDisplay codeResult={messageCodeResult.payload} />
+            </Card>
+          )}
       </div>
       <div className="w-full flex justify-center pr-0 animate-in duration-300 ease-in-out peer-[[data-state=open]]:xl:pr-[50%]">
         <ChatList chat={chat} userId={userId} />

components/chat/ChatMessage.tsx

@@ -75,17 +75,16 @@ export const ChatMessage: React.FC<ChatMessageProps> = ({
 }) => {
   const [messageId, setMessageId] = useAtom(selectedMessageId);
   const { id, mediaUrl, prompt, response, result, responseBody } = message;
-  const [formattedSections, finalResult, finalError] = useMemo(
+  const { formattedSections, finalResult, finalError } = useMemo(
     () =>
       formatStreamLogs(
-        responseBody ??
-          wipAssistantMessage ??
+        wipAssistantMessage ??
+          responseBody ??
           (response ? getParsedStreamLogs(response) : []),
+        result,
       ),
     [response, wipAssistantMessage, responseBody],
   );
-  // prioritize the result from the message over the WIP message
-  const codeResult = result?.payload ?? finalResult;
   return (
     <div
       className={cn(
@@ -143,7 +142,7 @@ export const ChatMessage: React.FC<ChatMessageProps> = ({
                       </TableCell>
                       <TableCell className="font-medium">
                         <ChunkTypeToText
-                          useTimer={!codeResult && !finalError}
+                          useTimer={!finalResult && !finalError}
                           chunk={section}
                         />
                       </TableCell>
@@ -154,15 +153,15 @@ export const ChatMessage: React.FC<ChatMessageProps> = ({
                   ))}
                 </TableBody>
               </Table>
-              {codeResult && (
+              {finalResult && (
                 <>
                   <div className="xl:hidden">
-                    <CodeResultDisplay codeResult={codeResult} />
+                    <CodeResultDisplay codeResult={finalResult} />
                   </div>
                   <p>✨ Coding complete</p>
                 </>
               )}
-              {!codeResult && finalError && (
+              {!finalResult && finalError && (
                 <>
                   <p>❌ {finalError.name}</p>
                   <div>

lib/db/functions.ts

@@ -183,19 +183,21 @@ export async function dbPostCreateMessage(
 export async function dbPostUpdateMessageResponse(
   messageId: string,
   messageResponse: MessageAssistantResponse,
+  shouldRevalidatePath = true,
 ) {
   await prisma.message.update({
     data: {
       response: messageResponse.response,
       result: messageResponse.result ?? undefined,
       responseBody: messageResponse.responseBody,
+      streamDuration: messageResponse.streamDuration,
     },
     where: {
       id: messageId,
     },
   });
 
-  revalidatePath('/chat');
+  shouldRevalidatePath && revalidatePath('/chat');
 }
 
 export async function dbDeleteChat(chatId: string) {

lib/db/prisma.ts

@@ -57,6 +57,8 @@ declare global {
       payload: StructuredError;
     }
 
+    type FinalResultBody = FinalCodeBody | FinalErrorBody;
+
     type MessageBody =
       | PlanAndToolsBody
       | CodeBody

lib/hooks/useVisionAgent.ts

@@ -1,7 +1,7 @@
 import { useChat } from 'ai/react';
 import { toast } from 'react-hot-toast';
 import { useEffect, useRef } from 'react';
-import { ChatWithMessages, MessageUI } from '../types';
+import { ChatWithMessages } from '../types';
 import { convertDBMessageToAPIMessage } from '../utils/message';
 import { useSetAtom } from 'jotai';
 import { selectedMessageId } from '@/state/chat';
@@ -25,7 +25,7 @@ const useVisionAgent = (chat: ChatWithMessages) => {
         toast.error(response.statusText);
       }
     },
-    onFinish: message => {
+    onFinish: () => {
       router.refresh();
       setMessageId(currMessageId.current);
     },
@@ -47,7 +47,11 @@ const useVisionAgent = (chat: ChatWithMessages) => {
    */
   const once = useRef(true);
   useEffect(() => {
-    if (!isLoading && !latestDbMessage.result && once.current) {
+    if (
+      !isLoading &&
+      !(latestDbMessage.response || latestDbMessage.responseBody) &&
+      once.current
+    ) {
       once.current = false;
       reload();
     }

lib/types.ts

@@ -5,9 +5,10 @@ export type ChatWithMessages = Chat & { messages: Message[] };
 
 export type MessageUserInput = Pick<Message, 'prompt' | 'mediaUrl'>;
 export type MessageAssistantResponse = {
-  result?: PrismaJson.FinalCodeBody;
+  result?: PrismaJson.FinalResultBody;
   response: string;
   responseBody: PrismaJson.AgentResponseBodies;
+  streamDuration: number;
 };
 
 export type MessageUI = Pick<MessageAI, 'role' | 'content' | 'id'>;

lib/utils/content.ts

@@ -32,42 +32,51 @@ export type WIPChunkBodyGroup = PrismaJson.MessageBody & {
  */
 export const formatStreamLogs = (
   content: WIPChunkBodyGroup[] | null,
-): [
-  WIPChunkBodyGroup[],
-  PrismaJson.FinalCodeBody['payload']?,
-  PrismaJson.StructuredError?,
-] => {
-  if (!content) return [[], undefined, undefined];
+  result: PrismaJson.FinalResultBody | null,
+): {
+  formattedSections: WIPChunkBodyGroup[];
+  finalResult?: PrismaJson.FinalCodeBody['payload'];
+  finalError?: PrismaJson.StructuredError;
+} => {
+  if (!content)
+    return {
+      formattedSections: [],
+    };
 
   // Merge consecutive logs of the same type to the latest status
-  const groupedSections = content.reduce((acc: WIPChunkBodyGroup[], curr) => {
-    const lastGroup = acc[acc.length - 1];
-    if (
-      acc.length > 0 &&
-      lastGroup.type === curr.type &&
-      curr.status !== 'started'
-    ) {
-      acc[acc.length - 1] = {
-        ...curr,
-        // always use the timestamp of the first log
-        timestamp: lastGroup?.timestamp,
-        // duration is the difference between the last log and the first log
-        duration:
-          lastGroup?.timestamp && curr.timestamp
-            ? Date.parse(curr.timestamp) - Date.parse(lastGroup.timestamp)
-            : undefined,
-      };
-    } else {
-      acc.push(curr);
-    }
-    return acc;
-  }, []);
+  const groupedSections = [...content, ...(result ? [result] : [])].reduce(
+    (acc: WIPChunkBodyGroup[], curr: WIPChunkBodyGroup) => {
+      const lastGroup = acc[acc.length - 1];
+      if (
+        acc.length > 0 &&
+        lastGroup.type === curr.type &&
+        curr.status !== 'started'
+      ) {
+        acc[acc.length - 1] = {
+          ...curr,
+          // always use the timestamp of the first log
+          timestamp: lastGroup?.timestamp,
+          // duration is the difference between the last log and the first log
+          duration:
+            lastGroup?.timestamp && curr.timestamp
+              ? Date.parse(curr.timestamp) - Date.parse(lastGroup.timestamp)
+              : undefined,
+        };
+      } else {
+        acc.push(curr);
+      }
+      return acc;
+    },
+    [],
+  );
 
-  return [
-    groupedSections.filter(section => WIPLogTypes.includes(section.type)),
-    groupedSections.find(section => section.type === 'final_code')
+  return {
+    formattedSections: groupedSections.filter(section =>
+      WIPLogTypes.includes(section.type),
+    ),
+    finalResult: groupedSections.find(section => section.type === 'final_code')
       ?.payload as PrismaJson.FinalCodeBody['payload'],
-    groupedSections.find(section => section.type === 'final_error')
+    finalError: groupedSections.find(section => section.type === 'final_error')
       ?.payload as PrismaJson.StructuredError,
-  ];
+  };
 };

lib/utils/message.ts

@@ -1,5 +1,5 @@
 import { Message } from '@prisma/client';
-import { MessageAssistantResponse, MessageUI } from '../types';
+import { MessageUI } from '../types';
 
 /**
  * The Message we saved to database consists of a prompt and a response
@@ -17,7 +17,7 @@ export const convertDBMessageToAPIMessage = (
         content: prompt,
       });
     }
-    if (result) {
+    if (result && result.type === 'final_code') {
       acc.push({
         id: id + '-assistant',
         role: 'assistant',

prisma/migrations/20240617015509_add_stream_duration_to_message/data-migration.mjs

@@ -0,0 +1,59 @@
+import { PrismaClient } from '@prisma/client';
+
+const prisma = new PrismaClient();
+
+async function main() {
+  await prisma.$transaction(
+    async tx => {
+      const messages = await tx.message.findMany();
+      for (const message of messages) {
+        const { responseBody, response, result } = message;
+        let newResponseBody = responseBody;
+        let newResult = result;
+        if (responseBody && result) {
+          continue;
+        }
+        if (!responseBody && response) {
+          console.log('update response body', message.id);
+          newResponseBody = response
+            .split('\n')
+            .filter(chunk => !!chunk.trim())
+            .map(chunk => JSON.parse(chunk))
+            .filter(
+              body => body.type !== 'final_error' && body.type !== 'final_code',
+            );
+        }
+        if (!result && response) {
+          console.log('update response result', message.id);
+          newResult = response
+            .split('\n')
+            .filter(chunk => !!chunk.trim())
+            .map(chunk => JSON.parse(chunk))
+            ?.find(
+              body => body.type === 'final_error' || body.type === 'final_code',
+            );
+        }
+
+        await tx.message.update({
+          where: { id: message.id },
+          data: {
+            ...(newResult ? { result: newResult } : {}),
+            ...(newResponseBody ? { responseBody: newResponseBody } : {}),
+          },
+        });
+      }
+    },
+    { timeout: 30000 },
+  );
+}
+
+main()
+  .catch(async e => {
+    console.error(e);
+    await prisma.$disconnect();
+    process.exit(1);
+  })
+  .finally(async () => {
+    console.log('finished');
+    await prisma.$disconnect();
+  });

prisma/migrations/20240617015509_add_stream_duration_to_message/migration.sql

@@ -0,0 +1,2 @@
+-- AlterTable
+ALTER TABLE "message" ADD COLUMN     "streamDuration" INTEGER;

prisma/schema.prisma

@@ -47,10 +47,11 @@ model Message {
   mediaUrl  String
   prompt    String
   response  String?
-  /// [FinalCodeBody]
+  /// [FinalResultBody]
   result    Json?
   /// [AgentResponseBodies]
   responseBody Json?
+  streamDuration Int?
   chat      Chat     @relation(fields: [chatId], references: [id], onDelete: Cascade)
   user      User?    @relation(fields: [userId], references: [id])
 

tests/e2e/index.spec.ts

@@ -83,3 +83,10 @@ test.describe('Page with example 1 can be opened', () => {
     await checkMainElementVisible(page);
   });
 });
+
+const TEST_FINAL_ERROR_CHAT_ID_1 = '6bLzMKS';
+test('Page with final error in response can be opened', async ({ page }) => {
+  await page.goto(`/chat/${TEST_FINAL_ERROR_CHAT_ID_1}`);
+
+  await expect(page.getByText('❌ StreamingError')).toBeVisible();
+});