Merge pull request #173 from SentienceAPI/improve_tracing

improve tracing with automated upload if opt-in; auto step_start

Author: rcholic

src/agent-runtime.ts

@@ -626,12 +626,25 @@ export class AgentRuntime {
    * Take a snapshot of the current page state.
    *
    * This updates lastSnapshot which is used as context for assertions.
+   * When emitTrace=true (default), automatically emits a 'snapshot' trace event
+   * with screenshot_base64 for Sentience Studio visualization.
    *
    * @param options - Options passed through to browser.snapshot()
+   * @param options.emitTrace - If true (default), emit a 'snapshot' trace event with screenshot.
+   *                           Set to false to disable automatic trace emission.
    * @returns Snapshot of current page state
+   *
+   * @example
+   * // Default: snapshot with auto-emit trace event
+   * const snapshot = await runtime.snapshot();
+   *
+   * // Disable auto-emit for manual control
+   * const snapshot = await runtime.snapshot({ emitTrace: false });
+   * // Later, manually emit if needed:
+   * tracer.emitSnapshot(snapshot, runtime.getStepId());
    */
   async snapshot(options?: Record<string, any>): Promise<Snapshot> {
-    const { _skipCaptchaHandling, ...snapshotOptions } = options || {};
+    const { _skipCaptchaHandling, emitTrace = true, ...snapshotOptions } = options || {};
     this.lastSnapshot = await this.browser.snapshot(this.page, snapshotOptions);
     if (this.lastSnapshot && !this.stepPreSnapshot) {
       this.stepPreSnapshot = this.lastSnapshot;
@@ -640,9 +653,32 @@ export class AgentRuntime {
     if (!_skipCaptchaHandling) {
       await this.handleCaptchaIfNeeded(this.lastSnapshot, 'gateway');
     }
+
+    // Auto-emit snapshot trace event for Studio visualization
+    if (emitTrace && this.lastSnapshot && this.tracer) {
+      this.emitSnapshotTrace(this.lastSnapshot);
+    }
+
     return this.lastSnapshot;
   }
 
+  /**
+   * Emit a snapshot trace event with screenshot for Studio visualization.
+   *
+   * This is called automatically by snapshot() when emitTrace=true.
+   */
+  private emitSnapshotTrace(snapshot: Snapshot): void {
+    if (!this.tracer) {
+      return;
+    }
+
+    try {
+      this.tracer.emitSnapshot(snapshot, this.stepId ?? undefined, this.stepIndex, 'jpeg');
+    } catch {
+      // Best-effort: don't let trace emission errors break snapshot
+    }
+  }
+
   /**
    * Evaluate JavaScript in the page context.
    */
@@ -1167,12 +1203,20 @@ export class AgentRuntime {
    * - Generates a new stepId
    * - Clears assertions from previous step
    * - Increments stepIndex (or uses provided value)
+   * - Emits step_start trace event (optional)
    *
    * @param goal - Description of what this step aims to achieve
    * @param stepIndex - Optional explicit step index (otherwise auto-increments)
+   * @param options - Optional settings: emitTrace (default true), preUrl
    * @returns Generated stepId in format 'step-N' where N is the step index
    */
-  beginStep(goal: string, stepIndex?: number): string {
+  beginStep(
+    goal: string,
+    stepIndex?: number,
+    options?: { emitTrace?: boolean; preUrl?: string }
+  ): string {
+    const { emitTrace = true, preUrl } = options || {};
+
     // Clear previous step state
     this.assertionsThisStep = [];
     this.stepPreSnapshot = null;
@@ -1190,6 +1234,16 @@ export class AgentRuntime {
     // Generate stepId in 'step-N' format for Studio compatibility
     this.stepId = `step-${this.stepIndex}`;
 
+    // Emit step_start trace event for Studio timeline display
+    if (emitTrace && this.tracer) {
+      try {
+        const url = preUrl || this.lastSnapshot?.url || this.page?.url?.() || '';
+        this.tracer.emitStepStart(this.stepId, this.stepIndex, goal, 0, url);
+      } catch {
+        // Tracing must be non-fatal
+      }
+    }
+
     return this.stepId;
   }
 

src/tracing/tracer-factory.ts

@@ -19,6 +19,35 @@ import { Tracer } from './tracer';
 import { CloudTraceSink, SentienceLogger } from './cloud-sink';
 import { JsonlTraceSink } from './jsonl-sink';
 
+/**
+ * Helper to emit run_start event with available metadata
+ */
+function emitRunStart(
+  tracer: Tracer,
+  agentType?: string,
+  llmModel?: string,
+  goal?: string,
+  startUrl?: string
+): void {
+  try {
+    const config: Record<string, string> = {};
+    if (goal) {
+      config.goal = goal;
+    }
+    if (startUrl) {
+      config.start_url = startUrl;
+    }
+
+    tracer.emitRunStart(
+      agentType || 'SentienceAgent',
+      llmModel,
+      Object.keys(config).length > 0 ? config : undefined
+    );
+  } catch {
+    // Tracing must be non-fatal
+  }
+}
+
 /**
  * Sentience API base URL (constant)
  */
@@ -198,6 +227,7 @@ function httpPost(
  * @param options.llmModel - LLM model used (e.g., "gpt-4-turbo", "claude-3-5-sonnet")
  * @param options.startUrl - Starting URL of the agent run (e.g., "https://amazon.com")
  * @param options.screenshotProcessor - Optional function to process screenshots before upload. Takes base64 string, returns processed base64 string. Useful for PII redaction or custom image processing.
+ * @param options.autoEmitRunStart - If true (default), automatically emit run_start event with provided metadata. This ensures traces have complete structure for Studio visualization.
  * @returns Tracer configured with appropriate sink
  *
  * @example
@@ -213,6 +243,7 @@ function httpPost(
  *   uploadTrace: true
  * });
  * // Returns: Tracer with CloudTraceSink
+ * // run_start event is automatically emitted
  *
  * // With screenshot processor for PII redaction
  * const redactPII = (screenshot: string): string => {
@@ -229,6 +260,10 @@ function httpPost(
  * const tracer = await createTracer({ apiKey: "sk_pro_xyz", runId: "demo", uploadTrace: false });
  * // Returns: Tracer with JsonlTraceSink (local-only)
  *
+ * // Disable auto-emit for manual control
+ * const tracer = await createTracer({ runId: "demo", autoEmitRunStart: false });
+ * tracer.emitRunStart("MyAgent", "gpt-4o"); // Manual emit
+ *
  * // Free tier user
  * const tracer = await createTracer({ runId: "demo" });
  * // Returns: Tracer with JsonlTraceSink (local-only)
@@ -250,6 +285,7 @@ export async function createTracer(options: {
   llmModel?: string;
   startUrl?: string;
   screenshotProcessor?: (screenshot: string) => string;
+  autoEmitRunStart?: boolean;
 }): Promise<Tracer> {
   const runId = options.runId || randomUUID();
   const apiUrl = options.apiUrl || SENTIENCE_API_URL;
@@ -303,11 +339,18 @@ export async function createTracer(options: {
 
         console.log('☁️  [Sentience] Cloud tracing enabled (Pro tier)');
         // PRODUCTION FIX: Pass runId for persistent cache naming
-        return new Tracer(
+        const tracer = new Tracer(
           runId,
           new CloudTraceSink(uploadUrl, runId, options.apiKey, apiUrl, options.logger),
           options.screenshotProcessor
         );
+
+        // Auto-emit run_start for complete trace structure
+        if (options.autoEmitRunStart !== false) {
+          emitRunStart(tracer, options.agentType, options.llmModel, options.goal, options.startUrl);
+        }
+
+        return tracer;
       } else if (response.status === 403) {
         console.log('⚠️  [Sentience] Cloud tracing requires Pro tier');
         console.log('   Falling back to local-only tracing');
@@ -338,7 +381,14 @@ export async function createTracer(options: {
   const localPath = path.join(tracesDir, `${runId}.jsonl`);
   console.log(`💾 [Sentience] Local tracing: ${localPath}`);
 
-  return new Tracer(runId, new JsonlTraceSink(localPath), options.screenshotProcessor);
+  const tracer = new Tracer(runId, new JsonlTraceSink(localPath), options.screenshotProcessor);
+
+  // Auto-emit run_start for complete trace structure
+  if (options.autoEmitRunStart !== false) {
+    emitRunStart(tracer, options.agentType, options.llmModel, options.goal, options.startUrl);
+  }
+
+  return tracer;
 }
 
 /**

src/tracing/tracer.ts

@@ -202,6 +202,93 @@ export class Tracer {
     this.emit('error', { step_id: stepId, error, attempt }, stepId);
   }
 
+  /**
+   * Emit snapshot event with screenshot for Studio visualization.
+   *
+   * This method builds and emits a 'snapshot' trace event that includes:
+   * - Page URL and element data
+   * - Screenshot (if present in snapshot)
+   * - Step correlation info
+   *
+   * Use this when you want screenshots to appear in the Sentience Studio timeline.
+   *
+   * @param snapshot - Snapshot object (must have 'screenshot' property for images)
+   * @param stepId - Step UUID (for correlating snapshot with a step)
+   * @param stepIndex - Step index (0-based) for Studio timeline ordering
+   * @param screenshotFormat - Format of screenshot ("jpeg" or "png", default: "jpeg")
+   *
+   * @example
+   * // After taking a snapshot with AgentRuntime
+   * const snapshot = await runtime.snapshot();
+   * tracer.emitSnapshot(snapshot, runtime.getStepId(), runtime.getStepIndex());
+   *
+   * // Or use auto-emit (default in AgentRuntime.snapshot())
+   * const snapshot = await runtime.snapshot();  // Auto-emits snapshot event
+   */
+  emitSnapshot(
+    snapshot: any,
+    stepId?: string,
+    stepIndex?: number,
+    screenshotFormat: string = 'jpeg'
+  ): void {
+    if (!snapshot) {
+      return;
+    }
+
+    try {
+      // Build the snapshot event data
+      const data: TraceEventData = {
+        url: snapshot.url,
+        element_count: snapshot.elements?.length || 0,
+        timestamp: snapshot.timestamp,
+      };
+
+      // Include step_index if provided (required for UUID step_ids)
+      if (stepIndex !== undefined) {
+        data.step_index = stepIndex;
+      }
+
+      // Include elements data (simplified for trace)
+      if (snapshot.elements && snapshot.elements.length > 0) {
+        // Normalize importance values to importance_score (0-1 range)
+        const importanceValues = snapshot.elements.map((el: any) => el.importance || 0);
+        const minImportance = Math.min(...importanceValues);
+        const maxImportance = Math.max(...importanceValues);
+        const importanceRange = maxImportance - minImportance;
+
+        data.elements = snapshot.elements.map((el: any) => {
+          const importanceScore =
+            importanceRange > 0 ? (el.importance - minImportance) / importanceRange : 0.5;
+          return {
+            ...el,
+            importance_score: importanceScore,
+          };
+        });
+      }
+
+      // Extract and add screenshot if present
+      const screenshotRaw = snapshot.screenshot;
+      if (screenshotRaw) {
+        // Extract base64 string from data URL if needed
+        // Format: "data:image/jpeg;base64,{base64_string}"
+        let screenshotBase64: string;
+        if (typeof screenshotRaw === 'string' && screenshotRaw.startsWith('data:image')) {
+          const commaIndex = screenshotRaw.indexOf(',');
+          screenshotBase64 =
+            commaIndex !== -1 ? screenshotRaw.slice(commaIndex + 1) : screenshotRaw;
+        } else {
+          screenshotBase64 = screenshotRaw;
+        }
+        data.screenshot_base64 = screenshotBase64;
+        data.screenshot_format = screenshotFormat;
+      }
+
+      this.emit('snapshot', data, stepId);
+    } catch {
+      // Best-effort: don't let trace emission errors break the caller
+    }
+  }
+
   /**
    * Automatically infer finalStatus from tracked step outcomes if not explicitly set.
    * This is called automatically in close() if finalStatus is still "unknown".

tests/tracing/tracer-factory.test.ts

@@ -305,6 +305,7 @@ describe('createTracer', () => {
     it('should work with agent workflow (Free tier)', async () => {
       const tracer = await createTracer({
         runId: 'agent-test',
+        autoEmitRunStart: false, // Disable auto-emit since we're manually emitting
       });
 
       tracer.emitRunStart('SentienceAgent', 'gpt-4');