diff --git a/src/main/services/agent-run-loop.ts b/src/main/services/agent-run-loop.ts index c147ac0b..c332312f 100644 --- a/src/main/services/agent-run-loop.ts +++ b/src/main/services/agent-run-loop.ts @@ -4,7 +4,7 @@ import { completeGoalTool } from "./agent-goal-tools"; import { executeWebTool, isWebTool, WEB_TOOLS } from "./agent-web-tools"; import { executeWorkspaceTool, isWorkspaceTool, WORKSPACE_TOOLS } from "./agent-workspace-tools"; import { canonicalMentionHandle } from "../../shared/mention-handles"; -import type { AgentTurnMessage, ToolCall, ToolDefinition } from "./provider-tool-turn"; +import type { AgentTurnMessage, AgentTurnResult, ToolCall, ToolDefinition } from "./provider-tool-turn"; import { isComposioTool, executeComposioTool } from "./composio/composio-tool-bridge"; import { capToolOutput, compactToolHistory } from "./run-context-compaction"; import { fitContextBudget } from "./context-budget"; @@ -25,8 +25,8 @@ export interface AgentRunLoopInput { externalTools?: ToolDefinition[]; /** Messages that arrived mid-run, or null. */ pollNewMessages?: () => string | null; - /** Called after each model turn the provider answered. */ - onModelTurn?: () => void; + /** Called after each model turn the provider answered, with that turn. */ + onModelTurn?: (turn: AgentTurnResult) => void; /** Called with every executed tool call and its full output. */ onToolCompleted?: (call: ToolCall, output: string) => void; /** Input tokens the provider's context window leaves after the output reserve. */ @@ -84,7 +84,7 @@ export async function runAgentLoop(deps: AgentRunLoopDeps, input: AgentRunLoopIn }); throwIfRunStopped(); lastRequestId = turn.requestId; - input.onModelTurn?.(); + input.onModelTurn?.(turn); if (turn.toolCalls.length === 0) { if (toolCallCount === 0 && !turn.text.trim()) { diff --git a/src/main/services/provider-cooldowns.ts b/src/main/services/provider-cooldowns.ts index 64d533ec..39976e43 100644 --- a/src/main/services/provider-cooldowns.ts +++ b/src/main/services/provider-cooldowns.ts @@ -3,9 +3,9 @@ import type { DatabaseSync } from "node:sqlite"; /** Without a provider-given wait, parks double from this base with each failure in a row, up to the cap. */ const BACKOFF_BASE_MS = 2 * 60_000; const BACKOFF_CAP_MS = 30 * 60_000; -/** A provider-given wait is honored within these bounds. */ +/** A provider-given wait (Retry-After, a usage window's reset) is honored within these bounds; the cap only guards against garbage. */ const RETRY_AFTER_MIN_MS = 30_000; -const RETRY_AFTER_MAX_MS = 60 * 60_000; +const RETRY_AFTER_MAX_MS = 24 * 60 * 60_000; interface CooldownRow { provider_id: string; diff --git a/src/main/services/provider-response-runner.ts b/src/main/services/provider-response-runner.ts index 8ca0e7cf..eb03510b 100644 --- a/src/main/services/provider-response-runner.ts +++ b/src/main/services/provider-response-runner.ts @@ -146,6 +146,7 @@ export class ProviderResponseRunner { try { const credential = await resolveTurnCredential({ authBroker: this.deps.authBroker, integrationKeys: this.deps.integrationKeys, orgId: input.orgId, provider }); const result = await turn({ credential, model, instructions: input.instructions, messages: input.messages, tools: input.tools }); + recordObservedRateLimit(this.deps.rateLimitRegistry, input.orgId, owner, model, result.rateLimitState ?? null); // Tool calls are output the model wrote, as much as its text. const outputChars = result.text.length + (result.toolCalls.length > 0 ? JSON.stringify(result.toolCalls).length : 0); this.deps.requestStore.markSucceeded({ diff --git a/src/main/services/provider-tool-turn.ts b/src/main/services/provider-tool-turn.ts index 970957de..34e80f46 100644 --- a/src/main/services/provider-tool-turn.ts +++ b/src/main/services/provider-tool-turn.ts @@ -1,3 +1,4 @@ +import type { ProviderObservedRateLimitState } from "../../shared/provider-runtime-domain"; /** * Shared contract for a single tool-use ("function calling") turn against any * provider backend. The agent loop maintains a normalized conversation @@ -30,6 +31,8 @@ export interface AgentTurnResult { text: string; toolCalls: ToolCall[]; rawResponseJson: string; + /** Rate-limit headers of the response, when the provider sent any. */ + rateLimitState?: ProviderObservedRateLimitState | null; } /** Resolved per-request credential, already refreshed if it was an OAuth token. */ diff --git a/src/main/services/provider-tool-turns/turns.ts b/src/main/services/provider-tool-turns/turns.ts index 1c7c6f94..8b7af38a 100644 --- a/src/main/services/provider-tool-turns/turns.ts +++ b/src/main/services/provider-tool-turns/turns.ts @@ -48,7 +48,7 @@ export async function runOpenAiResponsesTurn(input: ProviderToolTurnInput): Prom }; const response = await postJson("https://api.openai.com/v1/responses", { authorization: `Bearer ${requireApiKey(input.credential, "OpenAI")}` }, body); const parsed = parseResponsesOutput(response.data); - return { ...parsed, rawResponseJson: JSON.stringify(response.data) }; + return { ...parsed, rawResponseJson: JSON.stringify(response.data), rateLimitState: response.rateLimitState }; } // ---- OpenAI Chat Completions, streamed (xAI, Groq, Mistral, Cerebras, OpenRouter, Kimi, Kimi CN) ---- @@ -71,7 +71,7 @@ async function runChatCompletionsStream(endpoint: string, headers: Record Promise { @@ -90,7 +90,7 @@ export async function runAnthropicTurn(input: ProviderToolTurnInput): Promise { + onModelTurn: (turn) => { contextRecorder.commit(); // A served request ends the provider's failure streak, so the next park starts small again. this.cooldowns.clear(provider); + // A response that says no requests remain parks the provider until its reset, before the next call fails into it. + const limits = turn.rateLimitState; + const resetIn = limits?.resetAt ? Date.parse(limits.resetAt) - Date.now() : 0; + if (limits?.remainingRequests === 0 && resetIn > 0) this.cooldowns.park(provider, "no requests left until the provider's reset", resetIn); }, onToolCompleted: (call, output) => recordRunAction(this.db, { orgId, runId, taskId: task.id, call, output }), inputBudgetTokens: inputBudgetTokens(provider.context_window_tokens), @@ -265,7 +269,6 @@ export class TaskExecutionService { if (delivered) this.notifyDataChanged(); } this.completeRun(orgId, runId, result.lastRequestId, result.hitStepLimit); - this.cooldowns.clear(provider); if (task.origin === null) recordRunMemory(this.db, { orgId, runId, projectId: task.project_id, taskTitle: task.title, outcome: "completed", text: result.finalText }); this.audit.record({ orgId, diff --git a/tests/unit/provider-cooldown-backoff.test.ts b/tests/unit/provider-cooldown-backoff.test.ts index 0b32a2ce..29f68733 100644 --- a/tests/unit/provider-cooldown-backoff.test.ts +++ b/tests/unit/provider-cooldown-backoff.test.ts @@ -31,7 +31,9 @@ describe("provider cooldown sizing", () => { cooldowns.park(provider, "503", null); expect(parkedFor()).toBe(2 * 60_000); cooldowns.park(provider, "429", 5 * 3_600_000); - expect(parkedFor()).toBe(60 * 60_000); + expect(parkedFor()).toBe(5 * 3_600_000); + cooldowns.park(provider, "429", 30 * 3_600_000); + expect(parkedFor()).toBe(24 * 3_600_000); vi.restoreAllMocks(); connection.close(); }); diff --git a/tests/unit/rate-limit-observation.test.ts b/tests/unit/rate-limit-observation.test.ts new file mode 100644 index 00000000..51d31b26 --- /dev/null +++ b/tests/unit/rate-limit-observation.test.ts @@ -0,0 +1,49 @@ +import { randomUUID } from "node:crypto"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { describe, expect, test } from "vitest"; +import { DatabaseConnection } from "../../src/main/database/connection"; +import { OrganizationService } from "../../src/main/services/organization-service"; +import { ProviderService } from "../../src/main/services/provider-service"; +import { providerAdapters } from "../../src/main/services/provider-adapters"; +import { TaskExecutionService } from "../../src/main/services/task-execution-service"; +import { RISK_LEVEL, TASK_PRIORITY, TASK_STATUS } from "../../src/shared/domain"; +import type { AgentTurnResult } from "../../src/main/services/provider-tool-turn"; +import { claimNextQueuedRun } from "../../src/main/app/task-queue-worker"; + +const testVault = { + encryptSecret: (secret: string) => Buffer.from(secret, "utf8"), + decryptSecret: (ciphertext: Buffer) => ciphertext.toString("utf8"), +}; + +describe("rate limits seen on successful calls", () => { + test("a successful agent call records the provider's remaining quota, and an exhausted provider's work waits for its reset", async () => { + const resetAt = new Date(Date.now() + 10 * 60_000).toISOString(); + const turn = async (): Promise => ({ + text: "done", toolCalls: [], rawResponseJson: "{}", + rateLimitState: { requestLimit: 50, remainingRequests: 0, resetAt, retryAfterMs: null, rawStateJson: "{}" }, + }); + const connection = new DatabaseConnection(join(tmpdir(), `ac-ratelimit-${randomUUID()}.db`)); + const organization = new OrganizationService(connection.db); + const providers = new ProviderService(connection.db, providerAdapters, testVault, undefined, turn); + const execution = new TaskExecutionService(connection.db, providers); + let orgId = ""; + let providerId = ""; + connection.transaction(() => { + orgId = organization.createOrganization({ organizationName: "Acme", employeeName: "Dana", 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; + providers.saveCredential({ orgId, providerId, secret: "sk-test", model: "gpt-5.4" }); + const agentId = organization.createAgent({ orgId, name: "Ada", role: "Engineer", department: "Engineering", modelProviderId: providerId, modelId: "gpt-5.4", budgetCents: 1000, managerAgentId: null, heartbeatEnabled: true }); + organization.createTask({ orgId, title: "Work", description: "d", status: TASK_STATUS.BACKLOG, assigneeAgentId: agentId, risk: RISK_LEVEL.LOW, projectId: null, goalId: null, workspaceId: null, priority: TASK_PRIORITY.NORMAL, requiresFilesystem: false }); + }); + // The first session's response says no requests remain: it is recorded, and the provider is parked until the reset. + const [first] = connection.transaction(() => execution.enqueueHeartbeatRuns()); + connection.db.prepare("update task_runs set state = 'running' where id = ?").run(first); + await execution.executeRun(orgId, first); + expect(connection.db.prepare("select request_limit, remaining_requests, reset_at from provider_rate_limit_observations where provider_id = ?").get(providerId)) + .toEqual({ request_limit: 50, remaining_requests: 0, reset_at: resetAt }); + expect(execution.cooledProviderIds()).toEqual([providerId]); + expect(claimNextQueuedRun(connection, execution.cooledProviderIds())).toBeNull(); + connection.close(); + }); +});