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 <noreply@anthropic.com>
2026-01-19 14:21:30 +05:30
parent 9311c6fdf7
commit 2175118f9f
7 changed files with 227 additions and 52 deletions
--- 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<string>} 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<string>} 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()
+    };
 }

 /**
--- a/src/account-manager/index.js
+++ b/src/account-manager/index.js
@@ -297,7 +297,8 @@ export class AccountManager {
     * @returns {Promise<string>} 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());
    }

    /**
--- 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<string|null>} 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++) {
--- 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,
--- 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: {}
        };
--- 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();

--- 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'
                        });