Q17: budget every model request against the provider's context window

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

tests/unit/agent-goal-tools.test.ts

@@ -65,7 +65,7 @@ function call(args: Record<string, unknown>): ToolCall {
   return { id: `call_${randomUUID()}`, name: "complete_goal", arguments: args };
 }
 
-function argsFor(h: Harness, agentId: string, auditor: Auditor, toolCall: ToolCall): GoalCompletionArgs {
+function argsFor(h: Harness, agentId: string, auditor: Auditor, toolCall: ToolCall, inputBudgetTokens = 112_000): GoalCompletionArgs {
   const ctx: AgentToolContext = {
     orgId: h.orgId,
     agent: { id: agentId, modelProviderId: h.providerId, modelId: "gpt-5.4" },
@@ -73,7 +73,7 @@ function argsFor(h: Harness, agentId: string, auditor: Auditor, toolCall: ToolCa
     taskGoalId: h.goalId, taskId: null, sourceThreadId: null,
     workspace: { id: "wsp_audit", localPath: workspaceRoot, writeAccess: true },
   };
-  return { deps: h.deps, ctx, providers: { runToolTurn: auditor }, providerId: h.providerId, workspaceId: null, call: toolCall };
+  return { deps: h.deps, ctx, providers: { runToolTurn: auditor }, providerId: h.providerId, workspaceId: null, inputBudgetTokens, call: toolCall };
 }
 
 function goalStatus(h: Harness): string {
@@ -265,4 +265,18 @@ describe("company_status", () => {
     expect(output).toContain("Ask the CEO");
     h.connection.close();
   });
+
+  test("an audit that outgrows the context budget fails as an error instead of sending an oversized request", async () => {
+    const h = createHarness();
+    let calls = 0;
+    const auditor = async (): Promise<AuditorTurn> => {
+      calls += 1;
+      return { text: "", toolCalls: [], rawResponseJson: "{}", requestId: "r" };
+    };
+    const output = await completeGoalTool(argsFor(h, h.ceoId, auditor, call({ evidence: "证".repeat(5_000) }), 3_000));
+    expect(output).toMatch(/^ERROR: This run needs ~\d+ input tokens but the provider's context window leaves 3000/);
+    expect(calls).toBe(0);
+    expect(goalStatus(h)).toBe("active");
+    h.connection.close();
+  });
 });

tests/unit/agent-run-loop.test.ts

@@ -92,7 +92,7 @@ describe("agent run loop (real tool dispatch)", () => {
         tools: { db: connection.db, organization, chat: new ChatService(connection.db) },
         onToolResults: () => { toolResultNotifications += 1; },
       },
-      { orgId, providerId, workspaceId, instructions: "You are the CEO.", initialUserTurn: "Start the company.", toolContext },
+      { orgId, providerId, inputBudgetTokens: 112_000, workspaceId, instructions: "You are the CEO.", initialUserTurn: "Start the company.", toolContext },
     );
 
     expect(result.iterations).toBe(4);
@@ -147,7 +147,7 @@ describe("agent run loop (real tool dispatch)", () => {
     ];
     const result = await runAgentLoop(
       { providers: fakeProvider(script, connection.db), tools: { db: connection.db, organization, chat: new ChatService(connection.db) } },
-      { orgId, providerId, workspaceId: null, instructions: "i", initialUserTurn: "u", toolContext: ctx },
+      { orgId, providerId, inputBudgetTokens: 112_000, workspaceId: null, instructions: "i", initialUserTurn: "u", toolContext: ctx },
     );
     // The tool returns an error string to the model; nothing is written outside the root.
     expect(existsSync(join(workspaceRoot, "..", "escape.txt"))).toBe(false);
@@ -184,7 +184,7 @@ describe("agent run loop (real tool dispatch)", () => {
 
     await runAgentLoop(
       { providers: fakeProvider(script, connection.db), tools: { db: connection.db, organization, chat: new ChatService(connection.db) } },
-      { orgId, providerId, workspaceId: null, instructions: "i", initialUserTurn: "u", toolContext: ctx },
+      { orgId, providerId, inputBudgetTokens: 112_000, workspaceId: null, instructions: "i", initialUserTurn: "u", toolContext: ctx },
     );
 
     const alex = connection.db.prepare("select id from agents where org_id = ? and name = 'Alex Rivera'").get(orgId) as { id: string };
@@ -230,7 +230,7 @@ describe("agent run loop (real tool dispatch)", () => {
     };
     await expect(runAgentLoop(
       { providers: fakeProvider([{ text: "", toolCalls: [] }], connection.db), tools: { db: connection.db, organization, chat: new ChatService(connection.db) } },
-      { orgId, providerId, workspaceId: null, instructions: "i", initialUserTurn: "u", toolContext: ctx },
+      { orgId, providerId, inputBudgetTokens: 112_000, workspaceId: null, instructions: "i", initialUserTurn: "u", toolContext: ctx },
     )).rejects.toThrow(/empty turn/);
     connection.close();
   });

tests/unit/builtin-skills.test.ts

@@ -6,7 +6,10 @@ import { describe, expect, test } from "vitest";
 import { DatabaseConnection } from "../../src/main/database/connection";
 import { OrganizationService } from "../../src/main/services/organization-service";
 import { ChatService } from "../../src/main/services/chat-service";
-import { SkillService, MAX_COMPILED_PROMPT_CHARS } from "../../src/main/services/skill-service";
+import { SkillService } from "../../src/main/services/skill-service";
+import { skillBudgetTokens } from "../../src/main/services/context-budget";
+
+const SKILL_TOKENS = skillBudgetTokens(128_000);
 import { importStagedSkills, parseFrontmatter, stageBuiltinSkillsLibrary } from "../../src/main/services/builtin-skills";
 import { dispatchAgentTool, type AgentToolContext } from "../../src/main/services/agent-tools";
 
@@ -141,7 +144,7 @@ describe("built-in skills", () => {
     // A third, oversized flat skill to prove the compiled-prompt cap.
     writeFileSync(
       join(h.sourceRoot, "engineering", "engineering-epic-novelist.md"),
-      `---\nname: Epic Novelist\ndescription: too big to pack\n---\n\n${"x".repeat(MAX_COMPILED_PROMPT_CHARS + 10)}\n`,
+      `---\nname: Epic Novelist\ndescription: too big to pack\n---\n\n${"x".repeat(SKILL_TOKENS * 3 + 30)}\n`,
       "utf8",
     );
     const staged = stageBuiltinSkillsLibrary(h.sourceRoot, h.libraryRoot);
@@ -175,17 +178,17 @@ describe("built-in skills", () => {
     expect(hired).toContain("Unknown skills skipped: no-such-skill");
 
     const danaId = (h.connection.db.prepare("select id from agents where org_id = ? and name = 'Dana Lee'").get(h.orgId) as { id: string }).id;
-    const danaSnapshot = h.connection.transaction(() => h.skills.resolveRuntimeSnapshot({ orgId: h.orgId, agentId: danaId, runId: "run_dana" }));
+    const danaSnapshot = h.connection.transaction(() => h.skills.resolveRuntimeSnapshot({ orgId: h.orgId, agentId: danaId, runId: "run_dana", maxTokens: SKILL_TOKENS }));
     const danaPrompt = (h.connection.db.prepare("select compiled_prompt from skill_runtime_snapshots where id = ?").get(danaSnapshot) as { compiled_prompt: string }).compiled_prompt;
     expect(danaPrompt.indexOf("# Skill: humanizer")).toBe(0);
     expect(danaPrompt).toContain("# Skill: Backend Architect");
     expect(danaPrompt).toContain("Skills omitted for size");
     expect(danaPrompt).toContain("Epic Novelist");
-    expect(danaPrompt.length).toBeLessThan(MAX_COMPILED_PROMPT_CHARS + 5_000);
+    expect(danaPrompt.length).toBeLessThan(SKILL_TOKENS * 3 + 5_000);
 
     // The CEO carries humanizer too - forced by resolution, not by binding.
     expect(h.skills.hasRuntimeSkills(h.orgId, ceoId)).toBe(true);
-    const ceoSnapshot = h.connection.transaction(() => h.skills.resolveRuntimeSnapshot({ orgId: h.orgId, agentId: ceoId, runId: "run_ceo" }));
+    const ceoSnapshot = h.connection.transaction(() => h.skills.resolveRuntimeSnapshot({ orgId: h.orgId, agentId: ceoId, runId: "run_ceo", maxTokens: SKILL_TOKENS }));
     const ceoPrompt = (h.connection.db.prepare("select compiled_prompt from skill_runtime_snapshots where id = ?").get(ceoSnapshot) as { compiled_prompt: string }).compiled_prompt;
     expect(ceoPrompt).toContain("# Skill: humanizer");
     expect(ceoPrompt).not.toContain("# Skill: Backend Architect");
@@ -210,7 +213,7 @@ describe("built-in skills", () => {
       });
     });
     const promptFor = (runId: string) => {
-      const snapshotId = h.connection.transaction(() => h.skills.resolveRuntimeSnapshot({ orgId: h.orgId, agentId, runId }));
+      const snapshotId = h.connection.transaction(() => h.skills.resolveRuntimeSnapshot({ orgId: h.orgId, agentId, runId, maxTokens: SKILL_TOKENS }));
       return (h.connection.db.prepare("select compiled_prompt from skill_runtime_snapshots where id = ?").get(snapshotId) as { compiled_prompt: string }).compiled_prompt;
     };
 

tests/unit/context-budget.test.ts

@@ -0,0 +1,97 @@
+import { randomUUID } from "node:crypto";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
+import { describe, expect, test } from "vitest";
+import { estimateTokens, inputBudgetTokens, skillBudgetTokens } from "../../src/main/services/context-budget";
+import { runAgentLoop } from "../../src/main/services/agent-run-loop";
+import { DatabaseConnection } from "../../src/main/database/connection";
+import { OrganizationService } from "../../src/main/services/organization-service";
+import { ChatService } from "../../src/main/services/chat-service";
+import { IntegrationKeyStore } from "../../src/main/services/integration-key-store";
+import { ProviderService } from "../../src/main/services/provider-service";
+import type { AgentTurnMessage, AgentTurnResult } from "../../src/main/services/provider-tool-turn";
+import { AGENT_TOOLS } from "../../src/main/services/agent-tools";
+import { WORKSPACE_TOOLS } from "../../src/main/services/agent-workspace-tools";
+import { WEB_TOOLS } from "../../src/main/services/agent-web-tools";
+
+const FIXED_TOKENS = estimateTokens("i") + estimateTokens(JSON.stringify([...AGENT_TOOLS, ...WORKSPACE_TOOLS, ...WEB_TOOLS]));
+
+describe("token estimates", () => {
+  test("count CJK per character and other text per three characters", () => {
+    expect(estimateTokens("abcdef")).toBe(2);
+    expect(estimateTokens("abcdefg")).toBe(3);
+    expect(estimateTokens("技能规则")).toBe(4);
+    expect(estimateTokens("规则 rule")).toBe(4);
+    expect(inputBudgetTokens(128_000)).toBe(112_000);
+    expect(skillBudgetTokens(128_000)).toBe(44_800);
+  });
+});
+
+async function loop(budget: number, sent: AgentTurnMessage[][]) {
+  const connection = new DatabaseConnection(join(tmpdir(), `ac-budget-${randomUUID()}.db`));
+  const organization = new OrganizationService(connection.db);
+  let orgId = "";
+  let agentId = "";
+  connection.transaction(() => {
+    orgId = organization.createOrganization({ organizationName: "B", employeeName: "Y", mission: "m", goal: "g" });
+    agentId = organization.createAgent({ orgId, name: "Ana", role: "Analyst", department: "Operations", modelProviderId: null, modelId: null, budgetCents: 1000, managerAgentId: null });
+  });
+  const chat = new ChatService(connection.db);
+  const general = connection.db.prepare("select id from chat_spaces where org_id = ? and name = 'general'").get(orgId) as { id: string };
+  chat.sendMessage({ orgId, spaceId: general.id, threadId: null, senderType: "agent", senderId: agentId, bodyMd: "y".repeat(9000), attachmentFileIds: [] });
+  let turn = 0;
+  try {
+    return await runAgentLoop(
+      {
+        providers: {
+          integrationKeys: new IntegrationKeyStore(connection.db),
+          runToolTurn: async (input): Promise<AgentTurnResult & { requestId: string }> => {
+            sent.push(input.messages.map((message) => ({ ...message })));
+            turn += 1;
+            return turn <= 2
+              ? { text: "", toolCalls: [{ id: `c${turn}`, name: "read_messages", arguments: { channel: "general" } }], rawResponseJson: "{}", requestId: `r${turn}` }
+              : { text: "done", toolCalls: [], rawResponseJson: "{}", requestId: `r${turn}` };
+          },
+        },
+        tools: { db: connection.db, organization, chat },
+      },
+      {
+        orgId, providerId: "p", inputBudgetTokens: budget, workspaceId: null, instructions: "i", initialUserTurn: "go",
+        toolContext: { orgId, agent: { id: agentId, modelProviderId: null, modelId: null }, taskProjectId: null, taskGoalId: null, taskId: null, sourceThreadId: null, workspace: null },
+      },
+    );
+  } finally {
+    connection.close();
+  }
+}
+
+describe("run context budget", () => {
+  test("over budget, older tool results are elided before the next call", async () => {
+    const sent: AgentTurnMessage[][] = [];
+    const result = await loop(FIXED_TOKENS + 5_000, sent);
+    expect(result.finalText).toBe("done");
+    const outputs = sent[2].flatMap((message) => (message.role === "tool" ? message.results.map((item) => item.output) : []));
+    expect(outputs[0]).toMatch(/^\[Earlier result of read_messages elided/);
+    expect(outputs[1]).toContain("y".repeat(9000));
+  });
+
+  test("a request that cannot fit even after compaction fails loudly with the numbers", async () => {
+    await expect(loop(2_000, [])).rejects.toThrow(/This run needs ~\d+ input tokens but the provider's context window leaves 2000/);
+  });
+
+  test("the provider context window has a floor", () => {
+    const connection = new DatabaseConnection(join(tmpdir(), `ac-window-${randomUUID()}.db`));
+    const organization = new OrganizationService(connection.db);
+    let orgId = "";
+    connection.transaction(() => {
+      orgId = organization.createOrganization({ organizationName: "W", employeeName: "Y", mission: "m", goal: "g" });
+    });
+    const providerId = (connection.db.prepare("select id from providers where org_id = ? and provider_key = 'openai'").get(orgId) as { id: string }).id;
+    const providers = new ProviderService(connection.db);
+    const policy = { orgId, providerId, dataPolicy: "standard" as const, backgroundEnabled: false, monthlyRequestLimit: 100, workspaceScope: "organization" as const };
+    expect(() => providers.updatePolicy({ ...policy, contextWindowTokens: 8_000 })).toThrow("Context window must be at least 24000 tokens.");
+    providers.updatePolicy({ ...policy, contextWindowTokens: 200_000 });
+    expect(connection.db.prepare("select context_window_tokens from providers where id = ?").get(providerId)).toEqual({ context_window_tokens: 200_000 });
+    connection.close();
+  });
+});

tests/unit/domain-services.test.ts

@@ -24,7 +24,7 @@ describe("domain services", () => {
     const connection = createConnection();
 
     const rows = connection.db.prepare("select version from schema_migrations order by version asc").all() as Array<{ version: number }>;
-    expect(rows.map((row) => row.version)).toEqual([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61]);
+    expect(rows.map((row) => row.version)).toEqual([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62]);
     connection.close();
   });
 
@@ -186,7 +186,7 @@ describe("domain services", () => {
       skills.bind({ orgId, agentId, skillId: skill.id, revisionId });
       rmSync(skillDir, { recursive: true, force: true });
       skills.runTests({ orgId, revisionId });
-      skills.resolveRuntimeSnapshot({ orgId, agentId, runId: "run-review-1" });
+      skills.resolveRuntimeSnapshot({ orgId, agentId, runId: "run-review-1", maxTokens: 44_800 });
     });
 
     // Catalogue metadata ships in the snapshot; skill internals stay in the
@@ -247,10 +247,10 @@ describe("domain services", () => {
         .get(orgA, "review-routes") as { id: string }).id;
       expect(() => skills.bind({ orgId: orgA, agentId: agentB, skillId, revisionId })).toThrow("Agent was not found.");
       skills.bind({ orgId: orgA, agentId: agentA, skillId, revisionId });
-      const firstResolve = skills.resolveRuntimeSnapshot({ orgId: orgA, agentId: agentA, runId: "run-review-1" });
+      const firstResolve = skills.resolveRuntimeSnapshot({ orgId: orgA, agentId: agentA, runId: "run-review-1", maxTokens: 44_800 });
       // Idempotent per run: a recovered run reuses its frozen snapshot
       // instead of failing on the uniqueness guard.
-      expect(skills.resolveRuntimeSnapshot({ orgId: orgA, agentId: agentA, runId: "run-review-1" })).toBe(firstResolve);
+      expect(skills.resolveRuntimeSnapshot({ orgId: orgA, agentId: agentA, runId: "run-review-1", maxTokens: 44_800 })).toBe(firstResolve);
       expect(connection.db.prepare("select count(*) as n from skill_runtime_snapshots where run_id = 'run-review-1'").get())
         .toEqual({ n: 1 });
     });

tests/unit/operations-service.test.ts

@@ -85,7 +85,7 @@ describe("operations service", () => {
         dataPolicy: PROVIDER_DATA_POLICY.STANDARD,
         backgroundEnabled: true,
         monthlyRequestLimit: 10,
-        workspaceScope: "organization",
+        workspaceScope: "organization", contextWindowTokens: 128_000,
       });
       operations.setFeatureFlag({ orgId, flagId, enabled: false });
     });

tests/unit/provider-background-service.test.ts

@@ -52,7 +52,7 @@ describe("provider background service", () => {
         dataPolicy: PROVIDER_DATA_POLICY.STANDARD,
         backgroundEnabled: true,
         monthlyRequestLimit: 10,
-        workspaceScope: "organization",
+        workspaceScope: "organization", contextWindowTokens: 128_000,
       });
     });
 
@@ -115,7 +115,7 @@ describe("provider background service", () => {
         dataPolicy: PROVIDER_DATA_POLICY.STANDARD,
         backgroundEnabled: true,
         monthlyRequestLimit: 10,
-        workspaceScope: "organization",
+        workspaceScope: "organization", contextWindowTokens: 128_000,
       });
     });
 
@@ -152,7 +152,7 @@ describe("provider background service", () => {
         dataPolicy: PROVIDER_DATA_POLICY.STANDARD,
         backgroundEnabled: true,
         monthlyRequestLimit: 10,
-        workspaceScope: "organization",
+        workspaceScope: "organization", contextWindowTokens: 128_000,
       });
     });
 

