diff --git a/src/main/database/migrations/075-provider-cooldowns.ts b/src/main/database/migrations/075-provider-cooldowns.ts new file mode 100644 index 00000000..43935c85 --- /dev/null +++ b/src/main/database/migrations/075-provider-cooldowns.ts @@ -0,0 +1,12 @@ +export const providerCooldownsSchema = ` + create table if not exists provider_cooldowns ( + provider_id text primary key references providers(id) on delete cascade, + label text not null, + until_ms integer not null, + strikes integer not null, + interactive_probed_at_ms integer + ); + + insert into schema_migrations(version, name, applied_at) + values (75, 'provider_cooldowns', datetime('now')); +`; diff --git a/src/main/database/schema.ts b/src/main/database/schema.ts index bee8f601..4db3b6a8 100644 --- a/src/main/database/schema.ts +++ b/src/main/database/schema.ts @@ -73,6 +73,7 @@ import { backgroundPollFailuresSchema } from "./migrations/071-background-poll-f import { spendSchema } from "./migrations/072-spend"; import { usageWindowTokensSchema } from "./migrations/073-usage-window-tokens"; import { runFailureKindSchema } from "./migrations/074-run-failure-kind"; +import { providerCooldownsSchema } from "./migrations/075-provider-cooldowns"; interface Migration { version: number; @@ -155,6 +156,7 @@ const MIGRATIONS: Migration[] = [ { version: 72, name: "spend", sql: spendSchema }, { version: 73, name: "usage_window_tokens", sql: usageWindowTokensSchema }, { version: 74, name: "run_failure_kind", sql: runFailureKindSchema }, + { version: 75, name: "provider_cooldowns", sql: providerCooldownsSchema }, ]; const SCHEMA_VERSION = MIGRATIONS[MIGRATIONS.length - 1].version; diff --git a/src/main/services/provider-cooldowns.ts b/src/main/services/provider-cooldowns.ts index f825a14d..64d533ec 100644 --- a/src/main/services/provider-cooldowns.ts +++ b/src/main/services/provider-cooldowns.ts @@ -1,3 +1,5 @@ +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; @@ -5,78 +7,73 @@ const BACKOFF_CAP_MS = 30 * 60_000; const RETRY_AFTER_MIN_MS = 30_000; const RETRY_AFTER_MAX_MS = 60 * 60_000; -/** In-memory parks for providers that reported a transient failure. */ +interface CooldownRow { + provider_id: string; + label: string; + until_ms: number; + strikes: number; + interactive_probed_at_ms: number | null; +} + +/** Parks for providers that reported a transient failure, stored so a restart keeps waiting instead of hammering them. */ export class ProviderCooldowns { - private readonly parks = new Map(); - /** Transient failures in a row per provider, reset by a served request. */ - private readonly strikes = new Map(); + constructor(private readonly db: DatabaseSync, private readonly now: () => number) {} - constructor(private readonly now: () => number) {} - - /** - * Parks a provider after it reported a transient failure (usage limit, rate - * limit, overload), so the queue stops feeding it doomed runs every tick. - * The park is in-memory by design: after a restart the first probe run - * either succeeds or re-arms the cooldown at the cost of one request. - */ + /** Parks a provider after a transient failure, for the wait it gave or a growing backoff. */ park(provider: { id: string; label: string }, message: string, retryAfterMs: number | null): void { - const active = this.parks.get(provider.id); - const wasParked = active !== undefined && active.until > this.now(); - const strikes = (this.strikes.get(provider.id) ?? 0) + 1; - this.strikes.set(provider.id, strikes); + const row = this.row(provider.id); + const wasParked = row !== undefined && row.until_ms > this.now(); + const strikes = (row?.strikes ?? 0) + 1; const waitMs = retryAfterMs !== null ? Math.min(Math.max(retryAfterMs, RETRY_AFTER_MIN_MS), RETRY_AFTER_MAX_MS) : Math.min(BACKOFF_BASE_MS * 2 ** (strikes - 1), BACKOFF_CAP_MS); - this.parks.set(provider.id, { - until: this.now() + waitMs, - label: provider.label, - // A failed probe must not unlock endless probes: re-arming the park - // keeps the spent-probe mark until the window expires naturally. - interactiveProbedAt: wasParked ? active?.interactiveProbedAt : undefined, - }); + // A failed probe must not unlock endless probes: re-arming the park keeps the spent-probe mark. + const probedAt = wasParked ? row.interactive_probed_at_ms : null; + this.db.prepare(` + insert into provider_cooldowns(provider_id, label, until_ms, strikes, interactive_probed_at_ms) values (?, ?, ?, ?, ?) + on conflict(provider_id) do update set label = excluded.label, until_ms = excluded.until_ms, strikes = excluded.strikes, + interactive_probed_at_ms = excluded.interactive_probed_at_ms + `).run(provider.id, provider.label, this.now() + waitMs, strikes, probedAt); console.warn(`[task-queue] provider "${provider.label}" reported a transient failure ("${message}"); its tasks are parked for ${Math.round(waitMs / 1000)}s (failure ${strikes} in a row).`); } - /** - * Providers currently inside a transient-failure cooldown. Expired parks are - * pruned (with a resume log line) so the next tick probes the provider again. - */ + /** Providers currently parked; expired parks end (with a log line) but keep their strike count until a request is served. */ cooledProviderIds(): string[] { + const rows = this.db.prepare("select provider_id, label, until_ms from provider_cooldowns where until_ms > 0").all() as Array<{ provider_id: string; label: string; until_ms: number }>; const active: string[] = []; - for (const [id, park] of this.parks) { - if (park.until > this.now()) { - active.push(id); - } else { - this.parks.delete(id); - console.info(`[task-queue] provider "${park.label}" cooldown ended; its tasks resume.`); + for (const row of rows) { + if (row.until_ms > this.now()) { + active.push(row.provider_id); + continue; } + this.db.prepare("update provider_cooldowns set until_ms = 0, interactive_probed_at_ms = null where provider_id = ?").run(row.provider_id); + console.info(`[task-queue] provider "${row.label}" cooldown ended; its tasks resume.`); } return active; } - /** - * Cooled providers that even the interactive lane must skip: their one - * founder-probe for this park window is already spent. A cooled provider - * not yet probed stays interactively claimable — one human message is - * allowed through to learn whether the outage was a blip, instead of the - * founder sitting out the whole cooldown for a single failed background - * request. - */ + /** Parked providers whose one founder-probe for this park is already spent; the interactive lane skips only these. */ interactiveCooledProviderIds(): string[] { - return this.cooledProviderIds().filter((id) => this.parks.get(id)?.interactiveProbedAt !== undefined); + const active = new Set(this.cooledProviderIds()); + const probed = this.db.prepare("select provider_id from provider_cooldowns where interactive_probed_at_ms is not null").all() as Array<{ provider_id: string }>; + return probed.map((row) => row.provider_id).filter((id) => active.has(id)); } - /** Marks the cooled provider's founder-probe as spent (no-op when healthy). */ + /** Marks the parked provider's founder-probe as spent (no-op when healthy). */ noteInteractiveProbe(providerId: string): void { - const park = this.parks.get(providerId); - if (park) park.interactiveProbedAt = this.now(); + this.db.prepare("update provider_cooldowns set interactive_probed_at_ms = ? where provider_id = ? and until_ms > ?").run(this.now(), providerId, this.now()); } - /** A served request is proof of health: the park ends for everyone at once. */ + /** A served request is proof of health: the park and the failure streak end. */ clear(provider: { id: string; label: string }): void { - this.strikes.delete(provider.id); - if (this.parks.delete(provider.id)) { - console.info(`[task-queue] provider "${provider.label}" served a request; cooldown lifted early.`); - } + const row = this.row(provider.id); + if (!row) return; + this.db.prepare("delete from provider_cooldowns where provider_id = ?").run(provider.id); + if (row.until_ms > this.now()) console.info(`[task-queue] provider "${provider.label}" served a request; cooldown lifted early.`); + } + + private row(providerId: string): CooldownRow | undefined { + return this.db.prepare("select provider_id, label, until_ms, strikes, interactive_probed_at_ms from provider_cooldowns where provider_id = ?") + .get(providerId) as CooldownRow | undefined; } } diff --git a/src/main/services/task-execution-service.ts b/src/main/services/task-execution-service.ts index 560b243d..15f65e27 100644 --- a/src/main/services/task-execution-service.ts +++ b/src/main/services/task-execution-service.ts @@ -83,7 +83,7 @@ export class TaskExecutionService { this.audit = new AuditService(db); this.organization = new OrganizationService(db); this.chat = new ChatService(db); - this.cooldowns = new ProviderCooldowns(now); + this.cooldowns = new ProviderCooldowns(db, now); this.heartbeat = new HeartbeatScheduler(db, this.cooldowns, now, notifyDataChanged); } diff --git a/tests/unit/domain-services.test.ts b/tests/unit/domain-services.test.ts index fd1973ea..df2b53b8 100644 --- a/tests/unit/domain-services.test.ts +++ b/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, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74]); + 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, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75]); connection.close(); }); diff --git a/tests/unit/provider-cooldown-backoff.test.ts b/tests/unit/provider-cooldown-backoff.test.ts index 14fc3bbf..0b32a2ce 100644 --- a/tests/unit/provider-cooldown-backoff.test.ts +++ b/tests/unit/provider-cooldown-backoff.test.ts @@ -1,13 +1,18 @@ +import { randomUUID } from "node:crypto"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; import { describe, expect, test, vi } from "vitest"; +import { DatabaseConnection } from "../../src/main/database/connection"; +import { OrganizationService } from "../../src/main/services/organization-service"; import { ProviderCooldowns } from "../../src/main/services/provider-cooldowns"; describe("provider cooldown sizing", () => { test("a provider-given wait is honored; otherwise parks back off and reset once a request is served", () => { vi.spyOn(console, "warn").mockImplementation(() => {}); vi.spyOn(console, "info").mockImplementation(() => {}); + const { connection, provider } = company(); let now = 0; - const cooldowns = new ProviderCooldowns(() => now); - const provider = { id: "prv_1", label: "OpenAI" }; + const cooldowns = new ProviderCooldowns(connection.db, () => now); const parkedFor = (): number => { let elapsed = 0; while (cooldowns.cooledProviderIds().includes(provider.id)) { @@ -28,5 +33,36 @@ describe("provider cooldown sizing", () => { cooldowns.park(provider, "429", 5 * 3_600_000); expect(parkedFor()).toBe(60 * 60_000); vi.restoreAllMocks(); + connection.close(); + }); + + test("a park survives a restart, so the queue does not hammer a provider that is still down", () => { + vi.spyOn(console, "warn").mockImplementation(() => {}); + vi.spyOn(console, "info").mockImplementation(() => {}); + const { connection, provider } = company(); + let now = 1_000_000; + new ProviderCooldowns(connection.db, () => now).park(provider, "503", null); + new ProviderCooldowns(connection.db, () => now).park(provider, "503", null); + const afterRestart = new ProviderCooldowns(connection.db, () => now); + expect(afterRestart.cooledProviderIds()).toEqual([provider.id]); + now += 4 * 60_000 - 1; + expect(afterRestart.cooledProviderIds()).toEqual([provider.id]); + now += 1; + expect(afterRestart.cooledProviderIds()).toEqual([]); + afterRestart.park(provider, "503", null); + now += 8 * 60_000 - 1; + expect(new ProviderCooldowns(connection.db, () => now).cooledProviderIds()).toEqual([provider.id]); + vi.restoreAllMocks(); + connection.close(); }); }); + +function company() { + const connection = new DatabaseConnection(join(tmpdir(), `ac-cooldown-${randomUUID()}.db`)); + let orgId = ""; + connection.transaction(() => { + orgId = new OrganizationService(connection.db).createOrganization({ organizationName: "Acme", employeeName: "Dana", mission: "m", goal: "g" }); + }); + const provider = connection.db.prepare("select id, label from providers where org_id = ? and provider_key = 'openai'").get(orgId) as { id: string; label: string }; + return { connection, provider }; +}