@vellumai/assistant 0.9.0 → 0.10.0-staging.2

package/src/workflows/engine.test.ts

@@ -16,7 +16,11 @@ import type { TrustContext } from "../daemon/trust-context.js";
 import { getSqlite } from "../memory/db-connection.js";
 import { initializeDb } from "../memory/db-init.js";
 import type { ResolvedCapabilities } from "./capabilities.js";
-import { executeWorkflow, extractWorkflowMeta } from "./engine.js";
+import {
+  executeWorkflow,
+  extractWorkflowMeta,
+  type WorkflowLeafEvent,
+} from "./engine.js";
 import * as journalStore from "./journal-store.js";
 import type { LeafResult, RunLeafOptions } from "./leaf-runner.js";
 
@@ -677,6 +681,21 @@ describe("executeWorkflow — resume", () => {
     );
     expect(journalA.map((e) => e.seq)).toEqual([0, 1, 2]);
   });
+
+  test("a completed leaf persists its per-leaf token usage in the journal", async () => {
+    const runner = makeFakeRunner();
+    await execute("wf-leaf-tokens", `return agent("alpha");`, runner.runner);
+
+    const journal = journalStore.getJournal("wf-leaf-tokens");
+    expect(journal).toHaveLength(1);
+    // The fake runner reports 10/5 per leaf; the completed entry carries them.
+    expect(journal[0]).toMatchObject({
+      seq: 0,
+      status: "completed",
+      inputTokens: 10,
+      outputTokens: 5,
+    });
+  });
 });
 
 describe("executeWorkflow — schema vs tool leaf tool forwarding", () => {
@@ -1142,3 +1161,158 @@ describe("executeWorkflow — nested workflow()", () => {
     expect(second.prompts).toEqual([]);
   });
 });
