diff --git a/src/main/services/provider-background-service.ts b/src/main/services/provider-background-service.ts index 0bc544b3..f7d2c52d 100644 --- a/src/main/services/provider-background-service.ts +++ b/src/main/services/provider-background-service.ts @@ -9,7 +9,7 @@ import { type ProviderAuthType, type ProviderExecutionMode, } from "../../shared/domain"; -import { currentProviderUsageWindow, PROVIDER_MONTHLY_LIMIT_MESSAGE, type ProviderObservedRateLimitState } from "../../shared/provider-runtime-domain"; +import type { ProviderObservedRateLimitState } from "../../shared/provider-runtime-domain"; import { AuditService } from "./audit-service"; import { AuthBrokerService, type ProviderSecretVault } from "./auth-broker-service"; import { OperationsService } from "./operations-service"; @@ -249,14 +249,7 @@ export class ProviderBackgroundService { throw new Error("Provider background execution is disabled by feature flag."); } if (provider.background_enabled !== 1) throw new Error("Background execution is disabled for this provider policy."); - const window = currentProviderUsageWindow(); - const usage = this.db.prepare(` - select request_count - from provider_usage_windows - where org_id = ? and provider_id = ? and window_start = ? - `).get(orgId, provider.id, window.start) as { request_count: number } | undefined; - const requestCount = usage?.request_count ?? 0; - if (requestCount >= provider.monthly_request_limit) throw new Error(PROVIDER_MONTHLY_LIMIT_MESSAGE); + this.requestStore.assertMonthlyCapacity(orgId, provider); } private normalizeWorkspaceId(orgId: string, workspaceId: string | null): string | null { diff --git a/src/main/services/provider-request-store.ts b/src/main/services/provider-request-store.ts index b8d966eb..b68fd5da 100644 --- a/src/main/services/provider-request-store.ts +++ b/src/main/services/provider-request-store.ts @@ -5,7 +5,7 @@ import { providerPrices, readTokenUsage, requestCostCents, type TokenUsage } fro /** Providers already logged as not reporting token usage. */ const providersWithoutUsage = new Set(); -import { currentProviderUsageWindow } from "../../shared/provider-runtime-domain"; +import { currentProviderUsageWindow, PROVIDER_MONTHLY_LIMIT_MESSAGE } from "../../shared/provider-runtime-domain"; export interface ProviderRequestOwner { id: string; @@ -272,6 +272,19 @@ export class ProviderRequestStore { return request; } + /** Refuses a call once finished plus in-flight calls reach the limit; callers create the request right after, before any await, as the reservation. */ + assertMonthlyCapacity(orgId: string, provider: { id: string; monthly_request_limit: number }): void { + const window = currentProviderUsageWindow(); + const finished = this.db.prepare("select request_count from provider_usage_windows where org_id = ? and provider_id = ? and window_start = ?") + .get(orgId, provider.id, window.start) as { request_count: number } | undefined; + // Startup recovery fails requests a crash left open, so every open row here is a live reservation. + const inFlight = this.db.prepare(` + select count(*) as n from provider_requests + where org_id = ? and provider_id = ? and state in (?, ?) and created_at >= ? + `).get(orgId, provider.id, PROVIDER_REQUEST_STATE.PENDING, PROVIDER_REQUEST_STATE.STREAMING, window.start) as { n: number }; + if ((finished?.request_count ?? 0) + inFlight.n >= provider.monthly_request_limit) throw new Error(PROVIDER_MONTHLY_LIMIT_MESSAGE); + } + /** Settles a successful call: reads the provider's token usage, prices it, and adds it to the request and the month's totals. */ settleUsage(input: { orgId: string; requestId: string; provider: { id: string; provider_key: string }; rawResponseJson: string; inputCharCount: number; outputCharCount: number }): void { const usage = readTokenUsage(input.rawResponseJson); diff --git a/src/main/services/provider-response-runner.ts b/src/main/services/provider-response-runner.ts index abea2dce..8ca0e7cf 100644 --- a/src/main/services/provider-response-runner.ts +++ b/src/main/services/provider-response-runner.ts @@ -3,7 +3,6 @@ import type { DatabaseSync } from "node:sqlite"; import type { CancelProviderResponseInput, CreateProviderResponseInput, StreamProviderResponseInput } from "../../shared/commands"; import { FEATURE_FLAG_KEY, PROVIDER_AUTH_TYPE, PROVIDER_DATA_POLICY, PROVIDER_EXECUTION_MODE, PROVIDER_REQUEST_STATE, type ProviderExecutionMode } from "../../shared/domain"; import type { ProviderStreamRendererEvent } from "../../shared/provider-request-domain"; -import { currentProviderUsageWindow, PROVIDER_MONTHLY_LIMIT_MESSAGE } from "../../shared/provider-runtime-domain"; import type { AuditService } from "./audit-service"; import type { AuthBrokerService } from "./auth-broker-service"; import type { OperationsService } from "./operations-service"; @@ -128,7 +127,6 @@ export class ProviderResponseRunner { : null; const blocked = agent ? budgetBlock(this.deps.db, agent.id, isPriced(providerPrices(this.deps.db, provider.id)), new Date()) : null; if (blocked) throw new Error(blocked); - const credential = await resolveTurnCredential({ authBroker: this.deps.authBroker, integrationKeys: this.deps.integrationKeys, orgId: input.orgId, provider }); // Everything the model reads: instructions, the conversation, and the tool definitions. const turnInputChars = (input.instructions?.length ?? 0) + JSON.stringify(input.messages).length + JSON.stringify(input.tools).length; const requestId = this.deps.requestStore.create({ @@ -146,6 +144,7 @@ export class ProviderResponseRunner { }); const startedAt = Date.now(); 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 }); // 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); @@ -356,14 +355,7 @@ export class ProviderResponseRunner { if (executionMode === PROVIDER_EXECUTION_MODE.BACKGROUND && provider.background_enabled !== 1) { throw new Error("Background execution is disabled for this provider policy."); } - const window = currentProviderUsageWindow(); - const usage = this.deps.db.prepare(` - select request_count - from provider_usage_windows - where org_id = ? and provider_id = ? and window_start = ? - `).get(orgId, provider.id, window.start) as { request_count: number } | undefined; - const requestCount = usage?.request_count ?? 0; - if (requestCount >= provider.monthly_request_limit) throw new Error(PROVIDER_MONTHLY_LIMIT_MESSAGE); + this.deps.requestStore.assertMonthlyCapacity(orgId, provider); } private normalizeWorkspaceId(orgId: string, workspaceId: string | null): string | null { diff --git a/tests/unit/quota-reservation.test.ts b/tests/unit/quota-reservation.test.ts new file mode 100644 index 00000000..9963137b --- /dev/null +++ b/tests/unit/quota-reservation.test.ts @@ -0,0 +1,48 @@ +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 { PROVIDER_MONTHLY_LIMIT_MESSAGE } from "../../src/shared/provider-runtime-domain"; +import { PROVIDER_DATA_POLICY } from "../../src/shared/domain"; +import type { AgentTurnResult } from "../../src/main/services/provider-tool-turn"; + +const testVault = { + encryptSecret: (secret: string) => Buffer.from(secret, "utf8"), + decryptSecret: (ciphertext: Buffer) => ciphertext.toString("utf8"), +}; + +describe("monthly quota reservation", () => { + test("with one request left, concurrent calls cannot both take it", async () => { + let release = () => {}; + const gate = new Promise((resolve) => { release = resolve; }); + let started = 0; + const turn = async (): Promise => { + started += 1; + await gate; + return { text: "done", toolCalls: [], rawResponseJson: "{}" }; + }; + const connection = new DatabaseConnection(join(tmpdir(), `ac-quota-${randomUUID()}.db`)); + const providers = new ProviderService(connection.db, providerAdapters, testVault, undefined, turn); + let orgId = ""; + let providerId = ""; + connection.transaction(() => { + orgId = new OrganizationService(connection.db).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" }); + providers.updatePolicy({ orgId, providerId, dataPolicy: PROVIDER_DATA_POLICY.STANDARD, backgroundEnabled: false, monthlyRequestLimit: 1, workspaceScope: "organization", contextWindowTokens: 128_000 }); + }); + const call = () => providers.runToolTurn({ orgId, providerId, workspaceId: null, instructions: "i", messages: [{ role: "user", text: "u" }], tools: [] }); + const first = call(); + const second = call(); + await expect(second).rejects.toThrow(PROVIDER_MONTHLY_LIMIT_MESSAGE); + release(); + await first; + expect(started).toBe(1); + await expect(call()).rejects.toThrow(PROVIDER_MONTHLY_LIMIT_MESSAGE); + connection.close(); + }); +});