FIX Max Payload request error

cua2-core/src/cua2_core/services/agent_service.py

@@ -21,6 +21,7 @@ from cua2_core.services.agent_utils.function_parser import parse_function_call
 from cua2_core.services.agent_utils.get_model import get_model
 from cua2_core.services.archival_service import ArchivalService
 from cua2_core.services.sandbox_service import SandboxService
+from cua2_core.services.utils import compress_image_to_max_size
 from cua2_core.websocket.websocket_manager import WebSocketException, WebSocketManager
 from e2b_desktop import Sandbox, TimeoutException
 from fastapi import WebSocket
@@ -413,6 +414,7 @@ class AgentService:
             step_filename = f"{message_id}-{memory_step.step_number + 1}"
             screenshot_bytes = agent.desktop.screenshot()
             image = Image.open(BytesIO(screenshot_bytes))
+            image = compress_image_to_max_size(image, max_size_kb=500)
 
             for previous_memory_step in (
                 agent.memory.steps

cua2-core/src/cua2_core/services/utils.py

@@ -0,0 +1,83 @@
+"""Utility functions for services."""
+
+from io import BytesIO
+
+from PIL import Image
+
+
+def compress_image_to_max_size(
+    image: Image.Image, max_size_kb: int = 800
+) -> Image.Image:
+    """
+    Compress an image to ensure it doesn't exceed max_size_kb.
+    Always returns a PIL Image in PNG format.
+    Resizes the image if necessary to meet the size requirement.
+
+    Args:
+        image: PIL Image to compress
+        max_size_kb: Maximum size in KB (default: 800KB)
+
+    Returns:
+        PIL Image in PNG format that meets the size requirement
+    """
+    max_size_bytes = max_size_kb * 1024
+    current_image = image.copy()
+
+    # Try PNG compression first with different compression levels
+    for compress_level in [3, 6, 9]:
+        buffered = BytesIO()
+        current_image.save(buffered, format="PNG", compress_level=compress_level)
+        png_size = buffered.tell()
+        if png_size <= max_size_bytes:
+            # Size is acceptable, return the image as-is (it's already in PNG format when saved)
+            return current_image
+
+    # PNG is too large even with maximum compression, need to resize
+    # Binary search for the right size
+    scale_factor = 1.0
+    min_scale = 0.1
+    max_scale = 1.0
+    iterations = 0
+    max_iterations = 10
+    best_image = None
+
+    while iterations < max_iterations:
+        # Resize image
+        new_size = (
+            int(image.size[0] * scale_factor),
+            int(image.size[1] * scale_factor),
+        )
+        resized_image = image.resize(new_size, Image.Resampling.LANCZOS)
+
+        # Check PNG size with best compression
+        buffered = BytesIO()
+        resized_image.save(buffered, format="PNG", compress_level=9)
+        png_size = buffered.tell()
+
+        if png_size <= max_size_bytes:
+            # Found acceptable size
+            best_image = resized_image
+            # Try slightly larger for better quality
+            if scale_factor >= max_scale - 0.01:  # Close enough to max
+                return resized_image
+            # Try larger scale
+            min_scale = scale_factor
+            scale_factor = (scale_factor + max_scale) / 2
+        else:
+            # Too large, need smaller scale
+            max_scale = scale_factor
+            scale_factor = (min_scale + scale_factor) / 2
+
+        iterations += 1
+
+    # Return the best image found, or resize to minimum scale if none worked
+    if best_image is not None:
+        return best_image
+
+    # Final fallback: resize to minimum scale
+    final_size = (
+        int(image.size[0] * min_scale),
+        int(image.size[1] * min_scale),
+    )
+    final_image = image.resize(final_size, Image.Resampling.LANCZOS)
+    return final_image