feat: add llm cost tracking and spending limits

2026-08-14 00:48:09 +00:00 · 2026-04-11 15:35:00 -07:00
parent 05e34d741f
commit 373c447e39
4 changed files with 289 additions and 1 deletions
@@ -8,6 +8,7 @@ All notable changes to Lobster will be documented in this file.
 - Add per-step `on_error` workflow policies (`stop|continue|skip_rest`) for partial-failure recovery, with structured step error fields (`error`, `errorMessage`) for condition-based branching. Thanks to [@scottgl9](https://github.com/scottgl9) (PR [#72](https://github.com/openclaw/lobster/pull/72)).
 - Add per-step workflow `timeout_ms` handling, including timeout-triggered aborts, `SIGKILL` for timed shell steps, and dry-run annotations. Thanks to [@scottgl9](https://github.com/scottgl9) (PR [#74](https://github.com/openclaw/lobster/pull/74)).
 - Add workflow condition comparison operators `<`, `<=`, `>`, and `>=` with strict numeric semantics (booleans/null do not coerce), including mixed boolean-expression support with `&&`/`||`. Thanks to [@scottgl9](https://github.com/scottgl9) (PR [#71](https://github.com/openclaw/lobster/pull/71)).
+- Add workflow-level LLM cost tracking with `_meta.cost` summaries, per-step usage attribution, and optional `cost_limit` controls with `warn`/`stop` actions (plus custom pricing via `LOBSTER_LLM_PRICING_JSON`). Thanks to [@scottgl9](https://github.com/scottgl9) (PR [#70](https://github.com/openclaw/lobster/pull/70)).

 ## 2026.4.6

@@ -0,0 +1,101 @@
+export type StepCost = {
+  stepId: string;
+  model: string | null;
+  inputTokens: number;
+  outputTokens: number;
+  costUsd: number;
+};
+
+export type CostSummary = {
+  totalInputTokens: number;
+  totalOutputTokens: number;
+  estimatedCostUsd: number;
+  byStep: StepCost[];
+};
+
+export type CostLimit = {
+  max_usd: number;
+  action?: 'warn' | 'stop';
+};
+
+const DEFAULT_PRICING: Record<string, { input: number; output: number }> = {
+  'gpt-4o': { input: 2.50, output: 10.00 },
+  'gpt-4o-mini': { input: 0.15, output: 0.60 },
+  'gpt-4-turbo': { input: 10.00, output: 30.00 },
+  'gpt-3.5-turbo': { input: 0.50, output: 1.50 },
+  'claude-opus-4-20250514': { input: 15.00, output: 75.00 },
+  'claude-sonnet-4-5-20250514': { input: 3.00, output: 15.00 },
+  'claude-haiku-3-5': { input: 0.80, output: 4.00 },
+  'gemini-1.5-pro': { input: 1.25, output: 5.00 },
+  'gemini-1.5-flash': { input: 0.075, output: 0.30 },
+};
+
+function toTokenCount(value: unknown): number {
+  const parsed = Number(value ?? 0);
+  if (!Number.isFinite(parsed) || parsed < 0) return 0;
+  return Math.floor(parsed);
+}
+
+export class CostTracker {
+  private steps: StepCost[] = [];
+
+  private pricing: Record<string, { input: number; output: number }>;
+
+  constructor(customPricing?: Record<string, { input: number; output: number }>) {
+    this.pricing = { ...DEFAULT_PRICING, ...(customPricing ?? {}) };
+  }
+
+  recordUsage(stepId: string, model: string | null, usage: Record<string, unknown>) {
+    const inputTokens = toTokenCount(usage.inputTokens ?? usage.input_tokens ?? usage.prompt_tokens);
+    const outputTokens = toTokenCount(usage.outputTokens ?? usage.output_tokens ?? usage.completion_tokens);
+    const pricing = this.pricing[model ?? ''] ?? { input: 0, output: 0 };
+    const costUsd = (inputTokens * pricing.input + outputTokens * pricing.output) / 1_000_000;
+    this.steps.push({ stepId, model, inputTokens, outputTokens, costUsd });
+  }
+
+  getSummary(): CostSummary {
+    let totalInputTokens = 0;
+    let totalOutputTokens = 0;
+    let estimatedCostUsd = 0;
+
+    for (const step of this.steps) {
+      totalInputTokens += step.inputTokens;
+      totalOutputTokens += step.outputTokens;
+      estimatedCostUsd += step.costUsd;
+    }
+
+    return {
+      totalInputTokens,
+      totalOutputTokens,
+      estimatedCostUsd: Math.round(estimatedCostUsd * 1_000_000) / 1_000_000,
+      byStep: [...this.steps],
+    };
+  }
+
+  hasUsage() {
+    return this.steps.length > 0;
+  }
+
+  checkLimit(limit: CostLimit, stderr?: NodeJS.WritableStream) {
+    const summary = this.getSummary();
+    if (summary.estimatedCostUsd <= limit.max_usd) return;
+
+    if (limit.action === 'stop') {
+      throw new Error(`Cost limit exceeded: $${summary.estimatedCostUsd.toFixed(4)} > $${limit.max_usd.toFixed(2)} limit`);
+    }
+
+    if (stderr) {
+      stderr.write(`[WARN] Cost $${summary.estimatedCostUsd.toFixed(4)} exceeds limit $${limit.max_usd.toFixed(2)}\n`);
+    }
+  }
+
+  static parsePricingFromEnv(env: Record<string, string | undefined>): Record<string, { input: number; output: number }> | undefined {
+    const raw = env.LOBSTER_LLM_PRICING_JSON;
+    if (!raw) return undefined;
+    try {
+      return JSON.parse(raw);
+    } catch {
+      return undefined;
+    }
+  }
+}
@@ -13,6 +13,8 @@ import { createApprovalIndex, deleteStateJson, readStateJson, writeStateJson } f
 import { readLineFromStream } from '../read_line.js';
 import { resolveInlineShellCommand } from '../shell.js';
 import { sharedAjv } from '../validation.js';
+import { CostTracker } from '../core/cost_tracker.js';
+import type { CostLimit, CostSummary } from '../core/cost_tracker.js';

 export type WorkflowFile = {
  name?: string;
@@ -21,6 +23,7 @@ export type WorkflowFile = {
  env?: Record<string, string>;
  cwd?: string;
  steps: WorkflowStep[];
+  cost_limit?: CostLimit;
 };

 export type WorkflowStep = {
@@ -88,6 +91,9 @@ export type WorkflowRunResult = {
    subject?: unknown;
    resumeToken?: string;
  };
+  _meta?: {
+    cost?: CostSummary;
+  };
 };

 type RunContext = {
@@ -154,6 +160,19 @@ export async function loadWorkflowFile(filePath: string): Promise<WorkflowFile>
    throw new Error('Workflow file requires a non-empty steps array');
  }

+  const costLimit = (parsed as WorkflowFile).cost_limit;
+  if (costLimit !== undefined) {
+    if (!costLimit || typeof costLimit !== 'object' || Array.isArray(costLimit)) {
+      throw new Error('Workflow cost_limit must be an object');
+    }
+    if (!Number.isFinite(Number(costLimit.max_usd)) || Number(costLimit.max_usd) < 0) {
+      throw new Error('Workflow cost_limit.max_usd must be a non-negative number');
+    }
+    if (costLimit.action !== undefined && costLimit.action !== 'warn' && costLimit.action !== 'stop') {
+      throw new Error('Workflow cost_limit.action must be "warn" or "stop"');
+    }
+  }
+
  const seen = new Set<string>();
  for (const step of steps) {
    if (!step || typeof step !== 'object') {
@@ -368,6 +387,7 @@ export async function runWorkflowFile({
      return dryRunWorkflow({ steps, resolvedArgs, results, startIndex, ctx });
    }

+    const costTracker = new CostTracker(CostTracker.parsePricingFromEnv(ctx.env));
    let lastStepId: string | null = resumeState?.inputStepId ?? findLastCompletedStepId(steps, results);

    for (let idx = startIndex; idx < steps.length; idx++) {
@@ -580,6 +600,11 @@ export async function runWorkflowFile({
    results[step.id] = result;
    lastStepId = step.id;

+    trackStepCost(costTracker, step.id, result);
+    if (workflow.cost_limit) {
+      costTracker.checkLimit(workflow.cost_limit, ctx.stderr);
+    }
+
    if (isApprovalStep(step.approval)) {
      const approval = extractApprovalRequest(step, results[step.id]);

@@ -638,7 +663,11 @@ export async function runWorkflowFile({
    if (consumedResumeStateKey) {
      await deleteStateJson({ env: ctx.env, key: consumedResumeStateKey });
    }
-    return { status: 'ok', output };
+    const runResult: WorkflowRunResult = { status: 'ok', output };
+    if (costTracker.hasUsage()) {
+      runResult._meta = { cost: costTracker.getSummary() };
+    }
+    return runResult;
  } finally {
    ctx._activeWorkflows?.delete(canonicalFilePath);
  }
@@ -1044,6 +1073,21 @@ function extractApprovalRequest(step: WorkflowStep, result: WorkflowStepResult)
  };
 }

+function trackStepCost(costTracker: CostTracker, stepId: string, result: WorkflowStepResult) {
+  const json = result.json;
+  if (!json || typeof json !== 'object') return;
+
+  const items = Array.isArray(json) ? json : [json];
+  for (const item of items) {
+    if (!item || typeof item !== 'object') continue;
+    const usage = (item as Record<string, unknown>).usage;
+    if (!usage || typeof usage !== 'object') continue;
+    const modelValue = (item as Record<string, unknown>).model;
+    const model = typeof modelValue === 'string' ? modelValue : null;
+    costTracker.recordUsage(stepId, model, usage as Record<string, unknown>);
+  }
+}
+
 function parseJson(stdout: string) {
  const trimmed = stdout.trim();
  if (!trimmed) return undefined;
@@ -0,0 +1,142 @@
+import test from 'node:test';
+import assert from 'node:assert/strict';
+import { promises as fsp } from 'node:fs';
+import os from 'node:os';
+import path from 'node:path';
+import { PassThrough } from 'node:stream';
+
+import { CostTracker } from '../src/core/cost_tracker.js';
+import { runWorkflowFile } from '../src/workflows/file.js';
+
+test('CostTracker records usage and computes totals', () => {
+  const tracker = new CostTracker();
+  tracker.recordUsage('step1', 'gpt-4o', { inputTokens: 1000, outputTokens: 500 });
+  const summary = tracker.getSummary();
+  assert.equal(summary.totalInputTokens, 1000);
+  assert.equal(summary.totalOutputTokens, 500);
+  assert.equal(summary.estimatedCostUsd, 0.0075);
+  assert.equal(summary.byStep.length, 1);
+  assert.equal(summary.byStep[0].stepId, 'step1');
+});
+
+test('CostTracker handles OpenAI token field names', () => {
+  const tracker = new CostTracker();
+  tracker.recordUsage('step1', 'gpt-4o', { prompt_tokens: 1000, completion_tokens: 500 });
+  const summary = tracker.getSummary();
+  assert.equal(summary.totalInputTokens, 1000);
+  assert.equal(summary.totalOutputTokens, 500);
+});
+
+test('CostTracker uses zero cost for unknown models', () => {
+  const tracker = new CostTracker();
+  tracker.recordUsage('step1', 'unknown-model', { inputTokens: 1000, outputTokens: 500 });
+  const summary = tracker.getSummary();
+  assert.equal(summary.estimatedCostUsd, 0);
+});
+
+test('CostTracker supports custom pricing from env json', () => {
+  const pricing = CostTracker.parsePricingFromEnv({
+    LOBSTER_LLM_PRICING_JSON: '{"my-model":{"input":1.0,"output":2.0}}',
+  });
+  const tracker = new CostTracker(pricing);
+  tracker.recordUsage('step1', 'my-model', { inputTokens: 1_000_000, outputTokens: 1_000_000 });
+  assert.equal(tracker.getSummary().estimatedCostUsd, 3);
+});
+
+test('CostTracker checkLimit throws when action=stop and limit exceeded', () => {
+  const tracker = new CostTracker();
+  tracker.recordUsage('step1', 'gpt-4o', { inputTokens: 10_000_000, outputTokens: 10_000_000 });
+  assert.throws(
+    () => tracker.checkLimit({ max_usd: 0.01, action: 'stop' }),
+    /Cost limit exceeded/,
+  );
+});
+
+test('CostTracker checkLimit warns when action=warn and limit exceeded', () => {
+  const tracker = new CostTracker();
+  tracker.recordUsage('step1', 'gpt-4o', { inputTokens: 10_000_000, outputTokens: 10_000_000 });
+  const stderr = new PassThrough();
+  let out = '';
+  stderr.on('data', (d: Buffer | string) => { out += String(d); });
+  tracker.checkLimit({ max_usd: 0.01, action: 'warn' }, stderr);
+  assert.match(out, /\[WARN\] Cost/);
+});
+
+async function runWorkflow(workflow: unknown, envOverride?: Record<string, string>) {
+  const tmpDir = await fsp.mkdtemp(path.join(os.tmpdir(), 'lobster-cost-'));
+  const stateDir = path.join(tmpDir, 'state');
+  const filePath = path.join(tmpDir, 'workflow.lobster');
+  await fsp.writeFile(filePath, JSON.stringify(workflow, null, 2), 'utf8');
+  const stderr = new PassThrough();
+  let stderrOutput = '';
+  stderr.on('data', (d: Buffer | string) => { stderrOutput += String(d); });
+
+  const result = await runWorkflowFile({
+    filePath,
+    ctx: {
+      stdin: process.stdin,
+      stdout: process.stdout,
+      stderr,
+      env: { ...process.env, LOBSTER_STATE_DIR: stateDir, ...(envOverride ?? {}) },
+      mode: 'tool',
+    },
+  });
+
+  return { result, stderrOutput };
+}
+
+test('workflow result includes _meta.cost when usage is present', async () => {
+  const { result } = await runWorkflow({
+    steps: [
+      {
+        id: 'llm',
+        command: 'node -e "process.stdout.write(JSON.stringify({model:\'gpt-4o\',usage:{inputTokens:100,outputTokens:50},output:{text:\'hi\'}}))"',
+      },
+    ],
+  });
+
+  assert.equal(result.status, 'ok');
+  assert.ok(result._meta?.cost);
+  assert.equal(result._meta!.cost!.totalInputTokens, 100);
+  assert.equal(result._meta!.cost!.totalOutputTokens, 50);
+  assert.equal(result._meta!.cost!.byStep[0].model, 'gpt-4o');
+});
+
+test('workflow result omits _meta.cost when no usage exists', async () => {
+  const { result } = await runWorkflow({
+    steps: [{ id: 'plain', command: 'echo "hello"' }],
+  });
+  assert.equal(result.status, 'ok');
+  assert.equal(result._meta, undefined);
+});
+
+test('cost_limit warn logs warning and continues', async () => {
+  const { result, stderrOutput } = await runWorkflow({
+    cost_limit: { max_usd: 0.00001, action: 'warn' },
+    steps: [
+      {
+        id: 'llm',
+        command: 'node -e "process.stdout.write(JSON.stringify({model:\'gpt-4o\',usage:{inputTokens:1000,outputTokens:1000}}))"',
+      },
+      { id: 'after', command: 'echo done' },
+    ],
+  });
+  assert.equal(result.status, 'ok');
+  assert.match(stderrOutput, /\[WARN\] Cost/);
+  assert.deepEqual(result.output, ['done\n']);
+});
+
+test('cost_limit stop throws when exceeded', async () => {
+  await assert.rejects(
+    () => runWorkflow({
+      cost_limit: { max_usd: 0.00001, action: 'stop' },
+      steps: [
+        {
+          id: 'llm',
+          command: 'node -e "process.stdout.write(JSON.stringify({model:\'gpt-4o\',usage:{inputTokens:1000,outputTokens:1000}}))"',
+        },
+      ],
+    }).then((x) => x.result),
+    /Cost limit exceeded/,
+  );
+});