@stigmer/runner 3.0.1

package/src/budget/__tests__/tracker.test.ts

@@ -0,0 +1,293 @@
+import { describe, it, expect } from "vitest";
+import { BudgetTracker, extractCostFromOutput, type WorkflowBudget } from "../tracker.js";
+
+const START = 1_000_000;
+
+describe("BudgetTracker", () => {
+  describe("unlimited (no budget)", () => {
+    it("returns ok=true when budget is undefined", () => {
+      const t = new BudgetTracker(undefined, START);
+      t.record(999_999, 100_000, 50_000);
+      const result = t.check(START + 3_600_000);
+      expect(result.ok).toBe(true);
+      expect(result.exceeded).toBe(false);
+    });
+
+    it("costRemaining returns -1 when unlimited", () => {
+      const t = new BudgetTracker(undefined, START);
+      expect(t.costRemaining()).toBe(-1);
+    });
+
+    it("tokensRemaining returns -1 when unlimited", () => {
+      const t = new BudgetTracker(undefined, START);
+      expect(t.tokensRemaining()).toBe(-1);
+    });
+  });
+
+  describe("record", () => {
+    it("accumulates cost and token usage", () => {
+      const t = new BudgetTracker(undefined, START);
+      t.record(100, 50, 30);
+      t.record(200, 25, 10);
+      expect(t.costMicros).toBe(300);
+      expect(t.inputTokens).toBe(75);
+      expect(t.outputTokens).toBe(40);
+      expect(t.totalTokens()).toBe(115);
+    });
+  });
+
+  describe("cost limit", () => {
+    const budget: WorkflowBudget = { maxCostMicros: 1000 };
+
+    it("passes when under limit", () => {
+      const t = new BudgetTracker(budget, START);
+      t.record(500, 0, 0);
+      const result = t.check(START);
+      expect(result.ok).toBe(true);
+      expect(result.warningPct).toBeCloseTo(0.5);
+    });
+
+    it("exceeds when over limit", () => {
+      const t = new BudgetTracker(budget, START);
+      t.record(1001, 0, 0);
+      const result = t.check(START);
+      expect(result.ok).toBe(false);
+      expect(result.exceeded).toBe(true);
+      expect(result.exceededLimit).toBe("cost");
+      expect(result.exceededMessage).toContain("1001/1000");
+    });
+
+    it("costRemaining returns remaining budget", () => {
+      const t = new BudgetTracker(budget, START);
+      t.record(300, 0, 0);
+      expect(t.costRemaining()).toBe(700);
+    });
+
+    it("costRemaining returns 0 when exceeded", () => {
+      const t = new BudgetTracker(budget, START);
+      t.record(2000, 0, 0);
+      expect(t.costRemaining()).toBe(0);
+    });
+  });
+
+  describe("token limit", () => {
+    const budget: WorkflowBudget = { maxTotalTokens: 10_000 };
+
+    it("passes when under limit", () => {
+      const t = new BudgetTracker(budget, START);
+      t.record(0, 3000, 2000);
+      const result = t.check(START);
+      expect(result.ok).toBe(true);
+      expect(result.warningPct).toBeCloseTo(0.5);
+    });
+
+    it("exceeds when input + output exceeds limit", () => {
+      const t = new BudgetTracker(budget, START);
+      t.record(0, 6000, 5000);
+      const result = t.check(START);
+      expect(result.ok).toBe(false);
+      expect(result.exceeded).toBe(true);
+      expect(result.exceededLimit).toBe("tokens");
+      expect(result.exceededMessage).toContain("11000/10000");
+    });
+
+    it("tokensRemaining returns remaining budget", () => {
+      const t = new BudgetTracker(budget, START);
+      t.record(0, 3000, 1000);
+      expect(t.tokensRemaining()).toBe(6000);
+    });
+
+    it("tokensRemaining returns 0 when exceeded", () => {
+      const t = new BudgetTracker(budget, START);
+      t.record(0, 6000, 5000);
+      expect(t.tokensRemaining()).toBe(0);
+    });
+  });
+
+  describe("duration limit", () => {
+    const budget: WorkflowBudget = { maxDurationSeconds: 60 };
+
+    it("passes when under limit", () => {
+      const t = new BudgetTracker(budget, START);
+      const result = t.check(START + 30_000);
+      expect(result.ok).toBe(true);
+      expect(result.warningPct).toBeCloseTo(0.5);
+    });
+
+    it("exceeds when elapsed > limit", () => {
+      const t = new BudgetTracker(budget, START);
+      const result = t.check(START + 61_000);
+      expect(result.ok).toBe(false);
+      expect(result.exceeded).toBe(true);
+      expect(result.exceededLimit).toBe("duration");
+      expect(result.exceededMessage).toContain("61s/60s");
+    });
+  });
+
+  describe("multiple limits", () => {
+    it("checks cost before tokens before duration", () => {
+      const budget: WorkflowBudget = {
+        maxCostMicros: 100,
+        maxTotalTokens: 1000,
+        maxDurationSeconds: 300,
+      };
+      const t = new BudgetTracker(budget, START);
+      t.record(200, 500, 500);
+      const result = t.check(START + 100_000);
+      expect(result.exceededLimit).toBe("cost");
+    });
+
+    it("reports token limit when cost is within budget", () => {
+      const budget: WorkflowBudget = {
+        maxCostMicros: 10_000,
+        maxTotalTokens: 1000,
+      };
+      const t = new BudgetTracker(budget, START);
+      t.record(50, 600, 500);
+      const result = t.check(START);
+      expect(result.exceededLimit).toBe("tokens");
+    });
+
+    it("warningPct reflects the tightest limit when ok", () => {
+      const budget: WorkflowBudget = {
+        maxCostMicros: 1000,
+        maxTotalTokens: 10_000,
+      };
+      const t = new BudgetTracker(budget, START);
+      t.record(800, 1000, 500);
+      const result = t.check(START);
+      expect(result.ok).toBe(true);
+      expect(result.warningPct).toBeCloseTo(0.8);
+    });
+  });
+
+  describe("policy", () => {
+    it("defaults to terminate when unspecified", () => {
+      const budget: WorkflowBudget = { maxCostMicros: 100 };
+      const t = new BudgetTracker(budget, START);
+      t.record(200, 0, 0);
+      const result = t.check(START);
+      expect(result.policy).toBe("terminate");
+    });
+
+    it("uses configured policy", () => {
+      const budget: WorkflowBudget = {
+        maxCostMicros: 100,
+        onExceeded: "warn",
+      };
+      const t = new BudgetTracker(budget, START);
+      t.record(200, 0, 0);
+      const result = t.check(START);
+      expect(result.policy).toBe("warn");
+    });
+
+    it("returns policy in ok results", () => {
+      const budget: WorkflowBudget = {
+        maxCostMicros: 1000,
+        onExceeded: "human_review",
+      };
+      const t = new BudgetTracker(budget, START);
+      const result = t.check(START);
+      expect(result.ok).toBe(true);
+      expect(result.policy).toBe("human_review");
+    });
+  });
+});
+
+describe("extractCostFromOutput", () => {
+  it("extracts __stigmer_* prefixed keys", () => {
+    const output = {
+      __stigmer_cost_micros: 500,
+      __stigmer_input_tokens: 1000,
+      __stigmer_output_tokens: 200,
+      result: "some text",
+    };
+    const cost = extractCostFromOutput(output);
+    expect(cost.costMicros).toBe(500);
+    expect(cost.inputTokens).toBe(1000);
+    expect(cost.outputTokens).toBe(200);
+  });
+
+  it("falls back to unprefixed keys (LLM activity output)", () => {
+    const output = {
+      cost_micros: 0,
+      input_tokens: 800,
+      output_tokens: 150,
+      result: "text",
+    };
+    const cost = extractCostFromOutput(output);
+    expect(cost.inputTokens).toBe(800);
+    expect(cost.outputTokens).toBe(150);
+  });
+
+  it("prefers __stigmer_* over unprefixed when both present", () => {
+    const output = {
+      __stigmer_input_tokens: 100,
+      input_tokens: 999,
+    };
+    const cost = extractCostFromOutput(output);
+    expect(cost.inputTokens).toBe(100);
+  });
+
+  it("returns zeroes for null output", () => {
+    const cost = extractCostFromOutput(null);
+    expect(cost.costMicros).toBe(0);
+    expect(cost.inputTokens).toBe(0);
+    expect(cost.outputTokens).toBe(0);
+  });
+
+  it("returns zeroes for undefined output", () => {
+    const cost = extractCostFromOutput(undefined);
+    expect(cost.costMicros).toBe(0);
+  });
+
+  it("returns zeroes for non-object output", () => {
+    const cost = extractCostFromOutput("some string");
+    expect(cost.costMicros).toBe(0);
+  });
+
+  it("floors fractional values", () => {
+    const output = { input_tokens: 10.7, output_tokens: 5.3 };
+    const cost = extractCostFromOutput(output);
+    expect(cost.inputTokens).toBe(10);
+    expect(cost.outputTokens).toBe(5);
+  });
+
+  it("treats negative values as zero", () => {
+    const output = { input_tokens: -5, output_tokens: 10 };
+    const cost = extractCostFromOutput(output);
+    expect(cost.inputTokens).toBe(0);
+    expect(cost.outputTokens).toBe(10);
+  });
+
+  it("treats non-numeric values as zero", () => {
+    const output = { input_tokens: "lots", output_tokens: true };
+    const cost = extractCostFromOutput(output);
+    expect(cost.inputTokens).toBe(0);
+    expect(cost.outputTokens).toBe(0);
+  });
+
+  it("extracts tokens from nested usage object (normalized llm_call output)", () => {
+    const output = {
+      text: "hello",
+      model: "gpt-4o",
+      __stigmer_cost_micros: 150,
+      usage: { input_tokens: 400, output_tokens: 80 },
+    };
+    const cost = extractCostFromOutput(output);
+    expect(cost.costMicros).toBe(150);
+    expect(cost.inputTokens).toBe(400);
+    expect(cost.outputTokens).toBe(80);
+  });
+
+  it("prefers top-level tokens over nested usage", () => {
+    const output = {
+      input_tokens: 100,
+      output_tokens: 50,
+      usage: { input_tokens: 999, output_tokens: 999 },
+    };
+    const cost = extractCostFromOutput(output);
+    expect(cost.inputTokens).toBe(100);
+    expect(cost.outputTokens).toBe(50);
+  });
+});

