Q67: in-flight requests reserve monthly quota

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

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 {

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<string>();
-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);

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 {

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<void>((resolve) => { release = resolve; });
+    let started = 0;
+    const turn = async (): Promise<AgentTurnResult> => {
+      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();
+  });
+});