+
+describe("executeWorkflow — onLeaf observe-only hook", () => {
+  beforeEach(resetTables);
+
+  /** Run a workflow, collecting every onLeaf event. */
+  function executeWithLeafEvents(
+    runId: string,
+    scriptSource: string,
+    runner: typeof import("./leaf-runner.js").runLeaf,
+    config: WorkflowsConfig = configWith(),
+    signal?: AbortSignal,
+  ): {
+    promise: ReturnType<typeof executeWorkflow>;
+    events: WorkflowLeafEvent[];
+  } {
+    const events: WorkflowLeafEvent[] = [];
+    const promise = executeWorkflow({
+      runId,
+      scriptSource: META + scriptSource,
+      args: null,
+      capabilities: CAPABILITIES,
+      config,
+      journal: journalStore,
+      leafRunner: runner,
+      trustContext: TRUST,
+      onLeaf: (e) => events.push(e),
+      ...(signal ? { signal } : {}),
+    });
+    return { promise, events };
+  }
+
+  test("a single agent() emits leaf_started then leaf_finished:completed at seq 0", async () => {
+    const fake = makeFakeRunner();
+    const { promise, events } = executeWithLeafEvents(
+      "wf-leaf-single",
+      `return agent("solo");`,
+      fake.runner,
+    );
+    const res = await promise;
+    expect(res.status).toBe("completed");
+
+    expect(events).toEqual([
+      { type: "leaf_started", seq: 0, promptSummary: "solo" },
+      {
+        type: "leaf_finished",
+        seq: 0,
+        status: "completed",
+        inputTokens: 10,
+        outputTokens: 5,
+        resultSummary: "out:solo",
+      },
+    ]);
+  });
+
+  test("a 3-spec parallel() emits 3 started + 3 finished across seqs 0,1,2", async () => {
+    const fake = makeFakeRunner();
+    const { promise, events } = executeWithLeafEvents(
+      "wf-leaf-parallel",
+      `return parallel([leaf("a"), leaf("b"), leaf("c")]);`,
+      fake.runner,
+    );
+    const res = await promise;
+    expect(res.status).toBe("completed");
+
+    const started = events.filter((e) => e.type === "leaf_started");
+    const finished = events.filter((e) => e.type === "leaf_finished");
+    expect(started.map((e) => e.seq).sort((a, b) => a - b)).toEqual([0, 1, 2]);
+    expect(finished.map((e) => e.seq).sort((a, b) => a - b)).toEqual([0, 1, 2]);
+    expect(finished.every((e) => e.status === "completed")).toBe(true);
+  });
+
+  test("a throwing (non-abort) leaf emits leaf_started then leaf_finished:failed", async () => {
+    const fake = makeFakeRunner({ failOn: (p) => p === "boom" });
+    const { promise, events } = executeWithLeafEvents(
+      "wf-leaf-fail",
+      `return parallel([leaf("boom")]);`,
+      fake.runner,
+    );
+    await promise;
+
+    expect(events).toEqual([
+      { type: "leaf_started", seq: 0, promptSummary: "boom" },
+      {
+        type: "leaf_finished",
+        seq: 0,
+        status: "failed",
+        inputTokens: 0,
+        outputTokens: 0,
+        resultSummary: "leaf failed: boom",
+      },
+    ]);
+  });
+
+  test("an aborted in-flight leaf emits leaf_started but NO leaf_finished", async () => {
+    const controller = new AbortController();
+    let started = false;
+    const runner = (async (): Promise<
+      import("./leaf-runner.js").LeafResult
+    > => {
+      started = true;
+      await new Promise<void>((resolve) => {
+        if (controller.signal.aborted) return resolve();
+        controller.signal.addEventListener("abort", () => resolve(), {
+          once: true,
+        });
+      });
+      const err = new Error("The operation was aborted");
+      err.name = "AbortError";
+      throw err;
+    }) as unknown as typeof import("./leaf-runner.js").runLeaf;
+
+    const { promise, events } = executeWithLeafEvents(
+      "wf-leaf-abort",
+      `return agent("long-running");`,
+      runner,
+      configWith(),
+      controller.signal,
+    );
+
+    await waitUntil(() => started);
+    controller.abort();
+
+    const res = await promise;
+    expect(res.status).toBe("aborted");
+    expect(events).toEqual([
+      { type: "leaf_started", seq: 0, promptSummary: "long-running" },
+    ]);
+  });
+
+  test("a journal-replayed leaf emits neither leaf_started nor leaf_finished", async () => {
+    const script = `return agent("alpha");`;
+
+    // First run journals the leaf as completed.
+    const first = makeFakeRunner();
+    const res1 = await executeWithLeafEvents(
+      "wf-leaf-replay",
+      script,
+      first.runner,
+    ).promise;
+    expect(res1.status).toBe("completed");
+    expect(first.prompts).toEqual(["alpha"]);
+
+    // Resume: the leaf replays from the journal, so no onLeaf events fire.
+    const second = makeFakeRunner();
+    const { promise, events } = executeWithLeafEvents(
+      "wf-leaf-replay",
+      script,
+      second.runner,
+    );
+    const res2 = await promise;
+    expect(res2.status).toBe("completed");
+    expect(second.prompts).toEqual([]);
+    expect(events).toEqual([]);
+  });
+});

package/src/workflows/engine.ts

@@ -73,6 +73,30 @@ export type WorkflowProgressEvent =
   | { type: "phase"; title: string }
   | { type: "log"; message: string };
 
