hardening validation rule for cloud trace sink fields image

Author: SentienceDEV

src/tracing/cloud-sink.ts

@@ -773,6 +773,57 @@ export class CloudTraceSink extends TraceSink {
     });
   }
 
+  /**
+   * 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..."
+   *
+   * @param screenshotRaw - Raw screenshot data (data URL or base64)
+   * @param defaultFormat - Default format if not detected from data URL
+   * @returns Tuple of [base64String, formatString]
+   */
+  private _normalizeScreenshotData(
+    screenshotRaw: string,
+    defaultFormat: string = 'jpeg'
+  ): [string, string] {
+    if (!screenshotRaw) {
+      return ['', defaultFormat];
+    }
+
+    // Check if it's a data URL
+    if (screenshotRaw.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 (screenshotRaw.includes(',')) {
+          const [header, base64Data] = screenshotRaw.split(',', 2);
+          // Extract format from header: "data:image/jpeg;base64"
+          if (header.includes('/') && header.includes(';')) {
+            const formatPart = header.split('/')[1]?.split(';')[0];
+            if (formatPart === 'jpeg' || formatPart === 'jpg') {
+              return [base64Data, 'jpeg'];
+            } else if (formatPart === 'png') {
+              return [base64Data, 'png'];
+            }
+          }
+          return [base64Data, defaultFormat];
+        } else {
+          // Malformed data URL - return as-is with warning
+          this.logger?.warn('Malformed data URL in screenshot_base64 (missing comma)');
+          return [screenshotRaw, defaultFormat];
+        }
+      } catch (error: any) {
+        this.logger?.warn(`Error parsing screenshot data URL: ${error.message}`);
+        return [screenshotRaw, defaultFormat];
+      }
+    }
+
+    // Already pure base64
+    return [screenshotRaw, defaultFormat];
+  }
+
   /**
    * Extract screenshots from trace events.
    *
@@ -798,15 +849,24 @@ export class CloudTraceSink extends TraceSink {
           // Check if this is a snapshot event with screenshot
           if (event.type === 'snapshot') {
             const data = event.data || {};
-            const screenshotBase64 = data.screenshot_base64;
-
-            if (screenshotBase64) {
-              sequence += 1;
-              screenshots.set(sequence, {
-                base64: screenshotBase64,
-                format: data.screenshot_format || 'jpeg',
-                stepId: event.step_id,
-              });
+            const screenshotRaw = data.screenshot_base64;
+
+            if (screenshotRaw) {
+              // Normalize: extract base64 from data URL if needed
+              // Handles both "data:image/jpeg;base64,..." and pure base64
+              const [screenshotBase64, screenshotFormat] = this._normalizeScreenshotData(
+                screenshotRaw,
+                data.screenshot_format || 'jpeg'
+              );
+
+              if (screenshotBase64) {
+                sequence += 1;
+                screenshots.set(sequence, {
+                  base64: screenshotBase64,
+                  format: screenshotFormat,
+                  stepId: event.step_id,
+                });
+              }
             }
           }
         } catch {

tests/tracing/cloud-sink.test.ts

@@ -339,6 +339,96 @@ describe('CloudTraceSink', () => {
     });
   });
 
+  describe('Screenshot data URL handling', () => {
+    it('should normalize screenshot data URLs to pure base64', async () => {
+      const runId = 'test-run-' + Date.now();
+      const sink = new CloudTraceSink(uploadUrl, runId);
+
+      // Test the private _normalizeScreenshotData method via type casting
+      const sinkAny = sink as any;
+
+      // Test JPEG data URL
+      const [jpegBase64, jpegFormat] = sinkAny._normalizeScreenshotData(
+        'data:image/jpeg;base64,/9j/4AAQSkZJRg...',
+        'png'
+      );
+      expect(jpegBase64).toBe('/9j/4AAQSkZJRg...');
+      expect(jpegFormat).toBe('jpeg');
+
+      // Test PNG data URL
+      const [pngBase64, pngFormat] = sinkAny._normalizeScreenshotData(
+        'data:image/png;base64,iVBORw0KGgoAAAA...',
+        'jpeg'
+      );
+      expect(pngBase64).toBe('iVBORw0KGgoAAAA...');
+      expect(pngFormat).toBe('png');
+
+      // Test pure base64 (should pass through unchanged)
+      const [pureBase64, pureFormat] = sinkAny._normalizeScreenshotData(
+        '/9j/4AAQSkZJRg...',
+        'jpeg'
+      );
+      expect(pureBase64).toBe('/9j/4AAQSkZJRg...');
+      expect(pureFormat).toBe('jpeg');
+
+      // Test empty string
+      const [emptyBase64, emptyFormat] = sinkAny._normalizeScreenshotData('', 'jpeg');
+      expect(emptyBase64).toBe('');
+      expect(emptyFormat).toBe('jpeg');
+
+      // Clean up
+      await sink.close();
+    });
+
+    it('should handle data URL screenshots when extracting from trace', async () => {
+      const runId = 'test-run-' + Date.now();
+      const sink = new CloudTraceSink(uploadUrl, runId);
+
+      // Create test screenshot as a data URL (how some demos send it)
+      const testImageBase64 = 'iVBORw0KGgoAAAANSUhEUgAAAAE=';
+      const dataUrl = `data:image/png;base64,${testImageBase64}`;
+
+      // Emit snapshot event with data URL
+      sink.emit({
+        v: 1,
+        type: 'snapshot',
+        ts: new Date().toISOString(),
+        run_id: runId,
+        seq: 1,
+        step_id: 'step-1',
+        data: {
+          url: 'https://example.com',
+          element_count: 10,
+          screenshot_base64: dataUrl, // Data URL, not pure base64
+          screenshot_format: 'png',
+        },
+      });
+
+      // Close the write stream first
+      const sinkAny = sink as any;
+      if (sinkAny.writeStream && !sinkAny.writeStream.destroyed) {
+        await new Promise<void>(resolve => {
+          sinkAny.writeStream.end(() => resolve());
+        });
+      }
+
+      // Extract screenshots - should normalize data URL to pure base64
+      const screenshots = await sinkAny._extractScreenshotsFromTrace();
+
+      expect(screenshots.size).toBe(1);
+      expect(screenshots.has(1)).toBe(true);
+
+      const screenshot = screenshots.get(1);
+      // Verify the base64 was extracted from data URL (no "data:image" prefix)
+      expect(screenshot.base64).toBe(testImageBase64);
+      expect(screenshot.base64.startsWith('data:')).toBe(false);
+      expect(screenshot.format).toBe('png');
+
+      // Clean up
+      await sink.close();
+    });
+  });
+
   describe('Index upload', () => {
     let indexServer: http.Server;
     let indexServerPort: number;