fix cloud tracing index issue

Author: SentienceDEV

sentience/cloud_tracing.py

@@ -581,6 +581,56 @@ def _complete_trace(self) -> None:
             if self.logger:
                 self.logger.warning(f"Error reporting trace completion: {e}")
 
+    def _normalize_screenshot_data(
+        self, screenshot_raw: str, default_format: str = "jpeg"
+    ) -> tuple[str, str]:
+        """
+        Normalize screenshot data by extracting base64 from data URL if needed.
+
+        Handles both formats:
+        - Data URL: "data:image/jpeg;base64,/9j/4AAQ..."
+        - Pure base64: "/9j/4AAQ..."
+
+        Args:
+            screenshot_raw: Raw screenshot data (data URL or base64)
+            default_format: Default format if not detected from data URL
+
+        Returns:
+            Tuple of (base64_string, format_string)
+        """
+        if not screenshot_raw:
+            return "", default_format
+
+        # Check if it's a data URL
+        if screenshot_raw.startswith("data:image"):
+            # Extract format from "data:image/jpeg;base64,..." or "data:image/png;base64,..."
+            try:
+                # Split on comma to get the base64 part
+                if "," in screenshot_raw:
+                    header, base64_data = screenshot_raw.split(",", 1)
+                    # Extract format from header: "data:image/jpeg;base64"
+                    if "/" in header and ";" in header:
+                        format_part = header.split("/")[1].split(";")[0]
+                        if format_part in ("jpeg", "jpg"):
+                            return base64_data, "jpeg"
+                        elif format_part == "png":
+                            return base64_data, "png"
+                    return base64_data, default_format
+                else:
+                    # Malformed data URL - return as-is with warning
+                    if self.logger:
+                        self.logger.warning(
+                            "Malformed data URL in screenshot_base64 (missing comma)"
+                        )
+                    return screenshot_raw, default_format
+            except Exception as e:
+                if self.logger:
+                    self.logger.warning(f"Error parsing screenshot data URL: {e}")
+                return screenshot_raw, default_format
+
+        # Already pure base64
+        return screenshot_raw, default_format
+
     def _extract_screenshots_from_trace(self) -> dict[int, dict[str, Any]]:
         """
         Extract screenshots from trace events.
@@ -604,15 +654,22 @@ def _extract_screenshots_from_trace(self) -> dict[int, dict[str, Any]]:
                 # Check if this is a snapshot event with screenshot
                 if event.get("type") == "snapshot":
                     data = event.get("data", {})
-                    screenshot_base64 = data.get("screenshot_base64")
-
-                    if screenshot_base64:
-                        sequence += 1
-                        screenshots[sequence] = {
-                            "base64": screenshot_base64,
-                            "format": data.get("screenshot_format", "jpeg"),
-                            "step_id": event.get("step_id"),
-                        }
+                    screenshot_raw = data.get("screenshot_base64")
+
+                    if screenshot_raw:
+                        # Normalize: extract base64 from data URL if needed
+                        # Handles both "data:image/jpeg;base64,..." and pure base64
+                        screenshot_base64, screenshot_format = self._normalize_screenshot_data(
+                            screenshot_raw,
+                            data.get("screenshot_format", "jpeg"),
+                        )
+                        if screenshot_base64:
+                            sequence += 1
+                            screenshots[sequence] = {
+                                "base64": screenshot_base64,
+                                "format": screenshot_format,
+                                "step_id": event.get("step_id"),
+                            }
         except Exception as e:
             if self.logger:
                 self.logger.error(f"Error extracting screenshots: {e}")

tests/test_cloud_tracing.py

@@ -394,6 +394,92 @@ def post_side_effect(*args, **kwargs):
         if cleaned_trace_path.exists():
             os.remove(cleaned_trace_path)
 
+    def test_normalize_screenshot_data_handles_data_url(self):
+        """Test that _normalize_screenshot_data extracts base64 from data URLs."""
+        upload_url = "https://sentience.nyc3.digitaloceanspaces.com/user123/run456/trace.jsonl.gz"
+        run_id = f"test-run-{uuid.uuid4().hex[:8]}"
+
+        sink = CloudTraceSink(upload_url, run_id=run_id)
+
+        # Test JPEG data URL
+        jpeg_data_url = "data:image/jpeg;base64,/9j/4AAQSkZJRg..."
+        base64_str, fmt = sink._normalize_screenshot_data(jpeg_data_url)
+        assert base64_str == "/9j/4AAQSkZJRg..."
+        assert fmt == "jpeg"
+
+        # Test PNG data URL
+        png_data_url = "data:image/png;base64,iVBORw0KGgoAAAA..."
+        base64_str, fmt = sink._normalize_screenshot_data(png_data_url)
+        assert base64_str == "iVBORw0KGgoAAAA..."
+        assert fmt == "png"
+
+        # Test pure base64 (should pass through unchanged)
+        pure_base64 = "/9j/4AAQSkZJRg..."
+        base64_str, fmt = sink._normalize_screenshot_data(pure_base64, "jpeg")
+        assert base64_str == "/9j/4AAQSkZJRg..."
+        assert fmt == "jpeg"
+
+        # Test empty string
+        base64_str, fmt = sink._normalize_screenshot_data("")
+        assert base64_str == ""
+        assert fmt == "jpeg"
+
+        # Cleanup
+        cache_dir = Path.home() / ".sentience" / "traces" / "pending"
+        trace_path = cache_dir / f"{run_id}.jsonl"
+        if trace_path.exists():
+            os.remove(trace_path)
+
+    def test_cloud_trace_sink_handles_data_url_in_screenshot(self):
+        """Test that CloudTraceSink properly extracts screenshots from data URLs."""
+        upload_url = "https://sentience.nyc3.digitaloceanspaces.com/user123/run456/trace.jsonl.gz"
+        run_id = f"test-run-{uuid.uuid4().hex[:8]}"
+        api_key = "sk_test_123"
+
+        # Create test screenshot as a data URL (how langchain-debugging was sending it)
+        test_image_base64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg=="
+        data_url = f"data:image/png;base64,{test_image_base64}"
+
+        sink = CloudTraceSink(upload_url, run_id=run_id, api_key=api_key)
+
+        # Emit trace event with screenshot as data URL (not pure base64)
+        sink.emit(
+            {
+                "v": 1,
+                "type": "snapshot",
+                "ts": "2026-01-01T00:00:00.000Z",
+                "run_id": run_id,
+                "seq": 1,
+                "step_id": "step-1",
+                "data": {
+                    "url": "https://example.com",
+                    "element_count": 10,
+                    "screenshot_base64": data_url,  # Data URL, not pure base64
+                    "screenshot_format": "png",
+                },
+            }
+        )
+
+        # Extract screenshots - should normalize data URL to pure base64
+        screenshots = sink._extract_screenshots_from_trace()
+
+        assert len(screenshots) == 1
+        assert 1 in screenshots
+        # Verify the base64 was extracted from data URL (no "data:image" prefix)
+        assert screenshots[1]["base64"] == test_image_base64
+        assert not screenshots[1]["base64"].startswith("data:")
+        assert screenshots[1]["format"] == "png"
+
+        # Cleanup
+        sink.close()
+        cache_dir = Path.home() / ".sentience" / "traces" / "pending"
+        trace_path = cache_dir / f"{run_id}.jsonl"
+        cleaned_trace_path = cache_dir / f"{run_id}.cleaned.jsonl"
+        if trace_path.exists():
+            os.remove(trace_path)
+        if cleaned_trace_path.exists():
+            os.remove(cleaned_trace_path)
+
 
 class TestTracerFactory:
     """Test create_tracer factory function."""