package/src/budget/index.ts

@@ -0,0 +1,9 @@
+export {
+  BudgetTracker,
+  extractCostFromOutput,
+  type WorkflowBudget,
+  type BudgetExceededPolicy,
+  type CheckResult,
+  type LimitKind,
+  type TaskCostInfo,
+} from "./tracker.js";

package/src/budget/tracker.ts

@@ -0,0 +1,171 @@
+/**
+ * Budget tracker — accumulates cost, token, and duration usage against
+ * workflow-level budget limits.
+ *
+ * Port of Go's `pkg/budget/tracker.go`. The tracker is a pure, sandbox-safe
+ * class with zero external dependencies — it can run inside the Temporal
+ * deterministic isolate without importing Node.js or Temporal APIs.
+ *
+ * Budget config arrives via TemporalWorkflowInput (outside YAML) because
+ * budget is a Stigmer extension that doesn't survive YAML round-trip.
+ */
+
+export type LimitKind = "cost" | "tokens" | "duration";
+
+export type BudgetExceededPolicy = "terminate" | "warn" | "human_review";
+
+export interface WorkflowBudget {
+  readonly maxCostMicros?: number;
+  readonly maxTotalTokens?: number;
+  readonly maxDurationSeconds?: number;
+  readonly onExceeded?: BudgetExceededPolicy;
+}
+
+export interface CheckResult {
+  readonly ok: boolean;
+  readonly warningPct: number;
+  readonly exceeded: boolean;
+  readonly exceededLimit?: LimitKind;
+  readonly exceededMessage?: string;
+  readonly policy: BudgetExceededPolicy;
+}
+
+export class BudgetTracker {
+  private readonly budget: WorkflowBudget | undefined;
+  private readonly startedAtMs: number;
+
+  costMicros = 0;
+  inputTokens = 0;
+  outputTokens = 0;
+
+  constructor(budget: WorkflowBudget | undefined, startedAtMs: number) {
+    this.budget = budget;
+    this.startedAtMs = startedAtMs;
+  }
+
+  record(costMicros: number, inputTokens: number, outputTokens: number): void {
+    this.costMicros += costMicros;
+    this.inputTokens += inputTokens;
+    this.outputTokens += outputTokens;
+  }
+
+  totalTokens(): number {
+    return this.inputTokens + this.outputTokens;
+  }
+
+  check(nowMs: number): CheckResult {
+    const policy = this.resolvePolicy();
+
+    if (!this.budget) {
+      return { ok: true, warningPct: 0, exceeded: false, policy };
+    }
+
+    let highestPct = 0;
+
+    if (this.budget.maxCostMicros && this.budget.maxCostMicros > 0) {
+      const pct = this.costMicros / this.budget.maxCostMicros;
+      if (pct > highestPct) highestPct = pct;
+      if (this.costMicros > this.budget.maxCostMicros) {
+        return {
+          ok: false,
+          warningPct: pct,
+          exceeded: true,
+          exceededLimit: "cost",
+          exceededMessage: `cost budget exceeded: ${this.costMicros}/${this.budget.maxCostMicros} micro-USD`,
+          policy,
+        };
+      }
+    }
+
+    if (this.budget.maxTotalTokens && this.budget.maxTotalTokens > 0) {
+      const total = this.totalTokens();
+      const pct = total / this.budget.maxTotalTokens;
+      if (pct > highestPct) highestPct = pct;
+      if (total > this.budget.maxTotalTokens) {
+        return {
+          ok: false,
+          warningPct: pct,
+          exceeded: true,
+          exceededLimit: "tokens",
+          exceededMessage: `token budget exceeded: ${total}/${this.budget.maxTotalTokens} tokens`,
+          policy,
+        };
+      }
+    }
+
+    if (this.budget.maxDurationSeconds && this.budget.maxDurationSeconds > 0) {
+      const elapsedMs = nowMs - this.startedAtMs;
+      const maxMs = this.budget.maxDurationSeconds * 1000;
+      const pct = elapsedMs / maxMs;
+      if (pct > highestPct) highestPct = pct;
+      if (elapsedMs > maxMs) {
+        const elapsedSec = Math.round(elapsedMs / 1000);
+        return {
+          ok: false,
+          warningPct: pct,
+          exceeded: true,
+          exceededLimit: "duration",
+          exceededMessage: `duration budget exceeded: ${elapsedSec}s/${this.budget.maxDurationSeconds}s`,
+          policy,
+        };
+      }
+    }
+
+    return { ok: true, warningPct: highestPct, exceeded: false, policy };
+  }
+
+  costRemaining(): number {
+    if (!this.budget?.maxCostMicros) return -1;
+    const rem = this.budget.maxCostMicros - this.costMicros;
+    return rem < 0 ? 0 : rem;
+  }
+
+  tokensRemaining(): number {
+    if (!this.budget?.maxTotalTokens) return -1;
+    const rem = this.budget.maxTotalTokens - this.totalTokens();
+    return rem < 0 ? 0 : rem;
+  }
+
+  private resolvePolicy(): BudgetExceededPolicy {
+    return this.budget?.onExceeded ?? "terminate";
+  }
+}
+
+/**
+ * Extracts cost metadata from task output using the __stigmer_* convention.
+ * Activity outputs that contain LLM cost data should include these keys
+ * for the budget tracker to accumulate.
+ */
+export interface TaskCostInfo {
+  readonly costMicros: number;
+  readonly inputTokens: number;
+  readonly outputTokens: number;
+}
+
+export function extractCostFromOutput(output: unknown): TaskCostInfo {
+  if (output === null || output === undefined || typeof output !== "object") {
+    return { costMicros: 0, inputTokens: 0, outputTokens: 0 };
+  }
+
+  const o = output as Record<string, unknown>;
+  const usage = (typeof o.usage === "object" && o.usage !== null)
+    ? o.usage as Record<string, unknown>
+    : undefined;
+
+  return {
+    costMicros: toSafeInt(o.__stigmer_cost_micros ?? o.cost_micros),
+    inputTokens: toSafeInt(
+      o.__stigmer_input_tokens ?? o.input_tokens ?? usage?.input_tokens,
+    ),
+    outputTokens: toSafeInt(
+      o.__stigmer_output_tokens ?? o.output_tokens ?? usage?.output_tokens,
+    ),
+  };
+}
+
+function toSafeInt(value: unknown): number {
+  if (typeof value === "number" && Number.isFinite(value)) {
+    return Math.max(0, Math.floor(value));
+  }
+  return 0;
+}

