From 2175118f9fa4e212feea7f5c002ecd4fa41807ca Mon Sep 17 00:00:00 2001 From: Badri Narayanan S Date: Mon, 19 Jan 2026 14:21:30 +0530 Subject: [PATCH] feat: align project discovery with opencode-antigravity-auth reference - Store project IDs in composite refresh token format (refreshToken|projectId|managedProjectId) - Add parseRefreshParts() and formatRefreshParts() for token handling - Extract and persist subscription tier during project discovery - Fetch subscription in blocking mode when missing from cached accounts - Fix conditional duetProject setting to match reference implementation - Export parseTierId() for reuse across modules Co-Authored-By: Claude --- src/account-manager/credentials.js | 200 +++++++++++++++++++++++++---- src/account-manager/index.js | 3 +- src/account-manager/onboarding.js | 15 +-- src/auth/oauth.js | 40 +++++- src/cli/accounts.js | 16 +-- src/cloudcode/model-api.js | 2 +- src/webui/index.js | 3 +- 7 files changed, 227 insertions(+), 52 deletions(-) diff --git a/src/account-manager/credentials.js b/src/account-manager/credentials.js index 6e366fa..1b44225 100644 --- a/src/account-manager/credentials.js +++ b/src/account-manager/credentials.js @@ -11,11 +11,47 @@ import { LOAD_CODE_ASSIST_HEADERS, DEFAULT_PROJECT_ID } from '../constants.js'; -import { refreshAccessToken } from '../auth/oauth.js'; +import { refreshAccessToken, parseRefreshParts, formatRefreshParts } from '../auth/oauth.js'; import { getAuthStatus } from '../auth/database.js'; import { logger } from '../utils/logger.js'; import { isNetworkError } from '../utils/helpers.js'; import { onboardUser, getDefaultTierId } from './onboarding.js'; +import { parseTierId } from '../cloudcode/model-api.js'; + +// Track accounts currently fetching subscription to avoid duplicate calls +const subscriptionFetchInProgress = new Set(); + +/** + * Fetch subscription tier and save it (blocking) + * Used when we have a cached project but missing subscription data + * + * @param {string} token - OAuth access token + * @param {Object} account - Account object + * @param {Function} [onSave] - Callback to save account changes + */ +async function fetchAndSaveSubscription(token, account, onSave) { + // Avoid duplicate fetches for the same account + if (subscriptionFetchInProgress.has(account.email)) { + return; + } + subscriptionFetchInProgress.add(account.email); + + try { + // Call discoverProject just to get subscription info + const { subscription } = await discoverProject(token, account.projectId); + if (subscription && subscription.tier !== 'unknown') { + account.subscription = subscription; + if (onSave) { + await onSave(); + } + logger.info(`[AccountManager] Updated subscription tier for ${account.email}: ${subscription.tier}`); + } + } catch (e) { + logger.debug(`[AccountManager] Subscription fetch failed for ${account.email}: ${e.message}`); + } finally { + subscriptionFetchInProgress.delete(account.email); + } +} /** * Get OAuth token for an account @@ -82,27 +118,92 @@ export async function getTokenForAccount(account, tokenCache, onInvalid, onSave) /** * Get project ID for an account + * Aligned with opencode-antigravity-auth: parses refresh token for stored project IDs * * @param {Object} account - Account object * @param {string} token - OAuth access token * @param {Map} projectCache - Project cache map + * @param {Function} [onSave] - Callback to save account changes * @returns {Promise} Project ID */ -export async function getProjectForAccount(account, token, projectCache) { +export async function getProjectForAccount(account, token, projectCache, onSave = null) { // Check cache first const cached = projectCache.get(account.email); if (cached) { return cached; } - // OAuth or manual accounts may have projectId specified + // Parse refresh token to get stored project IDs (aligned with opencode-antigravity-auth) + const parts = account.refreshToken ? parseRefreshParts(account.refreshToken) : { refreshToken: null, projectId: undefined, managedProjectId: undefined }; + + // If we have a managedProjectId in the refresh token, use it + if (parts.managedProjectId) { + projectCache.set(account.email, parts.managedProjectId); + // If subscription is missing/unknown, fetch it now (blocking) + if (!account.subscription || account.subscription.tier === 'unknown') { + await fetchAndSaveSubscription(token, account, onSave); + } + return parts.managedProjectId; + } + + // Legacy: check account.projectId for backward compatibility if (account.projectId) { projectCache.set(account.email, account.projectId); + // If subscription is missing/unknown, fetch it now (blocking) + if (!account.subscription || account.subscription.tier === 'unknown') { + await fetchAndSaveSubscription(token, account, onSave); + } return account.projectId; } - // Discover project via loadCodeAssist API - const project = await discoverProject(token); + // Discover managed project, passing projectId for metadata.duetProject + // Reference: opencode-antigravity-auth - discoverProject handles fallback internally + const { project, subscription } = await discoverProject(token, parts.projectId); + + // Store managedProjectId back in refresh token (if we got a real project) + if (project && project !== DEFAULT_PROJECT_ID) { + let needsSave = false; + + if (account.refreshToken) { + // OAuth accounts: encode in refresh token + account.refreshToken = formatRefreshParts({ + refreshToken: parts.refreshToken, + projectId: parts.projectId, + managedProjectId: project, + }); + needsSave = true; + } else if (account.source === 'database' || account.source === 'manual') { + // Database/manual accounts: store in projectId field + account.projectId = project; + needsSave = true; + } + + // Save subscription tier if discovered + if (subscription) { + account.subscription = subscription; + needsSave = true; + } + + // Trigger save to persist the updated project and subscription + if (needsSave && onSave) { + try { + await onSave(); + } catch (e) { + logger.warn(`[AccountManager] Failed to save updated project: ${e.message}`); + } + } + } else if (subscription) { + // Even if no project discovered, save subscription if we got it + account.subscription = subscription; + if (onSave) { + try { + await onSave(); + } catch (e) { + logger.warn(`[AccountManager] Failed to save subscription: ${e.message}`); + } + } + } + projectCache.set(account.email, project); return project; } @@ -111,13 +212,27 @@ export async function getProjectForAccount(account, token, projectCache) { * Discover project ID via Cloud Code API * * @param {string} token - OAuth access token - * @returns {Promise} Project ID + * @param {string} [projectId] - Optional project ID from refresh token (for metadata.duetProject) + * @returns {Promise<{project: string, subscription: {tier: string, projectId: string|null, detectedAt: string}|null}>} Project and subscription info */ -export async function discoverProject(token) { +export async function discoverProject(token, projectId = undefined) { let lastError = null; let gotSuccessfulResponse = false; let loadCodeAssistData = null; + // Build metadata matching reference: only set duetProject if projectId is truthy + // For loadCodeAssist, reference passes projectId ?? fallbackProjectId + // Reference: opencode-antigravity-auth/src/plugin/project.ts buildMetadata() + const projectIdForLoad = projectId ?? DEFAULT_PROJECT_ID; + const metadata = { + ideType: 'IDE_UNSPECIFIED', + platform: 'PLATFORM_UNSPECIFIED', + pluginType: 'GEMINI' + }; + if (projectIdForLoad) { + metadata.duetProject = projectIdForLoad; + } + for (const endpoint of LOAD_CODE_ASSIST_ENDPOINTS) { try { const response = await fetch(`${endpoint}/v1internal:loadCodeAssist`, { @@ -127,14 +242,7 @@ export async function discoverProject(token) { 'Content-Type': 'application/json', ...LOAD_CODE_ASSIST_HEADERS }, - body: JSON.stringify({ - metadata: { - ideType: 'IDE_UNSPECIFIED', - platform: 'PLATFORM_UNSPECIFIED', - pluginType: 'GEMINI', - duetProject: DEFAULT_PROJECT_ID - } - }) + body: JSON.stringify({ metadata }) }); if (!response.ok) { @@ -150,13 +258,16 @@ export async function discoverProject(token) { logger.debug(`[AccountManager] loadCodeAssist response from ${endpoint}:`, JSON.stringify(data)); + // Extract subscription tier from response + const subscription = extractSubscriptionFromResponse(data); + if (typeof data.cloudaicompanionProject === 'string') { logger.success(`[AccountManager] Discovered project: ${data.cloudaicompanionProject}`); - return data.cloudaicompanionProject; + return { project: data.cloudaicompanionProject, subscription }; } if (data.cloudaicompanionProject?.id) { logger.success(`[AccountManager] Discovered project: ${data.cloudaicompanionProject.id}`); - return data.cloudaicompanionProject.id; + return { project: data.cloudaicompanionProject.id, subscription }; } // No project found - log tier data and try to onboard the user @@ -189,29 +300,66 @@ export async function discoverProject(token) { tierId = tierId || 'free-tier'; logger.info(`[AccountManager] Onboarding user with tier: ${tierId} (source: ${tierSource})`); - // Check if this is a free tier (raw API values contain 'free') - const isFree = tierId.toLowerCase().includes('free'); - - // For non-free tiers, pass DEFAULT_PROJECT_ID as the GCP project - // The API requires a project for paid tier onboarding + // Pass projectId for metadata.duetProject (without fallback, matching reference) + // Reference: opencode-antigravity-auth passes parts.projectId (not fallback) to onboardManagedProject const onboardedProject = await onboardUser( token, tierId, - isFree ? null : DEFAULT_PROJECT_ID + projectId // Original projectId without fallback ); if (onboardedProject) { logger.success(`[AccountManager] Successfully onboarded, project: ${onboardedProject}`); - return onboardedProject; + const subscription = extractSubscriptionFromResponse(loadCodeAssistData); + return { project: onboardedProject, subscription }; } - logger.warn(`[AccountManager] Onboarding failed, using default project: ${DEFAULT_PROJECT_ID}`); + logger.warn(`[AccountManager] Onboarding failed - account may not work correctly`); } // Only warn if all endpoints failed with errors (not just missing project) if (!gotSuccessfulResponse) { logger.warn(`[AccountManager] loadCodeAssist failed for all endpoints: ${lastError}`); } - return DEFAULT_PROJECT_ID; + + // Fallback: use projectId if available, otherwise use default + // Reference: opencode-antigravity-auth/src/plugin/project.ts + if (projectId) { + return { project: projectId, subscription: null }; + } + return { project: DEFAULT_PROJECT_ID, subscription: null }; +} + +/** + * Extract subscription tier from loadCodeAssist response + * + * @param {Object} data - loadCodeAssist response data + * @returns {{tier: string, projectId: string|null, detectedAt: string}|null} Subscription info + */ +function extractSubscriptionFromResponse(data) { + if (!data) return null; + + // Priority: paidTier > currentTier (consistent with model-api.js) + let tier = 'free'; + let cloudProject = null; + + if (data.paidTier?.id) { + tier = parseTierId(data.paidTier.id); + } else if (data.currentTier?.id) { + tier = parseTierId(data.currentTier.id); + } + + // Get project ID + if (typeof data.cloudaicompanionProject === 'string') { + cloudProject = data.cloudaicompanionProject; + } else if (data.cloudaicompanionProject?.id) { + cloudProject = data.cloudaicompanionProject.id; + } + + return { + tier, + projectId: cloudProject, + detectedAt: new Date().toISOString() + }; } /** diff --git a/src/account-manager/index.js b/src/account-manager/index.js index 8231344..e9988b3 100644 --- a/src/account-manager/index.js +++ b/src/account-manager/index.js @@ -297,7 +297,8 @@ export class AccountManager { * @returns {Promise} Project ID */ async getProjectForAccount(account, token) { - return fetchProject(account, token, this.#projectCache); + // Pass onSave callback to persist managedProjectId in refresh token + return fetchProject(account, token, this.#projectCache, () => this.saveToDisk()); } /** diff --git a/src/account-manager/onboarding.js b/src/account-manager/onboarding.js index fc2033e..d445b49 100644 --- a/src/account-manager/onboarding.js +++ b/src/account-manager/onboarding.js @@ -43,7 +43,7 @@ export function getDefaultTierId(allowedTiers) { * @param {number} [delayMs=5000] - Delay between polling attempts * @returns {Promise} Managed project ID or null if failed */ -export async function onboardUser(token, tierId, projectId = null, maxAttempts = 10, delayMs = 5000) { +export async function onboardUser(token, tierId, projectId = undefined, maxAttempts = 10, delayMs = 5000) { const metadata = { ideType: 'IDE_UNSPECIFIED', platform: 'PLATFORM_UNSPECIFIED', @@ -58,16 +58,11 @@ export async function onboardUser(token, tierId, projectId = null, maxAttempts = tierId, metadata }; + // Note: Do NOT add cloudaicompanionProject to requestBody + // Reference implementation only sets metadata.duetProject, not the body field + // Adding cloudaicompanionProject causes 400 errors for auto-provisioned tiers (g1-pro, g1-ultra) - // Check if this is a free tier (handles raw API values like 'free-tier') - const isFree = tierId.toLowerCase().includes('free'); - - // Non-free tiers require a cloudaicompanionProject - if (!isFree && projectId) { - requestBody.cloudaicompanionProject = projectId; - } - - logger.debug(`[Onboarding] Starting onboard with tierId: ${tierId}, projectId: ${projectId}, isFree: ${isFree}`); + logger.debug(`[Onboarding] Starting onboard with tierId: ${tierId}, projectId: ${projectId}`); for (const endpoint of ONBOARD_USER_ENDPOINTS) { for (let attempt = 0; attempt < maxAttempts; attempt++) { diff --git a/src/auth/oauth.js b/src/auth/oauth.js index fa9f2b3..ea7461a 100644 --- a/src/auth/oauth.js +++ b/src/auth/oauth.js @@ -17,6 +17,34 @@ import { import { logger } from '../utils/logger.js'; import { onboardUser, getDefaultTierId } from '../account-manager/onboarding.js'; +/** + * Parse refresh token parts (aligned with opencode-antigravity-auth) + * Format: refreshToken|projectId|managedProjectId + * + * @param {string} refresh - Composite refresh token string + * @returns {{refreshToken: string, projectId: string|undefined, managedProjectId: string|undefined}} + */ +export function parseRefreshParts(refresh) { + const [refreshToken = '', projectId = '', managedProjectId = ''] = (refresh ?? '').split('|'); + return { + refreshToken, + projectId: projectId || undefined, + managedProjectId: managedProjectId || undefined, + }; +} + +/** + * Format refresh token parts back into composite string + * + * @param {{refreshToken: string, projectId?: string|undefined, managedProjectId?: string|undefined}} parts + * @returns {string} Composite refresh token + */ +export function formatRefreshParts(parts) { + const projectSegment = parts.projectId ?? ''; + const base = `${parts.refreshToken}|${projectSegment}`; + return parts.managedProjectId ? `${base}|${parts.managedProjectId}` : base; +} + /** * Generate PKCE code verifier and challenge */ @@ -267,11 +295,15 @@ export async function exchangeCode(code, verifier) { /** * Refresh access token using refresh token + * Handles composite refresh tokens (refreshToken|projectId|managedProjectId) * - * @param {string} refreshToken - OAuth refresh token + * @param {string} compositeRefresh - OAuth refresh token (may be composite) * @returns {Promise<{accessToken: string, expiresIn: number}>} New access token */ -export async function refreshAccessToken(refreshToken) { +export async function refreshAccessToken(compositeRefresh) { + // Parse the composite refresh token to extract the actual OAuth token + const parts = parseRefreshParts(compositeRefresh); + const response = await fetch(OAUTH_CONFIG.tokenUrl, { method: 'POST', headers: { @@ -280,7 +312,7 @@ export async function refreshAccessToken(refreshToken) { body: new URLSearchParams({ client_id: OAUTH_CONFIG.clientId, client_secret: OAUTH_CONFIG.clientSecret, - refresh_token: refreshToken, + refresh_token: parts.refreshToken, // Use the actual OAuth token grant_type: 'refresh_token' }) }); @@ -408,6 +440,8 @@ export async function completeOAuthFlow(code, verifier) { } export default { + parseRefreshParts, + formatRefreshParts, getAuthorizationUrl, extractCodeFromInput, startCallbackServer, diff --git a/src/cli/accounts.js b/src/cli/accounts.js index bb283f8..ce37319 100644 --- a/src/cli/accounts.js +++ b/src/cli/accounts.js @@ -210,20 +210,18 @@ async function addAccount(existingAccounts) { if (existing) { console.log(`\n⚠ Account ${result.email} already exists. Updating tokens.`); existing.refreshToken = result.refreshToken; - existing.projectId = result.projectId; + // Note: projectId will be discovered and stored in refresh token on first use existing.addedAt = new Date().toISOString(); return null; // Don't add duplicate } console.log(`\n✓ Successfully authenticated: ${result.email}`); - if (result.projectId) { - console.log(` Project ID: ${result.projectId}`); - } + console.log(' Project will be discovered on first API request.'); return { email: result.email, refreshToken: result.refreshToken, - projectId: result.projectId, + // Note: projectId stored in refresh token, not as separate field addedAt: new Date().toISOString(), modelRateLimits: {} }; @@ -267,20 +265,18 @@ async function addAccountNoBrowser(existingAccounts, rl) { if (existing) { console.log(`\n⚠ Account ${result.email} already exists. Updating tokens.`); existing.refreshToken = result.refreshToken; - existing.projectId = result.projectId; + // Note: projectId will be discovered and stored in refresh token on first use existing.addedAt = new Date().toISOString(); return null; // Don't add duplicate } console.log(`\n✓ Successfully authenticated: ${result.email}`); - if (result.projectId) { - console.log(` Project ID: ${result.projectId}`); - } + console.log(' Project will be discovered on first API request.'); return { email: result.email, refreshToken: result.refreshToken, - projectId: result.projectId, + // Note: projectId stored in refresh token, not as separate field addedAt: new Date().toISOString(), modelRateLimits: {} }; diff --git a/src/cloudcode/model-api.js b/src/cloudcode/model-api.js index 9d32c56..0dc6cf4 100644 --- a/src/cloudcode/model-api.js +++ b/src/cloudcode/model-api.js @@ -128,7 +128,7 @@ export async function getModelQuotas(token, projectId = null) { * @param {string} tierId - The tier ID from the API * @returns {'free' | 'pro' | 'ultra' | 'unknown'} The subscription tier */ -function parseTierId(tierId) { +export function parseTierId(tierId) { if (!tierId) return 'unknown'; const lower = tierId.toLowerCase(); diff --git a/src/webui/index.js b/src/webui/index.js index d449297..548c4f7 100644 --- a/src/webui/index.js +++ b/src/webui/index.js @@ -668,10 +668,11 @@ export function mountWebUI(app, dirname, accountManager) { const accountData = await completeOAuthFlow(code, verifier); // Add or update the account + // Note: Don't set projectId here - it will be discovered and stored + // in the refresh token via getProjectForAccount() on first use await addAccount({ email: accountData.email, refreshToken: accountData.refreshToken, - projectId: accountData.projectId, source: 'oauth' });