+/**
+ * An observe-only lifecycle event for a single leaf — the source of a live leaf
+ * tree. Emitted once when a leaf STARTS (after the cap check + spawn increment)
+ * and once when it FINISHES (completed or failed). A journal-replayed leaf and
+ * an aborted in-flight leaf emit neither.
+ */
+export type WorkflowLeafEvent =
+  | {
+      type: "leaf_started";
+      seq: number;
+      label?: string;
+      phase?: string;
+      promptSummary: string;
+    }
+  | {
+      type: "leaf_finished";
+      seq: number;
+      status: "completed" | "failed";
+      label?: string;
+      inputTokens: number;
+      outputTokens: number;
+      resultSummary: string;
+    };
+
 /** The journal-store surface the engine depends on (injectable for tests). */
 export interface WorkflowJournal {
   appendJournalEntry: typeof JournalStore.appendJournalEntry;
@@ -108,6 +132,13 @@ export interface ExecuteWorkflowOptions {
   trustContext: TrustContext;
   /** Receives `phase`/`log` progress events from the script. */
   onProgress?: (event: WorkflowProgressEvent) => void;
+  /**
+   * Observe-only per-leaf lifecycle hook. Emits once on START (after the cap
+   * check + spawn increment; NEVER on a journal-replay short-circuit) and once
+   * on FINISH (completed or failed). An aborted in-flight leaf emits NEITHER a
+   * finish nor a spurious failed.
+   */
+  onLeaf?: (event: WorkflowLeafEvent) => void;
   /** Cooperative cancellation for the whole run. */
   signal?: AbortSignal;
 }
@@ -451,6 +482,26 @@ function skipRegexLiteral(source: string, start: number): number {
   return source.length - 1;
 }
 
+/**
+ * A bounded string summary of a leaf prompt or result, for an observe-only
+ * {@link WorkflowLeafEvent}. Strings pass through; everything else is
+ * JSON-serialized (falling back to `String(value)` if it cannot be). Truncated
+ * to `max` chars with a trailing ellipsis.
+ */
+function summarizeForLeafEvent(value: unknown, max = 200): string {
+  let text: string;
+  if (typeof value === "string") {
+    text = value;
+  } else {
+    try {
+      text = JSON.stringify(value);
+    } catch {
+      text = String(value);
+    }
+  }
+  return text.length > max ? `${text.slice(0, max)}…` : text;
+}
+
 function callHashOf(prompt: string, opts: LeafCallOptions): string {
   return createHash("sha256")
     .update(deterministicStringify({ prompt, opts }))
@@ -542,6 +593,7 @@ export async function executeWorkflow(
     leafRunner,
     trustContext,
     onProgress,
+    onLeaf,
     signal,
   } = opts;
 
@@ -630,6 +682,14 @@ export async function executeWorkflow(
     }
     agentsSpawned += 1;
 
+    onLeaf?.({
+      type: "leaf_started",
+      seq,
+      ...(leafOpts.label !== undefined ? { label: leafOpts.label } : {}),
+      ...(leafOpts.phase !== undefined ? { phase: leafOpts.phase } : {}),
+      promptSummary: summarizeForLeafEvent(prompt),
+    });
+
     try {
       // A leaf is EITHER a schema leaf (structured output via forced
       // tool-choice — `schema` set, no tools) OR a tool leaf (free-form with
@@ -666,6 +726,17 @@ export async function executeWorkflow(
         request: { prompt, opts: leafOpts },
         result: result.output,
         status: "completed",
+        inputTokens: result.inputTokens,
+        outputTokens: result.outputTokens,
+      });
+      onLeaf?.({
+        type: "leaf_finished",
+        seq,
+        status: "completed",
+        ...(leafOpts.label !== undefined ? { label: leafOpts.label } : {}),
+        inputTokens: result.inputTokens,
+        outputTokens: result.outputTokens,
+        resultSummary: summarizeForLeafEvent(result.output),
       });
       return { output: result.output, failed: false };
     } catch (err) {
@@ -693,6 +764,17 @@ export async function executeWorkflow(
         { err, runId, seq, label: leafOpts.label },
         "Workflow leaf failed",
       );
+      onLeaf?.({
+        type: "leaf_finished",
+        seq,
+        status: "failed",
+        ...(leafOpts.label !== undefined ? { label: leafOpts.label } : {}),
+        inputTokens: 0,
+        outputTokens: 0,
+        resultSummary: summarizeForLeafEvent(
+          err instanceof Error ? err.message : String(err),
+        ),
+      });
       return { output: null, failed: true };
     }
   };

package/src/workflows/journal-store.test.ts

@@ -197,6 +197,76 @@ describe("workflow journal store", () => {
     });
   });
 