tests/unit/provider-service.test.ts

@@ -448,7 +448,7 @@ describe("provider service", () => {
         dataPolicy: PROVIDER_DATA_POLICY.STANDARD,
         backgroundEnabled: false,
         monthlyRequestLimit: 1,
-        workspaceScope: "organization",
+        workspaceScope: "organization", contextWindowTokens: 128_000,
       });
     });
     await providers.createResponse({
@@ -524,7 +524,7 @@ describe("provider service", () => {
         dataPolicy: PROVIDER_DATA_POLICY.STANDARD,
         backgroundEnabled: false,
         monthlyRequestLimit: 10,
-        workspaceScope: "workspace",
+        workspaceScope: "workspace", contextWindowTokens: 128_000,
       });
     });
 
@@ -582,7 +582,7 @@ describe("provider service", () => {
         dataPolicy: PROVIDER_DATA_POLICY.ZERO_DATA_RETENTION,
         backgroundEnabled: false,
         monthlyRequestLimit: 10,
-        workspaceScope: "organization",
+        workspaceScope: "organization", contextWindowTokens: 128_000,
       });
     });
 

tests/unit/run-context-compaction.test.ts

@@ -72,7 +72,7 @@ describe("tool history compaction", () => {
         tools: { db: connection.db, organization, chat },
       },
       {
-        orgId, providerId: "p", workspaceId: null, instructions: "i", initialUserTurn: "go",
+        orgId, providerId: "p", inputBudgetTokens: 112_000, workspaceId: null, instructions: "i", initialUserTurn: "go",
         toolContext: { orgId, agent: { id: agentId, modelProviderId: null, modelId: null }, taskProjectId: null, taskGoalId: null, taskId: null, sourceThreadId: null, workspace: null },
       },
     );

tests/unit/task-execution-service.test.ts

@@ -408,7 +408,7 @@ describe("task execution service", () => {
         dataPolicy: PROVIDER_DATA_POLICY.STANDARD,
         backgroundEnabled: false,
         monthlyRequestLimit: 5,
-        workspaceScope: "organization",
+        workspaceScope: "organization", contextWindowTokens: 128_000,
       });
       const agentId = organization.createAgent({
         orgId,
@@ -452,7 +452,7 @@ describe("task execution service", () => {
       dataPolicy: PROVIDER_DATA_POLICY.STANDARD,
       backgroundEnabled: false,
       monthlyRequestLimit: 50,
-      workspaceScope: "organization",
+      workspaceScope: "organization", contextWindowTokens: 128_000,
     }));
     const enqueued = connection.transaction(() => execution.enqueueHeartbeatRuns());
     expect(enqueued).toHaveLength(1);