diff --git a/src/lib/db/providers.ts b/src/lib/db/providers.ts index 4d95a6f484..a59e178cad 100644 --- a/src/lib/db/providers.ts +++ b/src/lib/db/providers.ts @@ -572,6 +572,61 @@ export async function updateProviderConnection(id: string, data: JsonRecord) { ); } +/** + * Atomic conditional clear of recoverable error state on a connection row. + * + * Returns true when the row was cleared, false when a concurrent writer + * (markAccountUnavailable, connectionRecovery tick, test, etc.) changed the + * row between the caller's snapshot read and this UPDATE — in which case the + * clear is skipped to preserve the freshest error state. Closes the TOCTOU + * window in the quota-recovery path. + * + * CAS token = (test_status, last_error_at, rate_limited_until). + * markAccountUnavailable always bumps last_error_at on every cooldown/error + * write, so an unchanged last_error_at reliably indicates no concurrent write. + */ +export async function clearConnectionErrorIfUnchanged( + id: string, + expected: { + testStatus: string | null | undefined; + lastErrorAt: string | null | undefined; + rateLimitedUntil: string | null | undefined; + } +): Promise { + const db = getDbInstance() as unknown as DbLike; + const result = db.prepare( + ` + UPDATE provider_connections SET + test_status = 'active', + last_error = NULL, + last_error_at = NULL, + last_error_type = NULL, + last_error_source = NULL, + error_code = NULL, + rate_limited_until = NULL, + backoff_level = 0, + updated_at = ? + WHERE id = ? + AND IFNULL(test_status, '') = ? + AND IFNULL(last_error_at, '') = ? + AND IFNULL(rate_limited_until, '') = ? + ` + ).run( + new Date().toISOString(), + id, + expected.testStatus ?? "", + expected.lastErrorAt ?? "", + expected.rateLimitedUntil ?? "" + ); + const applied = (result.changes ?? 0) > 0; + if (applied) { + backupDbFile("pre-write"); + invalidateDbCache("connections"); + bumpProxyConfigGeneration(); + } + return applied; +} + export async function deleteProviderConnection(id: string) { const db = getDbInstance() as unknown as DbLike; const existing = db.prepare("SELECT provider FROM provider_connections WHERE id = ?").get(id); diff --git a/src/lib/localDb.ts b/src/lib/localDb.ts index 366d899ed5..d911599723 100755 --- a/src/lib/localDb.ts +++ b/src/lib/localDb.ts @@ -12,6 +12,7 @@ export { getProviderConnectionById, createProviderConnection, updateProviderConnection, + clearConnectionErrorIfUnchanged, deleteProviderConnection, deleteProviderConnections, deleteProviderConnectionsByProvider, diff --git a/src/lib/usage/providerLimits.ts b/src/lib/usage/providerLimits.ts index cf02868f2d..9b529ca8ba 100644 --- a/src/lib/usage/providerLimits.ts +++ b/src/lib/usage/providerLimits.ts @@ -440,21 +440,37 @@ export async function maybeClearRecoveredQuotaState( if (!hasTransientState) return connection; + let cleared = true; try { - await clearRecoveredProviderState({ - connectionId: connection.id, - testStatus: connection.testStatus, - lastError: connection.lastError ?? null, - rateLimitedUntil: connection.rateLimitedUntil ?? null, - errorCode: connection.errorCode ?? null, - lastErrorType: connection.lastErrorType ?? null, - lastErrorSource: connection.lastErrorSource ?? null, - }); + const result = await clearRecoveredProviderState( + { + connectionId: connection.id, + testStatus: connection.testStatus, + lastError: connection.lastError ?? null, + rateLimitedUntil: connection.rateLimitedUntil ?? null, + errorCode: connection.errorCode ?? null, + lastErrorType: connection.lastErrorType ?? null, + lastErrorSource: connection.lastErrorSource ?? null, + }, + { + testStatus: connection.testStatus ?? null, + lastErrorAt: connection.lastErrorAt ?? null, + rateLimitedUntil: connection.rateLimitedUntil ?? null, + } + ); + cleared = result.applied; } catch (dbError) { console.warn("[ProviderLimits] Failed to clear recovered quota state:", dbError); return connection; } + if (!cleared) { + // CAS miss — a concurrent writer (markAccountUnavailable, etc.) updated + // the row between our read and the clear. Return the original snapshot; + // the next read from DB will surface the fresh state. + return connection; + } + return { ...connection, testStatus: "active", diff --git a/src/sse/services/auth.ts b/src/sse/services/auth.ts index 0158be94d7..6c07cf7b66 100644 --- a/src/sse/services/auth.ts +++ b/src/sse/services/auth.ts @@ -4,6 +4,7 @@ import { getProviderNodes, validateApiKey, updateProviderConnection, + clearConnectionErrorIfUnchanged, getSettings, getCachedSettings, } from "@/lib/localDb"; @@ -2266,11 +2267,41 @@ export async function clearAccountError( log.info("AUTH", `Account ${connectionId.slice(0, 8)} error cleared`); } +/** + * Optional CAS token. When provided, the clear is performed via an atomic + * conditional UPDATE (clearConnectionErrorIfUnchanged) that aborts if the row + * was written by a concurrent path between the caller's snapshot read and this + * clear. Closes the TOCTOU window in the quota-recovery path. When omitted, + * the clear is unconditional (preserves existing post-success-call behavior). + */ +export interface RecoveredStateExpectation { + testStatus: string | null; + lastErrorAt: string | null; + rateLimitedUntil: string | null; +} + export async function clearRecoveredProviderState( - credentials: Partial | null -) { - if (!credentials?.connectionId) return; + credentials: Partial | null, + expectedState?: RecoveredStateExpectation +): Promise<{ applied: boolean }> { + if (!credentials?.connectionId) return { applied: false }; + if (expectedState) { + const applied = await clearConnectionErrorIfUnchanged( + credentials.connectionId, + expectedState + ); + if (!applied) { + log.info( + "AUTH", + `Skipped recovery clear for ${credentials.connectionId.slice(0, 8)} — state changed concurrently (CAS miss)` + ); + return { applied: false }; + } + log.info("AUTH", `Account ${credentials.connectionId.slice(0, 8)} error cleared (CAS)`); + return { applied: true }; + } await clearAccountError(credentials.connectionId, credentials); + return { applied: true }; } type AuthRequestHeaders = Headers | Record; diff --git a/tests/unit/provider-limits-recovery.test.ts b/tests/unit/provider-limits-recovery.test.ts index f719d19090..ae25b84ab5 100644 --- a/tests/unit/provider-limits-recovery.test.ts +++ b/tests/unit/provider-limits-recovery.test.ts @@ -201,3 +201,104 @@ test("error-only quota response does not clear transient state", async () => { assert.equal(updated.testStatus, "unavailable", "transient state should not be cleared on error"); assert.equal(updated.lastErrorType, "rate_limited"); }); + +test("CAS primitive clears when expected state matches", async () => { + const created = await createGlmConnectionWithTransientCooldown(); + const connectionId = (created as { id: string }).id; + const before = (await providersDb.getProviderConnectionById(connectionId)) as Record< + string, + unknown + >; + + const applied = await providersDb.clearConnectionErrorIfUnchanged(connectionId, { + testStatus: (before.testStatus as string) ?? null, + lastErrorAt: (before.lastErrorAt as string) ?? null, + rateLimitedUntil: (before.rateLimitedUntil as string) ?? null, + }); + + assert.equal(applied, true, "CAS UPDATE should apply when expected state matches"); + const after = (await providersDb.getProviderConnectionById(connectionId)) as Record< + string, + unknown + >; + assert.equal(after.testStatus, "active"); + assert.equal(after.rateLimitedUntil, undefined); + assert.equal(after.backoffLevel, 0); +}); + +test("CAS primitive aborts when state changed concurrently", async () => { + const created = await createGlmConnectionWithTransientCooldown(); + const connectionId = (created as { id: string }).id; + const before = (await providersDb.getProviderConnectionById(connectionId)) as Record< + string, + unknown + >; + + // Simulate a concurrent markAccountUnavailable writing a fresh error state. + const newLastErrorAt = new Date(Date.now() + 1000).toISOString(); + const newRateLimitedUntil = new Date(Date.now() + 120_000).toISOString(); + await providersDb.updateProviderConnection(connectionId, { + lastErrorAt: newLastErrorAt, + rateLimitedUntil: newRateLimitedUntil, + lastError: "fresh 429", + errorCode: 429, + backoffLevel: 3, + }); + + const applied = await providersDb.clearConnectionErrorIfUnchanged(connectionId, { + testStatus: (before.testStatus as string) ?? null, + lastErrorAt: (before.lastErrorAt as string) ?? null, + rateLimitedUntil: (before.rateLimitedUntil as string) ?? null, + }); + + assert.equal(applied, false, "CAS UPDATE should abort when state changed"); + const after = (await providersDb.getProviderConnectionById(connectionId)) as Record< + string, + unknown + >; + assert.equal(after.testStatus, "unavailable", "fresh mark should be preserved"); + assert.equal(after.backoffLevel, 3, "fresh backoff level should be preserved"); + assert.equal(after.lastError, "fresh 429"); +}); + +test("quota recovery path does NOT overwrite a concurrent mark (TOCTOU closed)", async () => { + const created = await createGlmConnectionWithTransientCooldown(); + const connectionId = (created as { id: string }).id; + const snapshotBeforeClear = (await providersDb.getProviderConnectionById( + connectionId + )) as Record; + const expectedLastErrorAt = (snapshotBeforeClear.lastErrorAt as string) ?? null; + + // Mock fetch so that DURING the quota fetch (between read and clear), a + // concurrent mark writes a fresh error state. This deterministically + // reproduces the TOCTOU window the CAS primitive is meant to close. + const concurrentMarkFetch = (() => { + // Simulate concurrent markAccountUnavailable writing fresh state. + providersDb.updateProviderConnection(connectionId, { + lastErrorAt: new Date(Date.now() + 1000).toISOString(), + rateLimitedUntil: new Date(Date.now() + 120_000).toISOString(), + lastError: "fresh concurrent 429", + errorCode: 429, + backoffLevel: 3, + }); + return glmQuotaResponse(); + }) as typeof fetch; + + await withMockedFetch(concurrentMarkFetch, async () => { + await providerLimits.fetchAndPersistProviderLimits(connectionId, "manual"); + }); + + const after = (await providersDb.getProviderConnectionById(connectionId)) as Record< + string, + unknown + >; + // Recovery should have aborted (CAS miss) — fresh mark must survive. + assert.notEqual( + after.lastErrorAt, + expectedLastErrorAt, + "fresh lastErrorAt must not be overwritten by recovery clear" + ); + assert.equal(after.testStatus, "unavailable", "fresh testStatus must survive"); + assert.equal(after.backoffLevel, 3, "fresh backoff level must survive"); + assert.equal(after.lastError, "fresh concurrent 429"); +});