package/src/claimcheck/compressor.ts

@@ -0,0 +1,9 @@
+import { gzipSync, gunzipSync } from "node:zlib";
+
+export function compress(data: Buffer): Buffer<ArrayBuffer> {
+  return gzipSync(data) as Buffer<ArrayBuffer>;
+}
+
+export function decompress(data: Buffer): Buffer<ArrayBuffer> {
+  return gunzipSync(data) as Buffer<ArrayBuffer>;
+}

package/src/claimcheck/config.ts

@@ -0,0 +1,20 @@
+export interface ClaimcheckConfig {
+  readonly enabled: boolean;
+  readonly thresholdBytes: number;
+  readonly compressionEnabled: boolean;
+  readonly keyPrefix: string;
+}
+
+const DEFAULT_THRESHOLD_BYTES = 128 * 1024; // 128KB
+
+export function loadClaimcheckConfig(): ClaimcheckConfig {
+  return {
+    enabled: process.env.CLAIMCHECK_ENABLED === "true",
+    thresholdBytes: parseInt(
+      process.env.CLAIMCHECK_THRESHOLD_BYTES ?? String(DEFAULT_THRESHOLD_BYTES),
+      10,
+    ),
+    compressionEnabled: process.env.CLAIMCHECK_COMPRESSION_ENABLED !== "false",
+    keyPrefix: process.env.CLAIMCHECK_KEY_PREFIX ?? "claimcheck/",
+  };
+}