+  test("appendJournalEntry round-trips per-leaf token usage", () => {
+    createRun({ id: "wf-tok", scriptSource: "s", scriptHash: "h" });
+
+    appendJournalEntry({
+      runId: "wf-tok",
+      seq: 0,
+      callHash: "c0",
+      kind: "agent",
+      request: { prompt: "with tokens" },
+      result: { text: "r" },
+      status: "completed",
+      inputTokens: 120,
+      outputTokens: 45,
+    });
+    // A leaf with no recorded usage stores NULL → reads back as undefined.
+    appendJournalEntry({
+      runId: "wf-tok",
+      seq: 1,
+      callHash: "c1",
+      kind: "agent",
+      request: { prompt: "no tokens" },
+      result: { error: "boom" },
+      status: "failed",
+    });
+
+    const journal = getJournal("wf-tok");
+    expect(journal[0]).toMatchObject({
+      seq: 0,
+      inputTokens: 120,
+      outputTokens: 45,
+    });
+    expect(journal[1].inputTokens).toBeUndefined();
+    expect(journal[1].outputTokens).toBeUndefined();
+  });
+
+  test("appendJournalEntry upserts token usage on a changed-hash re-run", () => {
+    createRun({ id: "wf-tok-resume", scriptSource: "s", scriptHash: "h" });
+
+    appendJournalEntry({
+      runId: "wf-tok-resume",
+      seq: 0,
+      callHash: "hash-old",
+      kind: "agent",
+      result: { text: "old" },
+      status: "completed",
+      inputTokens: 10,
+      outputTokens: 5,
+    });
+    // Resume re-runs the leaf with new usage; the upsert overwrites the tokens.
+    appendJournalEntry({
+      runId: "wf-tok-resume",
+      seq: 0,
+      callHash: "hash-new",
+      kind: "agent",
+      result: { text: "new" },
+      status: "completed",
+      inputTokens: 30,
+      outputTokens: 12,
+    });
+
+    const journal = getJournal("wf-tok-resume");
+    expect(journal).toHaveLength(1);
+    expect(journal[0]).toMatchObject({
+      seq: 0,
+      callHash: "hash-new",
+      inputTokens: 30,
+      outputTokens: 12,
+    });
+  });
+
   test("appendJournalEntry is idempotent on (run_id, seq)", () => {
     createRun({ id: "wf-dup", scriptSource: "s", scriptHash: "h" });
 

package/src/workflows/journal-store.ts

@@ -68,6 +68,9 @@ export interface WorkflowJournalEntry {
   request: unknown;
   result: unknown;
   status: string;
+  /** Per-leaf token usage for a completed agent leaf; undefined otherwise. */
+  inputTokens?: number;
+  outputTokens?: number;
   createdAt: number | null;
 }
 
@@ -111,6 +114,9 @@ export interface AppendJournalEntryInput {
   request?: unknown;
   result?: unknown;
   status: string;
+  /** Per-leaf token usage; set for a completed agent leaf, omitted otherwise. */
+  inputTokens?: number;
+  outputTokens?: number;
 }
 
 // ---------------------------------------------------------------------------
@@ -145,6 +151,8 @@ interface WorkflowJournalRow {
   request_json: string | null;
   result_json: string | null;
   status: string;
+  input_tokens: number | null;
+  output_tokens: number | null;
   created_at: number | null;
 }
 
@@ -194,6 +202,8 @@ function rowToJournalEntry(row: WorkflowJournalRow): WorkflowJournalEntry {
     request: parseJsonColumn(row.request_json),
     result: parseJsonColumn(row.result_json),
     status: row.status,
+    inputTokens: row.input_tokens ?? undefined,
+    outputTokens: row.output_tokens ?? undefined,
     createdAt: row.created_at,
   };
 }
@@ -318,14 +328,17 @@ export function appendJournalEntry(
   rawRun(
     /*sql*/ `
     INSERT INTO workflow_journal (
-      run_id, seq, call_hash, kind, request_json, result_json, status, created_at
-    ) VALUES (?, ?, ?, ?, ?, ?, ?, ?)
+      run_id, seq, call_hash, kind, request_json, result_json, status,
+      input_tokens, output_tokens, created_at
+    ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
     ON CONFLICT(run_id, seq) DO UPDATE SET
       call_hash = excluded.call_hash,
       kind = excluded.kind,
       request_json = excluded.request_json,
       result_json = excluded.result_json,
-      status = excluded.status
+      status = excluded.status,
+      input_tokens = excluded.input_tokens,
+      output_tokens = excluded.output_tokens
     `,
     input.runId,
     input.seq,
@@ -334,6 +347,8 @@ export function appendJournalEntry(
     serializeJsonColumn(input.request),
     serializeJsonColumn(input.result),
     input.status,
+    input.inputTokens ?? null,
+    input.outputTokens ?? null,
     Date.now(),
   );
   return getJournalEntry(input.runId, input.seq)!;