From 401e78e01c54b3b0fc69ae5137aebb1c8b6c42d6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Marc-Andr=C3=A9=20Bombeck?= Date: Mon, 6 Apr 2026 00:22:54 +0200 Subject: [PATCH] =?UTF-8?q?feat(ai):=20add=20provider=20resolution=20with?= =?UTF-8?q?=20codex=20=E2=86=92=20admin-key=20=E2=86=92=20none=20hierarchy?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-Authored-By: Claude Opus 4.6 (1M context) --- src/lib/ai/__tests__/provider.test.ts | 94 +++++++++++++++++++ src/lib/ai/provider.ts | 128 ++++++++++++++++++++++++++ 2 files changed, 222 insertions(+) create mode 100644 src/lib/ai/__tests__/provider.test.ts create mode 100644 src/lib/ai/provider.ts diff --git a/src/lib/ai/__tests__/provider.test.ts b/src/lib/ai/__tests__/provider.test.ts new file mode 100644 index 000000000..542073ecd --- /dev/null +++ b/src/lib/ai/__tests__/provider.test.ts @@ -0,0 +1,94 @@ +import { describe, it, expect, vi, beforeEach } from "vitest"; + +vi.mock("@/lib/db", () => ({ + prisma: { + user: { findUnique: vi.fn(), update: vi.fn() }, + appSettings: { findUnique: vi.fn() }, + }, +})); +vi.mock("@/lib/crypto", () => ({ + decrypt: vi.fn((v: string) => `decrypted:${v}`), + encrypt: vi.fn((v: string) => `encrypted:${v}`), +})); +vi.mock("@/lib/ai/codex-oauth", () => ({ + refreshAccessToken: vi.fn(), + encryptTokens: vi.fn((t: any) => ({ accessEncrypted: `enc:${t.accessToken}`, refreshEncrypted: `enc:${t.refreshToken}` })), + decryptTokens: vi.fn(), +})); + +import { resolveProvider } from "../provider"; +import { prisma } from "@/lib/db"; +import { decryptTokens } from "@/lib/ai/codex-oauth"; + +describe("resolveProvider", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("returns codex provider when user has valid tokens", async () => { + vi.mocked(decryptTokens).mockReturnValue({ + accessToken: "decrypted-access", + refreshToken: "decrypted-refresh", + }); + + vi.mocked(prisma.user.findUnique).mockResolvedValue({ + codexAccessTokenEncrypted: "enc-access", + codexRefreshTokenEncrypted: "enc-refresh", + codexTokenExpiresAt: new Date(Date.now() + 3600000), + codexConnectionStatus: "connected", + } as any); + + const provider = await resolveProvider("user-123"); + expect(provider.type).toBe("codex"); + }); + + it("returns admin-key when user has no codex but admin key exists", async () => { + vi.mocked(prisma.user.findUnique).mockResolvedValue({ + codexAccessTokenEncrypted: null, + codexRefreshTokenEncrypted: null, + codexTokenExpiresAt: null, + codexConnectionStatus: "disconnected", + } as any); + + vi.mocked(prisma.appSettings.findUnique).mockResolvedValue({ + adminAiKeyEncrypted: "enc-admin-key", + adminAiModel: "gpt-4o-mini", + adminAiBaseUrl: "https://api.openai.com/v1", + } as any); + + const provider = await resolveProvider("user-123"); + expect(provider.type).toBe("admin-key"); + }); + + it("returns none when nothing is configured", async () => { + vi.mocked(prisma.user.findUnique).mockResolvedValue({ + codexAccessTokenEncrypted: null, + codexConnectionStatus: "disconnected", + } as any); + + vi.mocked(prisma.appSettings.findUnique).mockResolvedValue({ + adminAiKeyEncrypted: null, + } as any); + + const provider = await resolveProvider("user-123"); + expect(provider.type).toBe("none"); + }); + + it("falls back to admin-key when codex status is expired", async () => { + vi.mocked(prisma.user.findUnique).mockResolvedValue({ + codexAccessTokenEncrypted: "enc-access", + codexRefreshTokenEncrypted: "enc-refresh", + codexTokenExpiresAt: new Date(Date.now() - 1000), + codexConnectionStatus: "expired", + } as any); + + vi.mocked(prisma.appSettings.findUnique).mockResolvedValue({ + adminAiKeyEncrypted: "enc-admin", + adminAiModel: "gpt-4o", + adminAiBaseUrl: "https://api.openai.com/v1", + } as any); + + const provider = await resolveProvider("user-123"); + expect(provider.type).toBe("admin-key"); + }); +}); diff --git a/src/lib/ai/provider.ts b/src/lib/ai/provider.ts new file mode 100644 index 000000000..db166f87b --- /dev/null +++ b/src/lib/ai/provider.ts @@ -0,0 +1,128 @@ +import type { AIProvider, CompletionParams, CompletionResult } from "./types"; +import { prisma } from "@/lib/db"; +import { decrypt } from "@/lib/crypto"; +import { CodexClient } from "./codex-client"; +import { OpenAIClient } from "./openai-client"; +import { refreshAccessToken, encryptTokens, decryptTokens } from "./codex-oauth"; + +const TOKEN_REFRESH_BUFFER_MS = 5 * 60 * 1000; // 5 minutes + +class NoProvider implements AIProvider { + readonly type = "none" as const; + + async generateCompletion(_params: CompletionParams): Promise { + throw new Error( + "No AI provider configured. Connect Codex or set an admin API key.", + ); + } +} + +async function resolveAdminProvider(): Promise { + const settings = await prisma.appSettings.findUnique({ where: { id: 1 } }); + + if (settings?.adminAiKeyEncrypted) { + return new OpenAIClient({ + apiKey: decrypt(settings.adminAiKeyEncrypted), + model: settings.adminAiModel ?? "gpt-4o-mini", + baseUrl: settings.adminAiBaseUrl ?? "https://api.openai.com/v1", + }); + } + + return new NoProvider(); +} + +export async function resolveProvider(userId: string): Promise { + const user = await prisma.user.findUnique({ + where: { id: userId }, + select: { + codexAccessTokenEncrypted: true, + codexRefreshTokenEncrypted: true, + codexTokenExpiresAt: true, + codexConnectionStatus: true, + }, + }); + + if ( + user?.codexConnectionStatus === "connected" && + user.codexAccessTokenEncrypted && + user.codexRefreshTokenEncrypted + ) { + const { accessToken, refreshToken } = decryptTokens({ + accessEncrypted: user.codexAccessTokenEncrypted, + refreshEncrypted: user.codexRefreshTokenEncrypted, + }); + + let currentAccessToken = accessToken; + + // Proactive refresh if token expires within 5 minutes + const expiresAt = user.codexTokenExpiresAt?.getTime() ?? 0; + if (expiresAt < Date.now() + TOKEN_REFRESH_BUFFER_MS) { + try { + const newTokens = await refreshAccessToken(refreshToken); + currentAccessToken = newTokens.access_token; + + const encrypted = encryptTokens({ + accessToken: newTokens.access_token, + refreshToken: newTokens.refresh_token, + }); + + await prisma.user.update({ + where: { id: userId }, + data: { + codexAccessTokenEncrypted: encrypted.accessEncrypted, + codexRefreshTokenEncrypted: encrypted.refreshEncrypted, + codexTokenExpiresAt: new Date( + Date.now() + newTokens.expires_in * 1000, + ), + }, + }); + } catch { + // If proactive refresh fails, still try with current token + // CodexClient will retry via onTokenRefresh on 401 + } + } + + return new CodexClient({ + accessToken: currentAccessToken, + onTokenRefresh: async () => { + const freshUser = await prisma.user.findUnique({ + where: { id: userId }, + select: { + codexRefreshTokenEncrypted: true, + }, + }); + + if (!freshUser?.codexRefreshTokenEncrypted) { + throw new Error("No refresh token available"); + } + + const { refreshToken: currentRefreshToken } = decryptTokens({ + accessEncrypted: "", + refreshEncrypted: freshUser.codexRefreshTokenEncrypted, + }); + + const newTokens = await refreshAccessToken(currentRefreshToken); + const encrypted = encryptTokens({ + accessToken: newTokens.access_token, + refreshToken: newTokens.refresh_token, + }); + + await prisma.user.update({ + where: { id: userId }, + data: { + codexAccessTokenEncrypted: encrypted.accessEncrypted, + codexRefreshTokenEncrypted: encrypted.refreshEncrypted, + codexTokenExpiresAt: new Date( + Date.now() + newTokens.expires_in * 1000, + ), + }, + }); + + return newTokens.access_token; + }, + }); + } + + // Codex not connected (disconnected, expired, or error) — fall back to admin key + return resolveAdminProvider(); +}