package/src/claimcheck/index.ts

@@ -0,0 +1,3 @@
+export { ClaimcheckPayloadCodec } from "./payload-codec.js";
+export { loadClaimcheckConfig, type ClaimcheckConfig } from "./config.js";
+export { compress, decompress } from "./compressor.js";

package/src/claimcheck/payload-codec.ts

@@ -0,0 +1,107 @@
+/**
+ * Temporal PayloadCodec that transparently offloads large payloads to
+ * external storage (ArtifactStorage). Payloads below the threshold pass
+ * through unchanged. Payloads at or above the threshold are compressed
+ * (optional), uploaded, and replaced with a small reference marker.
+ *
+ * On decode, markers are detected, the original payload is downloaded
+ * and decompressed, and the original bytes are restored — transparent
+ * to workflow/activity code.
+ */
+
+import { randomUUID } from "node:crypto";
+import type { Payload, PayloadCodec } from "@temporalio/common";
+import type { ArtifactStorage } from "../shared/artifact-storage.js";
+import type { ClaimcheckConfig } from "./config.js";
+import { compress, decompress } from "./compressor.js";
+
+const MARKER_METADATA_KEY = "encoding";
+const MARKER_ENCODING_VALUE = "binary/claimcheck";
+
+interface ClaimcheckMarker {
+  key: string;
+  size: number;
+  compressed: boolean;
+}
+
+export class ClaimcheckPayloadCodec implements PayloadCodec {
+  constructor(
+    private readonly storage: ArtifactStorage,
+    private readonly config: ClaimcheckConfig,
+  ) {}
+
+  async encode(payloads: Payload[]): Promise<Payload[]> {
+    return Promise.all(payloads.map((p) => this.encodePayload(p)));
+  }
+
+  async decode(payloads: Payload[]): Promise<Payload[]> {
+    return Promise.all(payloads.map((p) => this.decodePayload(p)));
+  }
+
+  private async encodePayload(payload: Payload): Promise<Payload> {
+    const data = payload.data;
+    if (!data || data.length < this.config.thresholdBytes) {
+      return payload;
+    }
+
+    const originalBuf = Buffer.from(data);
+    let uploadBuf = originalBuf;
+    let compressed = false;
+
+    if (this.config.compressionEnabled) {
+      const compressedBuf = compress(originalBuf);
+      if (compressedBuf.length < originalBuf.length) {
+        uploadBuf = compressedBuf;
+        compressed = true;
+      }
+    }
+
+    const key = `${this.config.keyPrefix}${randomUUID()}`;
+    await this.storage.upload(key, uploadBuf, "application/octet-stream");
+
+    const marker: ClaimcheckMarker = {
+      key,
+      size: data.length,
+      compressed,
+    };
+
+    return {
+      metadata: {
+        [MARKER_METADATA_KEY]: Buffer.from(MARKER_ENCODING_VALUE),
+      },
+      data: Buffer.from(JSON.stringify(marker)),
+    };
+  }
+
+  private async decodePayload(payload: Payload): Promise<Payload> {
+    if (!this.isClaimcheckPayload(payload)) {
+      return payload;
+    }
+
+    const marker: ClaimcheckMarker = JSON.parse(
+      Buffer.from(payload.data!).toString("utf-8"),
+    );
+
+    const downloadUrl = await this.storage.getDownloadUrl(marker.key);
+    const response = await fetch(downloadUrl);
+    if (!response.ok) {
+      throw new Error(
+        `Claimcheck retrieve failed for key ${marker.key}: HTTP ${response.status}`,
+      );
+    }
+
+    const rawBuf = Buffer.from(await response.arrayBuffer());
+    const dataBuf = marker.compressed ? decompress(rawBuf) : rawBuf;
+
+    return {
+      metadata: payload.metadata,
+      data: dataBuf,
+    };
+  }
+
+  private isClaimcheckPayload(payload: Payload): boolean {
+    const encoding = payload.metadata?.[MARKER_METADATA_KEY];
+    if (!encoding) return false;
+    return Buffer.from(encoding).toString("utf-8") === MARKER_ENCODING_VALUE;
+  }
+}