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/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;