feat: enforce per-task execution budgets

server/contracts.ts

@@ -76,6 +76,18 @@ export type RuntimeEvent = RuntimeEventBase &
       }
     | { type: "request.resolved"; behavior: string; source: string }
     | { type: "thread.token-usage.updated"; input: number; output: number }
+    | {
+        type: "task.budget.updated";
+        budget: Record<string, number>;
+        usage: Record<string, number | undefined>;
+        limit: string;
+      }
+    | {
+        type: "task.budget.exhausted";
+        budget: Record<string, number>;
+        usage: Record<string, number | undefined>;
+        limit: string;
+      }
     | { type: "runtime.error"; message: string }
   );
 

server/index.test.ts

@@ -143,6 +143,14 @@ describe("harness HTTP API", () => {
     expect(send.body.error).toContain("unavailable");
   });
 
+  it("validates an optional per-task budget before dispatch", async () => {
+    const { body } = await api("GET", "/api/bots");
+    const bot = body.bots[0];
+    const invalid = await api("POST", `/api/bots/${bot.id}/messages`, { text: "hello?", budget: { maxToolCalls: 0 } });
+    expect(invalid.status).toBe(400);
+    expect(invalid.body.error).toContain("budget");
+  });
+
   it("refuses to fork a message when the provider is unavailable, without mutating", async () => {
     const { body } = await api("GET", "/api/bots");
     const bot = body.bots[0];

server/index.ts

@@ -17,6 +17,7 @@ import { BUILT_IN_DRIVERS } from "./drivers/builtIn.ts";
 import { EventBus } from "./harness/bus.ts";
 import { ProviderRegistry } from "./harness/registry.ts";
 import { mentionedBots, Store, type Message } from "./store.ts";
+import { normalizeTaskBudget, TaskBudgetGuard, type TaskBudget } from "./task-budget.ts";
 
 const PORT = Number(process.env.OMB_PORT || process.env.OGB_PORT || 8799);
 const STATIC_DIR = process.env.OMB_STATIC_DIR || null;
@@ -131,13 +132,18 @@ function broadcast(payload: unknown) {
 // and every client view are projections of it.
 const toolMessageByItem = new Map<string, string>(); // itemId -> messageId
 const askMessageByRequest = new Map<string, string>(); // requestId -> messageId
+const taskBudgetGuards = new Map<string, TaskBudgetGuard>();
 
 // Group threads: the fold needs to know WHO is talking — the turn engine
 // records the active member here before dispatching its turn.
 const groupSpeakers = new Map<string, { botId: string; name: string; color: string }>();
 
 bus.subscribe((event: RuntimeEvent) => {
   broadcast({ kind: "runtime", event });
+  const budgetGuard = taskBudgetGuards.get(event.threadId);
+  if (event.type === "thread.token-usage.updated") budgetGuard?.noteTokenUsage(event.input, event.output);
+  if (event.type === "item.started" && event.itemType === "tool") budgetGuard?.noteToolStarted(event.title);
+  if (event.type === "item.completed" && event.itemType === "tool" && !event.ok) budgetGuard?.noteFailedTool();
   const bot = store.botByThread(event.threadId);
   const group = bot ? undefined : store.groupByThread(event.threadId);
   if (!bot && !group) return;
@@ -213,6 +219,8 @@ bus.subscribe((event: RuntimeEvent) => {
       pushMessage({ role: "bot", kind: "activity", tool: { name: `error: ${event.message.slice(0, 160)}`, ok: false } });
       break;
     case "turn.completed": {
+      taskBudgetGuards.get(event.threadId)?.dispose();
+      taskBudgetGuards.delete(event.threadId);
       if (bot) {
         // the last live frame becomes a settled inline screen message —
         // the screenshot-in-chat moment
@@ -312,7 +320,7 @@ function readCuaConnection(): { command: string; args: string[]; env: Record<str
 async function startTurn(
   botId: string,
   text: string,
-  opts?: { commsDepth?: number; userMessage?: Message },
+  opts?: { commsDepth?: number; userMessage?: Message; budget?: TaskBudget },
 ) {
   const bot = store.bot(botId);
   if (!bot) throw Object.assign(new Error("no such bot"), { status: 404 });
@@ -378,6 +386,31 @@ async function startTurn(
 
   void (async () => {
     try {
+      if (opts?.budget) {
+        const guard = new TaskBudgetGuard(opts.budget, (status) => {
+          bus.publish({
+            eventId: `budget-${Date.now().toString(36)}`,
+            provider: instance.driverKind,
+            providerInstanceId: instance.instanceId,
+            threadId: bot.threadId,
+            createdAt: new Date().toISOString(),
+            type: status.kind === "exhausted" ? "task.budget.exhausted" : "task.budget.updated",
+            budget: status.budget as Record<string, number>,
+            usage: { ...status.usage },
+            limit: status.limit,
+          });
+          if (status.kind !== "exhausted") return;
+          const message = store.appendMessage(bot.threadId, {
+            role: "bot",
+            kind: "activity",
+            tool: { name: `task budget reached: ${status.limit.replace(/_/g, " ")}`, ok: false },
+          });
+          broadcast({ kind: "message", threadId: bot.threadId, message });
+          void instance.adapter.interruptTurn(bot.threadId).catch(() => {});
+        });
+        taskBudgetGuards.set(bot.threadId, guard);
+        guard.start();
+      }
       const integrations: NonNullable<Parameters<typeof instance.adapter.sendTurn>[0]["integrations"]> = {};
       if (cfg.composio?.key) integrations.composio = { key: cfg.composio.key, url: cfg.composio.url };
       const wants = bot.computer; // 'cloud' | 'local' | 'off' | undefined(auto)
@@ -450,6 +483,8 @@ async function startTurn(
       if (rewound) store.patchBot(bot.id, { rewound: false, resumeCursors: {} });
       if (integrations.computer) startScreenPoller(bot.id);
     } catch (e) {
+      taskBudgetGuards.get(bot.threadId)?.dispose();
+      taskBudgetGuards.delete(bot.threadId);
       const message = e instanceof Error ? e.message : String(e);
       const failure = store.appendMessage(bot.threadId, {
         role: "bot",
@@ -930,8 +965,10 @@ const server = createServer(async (req, res) => {
       const body = await readBody(req);
       const text = String(body.text ?? "").trim();
       if (!text) return json(res, 400, { error: "text required" });
-      await startTurn(m[1], text);
-      return json(res, 202, { ok: true });
+      const budget = body.budget === undefined ? undefined : normalizeTaskBudget(body.budget);
+      if (body.budget !== undefined && !budget) return json(res, 400, { error: "budget needs one or more supported positive limits within documented bounds" });
+      await startTurn(m[1], text, { budget: budget ?? undefined });
+      return json(res, 202, { ok: true, ...(budget ? { budget } : {}) });
     }
 
     // edit a user message → fork the conversation there and rerun the turn.
@@ -1000,6 +1037,8 @@ const server = createServer(async (req, res) => {
       const bot = store.bot(m[1]);
       if (!bot) return json(res, 404, { error: "no such bot" });
       const instance = registry.get(bot.modelSelection.instanceId);
+      taskBudgetGuards.get(bot.threadId)?.dispose();
+      taskBudgetGuards.delete(bot.threadId);
       await instance?.adapter.interruptTurn(bot.threadId);
       return json(res, 200, { ok: true });
     }

server/task-budget.test.ts

@@ -0,0 +1,63 @@
+import { describe, expect, it, vi } from "vitest";
+import { isComputerAction, isDelegation, normalizeTaskBudget, TaskBudgetGuard } from "./task-budget.ts";
+
+describe("task budgets", () => {
+  it("validates bounded explicit budgets and leaves unconfigured tasks alone", () => {
+    expect(normalizeTaskBudget({ maxTotalTokens: 30_000, maxToolCalls: 4 })).toEqual({ maxTotalTokens: 30_000, maxToolCalls: 4 });
+    expect(normalizeTaskBudget({ maxDurationMs: 999 })).toBeNull();
+    expect(normalizeTaskBudget({ maxToolCalls: 0 })).toBeNull();
+    expect(normalizeTaskBudget({})).toBeNull();
+  });
+
+  it("uses reported token snapshots without inventing absent dimensions", () => {
+    const statuses: string[] = [];
+    const guard = new TaskBudgetGuard({ maxInputTokens: 10, maxTotalTokens: 15 }, (s) => statuses.push(s.limit));
+    guard.noteTokenUsage(8, undefined);
+    expect(guard.usage).toEqual({ inputTokens: 8, toolCalls: 0, computerActions: 0, retries: 0, delegations: 0 });
+    guard.noteTokenUsage(7, 7); // snapshots cannot reduce input usage
+    expect(statuses).toEqual(["input_tokens", "total_tokens"]);
+  });
+
+  it("enforces tool and computer-action boundaries once", () => {
+    const statuses: string[] = [];
+    const guard = new TaskBudgetGuard({ maxToolCalls: 3, maxComputerActions: 2 }, (s) => statuses.push(s.limit));
+    guard.noteToolStarted("mcp__computer__screenshot");
+    guard.noteToolStarted("mcp__computer__click");
+    guard.noteToolStarted("mcp__computer__type_text");
+    guard.noteToolStarted("another_tool");
+    expect(statuses.filter((x) => x === "computer_actions")).toHaveLength(1);
+    expect(guard.usage).toMatchObject({ toolCalls: 4, computerActions: 2 });
+  });
+
+  it("counts failed tools as observed retry attempts and delegates separately", () => {
+    const statuses: string[] = [];
+    const guard = new TaskBudgetGuard({ maxRetries: 2, maxDelegations: 1 }, (s) => statuses.push(s.limit));
+    guard.noteFailedTool();
+    guard.noteFailedTool();
+    guard.noteToolStarted("mcp__agents__ask_bot");
+    expect(statuses).toContain("retries");
+    expect(statuses).not.toContain("delegations"); // a task already stopped cannot emit another exhaustion
+    expect(isDelegation("mcp__agents__ask_bot")).toBe(true);
+    expect(isComputerAction("mcp__computer__click")).toBe(true);
+    expect(isComputerAction("mcp__computer__screenshot")).toBe(false);
+  });
+
+  it("warns before, then exhausts on a token boundary", () => {
+    const statuses: Array<{ kind: string; limit: string }> = [];
+    const guard = new TaskBudgetGuard({ maxOutputTokens: 10 }, (s) => statuses.push({ kind: s.kind, limit: s.limit }));
+    guard.noteTokenUsage(undefined, 8);
+    guard.noteTokenUsage(undefined, 10);
+    expect(statuses).toEqual([{ kind: "approaching", limit: "output_tokens" }, { kind: "exhausted", limit: "output_tokens" }]);
+  });
+
+  it("cancels the deadline timer when disposed", () => {
+    vi.useFakeTimers();
+    const exceeded = vi.fn();
+    const guard = new TaskBudgetGuard({ maxDurationMs: 5_000 }, exceeded);
+    guard.start();
+    guard.dispose();
+    vi.advanceTimersByTime(5_000);
+    expect(exceeded).not.toHaveBeenCalled();
+    vi.useRealTimers();
+  });
+});

server/task-budget.ts

@@ -0,0 +1,126 @@
+/** Central, provider-neutral budgets for one running task. Values are counted
+ * only from canonical runtime events; missing provider telemetry stays absent
+ * rather than becoming a guessed usage number. */
+export interface TaskBudget {
+  maxInputTokens?: number;
+  maxOutputTokens?: number;
+  maxTotalTokens?: number;
+  maxToolCalls?: number;
+  maxComputerActions?: number;
+  maxRetries?: number;
+  maxDelegations?: number;
+  maxDurationMs?: number;
+}
+
+export interface TaskBudgetUsage {
+  inputTokens?: number;
+  outputTokens?: number;
+  toolCalls: number;
+  computerActions: number;
+  retries: number;
+  delegations: number;
+}
+
+export type BudgetLimit = "input_tokens" | "output_tokens" | "total_tokens" | "tool_calls" | "computer_actions" | "retries" | "delegations" | "duration";
+export type BudgetStatus = { kind: "approaching" | "exhausted"; limit: BudgetLimit; usage: TaskBudgetUsage; budget: TaskBudget };
+
+const LIMITS: Array<[keyof TaskBudget, BudgetLimit, number, number]> = [
+  ["maxInputTokens", "input_tokens", 1, 10_000_000], ["maxOutputTokens", "output_tokens", 1, 10_000_000],
+  ["maxTotalTokens", "total_tokens", 1, 10_000_000], ["maxToolCalls", "tool_calls", 1, 10_000],
+  ["maxComputerActions", "computer_actions", 1, 10_000], ["maxRetries", "retries", 1, 10_000],
+  ["maxDelegations", "delegations", 1, 100], ["maxDurationMs", "duration", 1_000, 24 * 60 * 60_000],
+];
+
+export function normalizeTaskBudget(raw: unknown): TaskBudget | null {
+  if (!raw || typeof raw !== "object") return null;
+  const value = raw as Record<string, unknown>;
+  const result: TaskBudget = {};
+  for (const [key, , min, max] of LIMITS) {
+    if (value[key] === undefined) continue;
+    const number = Math.round(Number(value[key]));
+    if (!Number.isFinite(number) || number < min || number > max) return null;
+    result[key] = number;
+  }
+  return Object.keys(result).length ? result : null;
+}
+
+export function isComputerAction(title: string | undefined) {
+  return Boolean(title && /computer.*(click|type|press|scroll|exec|open_url)|mcp__computer__(click|type|press|scroll|exec|open_url)/i.test(title));
+}
+
+export function isDelegation(title: string | undefined) {
+  return Boolean(title && /ask_bot/i.test(title));
+}
+
+export class TaskBudgetGuard {
+  private readonly budget: TaskBudget;
+  private readonly onStatus: (status: BudgetStatus) => void;
+  private readonly usageValue: TaskBudgetUsage = { toolCalls: 0, computerActions: 0, retries: 0, delegations: 0 };
+  private readonly warned = new Set<BudgetLimit>();
+  private exhausted: BudgetLimit | null = null;
+  private timer: ReturnType<typeof setTimeout> | null = null;
+
+  constructor(budget: TaskBudget, onStatus: (status: BudgetStatus) => void) {
+    this.budget = budget;
+    this.onStatus = onStatus;
+  }
+
+  start() {
+    if (this.budget.maxDurationMs) this.timer = setTimeout(() => this.exceed("duration"), this.budget.maxDurationMs);
+    this.timer?.unref?.();
+  }
+
+  dispose() {
+    if (this.timer) clearTimeout(this.timer);
+    this.timer = null;
+  }
+
+  noteTokenUsage(input: number | undefined, output: number | undefined) {
+    // Providers may omit either dimension. Usage events are snapshots, so a
+    // later smaller value must not lower observed consumption or fabricate it.
+    if (Number.isFinite(input)) this.usageValue.inputTokens = Math.max(this.usageValue.inputTokens ?? 0, Math.max(0, input!));
+    if (Number.isFinite(output)) this.usageValue.outputTokens = Math.max(this.usageValue.outputTokens ?? 0, Math.max(0, output!));
+    this.check();
+  }
+
+  noteToolStarted(title?: string) {
+    this.usageValue.toolCalls += 1;
+    if (isComputerAction(title)) this.usageValue.computerActions += 1;
+    if (isDelegation(title)) this.usageValue.delegations += 1;
+    this.check();
+  }
+
+  noteFailedTool() {
+    this.usageValue.retries += 1;
+    this.check();
+  }
+
+  get usage(): TaskBudgetUsage { return { ...this.usageValue }; }
+
+  private check() {
+    const values: Array<[BudgetLimit, number | undefined, number | undefined]> = [
+      ["input_tokens", this.usageValue.inputTokens, this.budget.maxInputTokens],
+      ["output_tokens", this.usageValue.outputTokens, this.budget.maxOutputTokens],
+      ["total_tokens", this.usageValue.inputTokens === undefined && this.usageValue.outputTokens === undefined ? undefined : (this.usageValue.inputTokens ?? 0) + (this.usageValue.outputTokens ?? 0), this.budget.maxTotalTokens],
+      ["tool_calls", this.usageValue.toolCalls, this.budget.maxToolCalls],
+      ["computer_actions", this.usageValue.computerActions, this.budget.maxComputerActions],
+      ["retries", this.usageValue.retries, this.budget.maxRetries],
+      ["delegations", this.usageValue.delegations, this.budget.maxDelegations],
+    ];
+    for (const [limit, usage, max] of values) {
+      if (usage === undefined || max === undefined) continue;
+      if (usage >= max) return this.exceed(limit);
+      if (usage >= Math.max(1, Math.floor(max * 0.8)) && !this.warned.has(limit)) {
+        this.warned.add(limit);
+        this.onStatus({ kind: "approaching", limit, usage: this.usage, budget: this.budget });
+      }
+    }
+  }
+
+  private exceed(limit: BudgetLimit) {
+    if (this.exhausted) return;
+    this.exhausted = limit;
+    this.dispose();
+    this.onStatus({ kind: "exhausted", limit, usage: this.usage, budget: this.budget });
+  }
+}