From 03a1a726cefd1adbf6cbd299e7c40efb1deb4af9 Mon Sep 17 00:00:00 2001
From: Waishnav <86405648+Waishnav@users.noreply.github.com>
Date: Mon, 27 Jul 2026 20:59:36 +0000
Subject: [PATCH 1/2] feat(workflow): enforce per-run agent call budget

---
 src/workflow-api.ts       | 16 ++++++++++++++--
 src/workflow-contracts.ts |  1 +
 src/workflow-errors.ts    |  1 +
 src/workflow-types.ts     |  1 +
 4 files changed, 17 insertions(+), 2 deletions(-)

diff --git a/src/workflow-api.ts b/src/workflow-api.ts
index 4b6ce47e..d30f4df2 100644
--- a/src/workflow-api.ts
+++ b/src/workflow-api.ts
@@ -13,6 +13,7 @@ import type { JsonSchema, JsonValue } from "./json-types.js";
 import { jsonValueSchema } from "./json-types.js";
 import {
   WORKFLOW_LIMITS,
+  WORKFLOW_MAX_AGENT_CALLS,
   WORKFLOW_MAX_ITEMS,
   WORKFLOW_MAX_NEST_DEPTH,
   buildAgentCacheKeyInput,
@@ -293,8 +294,7 @@ export function createWorkflowApi(deps: WorkflowApiDeps): WorkflowApi {
     const phase = agentOpts.phase ?? phaseAls.getStore();
     const isolation: AgentIsolationMode =
       agentOpts.isolation === "worktree" ? "worktree" : "shared";
-    const index = runtime.callIndex;
-    runtime.callIndex += 1;
+    const index = allocateAgentCallIndex(runtime);
 
     const cacheKeyInput = buildAgentCacheKeyInput({
       prompt,
@@ -744,6 +744,18 @@ function assertMaxItems(count: number, label: string): void {
   }
 }
 
+function allocateAgentCallIndex(runtime: WorkflowApiRuntime): number {
+  if (runtime.callIndex >= WORKFLOW_MAX_AGENT_CALLS) {
+    throw new WorkflowEngineError(
+      "call_limit",
+      `Workflow exceeded the limit of ${WORKFLOW_MAX_AGENT_CALLS} agent calls`,
+    );
+  }
+  const index = runtime.callIndex;
+  runtime.callIndex += 1;
+  return index;
+}
+
 function throwIfCancelled(deps: WorkflowApiDeps): void {
   if (deps.signal.aborted || deps.journal.isCancelRequested(deps.runId)) {
     throw cancelledError();
diff --git a/src/workflow-contracts.ts b/src/workflow-contracts.ts
index b06841d0..8c15c75a 100644
--- a/src/workflow-contracts.ts
+++ b/src/workflow-contracts.ts
@@ -131,6 +131,7 @@ export const workflowErrorKindSchema = z.enum([
   "heartbeat",
   "worktree",
   "nest_depth",
+  "call_limit",
   "path",
   "result_too_large",
   "args_too_large",
diff --git a/src/workflow-errors.ts b/src/workflow-errors.ts
index 26eca02e..262ab3e2 100644
--- a/src/workflow-errors.ts
+++ b/src/workflow-errors.ts
@@ -19,6 +19,7 @@ export class WorkflowEngineError extends Error {
       | "no_provider"
       | "profile"
       | "nest_depth"
+      | "call_limit"
       | "worktree"
       | "schema"
       | "path"
diff --git a/src/workflow-types.ts b/src/workflow-types.ts
index 2e3996b7..75286338 100644
--- a/src/workflow-types.ts
+++ b/src/workflow-types.ts
@@ -48,6 +48,7 @@ export type {
 // ---------------------------------------------------------------------------
 
 export const WORKFLOW_MAX_ITEMS = 4096;
+export const WORKFLOW_MAX_AGENT_CALLS = 256;
 export const WORKFLOW_MAX_NEST_DEPTH = 1;
 export const WORKFLOW_MAX_SCHEMA_RETRIES = 2;
 export const WORKFLOW_HEARTBEAT_MS = 5_000;

From 9860cbd84b550e051fd5a02175ab1517555f41bf Mon Sep 17 00:00:00 2001
From: Waishnav <86405648+Waishnav@users.noreply.github.com>
Date: Mon, 27 Jul 2026 20:59:36 +0000
Subject: [PATCH 2/2] test(workflow): cover agent call budget boundary

---
 src/workflow-engine.test.ts | 62 ++++++++++++++++++++++++++++++++++++-
 src/workflow-types.test.ts  |  2 ++
 2 files changed, 63 insertions(+), 1 deletion(-)

diff --git a/src/workflow-engine.test.ts b/src/workflow-engine.test.ts
index 8279f0b7..8fd1bee8 100644
--- a/src/workflow-engine.test.ts
+++ b/src/workflow-engine.test.ts
@@ -6,13 +6,18 @@ import { WorkflowStore } from "./workflow-store.js";
 import { executeWorkflow } from "./workflow-engine.js";
 import {
   createWorkflowApi,
+  createWorkflowApiRuntime,
   WorkflowEngineError,
   WorkflowSemaphore,
   getCurrentWorkflowPhase,
   type WorkflowProviderRunInput,
   type CreateAgentWorktree,
 } from "./workflow-api.js";
-import { createStubBudget, WORKFLOW_LIMITS } from "./workflow-types.js";
+import {
+  createStubBudget,
+  WORKFLOW_LIMITS,
+  WORKFLOW_MAX_AGENT_CALLS,
+} from "./workflow-types.js";
 import type { LocalAgentProfile } from "./local-agent-profiles.js";
 
 // ---------------------------------------------------------------------------
@@ -35,6 +40,61 @@ import type { LocalAgentProfile } from "./local-agent-profiles.js";
   assert.equal(maxConcurrent, 2);
 }
 
+// ---------------------------------------------------------------------------
+// Per-run agent call budget
+// ---------------------------------------------------------------------------
+{
+  const dir = await mkdtemp(join(tmpdir(), "wf-call-limit-"));
+  const store = new WorkflowStore(dir);
+  const run = store.createRun({
+    name: "call-limit",
+    source: "inline",
+    scriptPath: "inline",
+    scriptHash: "h",
+    workspaceRoot: dir,
+  });
+  const runtime = createWorkflowApiRuntime(2);
+  runtime.callIndex = WORKFLOW_MAX_AGENT_CALLS - 1;
+  let providerCalls = 0;
+  const api = createWorkflowApi({
+    runId: run.id,
+    journal: store,
+    meta: { name: "call-limit", description: "d" },
+    args: undefined,
+    concurrency: 2,
+    signal: new AbortController().signal,
+    workspaceRoot: dir,
+    availableProviders: ["codex"],
+    runtime,
+    runProvider: async (input) => {
+      providerCalls += 1;
+      return { finalResponse: `ok:${input.prompt}` };
+    },
+  });
+
+  const [lastAllowed, overflow] = await Promise.allSettled([
+    api.agent("last allowed"),
+    api.agent("overflow"),
+  ]);
+  assert.equal(lastAllowed.status, "fulfilled");
+  assert.equal(overflow.status, "rejected");
+  assert.ok(
+    overflow.status === "rejected" &&
+      overflow.reason instanceof WorkflowEngineError &&
+      overflow.reason.kind === "call_limit",
+  );
+  assert.equal(providerCalls, 1);
+  assert.equal(api.getCallCount(), WORKFLOW_MAX_AGENT_CALLS);
+  assert.equal(store.listAgentCalls(run.id).length, 1);
+  assert.equal(
+    store.getAgentCall(run.id, WORKFLOW_MAX_AGENT_CALLS - 1)?.status,
+    "completed",
+  );
+
+  store.close();
+  await rm(dir, { recursive: true, force: true });
+}
+
 // ---------------------------------------------------------------------------
 // parallel → null on throw; barrier
 // ---------------------------------------------------------------------------
diff --git a/src/workflow-types.test.ts b/src/workflow-types.test.ts
index 55e1abc8..b8b06d7a 100644
--- a/src/workflow-types.test.ts
+++ b/src/workflow-types.test.ts
@@ -1,5 +1,6 @@
 import assert from "node:assert/strict";
 import {
+  WORKFLOW_MAX_AGENT_CALLS,
   WORKFLOW_MAX_ITEMS,
   WORKFLOW_MAX_NEST_DEPTH,
   buildAgentCacheKeyInput,
@@ -9,6 +10,7 @@ import {
 } from "./workflow-types.js";
 
 assert.equal(WORKFLOW_MAX_ITEMS, 4096);
+assert.equal(WORKFLOW_MAX_AGENT_CALLS, 256);
 assert.equal(WORKFLOW_MAX_NEST_DEPTH, 1);
 
 assert.deepEqual(