Automatically screenshot after computer tool use

Author: atupem

packages/bytebot-agent/src/agent/agent.computer-use.ts

@@ -180,8 +180,18 @@ export async function handleComputerToolUse(
         };
       }
     }
+
+    let image: string | null = null;
+    try {
+      logger.debug('Taking screenshot');
+      image = await screenshot();
+      logger.debug('Screenshot captured successfully');
+    } catch (error) {
+      logger.error('Failed to take screenshot', error);
+    }
+
     logger.debug(`Tool execution successful for tool_use_id: ${block.id}`);
-    return {
+    const toolResult: ToolResultContentBlock = {
       type: MessageContentType.ToolResult,
       tool_use_id: block.id,
       content: [
@@ -191,6 +201,19 @@ export async function handleComputerToolUse(
         },
       ],
     };
+
+    if (image) {
+      toolResult.content.push({
+        type: MessageContentType.Image,
+        source: {
+          data: image,
+          media_type: 'image/png',
+          type: 'base64',
+        },
+      });
+    }
+
+    return toolResult;
   } catch (error) {
     logger.error(
       `Error executing ${block.name} tool: ${error.message}`,
@@ -591,7 +614,7 @@ export async function writeFile(input: {
   try {
     // Content is always base64 encoded
     const base64Data = content;
-    
+
     const response = await fetch(`${BYTEBOT_DESKTOP_BASE_URL}/computer-use`, {
       method: 'POST',
       headers: { 'Content-Type': 'application/json' },

packages/bytebot-ui/src/components/messages/AssistantMessage.tsx

@@ -44,22 +44,23 @@ export function AssistantMessage({
                     block.content &&
                     block.content.length > 0
                   ) {
-                    const imageBlock = block.content[0];
-                    if (isImageContentBlock(imageBlock)) {
-                      return (
-                        <div
-                          key={blockIndex}
-                          data-message-index={messageIdToIndex[message.id]}
-                          data-block-index={blockIndex}
-                          style={{
-                            position: "absolute",
-                            width: 0,
-                            height: 0,
-                            overflow: "hidden",
-                          }}
-                        />
-                      );
-                    }
+                    block.content.map((contentBlock) => {
+                      if (isImageContentBlock(contentBlock)) {
+                        return (
+                          <div
+                            key={blockIndex}
+                            data-message-index={messageIdToIndex[message.id]}
+                            data-block-index={blockIndex}
+                            style={{
+                              position: "absolute",
+                              width: 0,
+                              height: 0,
+                              overflow: "hidden",
+                            }}
+                          />
+                        );
+                      }
+                    });
                   }
                   return null;
                 })}

packages/bytebot-ui/src/components/messages/content/MessageContent.tsx

@@ -25,7 +25,8 @@ export function MessageContent({
     // Filter logic from the original code
     if (
       isToolResultContentBlock(block) &&
-      isImageContentBlock(block.content?.[0])
+      block.content &&
+      block.content.some((contentBlock) => isImageContentBlock(contentBlock))
     ) {
       return true;
     }
@@ -50,9 +51,16 @@ export function MessageContent({
         <div key={index}>
           {isTextContentBlock(block) && <TextContent block={block} />}
 
-          {isImageContentBlock(block.content?.[0]) && (
-            <ImageContent block={block.content[0]} />
-          )}
+          {isToolResultContentBlock(block) &&
+            !block.is_error &&
+            block.content.map((contentBlock, contentBlockIndex) => {
+              if (isImageContentBlock(contentBlock)) {
+                return (
+                  <ImageContent key={contentBlockIndex} block={contentBlock} />
+                );
+              }
+              return null;
+            })}
 
           {isComputerToolUseContentBlock(block) && (
             <ComputerToolContent block={block} isTakeOver={isTakeOver} />