upload screenshots along with traces

Author: rcholic

sentience/agent.py

@@ -9,6 +9,7 @@
 from typing import TYPE_CHECKING, Any, Optional
 
 from .actions import click, click_async, press, press_async, type_text, type_text_async
+from .agent_config import AgentConfig
 from .base_agent import BaseAgent, BaseAgentAsync
 from .browser import AsyncSentienceBrowser, SentienceBrowser
 from .llm_provider import LLMProvider, LLMResponse
@@ -25,7 +26,6 @@
 from .snapshot import snapshot, snapshot_async
 
 if TYPE_CHECKING:
-    from .agent_config import AgentConfig
     from .tracing import Tracer
 
 
@@ -78,7 +78,10 @@ def __init__(
         self.default_snapshot_limit = default_snapshot_limit
         self.verbose = verbose
         self.tracer = tracer
-        self.config = config
+        self.config = config or AgentConfig()
+
+        # Screenshot sequence counter
+        self._screenshot_sequence = 0
 
         # Execution history
         self.history: list[dict[str, Any]] = []
@@ -150,12 +153,42 @@ def act(  # noqa: C901
                 if snap_opts.goal is None:
                     snap_opts.goal = goal
 
+                # Apply AgentConfig screenshot settings if not overridden by snapshot_options
+                if snapshot_options is None and self.config:
+                    if self.config.capture_screenshots:
+                        # Create ScreenshotConfig from AgentConfig
+                        snap_opts.screenshot = ScreenshotConfig(
+                            format=self.config.screenshot_format,
+                            quality=(
+                                self.config.screenshot_quality
+                                if self.config.screenshot_format == "jpeg"
+                                else None
+                            ),
+                        )
+                    else:
+                        snap_opts.screenshot = False
+
                 # Call snapshot with options object (matches TypeScript API)
                 snap = snapshot(self.browser, snap_opts)
 
                 if snap.status != "success":
                     raise RuntimeError(f"Snapshot failed: {snap.error}")
 
+                # Store screenshot if captured
+                if snap.screenshot and self.tracer:
+                    self._screenshot_sequence += 1
+                    seq = self._screenshot_sequence
+
+                    # Store screenshot in CloudTraceSink if available
+                    if hasattr(self.tracer.sink, "store_screenshot"):
+                        self.tracer.sink.store_screenshot(
+                            sequence=seq,
+                            screenshot_data=snap.screenshot,
+                            format=snap.screenshot_format
+                            or (self.config.screenshot_format if self.config else "jpeg"),
+                            step_id=step_id,
+                        )
+
                 # Apply element filtering based on goal
                 filtered_elements = self.filter_elements(snap, goal)
 
@@ -721,7 +754,10 @@ def __init__(
         self.default_snapshot_limit = default_snapshot_limit
         self.verbose = verbose
         self.tracer = tracer
-        self.config = config
+        self.config = config or AgentConfig()
+
+        # Screenshot sequence counter
+        self._screenshot_sequence = 0
 
         # Execution history
         self.history: list[dict[str, Any]] = []
@@ -790,12 +826,42 @@ async def act(  # noqa: C901
                 if snap_opts.goal is None:
                     snap_opts.goal = goal
 
+                # Apply AgentConfig screenshot settings if not overridden by snapshot_options
+                if snapshot_options is None and self.config:
+                    if self.config.capture_screenshots:
+                        # Create ScreenshotConfig from AgentConfig
+                        snap_opts.screenshot = ScreenshotConfig(
+                            format=self.config.screenshot_format,
+                            quality=(
+                                self.config.screenshot_quality
+                                if self.config.screenshot_format == "jpeg"
+                                else None
+                            ),
+                        )
+                    else:
+                        snap_opts.screenshot = False
+
                 # Call snapshot with options object (matches TypeScript API)
                 snap = await snapshot_async(self.browser, snap_opts)
 
                 if snap.status != "success":
                     raise RuntimeError(f"Snapshot failed: {snap.error}")
 
+                # Store screenshot if captured
+                if snap.screenshot and self.tracer:
+                    self._screenshot_sequence += 1
+                    seq = self._screenshot_sequence
+
+                    # Store screenshot in CloudTraceSink if available
+                    if hasattr(self.tracer.sink, "store_screenshot"):
+                        self.tracer.sink.store_screenshot(
+                            sequence=seq,
+                            screenshot_data=snap.screenshot,
+                            format=snap.screenshot_format
+                            or (self.config.screenshot_format if self.config else "jpeg"),
+                            step_id=step_id,
+                        )
+
                 # Apply element filtering based on goal
                 filtered_elements = self.filter_elements(snap, goal)
 

sentience/cloud_tracing.py

@@ -4,16 +4,19 @@
 Implements "Local Write, Batch Upload" pattern for enterprise cloud tracing.
 """
 
+import base64
 import gzip
 import json
 import os
 import threading
 from collections.abc import Callable
+from concurrent.futures import ThreadPoolExecutor, as_completed
 from pathlib import Path
 from typing import Any, Protocol
 
 import requests
 
+from sentience.models import ScreenshotMetadata
 from sentience.tracing import TraceSink
 
 
@@ -103,10 +106,17 @@ def __init__(
         self._closed = False
         self._upload_successful = False
 
-        # File size tracking (NEW)
+        # File size tracking
         self.trace_file_size_bytes = 0
         self.screenshot_total_size_bytes = 0
 
+        # Screenshot storage directory
+        self._screenshot_dir = cache_dir / f"{run_id}_screenshots"
+        self._screenshot_dir.mkdir(exist_ok=True)
+
+        # Screenshot metadata tracking (sequence -> ScreenshotMetadata)
+        self._screenshot_metadata: dict[int, ScreenshotMetadata] = {}
+
     def emit(self, event: dict[str, Any]) -> None:
         """
         Write event to local persistent file (Fast, non-blocking).
@@ -213,18 +223,21 @@ def _do_upload(self, on_progress: Callable[[int, int], None] | None = None) -> N
                 if on_progress:
                     on_progress(compressed_size, compressed_size)
 
+                # Upload screenshots after trace upload succeeds
+                if self._screenshot_metadata:
+                    print(
+                        f"📸 [Sentience] Uploading {len(self._screenshot_metadata)} screenshots..."
+                    )
+                    self._upload_screenshots(on_progress)
+
                 # Upload trace index file
                 self._upload_index()
 
                 # Call /v1/traces/complete to report file sizes
                 self._complete_trace()
 
-                # Delete file only on successful upload
-                if os.path.exists(self._path):
-                    try:
-                        os.remove(self._path)
-                    except Exception:
-                        pass  # Ignore cleanup errors
+                # Delete files only on successful upload
+                self._cleanup_files()
             else:
                 self._upload_successful = False
                 print(f"❌ [Sentience] Upload failed: HTTP {response.status_code}")
@@ -353,6 +366,7 @@ def _complete_trace(self) -> None:
                     "stats": {
                         "trace_file_size_bytes": self.trace_file_size_bytes,
                         "screenshot_total_size_bytes": self.screenshot_total_size_bytes,
+                        "screenshot_count": len(self._screenshot_metadata),
                     },
                 },
                 timeout=10,

sentience/extension/background.js

@@ -144,13 +144,13 @@ async function handleScreenshotCapture(_tabId, options = {}) {
 async function handleSnapshotProcessing(rawData, options = {}) {
     const MAX_ELEMENTS = 10000; // Safety limit to prevent hangs
     const startTime = performance.now();
-    
+
     try {
         // Safety check: limit element count to prevent hangs
         if (!Array.isArray(rawData)) {
             throw new Error('rawData must be an array');
         }
-        
+
         if (rawData.length > MAX_ELEMENTS) {
             console.warn(`[Sentience Background] ⚠️ Large dataset: ${rawData.length} elements. Limiting to ${MAX_ELEMENTS} to prevent hangs.`);
             rawData = rawData.slice(0, MAX_ELEMENTS);
@@ -186,7 +186,7 @@ async function handleSnapshotProcessing(rawData, options = {}) {
             // Add timeout protection (18 seconds - less than content.js timeout)
             analyzedElements = await Promise.race([
                 wasmPromise,
-                new Promise((_, reject) => 
+                new Promise((_, reject) =>
                     setTimeout(() => reject(new Error('WASM processing timeout (>18s)')), 18000)
                 )
             ]);

sentience/extension/content.js

@@ -92,7 +92,7 @@ function handleSnapshotRequest(data) {
                 if (responded) return; // Already responded via timeout
                 responded = true;
                 clearTimeout(timeoutId);
-                
+
                 const duration = performance.now() - startTime;
 
                 // Handle Chrome extension errors (e.g., background script crashed)