diff --git a/src/main/services/agent-goal-tools.ts b/src/main/services/agent-goal-tools.ts index 6ee02394..208029e4 100644 --- a/src/main/services/agent-goal-tools.ts +++ b/src/main/services/agent-goal-tools.ts @@ -6,7 +6,7 @@ import type { AgentToolContext, AgentToolDeps } from "./agent-tools"; import type { JsonRecord } from "./provider-http"; import type { ProviderService } from "./provider-service"; import type { AgentTurnMessage, ToolCall, ToolDefinition } from "./provider-tool-turn"; -import { throwIfRunStopped } from "./run-control"; +import { MAX_TOOL_CALLS_PER_TURN, throwIfRunStopped, untilStopped } from "./run-control"; const AUDIT_MAX_ITERATIONS = 6; const STATUS_TASK_LINES = 8; @@ -134,19 +134,24 @@ async function auditGoalCompletion( for (let iteration = 1; iteration <= AUDIT_MAX_ITERATIONS; iteration += 1) { fitContextBudget(instructions, messages, auditTools, args.inputBudgetTokens, iteration); - const turn = await args.providers.runToolTurn({ + throwIfRunStopped(); + const turn = await untilStopped(args.providers.runToolTurn({ orgId: args.ctx.orgId, providerId: args.providerId, workspaceId: args.workspaceId, instructions, messages, tools: auditTools, - }); + })); if (turn.toolCalls.length === 0) return parseVerdict(turn.text); messages.push({ role: "assistant", text: turn.text, toolCalls: turn.toolCalls }); const results: Array<{ toolCallId: string; toolName: string; output: string }> = []; - for (const toolCall of turn.toolCalls) { - results.push({ toolCallId: toolCall.id, toolName: toolCall.name, output: await executeWorkspaceTool(auditCtx, toolCall) }); + for (const [index, toolCall] of turn.toolCalls.entries()) { + throwIfRunStopped(); + const output = index < MAX_TOOL_CALLS_PER_TURN + ? await untilStopped(executeWorkspaceTool(auditCtx, toolCall)) + : `ERROR: not run: this turn asked for ${turn.toolCalls.length} tool calls; the limit is ${MAX_TOOL_CALLS_PER_TURN}.`; + results.push({ toolCallId: toolCall.id, toolName: toolCall.name, output }); } messages.push({ role: "tool", results }); } diff --git a/src/main/services/agent-run-loop.ts b/src/main/services/agent-run-loop.ts index d378201a..c147ac0b 100644 --- a/src/main/services/agent-run-loop.ts +++ b/src/main/services/agent-run-loop.ts @@ -8,9 +8,11 @@ import type { AgentTurnMessage, ToolCall, ToolDefinition } from "./provider-tool import { isComposioTool, executeComposioTool } from "./composio/composio-tool-bridge"; import { capToolOutput, compactToolHistory } from "./run-context-compaction"; import { fitContextBudget } from "./context-budget"; -import { throwIfRunStopped, untilStopped } from "./run-control"; +import { MAX_TOOL_CALLS_PER_TURN, throwIfRunStopped, untilStopped } from "./run-control"; const MAX_ITERATIONS = 12; +/** Each goal audit runs its own model loop, so a session gets only a couple. */ +const MAX_GOAL_AUDITS_PER_SESSION = 2; const STEP_LIMIT_TEXT = "I used this session's steps and will continue in the next one."; export interface AgentRunLoopInput { @@ -52,18 +54,14 @@ export function runTools(externalTools: ToolDefinition[]): ToolDefinition[] { return [...AGENT_TOOLS, ...WORKSPACE_TOOLS, ...WEB_TOOLS, ...externalTools]; } -/** - * Drives a real tool-use loop against the live provider: prompt the model with - * the agent tools, execute every tool call it makes, feed results back, and - * repeat until the model returns a turn with no tool calls (it is done) or the - * iteration cap is hit (which throws rather than silently truncating). - */ +/** Runs the model with the agent tools until it stops calling them; step, time, tool-call and audit limits end it as a checkpoint or refuse the excess. */ export async function runAgentLoop(deps: AgentRunLoopDeps, input: AgentRunLoopInput): Promise { const allTools = runTools(input.externalTools ?? []); const messages: AgentTurnMessage[] = [{ role: "user", text: input.initialUserTurn }]; let toolCallCount = 0; let lastRequestId = ""; let lastText = ""; + let goalAudits = 0; for (let iteration = 1; iteration <= MAX_ITERATIONS; iteration += 1) { const compaction = compactToolHistory(messages); @@ -106,8 +104,17 @@ export async function runAgentLoop(deps: AgentRunLoopDeps, input: AgentRunLoopIn lastText = turn.text.trim() || lastText; messages.push({ role: "assistant", text: turn.text, toolCalls: turn.toolCalls }); const results: Array<{ toolCallId: string; toolName: string; output: string }> = []; - for (const call of turn.toolCalls) { + for (const [index, call] of turn.toolCalls.entries()) { throwIfRunStopped(); + if (index >= MAX_TOOL_CALLS_PER_TURN) { + results.push({ toolCallId: call.id, toolName: call.name, output: `ERROR: not run: this turn asked for ${turn.toolCalls.length} tool calls; the limit is ${MAX_TOOL_CALLS_PER_TURN}. Repeat the ones you still need.` }); + continue; + } + if (call.name === "complete_goal" && goalAudits >= MAX_GOAL_AUDITS_PER_SESSION) { + results.push({ toolCallId: call.id, toolName: call.name, output: `ERROR: this session already ran ${MAX_GOAL_AUDITS_PER_SESSION} goal audits. Close the gaps they named, then try again next session.` }); + continue; + } + if (call.name === "complete_goal") goalAudits += 1; toolCallCount += 1; let output: string; if (isComposioTool(call.name)) { diff --git a/src/main/services/run-control.ts b/src/main/services/run-control.ts index c2fd66ba..bdb7545d 100644 --- a/src/main/services/run-control.ts +++ b/src/main/services/run-control.ts @@ -2,20 +2,30 @@ import { AsyncLocalStorage } from "node:async_hooks"; /** The error message a run ends with when the operator stops it. */ export const RUN_STOPPED_MESSAGE = "Stopped by the operator."; +/** The reason a session hits its wall-clock budget; it ends as a checkpoint, not a failure. */ +export const SESSION_TIME_LIMIT_MESSAGE = "The session reached its time limit."; +export const SESSION_TIME_LIMIT_MS = 30 * 60_000; +/** Tool calls one model turn may run, in a session or a nested audit. */ +export const MAX_TOOL_CALLS_PER_TURN = 16; const runScope = new AsyncLocalStorage<{ runId: string; signal: AbortSignal }>(); const controllers = new Map(); const settling = new Map>(); -/** Runs a session under a stop signal that every provider request and tool inside it observes. */ -export async function withRunControl(runId: string, operation: () => Promise): Promise { +/** Runs a session under a stop signal that every provider request and tool inside it observes; the signal also fires at the session's deadline. */ +export async function withRunControl(runId: string, operation: () => Promise, timeLimitMs: number = SESSION_TIME_LIMIT_MS): Promise { const controller = new AbortController(); controllers.set(runId, controller); + const deadline = setTimeout(() => { + console.info(`[run-control] run ${runId} reached its ${timeLimitMs / 60_000}-minute limit`); + controller.abort(new Error(SESSION_TIME_LIMIT_MESSAGE)); + }, timeLimitMs); const running = runScope.run({ runId, signal: controller.signal }, operation); settling.set(runId, running.catch(() => undefined)); try { return await running; } finally { + clearTimeout(deadline); controllers.delete(runId); settling.delete(runId); } diff --git a/src/main/services/task-execution-service.ts b/src/main/services/task-execution-service.ts index bb79b0b0..9237d17e 100644 --- a/src/main/services/task-execution-service.ts +++ b/src/main/services/task-execution-service.ts @@ -13,7 +13,7 @@ import { runAgentLoop, runTools } from "./agent-run-loop"; import type { RunModel } from "./skill-runtime"; import { combineMidRunPolls, createGoalChangePoller } from "./goal-change-poller"; import { recordRunAction, runJournalLines } from "./run-journal"; -import { runStopReason, withRunControl } from "./run-control"; +import { runStopReason, SESSION_TIME_LIMIT_MESSAGE, SESSION_TIME_LIMIT_MS, withRunControl } from "./run-control"; import type { AgentToolContext } from "./agent-tools"; import { buildAgentInstructions, @@ -77,6 +77,8 @@ export class TaskExecutionService { private readonly notifyDataChanged: () => void = () => undefined, /** Injectable clock so tests can walk providers out of cooldown. */ now: () => number = () => Date.now(), + /** Wall-clock budget of one session. */ + private readonly sessionTimeLimitMs: number = SESSION_TIME_LIMIT_MS, ) { this.audit = new AuditService(db); this.organization = new OrganizationService(db); @@ -99,7 +101,7 @@ export class TaskExecutionService { const workspace = loadWorkspace(this.db, task.workspace_id); const isolation = createIsolationContext(this.db, input.orgId, task, agent, provider, workspace); const runId = this.createRun(input.orgId, task.id, agent.id, provider.id, isolation); - return withRunControl(runId, () => this.executeAgentRun({ orgId: input.orgId, runId, task, agent, provider, workspace, isolation })); + return withRunControl(runId, () => this.executeAgentRun({ orgId: input.orgId, runId, task, agent, provider, workspace, isolation }), this.sessionTimeLimitMs); } /** Executes a claimed run after re-checking that the task still matches it; a mismatch supersedes the run. */ @@ -157,7 +159,7 @@ export class TaskExecutionService { if (task && task.origin === null) recordRunMemory(this.db, { orgId, runId, projectId: task.project_id, taskTitle: task.title, outcome: "failed", text: message }); throw error; } - return withRunControl(runId, () => this.executeAgentRun(prepared)); + return withRunControl(runId, () => this.executeAgentRun(prepared), this.sessionTimeLimitMs); } private async executeAgentRun(args: { @@ -286,6 +288,12 @@ export class TaskExecutionService { const message = error instanceof Error ? error.message : String(error); // A stop wins over whatever error the interrupted step raised. const stopReason = runStopReason(); + // Out of time is not a failure: the session ends as a checkpoint and the next one continues from the journal. + if (stopReason === SESSION_TIME_LIMIT_MESSAGE) { + this.completeRun(orgId, runId, providerRequestId, true); + this.audit.record({ orgId, entityType: "task_run", entityId: runId, actorType: "system", actorId: null, action: "completed", payload: { taskId: task.id, agentId: agent.id, timeLimit: true } }); + return runId; + } if (stopReason) { this.markStopped(orgId, runId, providerRequestId, stopReason, { taskId: task.id, agentId: agent.id, providerId: provider.id }); throw new Error(stopReason); @@ -400,10 +408,10 @@ export class TaskExecutionService { return row?.language === "zh" ? "zh" : "en"; } - private completeRun(orgId: string, runId: string, providerRequestId: string, checkpoint: boolean): void { + private completeRun(orgId: string, runId: string, providerRequestId: string | null, checkpoint: boolean): void { this.db.prepare(` update task_runs - set state = ?, provider_request_id = ?, completed_at = ?, error_message = null + set state = ?, provider_request_id = coalesce(?, provider_request_id), completed_at = ?, error_message = null where id = ? and org_id = ? `).run(TASK_RUN_STATE.COMPLETED, providerRequestId, new Date().toISOString(), runId, orgId); if (!checkpoint) this.closeChatWakeTask(orgId, runId); diff --git a/tests/integration/run-stop.test.ts b/tests/integration/run-stop.test.ts index 393bb4da..92b38724 100644 --- a/tests/integration/run-stop.test.ts +++ b/tests/integration/run-stop.test.ts @@ -108,6 +108,32 @@ describe("stopping runs", () => { connection.close(); }); + test("a session that runs out of time ends as a checkpoint, even mid-request", async () => { + let stalls = true; + const turn = async (): Promise => { + if (stalls) await new Promise(() => {}); + return { text: "done", toolCalls: [], rawResponseJson: "{}" }; + }; + const connection = new DatabaseConnection(join(tmpdir(), `ac-deadline-${randomUUID()}.db`)); + const providers = new ProviderService(connection.db, providerAdapters, testVault, undefined, turn); + const onboarding = new OnboardingService({ db: connection.db, saveProviderCredential: (input) => providers.saveCredential(input) }); + const execution = new TaskExecutionService(connection.db, providers, undefined, undefined, undefined, 200); + let orgId = ""; + connection.transaction(() => { + orgId = onboarding.bootstrap({ + companyName: "FlowPilot", companyDescription: "CSM analytics.", employeeName: "Yang", workspacePath: "", + executionMode: "ai_ceo", autoApproveHires: true, reviewDeliverables: false, + goalTitle: "Ship the MVP.", goalDescription: "A working product.", + providerKey: "openai", providerSecret: "sk-test-key", providerModel: "gpt-5.4", + }).orgId; + }); + const runId = claimNextQueuedRun(connection)!.id; + await expect(execution.executeRun(orgId, runId)).resolves.toBe(runId); + expect(connection.db.prepare("select state, error_message from task_runs where id = ?").get(runId)).toEqual({ state: "completed", error_message: null }); + stalls = false; + connection.close(); + }); + test("stopping a run aborts its provider request in flight", async () => { const server = createServer(() => { /* never answers */ }); await new Promise((resolve) => server.listen(0, "127.0.0.1", resolve)); diff --git a/tests/unit/agent-run-loop.test.ts b/tests/unit/agent-run-loop.test.ts index b124c62e..36b30413 100644 --- a/tests/unit/agent-run-loop.test.ts +++ b/tests/unit/agent-run-loop.test.ts @@ -258,4 +258,44 @@ describe("agent run loop (real tool dispatch)", () => { expect(result).toMatchObject({ iterations: 12, toolCallCount: 12, hitStepLimit: true, finalText: "Still building." }); connection.close(); }); + + test("a session is bounded in tool calls per turn and in goal audits", async () => { + const connection = new DatabaseConnection(join(tmpdir(), `ac-loop5-${randomUUID()}.db`)); + const organization = new OrganizationService(connection.db); + let orgId = ""; + let ceoId = ""; + let providerId = ""; + connection.transaction(() => { + orgId = organization.createOrganization({ organizationName: "OrbitOps", employeeName: "Yang", mission: "m", goal: "g" }); + providerId = (connection.db.prepare("select id from providers where org_id = ? and provider_key = 'openai'").get(orgId) as { id: string }).id; + ceoId = organization.createAgent({ orgId, name: "CEO", role: "CEO", department: "Executive", modelProviderId: providerId, modelId: "gpt-5.4", budgetCents: 25000, managerAgentId: null }); + }); + const ctx: AgentToolContext = { + orgId, agent: { id: ceoId, modelProviderId: providerId, modelId: "gpt-5.4" }, + taskProjectId: null, taskGoalId: null, taskId: null, sourceThreadId: null, workspace: null, + }; + const tools = { db: connection.db, organization, chat: new ChatService(connection.db) }; + const loopInput = { orgId, providerId, inputBudgetTokens: 112_000, workspaceId: null, instructions: "i", initialUserTurn: "u", toolContext: ctx }; + + // 20 calls in one turn: 16 run, the rest are refused by name. + const flood = [{ text: "Many reads.", toolCalls: Array.from({ length: 20 }, () => call("read_messages", {})) }]; + let lastToolMessage: unknown = null; + const watching = fakeProvider(flood, connection.db); + const recorded = await runAgentLoop({ ...{ providers: { ...watching, runToolTurn: async (input: { messages: unknown[] }) => { lastToolMessage = input.messages.at(-1); return watching.runToolTurn(); } } }, tools } as never, loopInput); + expect(recorded.toolCallCount).toBe(16); + const results = (lastToolMessage as { results: Array<{ output: string }> }).results; + expect(results.filter((result) => result.output.startsWith("ERROR: not run: this turn asked for 20 tool calls; the limit is 16"))).toHaveLength(4); + + // A third goal audit in one session is refused before it runs. + const audits = [{ text: "Closing.", toolCalls: [call("complete_goal", { evidence: "a" }), call("complete_goal", { evidence: "b" }), call("complete_goal", { evidence: "c" })] }]; + let auditResults: Array<{ output: string }> = []; + const auditing = fakeProvider(audits, connection.db); + await runAgentLoop({ providers: { ...auditing, runToolTurn: async (input: { messages: Array<{ role: string; results?: Array<{ output: string }> }> }) => { + const last = input.messages.at(-1); + if (last?.role === "tool") auditResults = last.results ?? []; + return auditing.runToolTurn(); + } }, tools } as never, loopInput); + expect(auditResults.map((result) => result.output.startsWith("ERROR: this session already ran 2 goal audits"))).toEqual([false, false, true]); + connection.close(); + }); });