feat: Implement TaskMaster AI Gateway integration with enhanced UX

- Fix Zod schema conversion, update headers, add premium telemetry display, improve user auth flow, and standardize email fields Functionally complete on this end, mostly polish around user experience and need to add in profile, upgrade/downgrade, etc. But the AI commands are working off the gateway.
2025-06-01 19:37:12 -04:00
parent 9b87dd23de
commit 2819be51d3
11 changed files with 456 additions and 246 deletions
--- a/scripts/init.js
+++ b/scripts/init.js
@@ -385,51 +385,42 @@ async function initializeProject(options = {}) {
      };
    }

-    // NON-INTERACTIVE MODE - Use proper auth/init flow
-    let userSetupResult;
+    // NON-INTERACTIVE MODE - Try auth/init gracefully
+    let userSetupResult = null;
+    let isGatewayAvailable = false;

+    // Try to initialize user, but don't throw errors if it fails
    try {
-      // Check if existing config has userId
-      const existingConfigPath = path.join(process.cwd(), ".taskmasterconfig");
-      let existingUserId = null;
-
-      if (fs.existsSync(existingConfigPath)) {
-        const existingConfig = JSON.parse(
-          fs.readFileSync(existingConfigPath, "utf8")
-        );
-        existingUserId = existingConfig.account?.userId;
-      }
-
-      if (existingUserId) {
-        // Validate existing userId through auth/init
-        userSetupResult = await registerUserWithGateway(null, process.cwd());
-        if (!userSetupResult.success) {
-          throw new Error(
-            `Failed to validate existing user: ${userSetupResult.error}`
-          );
+      userSetupResult = await initializeUser(process.cwd());
+      if (userSetupResult.success) {
+        isGatewayAvailable = true;
+        if (!isSilentMode()) {
+          log("info", "Gateway connection successful");
        }
      } else {
-        // Create new user through auth/init
-        userSetupResult = await initializeUser(process.cwd());
-        if (!userSetupResult.success) {
-          throw new Error(
-            `Failed to initialize user: ${userSetupResult.error}`
-          );
+        if (!isSilentMode()) {
+          log("info", "Gateway not available, using BYOK mode");
        }
+        isGatewayAvailable = false;
      }
    } catch (error) {
-      log("error", `User initialization failed: ${error.message}`);
-      throw error; // Don't fall back to random userId!
+      // Silent failure - gateway not available
+      if (!isSilentMode()) {
+        log("info", "Gateway not available, using BYOK mode");
+      }
+      isGatewayAvailable = false;
+      userSetupResult = null;
    }

-    // Create project structure with properly authenticated userId
+    // Create project structure - always use BYOK for non-interactive mode
+    // since we don't want to prompt for mode selection
    createProjectStructure(
      addAliases,
      dryRun,
-      userSetupResult, // Pass the full auth result
-      "byok", // or determine from result
+      userSetupResult, // Pass the auth result (may be null)
+      "byok", // Always use BYOK for non-interactive
      null,
-      userSetupResult.userId || null
+      userSetupResult?.userId || null
    );
  } else {
    // Interactive logic - NEW FLOW STARTS HERE
@@ -440,7 +431,7 @@ async function initializeProject(options = {}) {
    });

    try {
-      // STEP 1: Create/find userId first
+      // STEP 1: Welcome message
      console.log(
        boxen(
          chalk.blue.bold("🚀 Welcome to Taskmaster AI") +
@@ -455,101 +446,119 @@ async function initializeProject(options = {}) {
        )
      );

-      // INTERACTIVE MODE - Also use proper auth/init flow
-      // STEP 1: Proper user setup
-      let userSetupResult;
+      // STEP 2: Try auth/init gracefully to detect gateway availability
+      let userSetupResult = null;
+      let isGatewayAvailable = false;

      try {
-        // Same logic as non-interactive mode
-        const existingConfigPath = path.join(
-          process.cwd(),
-          ".taskmasterconfig"
-        );
-        let existingUserId = null;
-
-        if (fs.existsSync(existingConfigPath)) {
-          const existingConfig = JSON.parse(
-            fs.readFileSync(existingConfigPath, "utf8")
+        userSetupResult = await initializeUser(process.cwd());
+        if (userSetupResult.success) {
+          isGatewayAvailable = true;
+          console.log(
+            boxen(
+              chalk.green("✅ Gateway Connection Successful") +
+                "\n\n" +
+                chalk.white("TaskMaster AI Gateway is available.") +
+                "\n" +
+                chalk.white("You can choose between BYOK or Hosted mode."),
+              {
+                padding: 1,
+                margin: { top: 1, bottom: 1 },
+                borderStyle: "round",
+                borderColor: "green",
+              }
+            )
          );
-          existingUserId = existingConfig.account?.userId;
-        }
-
-        if (existingUserId) {
-          userSetupResult = await registerUserWithGateway(null, process.cwd());
-          if (!userSetupResult.success) {
-            throw new Error(
-              `Failed to validate existing user: ${userSetupResult.error}`
-            );
-          }
        } else {
-          userSetupResult = await initializeUser(process.cwd());
-          if (!userSetupResult.success) {
-            throw new Error(
-              `Failed to initialize user: ${userSetupResult.error}`
-            );
-          }
+          // Silent failure - gateway not available
+          isGatewayAvailable = false;
        }
      } catch (error) {
-        log("error", `User initialization failed: ${error.message}`);
-        // Don't fall back to random userId - exit or prompt user
-        throw error;
+        // Silent failure - gateway not available
+        isGatewayAvailable = false;
+        userSetupResult = null;
      }

-      // STEP 2: Choose AI access method using inquirer
-      const modeResponse = await inquirer.prompt([
-        {
-          type: "list",
-          name: "accessMode",
-          message: "Choose Your AI Access Method:",
-          choices: [
-            {
-              name: "🔑 BYOK - Bring Your Own API Keys (You manage API keys & billing)",
-              value: "byok",
-            },
-            {
-              name: "🎯 Hosted API Gateway - All models, no keys needed (Recommended)",
-              value: "hosted",
-            },
-          ],
-          default: "hosted",
-        },
-      ]);
-
-      const selectedMode = modeResponse.accessMode;
-
-      console.log(
-        boxen(
-          selectedMode === "byok"
-            ? chalk.blue.bold("🔑 BYOK Mode Selected") +
-                "\n\n" +
-                chalk.white("You'll manage your own API keys and billing.") +
-                "\n" +
-                chalk.white("After setup, add your API keys to ") +
-                chalk.cyan(".env") +
-                chalk.white(" file.")
-            : chalk.green.bold("🎯 Hosted API Gateway Selected") +
-                "\n\n" +
-                chalk.white(
-                  "All AI models available instantly - no API keys needed!"
-                ) +
-                "\n" +
-                chalk.dim("Let's set up your subscription plan..."),
-          {
-            padding: 1,
-            margin: { top: 1, bottom: 1 },
-            borderStyle: "round",
-            borderColor: selectedMode === "byok" ? "blue" : "green",
-          }
-        )
-      );
-
-      // STEP 3: If hosted mode, handle subscription plan with Stripe simulation
+      // STEP 3: Choose AI access method (conditional based on gateway availability)
+      let selectedMode = "byok"; // Default to BYOK
      let selectedPlan = null;
-      if (selectedMode === "hosted") {
-        selectedPlan = await handleHostedSubscription();
+
+      if (isGatewayAvailable) {
+        // Gateway is available, show both options
+        const modeResponse = await inquirer.prompt([
+          {
+            type: "list",
+            name: "accessMode",
+            message: "Choose Your AI Access Method:",
+            choices: [
+              {
+                name: "🔑 BYOK - Bring Your Own API Keys (You manage API keys & billing)",
+                value: "byok",
+              },
+              {
+                name: "🎯 Hosted API Gateway - All models, no keys needed (Recommended)",
+                value: "hosted",
+              },
+            ],
+            default: "hosted",
+          },
+        ]);
+        selectedMode = modeResponse.accessMode;
+
+        console.log(
+          boxen(
+            selectedMode === "byok"
+              ? chalk.blue.bold("🔑 BYOK Mode Selected") +
+                  "\n\n" +
+                  chalk.white("You'll manage your own API keys and billing.") +
+                  "\n" +
+                  chalk.white("After setup, add your API keys to ") +
+                  chalk.cyan(".env") +
+                  chalk.white(" file.")
+              : chalk.green.bold("🎯 Hosted API Gateway Selected") +
+                  "\n\n" +
+                  chalk.white(
+                    "All AI models available instantly - no API keys needed!"
+                  ) +
+                  "\n" +
+                  chalk.dim("Let's set up your subscription plan..."),
+            {
+              padding: 1,
+              margin: { top: 1, bottom: 1 },
+              borderStyle: "round",
+              borderColor: selectedMode === "byok" ? "blue" : "green",
+            }
+          )
+        );
+
+        // If hosted mode selected, handle subscription plan
+        if (selectedMode === "hosted") {
+          selectedPlan = await handleHostedSubscription();
+        }
+      } else {
+        // Gateway not available, silently proceed with BYOK mode
+        // Show standard BYOK mode message without mentioning gateway failure
+        console.log(
+          boxen(
+            chalk.blue.bold("🔑 BYOK Mode") +
+              "\n\n" +
+              chalk.white("You'll manage your own API keys and billing.") +
+              "\n" +
+              chalk.white("After setup, add your API keys to ") +
+              chalk.cyan(".env") +
+              chalk.white(" file."),
+            {
+              padding: 1,
+              margin: { top: 1, bottom: 1 },
+              borderStyle: "round",
+              borderColor: "blue",
+            }
+          )
+        );
+        selectedMode = "byok";
      }

-      // STEP 4: Continue with aliases (this fixes the hanging issue)
+      // STEP 4: Continue with aliases
      const aliasResponse = await inquirer.prompt([
        {
          type: "confirm",
@@ -560,7 +569,6 @@ async function initializeProject(options = {}) {
      ]);

      const addAliases = aliasResponse.addAliases;
-
      const dryRun = options.dryRun || false;

      // STEP 5: Show overview and continue with project creation
@@ -571,7 +579,7 @@ async function initializeProject(options = {}) {
        userSetupResult,
        selectedMode,
        selectedPlan,
-        userSetupResult.userId
+        userSetupResult?.userId || null
      );
    } catch (error) {
      rl.close();
@@ -856,7 +864,7 @@ function configureTaskmasterConfig(
    // Store account-specific configuration
    config.account.mode = selectedMode;
    config.account.userId = userId || null;
-    config.account.userEmail = gatewayRegistration?.email || "";
+    config.account.email = gatewayRegistration?.email || "";
    config.account.telemetryEnabled = selectedMode === "hosted";

    // Store remaining global config items
--- a/scripts/modules/ai-services-unified.js
+++ b/scripts/modules/ai-services-unified.js
@@ -43,6 +43,8 @@ import {
  VertexAIProvider,
 } from "../../src/ai-providers/index.js";

+import { zodToJsonSchema } from "zod-to-json-schema";
+
 // Create provider instances
 const PROVIDERS = {
  anthropic: new AnthropicAIProvider(),
@@ -293,11 +295,11 @@ async function _callGatewayAI(
  initialRole
 ) {
  // Hard-code service-level constants
-  const gatewayUrl = "http://localhost:4444"; // or your production URL
-  const serviceApiKey = "339a81c9-5b9c-4d60-92d8-cba2ee2a8cc3"; // Hardcoded service key -- if you change this, the Hosted Gateway will not work
+  const gatewayUrl = "http://localhost:4444";
+  const serviceId = "98fb3198-2dfc-42d1-af53-07b99e4f3bde"; // Hardcoded service ID -- if you change this, the Hosted Gateway will not work

  // Get user auth info for headers
-  const userMgmt = require("./user-management.js");
+  const userMgmt = await import("./user-management.js");
  const userToken = await userMgmt.getUserToken(projectRoot);
  const userEmail = await userMgmt.getUserEmail(projectRoot);

@@ -316,6 +318,8 @@ async function _callGatewayAI(
    callParams.messages?.find((m) => m.role === "user")?.content || "";

  const requestBody = {
+    provider: providerName,
+    serviceType,
    role: initialRole,
    messages: callParams.messages,
    modelId,
@@ -326,14 +330,14 @@ async function _callGatewayAI(
      temperature: callParams.temperature,
    },
    ...(serviceType === "generateObject" && {
-      schema: callParams.schema,
+      schema: zodToJsonSchema(callParams.schema),
      objectName: callParams.objectName,
    }),
  };

  const headers = {
    "Content-Type": "application/json",
-    "X-TaskMaster-API-Key": serviceApiKey, // Service-level auth (hardcoded)
+    "X-TaskMaster-Service-ID": serviceId, // TaskMaster service ID for instance auth
    Authorization: `Bearer ${userToken}`, // User-level auth
  };

@@ -405,6 +409,10 @@ async function _unifiedServiceRunner(serviceType, params) {
      outputType,
      projectRoot,
    });
+
+    if (isHostedMode(projectRoot)) {
+      log("info", "Communicating with Taskmaster Gateway");
+    }
  }

  const effectiveProjectRoot = projectRoot || findProjectRoot();
@@ -426,10 +434,12 @@ async function _unifiedServiceRunner(serviceType, params) {

        log("info", "User successfully authenticated with gateway");
      } else {
-        log("warn", `Silent auth/init failed: ${initResult.error}`);
+        // Silent failure - only log at debug level during init sequence
+        log("debug", `Silent auth/init failed: ${initResult.error}`);
      }
    } catch (error) {
-      log("warn", `Silent auth/init attempt failed: ${error.message}`);
+      // Silent failure - only log at debug level during init sequence
+      log("debug", `Silent auth/init attempt failed: ${error.message}`);
    }
  }

--- a/scripts/modules/config-manager.js
+++ b/scripts/modules/config-manager.js
@@ -66,9 +66,9 @@ const defaultConfig = {
  },
  account: {
    userId: "1234567890", // Placeholder that triggers auth/init
-    userEmail: "",
+    email: "",
    mode: "byok",
-    telemetryEnabled: false,
+    telemetryEnabled: true,
  },
 };

@@ -129,7 +129,6 @@ function _loadAndValidateConfig(explicitRoot = null) {
              : { ...defaults.models.fallback },
        },
        global: { ...defaults.global, ...parsedConfig?.global },
-        ai: { ...defaults.ai, ...parsedConfig?.ai },
        account: { ...defaults.account, ...parsedConfig?.account },
      };
      configSource = `file (${configPath})`; // Update source info
@@ -756,7 +755,7 @@ function getTelemetryEnabled(explicitRoot = null) {
 // Update getUserEmail to use account
 function getUserEmail(explicitRoot = null) {
  const config = getConfig(explicitRoot);
-  return config.account?.userEmail || "";
+  return config.account?.email || "";
 }

 // Update getMode function to use account
--- a/scripts/modules/telemetry-submission.js
+++ b/scripts/modules/telemetry-submission.js
@@ -32,31 +32,24 @@ const MAX_RETRIES = 3;
 const RETRY_DELAY = 1000; // 1 second

 /**
- * Get telemetry configuration from environment variables only
- * @returns {Object} Configuration object with apiKey, userId, and email
+ * Get telemetry configuration from hardcoded service ID, user token, and config
+ * @returns {Object} Configuration object with serviceId, apiKey, userId, and email
 */
 function getTelemetryConfig() {
-  // Try environment variables first (includes .env file via resolveEnvVariable)
-  const envApiKey =
-    resolveEnvVariable("TASKMASTER_API_KEY") ||
-    resolveEnvVariable("GATEWAY_API_KEY") ||
-    resolveEnvVariable("TELEMETRY_API_KEY");
-  const envUserId =
-    resolveEnvVariable("TASKMASTER_USER_ID") ||
-    resolveEnvVariable("GATEWAY_USER_ID") ||
-    resolveEnvVariable("TELEMETRY_USER_ID");
-  const envEmail =
-    resolveEnvVariable("TASKMASTER_USER_EMAIL") ||
-    resolveEnvVariable("GATEWAY_USER_EMAIL") ||
-    resolveEnvVariable("TELEMETRY_USER_EMAIL");
-
-  // Get the config (which might contain userId)
+  // Get the config which contains userId and email
  const config = getConfig();

+  // Hardcoded service ID for TaskMaster telemetry service
+  const hardcodedServiceId = "98fb3198-2dfc-42d1-af53-07b99e4f3bde";
+
+  // Get user's API token from .env (managed by user-management.js)
+  const userApiKey = resolveEnvVariable("TASKMASTER_API_KEY");
+
  return {
-    apiKey: envApiKey || null, // API key should only come from environment
-    userId: envUserId || config?.account?.userId || null,
-    email: envEmail || null,
+    serviceId: hardcodedServiceId, // Hardcoded service identifier
+    apiKey: userApiKey || null, // User's Bearer token from .env
+    userId: config?.account?.userId || null, // From config
+    email: config?.account?.email || null, // From config
  };
 }

@@ -119,8 +112,11 @@ export async function registerUserWithGateway(email = null, userId = null) {
 */
 export async function submitTelemetryData(telemetryData) {
  try {
-    // Check user opt-out preferences first
-    if (!getTelemetryEnabled()) {
+    // Check user opt-out preferences first, but hosted mode always sends telemetry
+    const config = getConfig();
+    const isHostedMode = config?.account?.mode === "hosted";
+
+    if (!isHostedMode && !getTelemetryEnabled()) {
      return {
        success: true,
        skipped: true,
@@ -138,7 +134,7 @@ export async function submitTelemetryData(telemetryData) {
      return {
        success: false,
        error:
-          "Telemetry configuration incomplete. Run 'task-master init' and select hosted gateway option, or manually set TASKMASTER_API_KEY, TASKMASTER_USER_ID, and TASKMASTER_USER_EMAIL environment variables",
+          "Telemetry configuration incomplete. Please ensure you have completed 'task-master init' to set up your user account.",
      };
    }

@@ -167,8 +163,9 @@ export async function submitTelemetryData(telemetryData) {
          method: "POST",
          headers: {
            "Content-Type": "application/json",
-            Authorization: `Bearer ${telemetryConfig.apiKey}`, // Use Bearer token format
-            "X-User-Email": telemetryConfig.email, // Add required email header
+            "X-Service-ID": telemetryConfig.serviceId, // Hardcoded service ID
+            Authorization: `Bearer ${telemetryConfig.apiKey}`, // User's Bearer token
+            "X-User-Email": telemetryConfig.email, // User's email from config
          },
          body: JSON.stringify(completeTelemetryData),
        });
--- a/scripts/modules/ui.js
+++ b/scripts/modules/ui.js
@@ -1513,18 +1513,16 @@ async function displayComplexityReport(reportPath) {
      )
    );

-    const readline = require("readline").createInterface({
+    const readline = await import("readline");
+    const rl = readline.default.createInterface({
      input: process.stdin,
      output: process.stdout,
    });

    const answer = await new Promise((resolve) => {
-      readline.question(
-        chalk.cyan("Generate complexity report? (y/n): "),
-        resolve
-      );
+      rl.question(chalk.cyan("Generate complexity report? (y/n): "), resolve);
    });
-    readline.close();
+    rl.close();

    if (answer.toLowerCase() === "y" || answer.toLowerCase() === "yes") {
      // Call the analyze-complexity command
@@ -2029,43 +2027,114 @@ function displayAvailableModels(availableModels) {
 * @param {string} outputType - 'cli' or 'mcp' (though typically only called for 'cli').
 */
 function displayAiUsageSummary(telemetryData, outputType = "cli") {
-  if (
-    (outputType !== "cli" && outputType !== "text") ||
-    !telemetryData ||
-    isSilentMode()
-  ) {
-    return; // Only display for CLI and if data exists and not in silent mode
-  }
+  if (!telemetryData || outputType !== "cli") return;

  const {
-    modelUsed,
-    providerName,
-    inputTokens,
-    outputTokens,
-    totalTokens,
-    totalCost,
    commandName,
+    providerName,
+    modelUsed,
+    inputTokens = 0,
+    outputTokens = 0,
+    totalTokens = 0,
+    totalCost = 0,
+    accountInfo,
  } = telemetryData;

-  let summary = chalk.bold.blue("AI Usage Summary:") + "\n";
-  summary += chalk.gray(`  Command: ${commandName}\n`);
-  summary += chalk.gray(`  Provider: ${providerName}\n`);
-  summary += chalk.gray(`  Model: ${modelUsed}\n`);
-  summary += chalk.gray(
-    `  Tokens: ${totalTokens} (Input: ${inputTokens}, Output: ${outputTokens})\n`
-  );
-  summary += chalk.gray(`  Est. Cost: $${totalCost.toFixed(6)}`);
+  const isHostedMode = !!accountInfo;

-  console.log(
-    boxen(summary, {
-      padding: 1,
-      margin: { top: 1 },
-      borderColor: "blue",
-      borderStyle: "round",
-      title: "💡 Telemetry",
-      titleAlignment: "center",
-    })
-  );
+  // Build the core telemetry content (same for both modes)
+  let content =
+    chalk.cyan.bold("AI Usage Summary:") +
+    "\n" +
+    chalk.gray("  Command: ") +
+    chalk.white(commandName) +
+    "\n" +
+    chalk.gray("  Provider: ") +
+    chalk.white(providerName) +
+    "\n" +
+    chalk.gray("  Model: ") +
+    chalk.white(modelUsed) +
+    "\n" +
+    chalk.gray("  Tokens: ") +
+    chalk.white(
+      `${totalTokens.toLocaleString()} (Input: ${inputTokens.toLocaleString()}, Output: ${outputTokens.toLocaleString()})`
+    );
+
+  if (isHostedMode) {
+    // Hosted mode - add premium credit information
+    const creditsUsed = accountInfo.creditsUsed || 0;
+    const remainingCredits = accountInfo.remainingCredits || 0;
+    const totalCredits = creditsUsed + remainingCredits;
+    const remainingPercentage =
+      totalCredits > 0 ? (remainingCredits / totalCredits) * 100 : 0;
+
+    // Create credit bar showing REMAINING credits (goes down as credits are consumed)
+    const barLength = 20;
+    const filledLength = Math.round((remainingPercentage / 100) * barLength);
+    const emptyLength = barLength - filledLength;
+    const creditBar =
+      chalk.green("█".repeat(filledLength)) +
+      chalk.gray("█".repeat(emptyLength));
+
+    // Determine footer message based on remaining credits
+    let footerMessage;
+    if (remainingPercentage <= 20) {
+      footerMessage = chalk.yellow(
+        "⚠️  Running low on credits - consider topping up soon"
+      );
+    } else {
+      footerMessage = chalk.dim("🚀 Powered by hosted AI infrastructure");
+    }
+
+    content +=
+      "\n" +
+      chalk.gray("  Credits Used: ") +
+      chalk.yellow.bold(creditsUsed.toLocaleString()) +
+      "\n" +
+      chalk.gray("  Remaining: ") +
+      chalk.green.bold(remainingCredits.toLocaleString()) +
+      "\n" +
+      chalk.gray("  Usage: ") +
+      `[${creditBar}] ${remainingPercentage.toFixed(1)}%` +
+      "\n\n" +
+      footerMessage;
+
+    // Premium double border
+    console.log(
+      boxen(content, {
+        padding: { top: 1, bottom: 1, left: 2, right: 2 },
+        margin: { top: 1, bottom: 1 },
+        borderStyle: "double",
+        borderColor: "cyan",
+        title: chalk.yellow.bold("💎 PREMIUM TELEMETRY"),
+        titleAlignment: "center",
+      })
+    );
+  } else {
+    // BYOK mode - add cost and upgrade CTA
+    content +=
+      "\n" +
+      chalk.gray("  Est. Cost: ") +
+      chalk.white(`$${totalCost.toFixed(6)}`) +
+      "\n\n" +
+      chalk.cyan(
+        "⚡ Upgrade to TaskMaster Premium for instant access to all AI models"
+      ) +
+      "\n" +
+      chalk.dim("Learn more: https://taskmaster.ai/premium");
+
+    // Regular single border
+    console.log(
+      boxen(content, {
+        padding: { top: 1, bottom: 1, left: 2, right: 2 },
+        margin: { top: 1, bottom: 1 },
+        borderStyle: "round",
+        borderColor: "blue",
+        title: chalk.blue.bold("💡 TELEMETRY"),
+        titleAlignment: "center",
+      })
+    );
+  }
 }

 /**
--- a/scripts/modules/user-management.js
+++ b/scripts/modules/user-management.js
@@ -1,7 +1,7 @@
 import fs from "fs";
 import path from "path";
 import { log, findProjectRoot } from "./utils.js";
-import { getConfig, writeConfig } from "./config-manager.js";
+import { getConfig, writeConfig, getUserId } from "./config-manager.js";

 /**
 * Registers or finds a user via the gateway's /auth/init endpoint
@@ -14,8 +14,18 @@ async function registerUserWithGateway(email = null, explicitRoot = null) {
    const gatewayUrl =
      process.env.TASKMASTER_GATEWAY_URL || "http://localhost:4444";

-    // Email is optional - only send if provided
-    const requestBody = email ? { email } : {};
+    // Check for existing userId and email to pass to gateway
+    const existingUserId = getUserId(explicitRoot);
+    const existingEmail = email || getUserEmail(explicitRoot);
+
+    // Build request body with existing values (gateway can handle userId for existing users)
+    const requestBody = {};
+    if (existingUserId && existingUserId !== "1234567890") {
+      requestBody.userId = existingUserId;
+    }
+    if (existingEmail) {
+      requestBody.email = existingEmail;
+    }

    const response = await fetch(`${gatewayUrl}/auth/init`, {
      method: "POST",
@@ -70,10 +80,17 @@ async function registerUserWithGateway(email = null, explicitRoot = null) {
 * @param {string} userId - User ID from gateway
 * @param {string} token - User authentication token from gateway (stored in .env)
 * @param {string} mode - User mode ('byok' or 'hosted')
+ * @param {string|null} email - Optional user email to save
 * @param {string|null} explicitRoot - Optional explicit project root path
 * @returns {boolean} Success status
 */
-function updateUserConfig(userId, token, mode, explicitRoot = null) {
+function updateUserConfig(
+  userId,
+  token,
+  mode,
+  email = null,
+  explicitRoot = null
+) {
  try {
    const config = getConfig(explicitRoot);

@@ -82,10 +99,20 @@ function updateUserConfig(userId, token, mode, explicitRoot = null) {
      config.account = {};
    }

+    // Ensure global section exists for email
+    if (!config.global) {
+      config.global = {};
+    }
+
    // Update user configuration in account section
    config.account.userId = userId;
    config.account.mode = mode; // 'byok' or 'hosted'

+    // Save email if provided
+    if (email) {
+      config.account.email = email;
+    }
+
    // Write user authentication token to .env file (not config)
    if (token) {
      writeApiKeyToEnv(token, explicitRoot);
@@ -94,7 +121,11 @@ function updateUserConfig(userId, token, mode, explicitRoot = null) {
    // Save updated config
    const success = writeConfig(config, explicitRoot);
    if (success) {
-      log("info", `User configuration updated: userId=${userId}, mode=${mode}`);
+      const emailInfo = email ? `, email=${email}` : "";
+      log(
+        "info",
+        `User configuration updated: userId=${userId}, mode=${mode}${emailInfo}`
+      );
    } else {
      log("error", "Failed to write updated user configuration");
    }
@@ -219,10 +250,12 @@ async function setupUser(email = null, mode = "hosted", explicitRoot = null) {
      };
    }

-    // Step 2: Update config with userId and mode
-    const configResult = await updateUserConfig(
+    // Step 2: Update config with userId, mode, and email
+    const configResult = updateUserConfig(
      registrationResult.userId,
+      registrationResult.token,
      mode,
+      email,
      explicitRoot
    );

@@ -261,39 +294,58 @@ async function setupUser(email = null, mode = "hosted", explicitRoot = null) {
 */
 async function initializeUser(explicitRoot = null) {
  try {
-    // Register with gateway without email
+    // Try to register with gateway without email
    const result = await registerUserWithGateway(null, explicitRoot);

-    if (!result.success) {
-      return {
-        success: false,
-        userId: "",
-        error: result.error,
-      };
-    }
+    // If gateway call succeeded, use the returned values
+    if (result.success) {
+      // Update config with userId, token, and preserve existing mode (or default)
+      const existingMode = getUserMode(explicitRoot);
+      const modeToUse = existingMode !== "unknown" ? existingMode : "byok";

-    // Update config with userId, token, and default hosted mode
-    const configResult = updateUserConfig(
-      result.userId,
-      result.token, // Include the token parameter
-      "hosted", // Default to hosted mode until user chooses plan
-      explicitRoot
-    );
+      const configResult = updateUserConfig(
+        result.userId,
+        result.token,
+        modeToUse,
+        null,
+        explicitRoot
+      );
+
+      if (!configResult) {
+        return {
+          success: false,
+          userId: result.userId,
+          error: "Failed to update user configuration",
+        };
+      }

-    if (!configResult) {
      return {
-        success: false,
+        success: true,
        userId: result.userId,
-        error: "Failed to update user configuration",
+        message: result.isNewUser
+          ? "New user registered with gateway"
+          : "Existing user found in gateway",
      };
    }

+    // Gateway call failed - check if we have existing credentials to use
+    const existingUserId = getUserId(explicitRoot);
+    const existingToken = getUserToken(explicitRoot);
+
+    if (existingUserId && existingUserId !== "1234567890" && existingToken) {
+      // We have existing credentials, use them (gateway unavailable scenario)
+      return {
+        success: true,
+        userId: existingUserId,
+        message: "Gateway unavailable, using existing user credentials",
+      };
+    }
+
+    // No existing credentials and gateway failed
    return {
-      success: true,
-      userId: result.userId,
-      message: result.isNewUser
-        ? "New user registered with gateway"
-        : "Existing user found in gateway",
+      success: false,
+      userId: "",
+      error: result.error,
    };
  } catch (error) {
    return {
@@ -351,7 +403,7 @@ function getUserToken(explicitRoot = null) {
 function getUserEmail(explicitRoot = null) {
  try {
    const config = getConfig(explicitRoot);
-    return config?.global?.email || null;
+    return config?.account?.email || null;
  } catch (error) {
    log("error", `Error getting user email: ${error.message}`);
    return null;