diff --git a/packages/web-shared/src/lib/ai-gateway/openai-chatgpt/refresh-terminal-race.test.ts b/packages/web-shared/src/lib/ai-gateway/openai-chatgpt/refresh-terminal-race.test.ts new file mode 100644 index 00000000000..1c5a2574265 --- /dev/null +++ b/packages/web-shared/src/lib/ai-gateway/openai-chatgpt/refresh-terminal-race.test.ts @@ -0,0 +1,134 @@ +jest.mock('@kilocode/web-shared/lib/config.server', () => ({ + BYOK_ENCRYPTION_KEY: process.env.BYOK_ENCRYPTION_KEY, + OPENAI_CLIENT_ID: 'test-client-id', + OPENAI_CLIENT_SECRET: 'test-client-secret', +})); + +import { afterEach, beforeAll, beforeEach, describe, expect, it } from '@jest/globals'; +import { randomUUID } from 'crypto'; +import { insertTestUser } from '@kilocode/web-shared/tests/helpers/user.helper'; +import { OPENAI_CHATGPT_RECONNECT_MESSAGE, resolveOpenAiChatGptAccessToken } from './refresh'; +import { + getOpenAiChatGptStoredConnection, + saveOpenAiChatGptConnection, + type OpenAiChatGptOwner, +} from './store'; +import type { OpenAiChatGptConnection } from './types'; + +function nowSeconds(): number { + return Math.floor(Date.now() / 1000); +} + +function connection(overrides: Partial = {}): OpenAiChatGptConnection { + return { + access_token: 'stale-access-token', + refresh_token: 'stale-refresh-token', + expires_at: nowSeconds() - 10, + scope: 'openid profile email offline_access', + token_type: 'Bearer', + issuer: 'https://auth.openai.com', + client_id: 'test-client-id', + subject: 'subject-1', + email: 'user@example.com', + connected_at: '2026-09-16T00:00:00.000Z', + status: 'connected', + ...overrides, + }; +} + +const TERMINAL_REFRESH_RESPONSE = { + ok: false, + status: 400, + headers: new Headers(), + json: async () => ({ error: 'invalid_grant' }), +} as unknown as Response; + +describe('terminal refresh failure against a real database', () => { + let owner: OpenAiChatGptOwner; + const fetchMock = jest.fn(); + const originalFetch = global.fetch; + + beforeAll(async () => { + const user = await insertTestUser({ + google_user_email: `chatgpt-refresh-terminal-${randomUUID()}@example.com`, + }); + owner = { kiloUserId: user.id, organizationId: null }; + }); + + beforeEach(() => { + fetchMock.mockReset(); + global.fetch = fetchMock as unknown as typeof fetch; + }); + + afterEach(() => { + global.fetch = originalFetch; + }); + + it('disables and clears the connection whose refresh token OpenAI rejected', async () => { + await saveOpenAiChatGptConnection(owner, connection(), owner.kiloUserId as string); + fetchMock.mockResolvedValue(TERMINAL_REFRESH_RESPONSE); + + await expect(resolveOpenAiChatGptAccessToken(owner)).resolves.toEqual({ kind: 'terminal' }); + + const request = fetchMock.mock.calls[0][1] as RequestInit; + expect(String(request.body)).toContain('stale-refresh-token'); + + const stored = await getOpenAiChatGptStoredConnection(owner); + expect(stored?.isEnabled).toBe(false); + expect(stored?.connection.access_token).toBe(''); + expect(stored?.connection.refresh_token).toBeUndefined(); + expect(stored?.connection.status).toBe('error'); + expect(stored?.connection.error_message).toBe(OPENAI_CHATGPT_RECONNECT_MESSAGE); + }); + + it('does not erase a connection reconnected while the failing refresh held the row lock', async () => { + await saveOpenAiChatGptConnection(owner, connection(), owner.kiloUserId as string); + + const fetchGate = Promise.withResolvers(); + const fetchCalled = Promise.withResolvers(); + + // The token exchange runs inside the transaction that holds the row lock. + // Parking it there keeps the lock held while the reconnect below is issued. + fetchMock.mockImplementation(async () => { + fetchCalled.resolve(); + await fetchGate.promise; + return TERMINAL_REFRESH_RESPONSE; + }); + + const refresh = resolveOpenAiChatGptAccessToken(owner); + await fetchCalled.promise; + + let reconnected = false; + const reconnect = saveOpenAiChatGptConnection( + owner, + connection({ + access_token: 'reconnected-access-token', + refresh_token: 'reconnected-refresh-token', + expires_at: nowSeconds() + 3600, + }), + owner.kiloUserId as string + ).then(() => { + reconnected = true; + }); + + // A blocked row-lock wait exposes no promise to await, and deterministic + // time control cannot apply to a real Postgres lock, so this one bounded + // wait observes the wait state directly. The lock holds it well past this. + await new Promise(resolve => setTimeout(resolve, 250)); + // The row lock keeps the reconnect from landing between the failing request + // and the terminal write, so the write cannot erase it. + expect(reconnected).toBe(false); + + fetchGate.resolve(); + await expect(refresh).resolves.toEqual({ kind: 'terminal' }); + await reconnect; + + const stored = await getOpenAiChatGptStoredConnection(owner); + expect(stored?.isEnabled).toBe(true); + expect(stored?.connection).toMatchObject({ + access_token: 'reconnected-access-token', + refresh_token: 'reconnected-refresh-token', + status: 'connected', + }); + }); +}); diff --git a/packages/web-shared/src/lib/ai-gateway/openai-chatgpt/store.test.ts b/packages/web-shared/src/lib/ai-gateway/openai-chatgpt/store.test.ts index 9660316e131..4cb9daef114 100644 --- a/packages/web-shared/src/lib/ai-gateway/openai-chatgpt/store.test.ts +++ b/packages/web-shared/src/lib/ai-gateway/openai-chatgpt/store.test.ts @@ -21,7 +21,6 @@ import { clearOpenAiChatGptConnection, getOpenAiChatGptConnection, getOpenAiChatGptStoredConnection, - markOpenAiChatGptError, saveOpenAiChatGptConnection, type OpenAiChatGptOwner, } from './store'; @@ -242,35 +241,22 @@ describe('openai-chatgpt connection store', () => { await expect(getOpenAiChatGptConnection(USER_OWNER)).resolves.toBeNull(); }); - it('markOpenAiChatGptError clears the token set and disables the connection', async () => { - await saveOpenAiChatGptConnection(USER_OWNER, buildConnection(), TEST_USER_ID); - - await markOpenAiChatGptError(USER_OWNER, 'Your ChatGPT connection has expired.'); - - const stored = await getOpenAiChatGptConnection(USER_OWNER); - expect(stored).toMatchObject({ + it('clears a previous error state when the connection is saved again', async () => { + const errored = { + ...buildConnection(), access_token: '', expires_at: 0, - status: 'error', - error_message: 'Your ChatGPT connection has expired.', + status: 'error' as const, + error_message: 'expired', + error_at: new Date().toISOString(), + }; + rows.push({ + kilo_user_id: TEST_USER_ID, + organization_id: null, + encrypted_connection: encryptApiKey(JSON.stringify(errored), BYOK_ENCRYPTION_KEY), + is_enabled: false, + created_by: TEST_USER_ID, }); - expect(stored?.refresh_token).toBeUndefined(); - // The identity the card renders survives the cleared credential. - expect(stored?.email).toBe('user@example.com'); - expect(stored?.error_at).toBeTruthy(); - expect(rows[0].is_enabled).toBe(false); - }); - - it('markOpenAiChatGptError is a no-op when no row exists', async () => { - await markOpenAiChatGptError(USER_OWNER, 'Your ChatGPT connection has expired.'); - - expect(rows).toHaveLength(0); - expect(mockDb.update).not.toHaveBeenCalled(); - }); - - it('clears a previous error state when the connection is saved again', async () => { - await saveOpenAiChatGptConnection(USER_OWNER, buildConnection(), TEST_USER_ID); - await markOpenAiChatGptError(USER_OWNER, 'expired'); await saveOpenAiChatGptConnection(USER_OWNER, buildConnection(), TEST_USER_ID); diff --git a/packages/web-shared/src/lib/ai-gateway/openai-chatgpt/store.ts b/packages/web-shared/src/lib/ai-gateway/openai-chatgpt/store.ts index 866c60a16eb..3cf9ed11294 100644 --- a/packages/web-shared/src/lib/ai-gateway/openai-chatgpt/store.ts +++ b/packages/web-shared/src/lib/ai-gateway/openai-chatgpt/store.ts @@ -322,21 +322,6 @@ export async function markOpenAiChatGptConnectionErrored( .where(openAiChatGptOwnerWhere(owner)); } -/** - * Terminal-failure write for callers outside a transaction. A missing row is a - * no-op. The refresh path uses `markOpenAiChatGptConnectionErrored` directly so - * the write stays inside the row lock. - */ -export async function markOpenAiChatGptError( - owner: OpenAiChatGptOwner, - message: string -): Promise { - const connection = await getOpenAiChatGptConnection(owner); - if (!connection) return; - - await markOpenAiChatGptConnectionErrored(db, owner, connection, message); -} - /** * Records OpenAI's refusal of the connected account. Only the connection whose * access token served the refused request is marked: a reconnect that landed