From 0610909116fb0f5e2e7ab41fabf0ef5da8ed1609 Mon Sep 17 00:00:00 2001 From: nyatoru Date: Mon, 23 Feb 2026 01:11:53 +0700 Subject: [PATCH] feat(codex): add workspace binding via chatgpt-account-id header --- open-sse/executors/codex.ts | 14 +- open-sse/services/usage.ts | 108 ++++++------- .../usage/components/ProviderLimits/index.tsx | 2 + .../usage/components/ProviderLimits/utils.tsx | 16 +- src/lib/db/providers.ts | 28 +++- src/lib/oauth/providers/codex.ts | 145 +++++++++++++++++- src/lib/usage/fetcher.ts | 14 +- 7 files changed, 244 insertions(+), 83 deletions(-) diff --git a/open-sse/executors/codex.ts b/open-sse/executors/codex.ts index 52f8dd6c56..6e7a861b77 100644 --- a/open-sse/executors/codex.ts +++ b/open-sse/executors/codex.ts @@ -4,7 +4,8 @@ import { PROVIDERS } from "../config/constants.ts"; /** * Codex Executor - handles OpenAI Codex API (Responses API format) - * Automatically injects default instructions if missing + * Automatically injects default instructions if missing. + * IMPORTANT: Includes chatgpt-account-id header for workspace binding. */ export class CodexExecutor extends BaseExecutor { constructor() { @@ -14,9 +15,18 @@ export class CodexExecutor extends BaseExecutor { /** * Codex Responses endpoint is SSE-first. * Always request event-stream from upstream, even when client requested stream=false. + * Includes chatgpt-account-id header for strict workspace binding. */ buildHeaders(credentials, stream = true) { - return super.buildHeaders(credentials, true); + const headers = super.buildHeaders(credentials, true); + + // Add workspace binding header if workspaceId is persisted + const workspaceId = credentials?.providerSpecificData?.workspaceId; + if (workspaceId) { + headers["chatgpt-account-id"] = workspaceId; + } + + return headers; } /** diff --git a/open-sse/services/usage.ts b/open-sse/services/usage.ts index 406d0f7192..060f0d9cb7 100644 --- a/open-sse/services/usage.ts +++ b/open-sse/services/usage.ts @@ -53,7 +53,7 @@ export async function getUsageForProvider(connection) { case "claude": return await getClaudeUsage(accessToken); case "codex": - return await getCodexUsage(accessToken); + return await getCodexUsage(accessToken, providerSpecificData); case "kiro": return await getKiroUsage(accessToken, providerSpecificData); case "qwen": @@ -483,38 +483,17 @@ async function getClaudeUsage(accessToken) { /** * Codex (OpenAI) Usage - Fetch from ChatGPT backend API + * IMPORTANT: Uses persisted workspaceId from OAuth to ensure correct workspace binding. + * No fallback to other workspaces - strict binding to user's selected workspace. */ -async function getCodexUsage(accessToken) { +async function getCodexUsage(accessToken, providerSpecificData: Record = {}) { try { - let accountId = null; - try { - const accountsRes = await fetch("https://chatgpt.com/backend-api/accounts/check/v4", { - method: "GET", - headers: { - Authorization: `Bearer ${accessToken}`, - Accept: "application/json", - }, - }); - if (accountsRes.ok) { - const accountsData = await accountsRes.json(); - if (accountsData.accounts) { - const accountsArray = Object.values(accountsData.accounts) as any[]; - const targetWorkspace = - accountsArray.find((a) => a.account?.plan_type === "biz") || - accountsArray.find((a) => a.account?.plan_type !== "free") || - accountsArray.find((a) => a.is_default) || - accountsArray[0]; - if (targetWorkspace && targetWorkspace.account?.id) { - accountId = targetWorkspace.account.id; - } - } - } - } catch (err) { - console.warn("Could not fetch ChatGPT accounts for quota:", err); - } + // Use persisted workspace ID from OAuth - NO FALLBACK + const accountId = providerSpecificData?.workspaceId || null; const headers: Record = { Authorization: `Bearer ${accessToken}`, + "Content-Type": "application/json", Accept: "application/json", }; if (accountId) { @@ -532,38 +511,53 @@ async function getCodexUsage(accessToken) { const data = await response.json(); - // Parse rate limit info - const rateLimit = data.rate_limit || {}; - const primaryWindow = rateLimit.primary_window || {}; - const secondaryWindow = rateLimit.secondary_window || {}; + // Helper to get field with snake_case/camelCase fallback + const getField = (obj: any, snakeKey: string, camelKey: string) => + obj?.[snakeKey] ?? obj?.[camelKey] ?? null; - // Parse reset dates (reset_at is Unix timestamp in seconds, multiply by 1000 for ms) - const sessionResetAt = parseResetTime( - primaryWindow.reset_at ? primaryWindow.reset_at * 1000 : null - ); - const weeklyResetAt = parseResetTime( - secondaryWindow.reset_at ? secondaryWindow.reset_at * 1000 : null - ); + // Parse rate limit info (supports both snake_case and camelCase) + const rateLimit = getField(data, "rate_limit", "rateLimit") || {}; + const primaryWindow = getField(rateLimit, "primary_window", "primaryWindow") || {}; + const secondaryWindow = getField(rateLimit, "secondary_window", "secondaryWindow") || {}; + + // Parse reset times (reset_at is Unix timestamp in seconds) + const parseWindowReset = (window: any) => { + const resetAt = getField(window, "reset_at", "resetAt"); + const resetAfterSeconds = getField(window, "reset_after_seconds", "resetAfterSeconds"); + if (resetAt) return parseResetTime(resetAt * 1000); + if (resetAfterSeconds) return parseResetTime(Date.now() + resetAfterSeconds * 1000); + return null; + }; + + // Build quota windows + const quotas: Record = {}; + + // Primary window (5-hour) + if (Object.keys(primaryWindow).length > 0) { + quotas.session = { + used: getField(primaryWindow, "used_percent", "usedPercent") || 0, + total: 100, + remaining: 100 - (getField(primaryWindow, "used_percent", "usedPercent") || 0), + resetAt: parseWindowReset(primaryWindow), + unlimited: false, + }; + } + + // Secondary window (weekly) + if (Object.keys(secondaryWindow).length > 0) { + quotas.weekly = { + used: getField(secondaryWindow, "used_percent", "usedPercent") || 0, + total: 100, + remaining: 100 - (getField(secondaryWindow, "used_percent", "usedPercent") || 0), + resetAt: parseWindowReset(secondaryWindow), + unlimited: false, + }; + } return { - plan: data.plan_type || "unknown", - limitReached: rateLimit.limit_reached || false, - quotas: { - session: { - used: primaryWindow.used_percent || 0, - total: 100, - remaining: 100 - (primaryWindow.used_percent || 0), - resetAt: sessionResetAt, - unlimited: false, - }, - weekly: { - used: secondaryWindow.used_percent || 0, - total: 100, - remaining: 100 - (secondaryWindow.used_percent || 0), - resetAt: weeklyResetAt, - unlimited: false, - }, - }, + plan: getField(data, "plan_type", "planType") || "unknown", + limitReached: getField(rateLimit, "limit_reached", "limitReached") || false, + quotas, }; } catch (error) { throw new Error(`Failed to fetch Codex usage: ${error.message}`); diff --git a/src/app/(dashboard)/dashboard/usage/components/ProviderLimits/index.tsx b/src/app/(dashboard)/dashboard/usage/components/ProviderLimits/index.tsx index e66e17d9b8..7f0206fd46 100644 --- a/src/app/(dashboard)/dashboard/usage/components/ProviderLimits/index.tsx +++ b/src/app/(dashboard)/dashboard/usage/components/ProviderLimits/index.tsx @@ -23,6 +23,7 @@ const PROVIDER_CONFIG = { const TIER_FILTERS = [ { key: "all", label: "All" }, { key: "enterprise", label: "Enterprise" }, + { key: "team", label: "Team" }, { key: "business", label: "Business" }, { key: "ultra", label: "Ultra" }, { key: "pro", label: "Pro" }, @@ -249,6 +250,7 @@ export default function ProviderLimits() { const counts = { all: sortedConnections.length, enterprise: 0, + team: 0, business: 0, ultra: 0, pro: 0, diff --git a/src/app/(dashboard)/dashboard/usage/components/ProviderLimits/utils.tsx b/src/app/(dashboard)/dashboard/usage/components/ProviderLimits/utils.tsx index 515dc9cc1f..7ae3207724 100644 --- a/src/app/(dashboard)/dashboard/usage/components/ProviderLimits/utils.tsx +++ b/src/app/(dashboard)/dashboard/usage/components/ProviderLimits/utils.tsx @@ -202,7 +202,7 @@ export function parseQuotaData(provider, data) { /** * Normalize provider-specific plan labels into a shared tier taxonomy. - * Supported tiers: enterprise, business, ultra, pro, free, unknown. + * Supported tiers: enterprise, business, team, ultra, pro, free, unknown. */ export function normalizePlanTier(plan) { const raw = typeof plan === "string" ? plan.trim() : ""; @@ -213,15 +213,15 @@ export function normalizePlanTier(plan) { const upper = raw.toUpperCase(); if (upper.includes("ENTERPRISE") || upper.includes("CORP") || upper.includes("ORG")) { - return { key: "enterprise", label: "Enterprise", variant: "info", rank: 6, raw }; + return { key: "enterprise", label: "Enterprise", variant: "info", rank: 7, raw }; } - if ( - upper.includes("BUSINESS") || - upper.includes("TEAM") || - upper.includes("STANDARD") || - upper.includes("BIZ") - ) { + // Team plan (e.g., ChatGPT Team, GitHub Team) + if (upper.includes("TEAM") || upper.includes("CHATGPTTEAM")) { + return { key: "team", label: "Team", variant: "info", rank: 6, raw }; + } + + if (upper.includes("BUSINESS") || upper.includes("STANDARD") || upper.includes("BIZ")) { return { key: "business", label: "Business", variant: "warning", rank: 5, raw }; } diff --git a/src/lib/db/providers.ts b/src/lib/db/providers.ts index 4cad15fffd..3a7c173645 100644 --- a/src/lib/db/providers.ts +++ b/src/lib/db/providers.ts @@ -44,13 +44,31 @@ export async function createProviderConnection(data: any) { const now = new Date().toISOString(); // Upsert check + // For Codex/OpenAI, a single email can have multiple workspaces (Team + Personal) + // We need to check for workspace uniqueness, not just email let existing = null; + if (data.authType === "oauth" && data.email) { - existing = db - .prepare( - "SELECT * FROM provider_connections WHERE provider = ? AND auth_type = 'oauth' AND email = ?" - ) - .get(data.provider, data.email); + // For Codex, check for existing connection with same workspace + const workspaceId = data.providerSpecificData?.workspaceId; + if (data.provider === "codex" && workspaceId) { + // Check for existing connection with same provider + workspace ONLY + // Do NOT fall back to email check - this allows multiple workspaces per email + existing = db + .prepare( + "SELECT * FROM provider_connections WHERE provider = ? AND auth_type = 'oauth' AND json_extract(provider_specific_data, '$.workspaceId') = ?" + ) + .get(data.provider, workspaceId); + // For Codex with workspaceId, don't fall back to email check + // This allows creating new connections for different workspaces + } else { + // For other providers (or Codex without workspaceId), use email check + existing = db + .prepare( + "SELECT * FROM provider_connections WHERE provider = ? AND auth_type = 'oauth' AND email = ?" + ) + .get(data.provider, data.email); + } } else if (data.authType === "apikey" && data.name) { existing = db .prepare( diff --git a/src/lib/oauth/providers/codex.ts b/src/lib/oauth/providers/codex.ts index d856bf9895..4f01d9508a 100644 --- a/src/lib/oauth/providers/codex.ts +++ b/src/lib/oauth/providers/codex.ts @@ -1,10 +1,70 @@ import { CODEX_CONFIG } from "../constants/oauth"; +/** + * OpenAI Codex Auth Info embedded in id_token JWT + * The JWT claims contain a custom claim at "https://api.openai.com/auth" + */ +interface CodexAuthInfo { + chatgpt_account_id: string; + chatgpt_plan_type: string; + chatgpt_user_id: string; + user_id: string; + organizations: Array<{ + id: string; + is_default: boolean; + role: string; + title: string; + }>; +} + +/** + * Parse the id_token JWT to extract Codex-specific auth information. + * The workspace selection is embedded in the JWT after OAuth completion. + * + * IMPORTANT: The workspace selected by the user during OAuth is stored in + * chatgpt_account_id - this is the workspace we MUST use, not a fallback. + */ +function parseIdToken(idToken: string): { email: string | null; authInfo: CodexAuthInfo | null } { + try { + const parts = idToken.split("."); + if (parts.length !== 3) { + return { email: null, authInfo: null }; + } + + // Base64 URL decode the payload + let payload = parts[1]; + // Add padding if necessary + switch (payload.length % 4) { + case 2: + payload += "=="; + break; + case 3: + payload += "="; + break; + } + + // Use atob for base64 decoding (works in both browser and Node.js) + // Replace URL-safe characters with standard base64 characters + const base64 = payload.replace(/-/g, "+").replace(/_/g, "/"); + const decoded = JSON.parse(atob(base64)); + + const email = decoded.email || null; + + // Extract Codex auth info from custom claim + const authInfo = decoded["https://api.openai.com/auth"] || null; + + return { email, authInfo }; + } catch (e) { + return { email: null, authInfo: null }; + } +} + export const codex = { config: CODEX_CONFIG, flowType: "authorization_code_pkce", fixedPort: 1455, callbackPath: "/auth/callback", + buildAuthUrl: (config, redirectUri, state, codeChallenge) => { const params = { response_type: "code", @@ -21,6 +81,7 @@ export const codex = { .join("&"); return `${config.authorizeUrl}?${queryString}`; }, + exchangeToken: async (config, code, redirectUri, codeVerifier) => { const response = await fetch(config.tokenUrl, { method: "POST", @@ -44,24 +105,92 @@ export const codex = { return await response.json(); }, - mapTokens: (tokens) => { - // Extract email from id_token JWT to distinguish between accounts + + /** + * Post-exchange hook: Parse id_token to extract workspace info. + * The workspace selected by the user during OAuth is embedded in the id_token. + */ + postExchange: async (tokens) => { + if (!tokens.id_token) { + return { authInfo: null }; + } + + const { authInfo } = parseIdToken(tokens.id_token); + return { authInfo }; + }, + + mapTokens: (tokens, extra) => { + // Parse id_token for email and auth info let email = null; + let authInfo = extra?.authInfo || null; + if (tokens.id_token) { - try { - const payload = tokens.id_token.split(".")[1]; - const decoded = JSON.parse(Buffer.from(payload, "base64").toString()); - email = decoded.email || null; - } catch { - // Ignore JWT parsing errors + const parsed = parseIdToken(tokens.id_token); + email = parsed.email; + // Use authInfo from postExchange if available, otherwise from parsing + if (!authInfo && parsed.authInfo) { + authInfo = parsed.authInfo; } } + + // Determine the correct workspace to use + // + // IMPORTANT: A user can have both Team and Personal workspaces. + // The JWT's chatgpt_account_id may not always reflect the workspace + // the user selected during OAuth. We need to be smart about selection. + // + // Selection logic: + // 1. If plan_type indicates team/business, use chatgpt_account_id + // 2. If plan_type is "free" but organizations has team workspace, use team + // 3. Otherwise use chatgpt_account_id as fallback + let workspaceId = authInfo?.chatgpt_account_id || null; + let planType = (authInfo?.chatgpt_plan_type || "").toLowerCase(); + + // Check if we should use a team workspace instead + const organizations = authInfo?.organizations || []; + if (organizations.length > 0) { + // Find team/business workspace (non-default usually means team) + const teamOrg = organizations.find((org) => { + const title = (org.title || "").toLowerCase(); + const role = (org.role || "").toLowerCase(); + // Team workspaces typically have role like "member" or "admin" and non-personal titles + return ( + !org.is_default && + (title.includes("team") || + title.includes("business") || + title.includes("workspace") || + title.includes("org") || + role === "admin" || + role === "member") + ); + }); + + // If user's plan_type is "team" or we found a team org, prefer it + if (planType.includes("team") || planType.includes("chatgptteam")) { + // User authenticated via Team, use the chatgpt_account_id from JWT + } else if (teamOrg && (planType === "free" || planType === "")) { + // User has a team org but plan_type shows free - use team org instead + workspaceId = teamOrg.id; + planType = "team"; + } + } + + const providerSpecificData = { + workspaceId, + workspacePlanType: planType, + // Also store the full authInfo for future reference + chatgptUserId: authInfo?.chatgpt_user_id || null, + organizations: organizations.length > 0 ? organizations : null, + }; + return { accessToken: tokens.access_token, refreshToken: tokens.refresh_token, idToken: tokens.id_token, expiresIn: tokens.expires_in, email, + // Persist workspace binding to prevent fallback to wrong workspace + providerSpecificData, }; }, }; diff --git a/src/lib/usage/fetcher.ts b/src/lib/usage/fetcher.ts index a87f93d8d5..0d8f0c308a 100644 --- a/src/lib/usage/fetcher.ts +++ b/src/lib/usage/fetcher.ts @@ -22,7 +22,7 @@ export async function getUsageForProvider(connection) { case "claude": return await getClaudeUsage(accessToken); case "codex": - return await getCodexUsage(accessToken); + return await getCodexUsage(accessToken, providerSpecificData); case "qwen": return await getQwenUsage(accessToken, providerSpecificData); case "iflow": @@ -169,10 +169,18 @@ async function getClaudeUsage(accessToken) { /** * Codex (OpenAI) Usage + * Note: Actual quota tracking is handled by open-sse/services/usage.ts + * This fallback returns a message directing users to the dashboard. */ -async function getCodexUsage(accessToken) { +async function getCodexUsage(accessToken, providerSpecificData: Record = {}) { try { - // OpenAI usage requires organization API access + // Check if workspace is bound + const workspaceId = providerSpecificData?.workspaceId; + if (workspaceId) { + return { + message: `Codex connected (workspace: ${workspaceId.slice(0, 8)}...). Check dashboard for quota.`, + }; + } return { message: "Codex connected. Check OpenAI dashboard for usage." }; } catch (error) { return { message: "Unable to fetch Codex usage." };