@unifiedmemory/cli 1.3.13 → 1.3.14

package/commands/usage.js

@@ -0,0 +1,163 @@
+import chalk from 'chalk';
+import { loadAndRefreshToken } from '../lib/token-validation.js';
+import { config } from '../lib/config.js';
+
+/**
+ * Show current usage and quota allowances
+ */
+export async function usage() {
+  // Load token and refresh if expired
+  const tokenData = await loadAndRefreshToken();
+
+  const accessToken = tokenData.idToken || tokenData.accessToken;
+
+  if (!accessToken) {
+    console.error(chalk.red('❌ Not authenticated'));
+    console.log(chalk.gray('Run `um login` to authenticate'));
+    process.exit(1);
+  }
+
+  console.log(chalk.blue('\n📊 Usage & Quota\n'));
+
+  try {
+    // Build headers with org context (selectedOrgId fallback for when JWT has no org claims)
+    const userId = tokenData.decoded?.sub;
+    const orgId = tokenData.decoded?.o?.o_id || tokenData.selectedOrgId;
+
+    const headers = {
+      'Authorization': `Bearer ${accessToken}`,
+    };
+    if (userId) headers['X-User-Id'] = userId;
+    if (orgId) headers['X-Org-Id'] = orgId;
+
+    // Fetch quota usage from API gateway
+    const response = await fetch(`${config.apiEndpoint}/v1/quota/usage`, { headers });
+
+    if (!response.ok) {
+      if (response.status === 401) {
+        console.error(chalk.red('❌ Authentication failed'));
+        console.log(chalk.gray('Run `um login` to re-authenticate'));
+        process.exit(1);
+      }
+      throw new Error(`API returned ${response.status}: ${response.statusText}`);
+    }
+
+    const data = await response.json();
+    displayUsage(data, tokenData);
+
+  } catch (error) {
+    console.error(chalk.red('❌ Failed to fetch usage data:'), error.message);
+    process.exit(1);
+  }
+}
+
+/**
+ * Display formatted usage information
+ */
+function displayUsage(quotaData, tokenData) {
+  const { quota_type, usage, period } = quotaData;
+  const { used, total, percentage_used, remaining } = usage;
+
+  // Display account context (use selectedOrgName fallback for when JWT has no org claims)
+  console.log(chalk.yellow('Account:'));
+  if (quota_type === 'organization') {
+    // Get org name from JWT claims, fallback to selectedOrgName from login state
+    const orgName = tokenData.decoded?.o?.o_name || tokenData.selectedOrgName || 'Organization';
+    console.log(chalk.white(`  Organization: ${orgName}`));
+  } else {
+    console.log(chalk.white('  Personal Account'));
+  }
+  console.log('');
+
+  // Handle special cases
+  if (total === 0) {
+    console.log(chalk.yellow('Monthly Queries:'));
+    console.log(chalk.yellow('  No active subscription'));
+    console.log(chalk.gray('  Start at: https://unifiedmemory.ai/pricing'));
+    console.log('');
+    return;
+  }
+
+  if (total < 0 || total > 1000000) {
+    console.log(chalk.yellow('Monthly Queries:'));
+    console.log(chalk.green('  Unlimited ♾️'));
+    console.log('');
+    return;
+  }
+
+  // Determine color based on usage percentage
+  const color = getColorForUsage(percentage_used);
+  const colorFn = getChalkFunction(color);
+
+  // Display usage count and percentage
+  console.log(chalk.yellow('Monthly Queries:'));
+  console.log(colorFn(`  ${used.toLocaleString()} / ${total.toLocaleString()} (${percentage_used.toFixed(1)}%)`));
+
+  // Display progress bar
+  const progressBar = buildProgressBar(percentage_used, 20);
+  console.log(colorFn(`  [${progressBar}]`));
+
+  // Display remaining quota
+  console.log(colorFn(`  Remaining: ${remaining.toLocaleString()} queries`));
+
+  // Display reset date if available
+  if (period?.next_reset_date) {
+    const resetDate = new Date(period.next_reset_date);
+    const formattedDate = resetDate.toLocaleDateString('en-US', {
+      month: 'numeric',
+      day: 'numeric',
+      year: 'numeric'
+    });
+    console.log(chalk.gray(`  Resets: ${formattedDate}`));
+  }
+
+  console.log('');
+
+  // Show upgrade prompt if at or above threshold
+  if (percentage_used >= 90) {
+    console.log(chalk.red('⚠️  Running low on quota!'));
+    console.log(chalk.gray('Upgrade your plan: https://unifiedmemory.ai/pricing'));
+    console.log('');
+  }
+}
+
+/**
+ * Get color category based on usage percentage
+ * @param {number} percentage - Usage percentage (0-100)
+ * @returns {string} - Color category: 'green', 'yellow', or 'red'
+ */
+function getColorForUsage(percentage) {
+  if (percentage >= 90) return 'red';
+  if (percentage >= 70) return 'yellow';
+  return 'green';
+}
+
+/**
+ * Get chalk color function based on color name
+ * @param {string} color - Color name
+ * @returns {Function} - Chalk color function
+ */
+function getChalkFunction(color) {
+  switch (color) {
+    case 'red':
+      return chalk.red;
+    case 'yellow':
+      return chalk.yellow;
+    case 'green':
+      return chalk.green;
+    default:
+      return chalk.white;
+  }
+}
+
+/**
+ * Build a text-based progress bar
+ * @param {number} percentage - Usage percentage (0-100)
+ * @param {number} length - Total length of the bar in characters
+ * @returns {string} - Progress bar string
+ */
+function buildProgressBar(percentage, length) {
+  const filled = Math.round((percentage / 100) * length);
+  const empty = length - filled;
+  return '█'.repeat(filled) + '░'.repeat(empty);
+}

package/index.js

@@ -11,8 +11,10 @@ import { login } from './commands/login.js';
 import { init } from './commands/init.js';
 import { switchOrg, showOrg, fixOrg } from './commands/org.js';
 import { record } from './commands/record.js';
+import { usage } from './commands/usage.js';
+import { debugAuth } from './commands/debug.js';
 import { config } from './lib/config.js';
-import { getSelectedOrg } from './lib/token-storage.js';
+import { getOrgContext } from './lib/token-storage.js';
 import { loadAndRefreshToken } from './lib/token-validation.js';
 import { showWelcome } from './lib/welcome.js';
 
@@ -74,7 +76,8 @@ program
     try {
       // Try to load and refresh token if expired
       const tokenData = await loadAndRefreshToken(false);
-      const selectedOrg = getSelectedOrg();
+      // Get org context from JWT claims (single source of truth)
+      const orgContext = getOrgContext();
 
       console.log(chalk.blue('\n📋 UnifiedMemory Status\n'));
 
@@ -94,10 +97,10 @@ program
       }
 
       console.log(chalk.yellow('\nOrganization Context:'));
-      if (selectedOrg) {
-        console.log(chalk.green(`  ${selectedOrg.name} (${selectedOrg.slug})`));
-        console.log(chalk.gray(`  ID: ${selectedOrg.id}`));
-        console.log(chalk.gray(`  Role: ${selectedOrg.role}`));
+      if (orgContext) {
+        console.log(chalk.green(`  ${orgContext.name}${orgContext.slug ? ` (${orgContext.slug})` : ''}`));
+        console.log(chalk.gray(`  ID: ${orgContext.id}`));
+        console.log(chalk.gray(`  Role: ${orgContext.role || 'member'}`));
       } else {
         console.log(chalk.cyan('  Personal Account'));
       }
@@ -128,6 +131,34 @@ program
     }
   });
 
+// Usage command
+program
+  .command('usage')
+  .description('Show current usage and quota allowances')
+  .action(async () => {
+    try {
+      await usage();
+      process.exit(0);
+    } catch (error) {
+      console.error(chalk.red('Failed to show usage:'), error.message);
+      process.exit(1);
+    }
+  });
+
+// Debug command
+program
+  .command('debug')
+  .description('Debug authentication and token issues')
+  .action(async () => {
+    try {
+      await debugAuth();
+      process.exit(0);
+    } catch (error) {
+      console.error(chalk.red('Debug failed:'), error.message);
+      process.exit(1);
+    }
+  });
+
 // Organization management
 const orgCommand = program
   .command('org')

package/lib/clerk-api.js

@@ -121,30 +121,27 @@ export function formatOrganization(membership) {
 }
 
 /**
- * Get an organization-scoped JWT token via backend API
+ * Get a JWT token with specific org context via backend API
  * @param {string} sessionId - Session ID from JWT sid claim
- * @param {string} orgId - Organization ID to set as active
+ * @param {string|null} orgId - Organization ID to set as active, or null for personal context
  * @param {string} currentToken - Current session token (for authentication)
- * @returns {Promise<{jwt: string, org_id: string}>} New token with org context
+ * @returns {Promise<{jwt: string, org_id: string|null}>} New token with specified context
  */
 export async function getOrgScopedToken(sessionId, orgId, currentToken) {
   if (!sessionId) {
     throw new Error('Session ID is required');
   }
-  if (!orgId) {
-    throw new Error('Organization ID is required');
-  }
   if (!currentToken) {
     throw new Error('Current token is required');
   }
+  // orgId can be null for personal context
 
-  console.log(chalk.gray(`Requesting org-scoped token from backend...`));
+  const contextDesc = orgId ? `org ${orgId}` : 'personal context';
+  console.log(chalk.gray(`Requesting token for ${contextDesc}...`));
   console.log(chalk.gray(`  Session: ${sessionId}`));
-  console.log(chalk.gray(`  Org: ${orgId}`));
 
   // Call backend API which proxies to Clerk Backend API
   const apiUrl = `${config.apiEndpoint}/v1/auth/org-token`;
-  console.log(chalk.gray(`Calling backend API: ${apiUrl}`));
 
   const response = await fetch(apiUrl, {
     method: 'POST',
@@ -154,7 +151,7 @@ export async function getOrgScopedToken(sessionId, orgId, currentToken) {
     },
     body: JSON.stringify({
       session_id: sessionId,
-      org_id: orgId
+      org_id: orgId  // null for personal context
     })
   });
 
@@ -162,11 +159,11 @@ export async function getOrgScopedToken(sessionId, orgId, currentToken) {
     const errorText = await response.text();
     console.error(chalk.red(`Backend API failed: ${response.status}`));
     console.error(chalk.gray(errorText));
-    throw new Error(`Failed to get org-scoped token: ${response.status} - ${errorText}`);
+    throw new Error(`Failed to get token: ${response.status} - ${errorText}`);
   }
 
   const tokenData = await response.json();
-  console.log(chalk.gray('✓ Received org-scoped token from backend'));
+  console.log(chalk.gray(`✓ Received token for ${contextDesc}`));
 
   return tokenData;
 }

package/lib/config.js

@@ -22,7 +22,10 @@ export const config = {
   port: parseInt(process.env.PORT || '3333', 10),
 
   // OAuth success page configuration (optional website link for account management)
-  oauthSuccessUrl: process.env.OAUTH_SUCCESS_URL || 'https://unifiedmemory.ai/oauth/callback'
+  oauthSuccessUrl: process.env.OAUTH_SUCCESS_URL || 'https://unifiedmemory.ai/oauth/callback',
+
+  // Frontend URL for CLI auth page (org picker)
+  frontendUrl: process.env.FRONTEND_URL || 'https://unifiedmemory.ai'
 };
 
 // Validation function - validates configuration values
@@ -51,5 +54,12 @@ export function validateConfig() {
     throw new Error(`OAUTH_SUCCESS_URL must be a valid URL (got: ${config.oauthSuccessUrl})`);
   }
 
+  // Validate frontendUrl format
+  try {
+    new URL(config.frontendUrl);
+  } catch (e) {
+    throw new Error(`FRONTEND_URL must be a valid URL (got: ${config.frontendUrl})`);
+  }
+
   return true;
 }

package/lib/mcp-proxy.js

@@ -82,8 +82,9 @@ function transformToolSchema(tool) {
 
 /**
  * Inject context parameters into tool arguments
+ * JWT is the single source of truth for org context
  * @param {Object} args - Tool arguments from AI agent
- * @param {Object} authContext - Auth context with user/org info
+ * @param {Object} authContext - Auth context with user/org info (from JWT claims)
  * @param {Object|null} projectContext - Project config from .um/config.json
  * @returns {Object} - Arguments with injected context params
  */
@@ -101,21 +102,19 @@ function injectContextParams(args, authContext, projectContext) {
     injected.headers = {};
   }
 
-  // Inject user ID from decoded JWT
+  // Inject user ID from decoded JWT (always present)
   if (authContext?.decoded?.sub) {
     injected.pathParams.user = authContext.decoded.sub;
     injected.headers['X-User-Id'] = authContext.decoded.sub;
   }
 
-  // Inject org ID (from selectedOrg or fallback to user ID for personal account)
-  if (authContext?.selectedOrg?.id) {
-    injected.pathParams.org = authContext.selectedOrg.id;
-    injected.headers['X-Org-Id'] = authContext.selectedOrg.id;
-  } else if (authContext?.decoded?.sub) {
-    // Fallback to user ID for personal account context
-    injected.pathParams.org = authContext.decoded.sub;
-    injected.headers['X-Org-Id'] = authContext.decoded.sub;
+  // ONLY inject org ID if JWT has org claims
+  // NO FALLBACK - gateway handles personal context by using userId
+  if (authContext?.orgContext?.id) {
+    injected.pathParams.org = authContext.orgContext.id;
+    injected.headers['X-Org-Id'] = authContext.orgContext.id;
   }
+  // For personal context: no X-Org-Id header set, gateway falls back to userId
 
   // Inject project ID from project config
   if (projectContext?.project_id) {

package/lib/mcp-server.js

@@ -11,28 +11,32 @@ import path from 'path';
 import os from 'os';
 import { loadAndRefreshToken } from './token-validation.js';
 import { fetchRemoteMCPTools, callRemoteMCPTool } from './mcp-proxy.js';
+import { config } from './config.js';
+
+/**
+ * Usage cache to avoid spamming quota API
+ * Map: orgId -> { data: QuotaResponse, timestamp: number }
+ */
+const usageCache = new Map();
+const CACHE_TTL_MS = 5 * 60 * 1000; // 5 minutes
+const UPGRADE_THRESHOLD = 90; // Show prompt at 90%+
 
 /**
  * Start the MCP server on stdio transport
  */
 export async function startMCPServer() {
   try {
-    // 1. Load and validate authentication
-    const authData = await loadAndValidateAuth();
+    // 1. Validate authentication at startup for immediate failure
+    await loadAndValidateAuth();
 
-    // 2. Load project context from current directory
+    // 2. Load project context from current directory (this doesn't change)
     const projectContext = loadProjectContext();
 
-    // 3. Build auth headers
-    const authHeaders = buildAuthHeaders(authData, projectContext);
+    // Note: We don't cache authData here anymore. Instead, we reload it
+    // on each request to ensure we always have the current token state,
+    // even after token refreshes.
 
-    // 4. Extract auth context for parameter injection
-    const authContext = {
-      decoded: authData.decoded,
-      selectedOrg: authData.selectedOrg
-    };
-
-    // 5. Create MCP server
+    // 3. Create MCP server
     const server = new Server(
       {
         name: "unifiedmemory",
@@ -46,22 +50,25 @@ export async function startMCPServer() {
       }
     );
 
-    // 6. Register handlers
+    // 4. Register handlers - reload auth on each request
     server.setRequestHandler(ListToolsRequestSchema, async () => {
+      const { authHeaders } = await loadFreshAuth(projectContext);
       return await handleListTools(authHeaders, projectContext);
     });
 
     server.setRequestHandler(CallToolRequestSchema, async (request) => {
+      const { authHeaders, authContext } = await loadFreshAuth(projectContext);
       return await handleCallTool(request, authHeaders, authContext, projectContext);
     });
 
     // Register resource handlers for authentication context
     server.setRequestHandler(ListResourcesRequestSchema, async () => {
+      // Reload auth for fresh context
+      const authData = await loadAndRefreshToken();
       const resources = [];
 
-      // Add org_id resource
-      const orgId = authData.selectedOrg?.id || authData.decoded?.sub;
-      if (orgId) {
+      // Add org_id resource only if JWT has org context
+      if (authData.decoded?.o?.o_id) {
         resources.push({
           uri: "um://context/org_id",
           name: "Organization ID",
@@ -95,14 +102,17 @@ export async function startMCPServer() {
     });
 
     server.setRequestHandler(ReadResourceRequestSchema, async (request) => {
+      // Reload auth for fresh context
+      const authData = await loadAndRefreshToken();
       const { uri } = request.params;
       console.error(`→ Reading resource: ${uri}`);
 
       switch(uri) {
         case "um://context/org_id": {
-          const orgId = authData.selectedOrg?.id || authData.decoded?.sub || "";
+          // Only return org_id if JWT has org claims
+          const orgId = authData.decoded?.o?.o_id;
           if (!orgId) {
-            throw new Error("Organization context not available");
+            throw new Error("Organization context not available (using personal account)");
           }
           return {
             contents: [{
@@ -193,6 +203,7 @@ function loadProjectContext() {
 
 /**
  * Build authentication headers for gateway requests
+ * Checks JWT claims first, falls back to selectedOrgId from login
  * @param {Object} authData - Token data
  * @param {Object|null} projectContext - Project config
  * @returns {Object} - Headers object
@@ -202,21 +213,22 @@ function buildAuthHeaders(authData, projectContext) {
     'Authorization': `Bearer ${authData.idToken || authData.accessToken}`,
   };
 
-  // Add org context
-  if (authData.selectedOrg) {
-    headers['X-Org-Id'] = authData.selectedOrg.id;
-    console.error(`✓ Organization: ${authData.selectedOrg.name} (${authData.selectedOrg.id})`);
-  } else if (authData.decoded?.sub) {
-    // Fallback to user ID for personal account
-    headers['X-Org-Id'] = authData.decoded.sub;
-    console.error(`✓ Using personal account (${authData.decoded.sub})`);
-  }
-
-  // Add user ID from JWT
+  // Add user ID from JWT (always present)
   if (authData.decoded?.sub) {
     headers['X-User-Id'] = authData.decoded.sub;
   }
 
+  // Add X-Org-Id - check JWT claims first, fall back to selectedOrgId
+  const orgId = authData.decoded?.o?.o_id || authData.selectedOrgId;
+  const orgName = authData.decoded?.o?.o_name || authData.selectedOrgName;
+
+  if (orgId) {
+    headers['X-Org-Id'] = orgId;
+    console.error(`✓ Organization: ${orgName || orgId}`);
+  } else {
+    console.error(`✓ Using personal account context`);
+  }
+
   // Add project context if available
   if (projectContext) {
     headers['X-Project-Id'] = projectContext.project_id;
@@ -226,6 +238,93 @@ function buildAuthHeaders(authData, projectContext) {
   return headers;
 }
 
+/**
+ * Load fresh authentication data and build headers
+ * Called on each MCP request to ensure current token state
+ * Checks JWT claims first, falls back to selectedOrgId from login
+ * @param {Object|null} projectContext - Project config
+ * @returns {Promise<{authHeaders: Object, authContext: Object}>}
+ */
+async function loadFreshAuth(projectContext) {
+  const authData = await loadAndRefreshToken();
+  const authHeaders = buildAuthHeaders(authData, projectContext);
+
+  // Build auth context - check JWT claims first, fall back to selectedOrgId
+  const jwtOrgId = authData.decoded?.o?.o_id;
+  const orgId = jwtOrgId || authData.selectedOrgId;
+  const orgName = authData.decoded?.o?.o_name || authData.selectedOrgName;
+  const orgRole = authData.decoded?.o?.o_role;
+
+  const authContext = {
+    decoded: authData.decoded,
+    orgContext: orgId ? {
+      id: orgId,
+      name: orgName,
+      role: orgRole,
+    } : null
+  };
+
+  return { authHeaders, authContext };
+}
+
+/**
+ * Check quota usage and return upgrade prompt if needed
+ * Uses caching to avoid API spam (5-min TTL)
+ */
+async function checkAndGetUpgradePrompt(authHeaders, authContext) {
+  try {
+    // Use org from JWT claims, or userId for personal context
+    const orgId = authContext?.orgContext?.id || authContext?.decoded?.sub;
+    if (!orgId) return null;
+
+    // Check cache
+    const cached = usageCache.get(orgId);
+    if (cached && Date.now() - cached.timestamp < CACHE_TTL_MS) {
+      return formatUpgradePrompt(cached.data, authContext);
+    }
+
+    // Fetch fresh data
+    const response = await fetch(`${config.apiEndpoint}/v1/quota/usage`, {
+      headers: authHeaders
+    });
+
+    if (!response.ok) {
+      console.error(`⚠️  Quota check failed: ${response.status}`);
+      return null; // Fail silently, don't break MCP
+    }
+
+    const data = await response.json();
+
+    // Update cache
+    usageCache.set(orgId, { data, timestamp: Date.now() });
+
+    return formatUpgradePrompt(data, authContext);
+
+  } catch (error) {
+    console.error(`⚠️  Quota check error: ${error.message}`);
+    return null; // Fail silently
+  }
+}
+
+/**
+ * Format upgrade prompt if usage >= threshold
+ */
+function formatUpgradePrompt(quotaData, authContext) {
+  const { percentage_used, used, total } = quotaData.usage;
+
+  if (percentage_used < UPGRADE_THRESHOLD) return null;
+
+  const isOrg = quotaData.quota_type === 'organization';
+  const contextName = isOrg
+    ? authContext?.orgContext?.name
+    : 'your account';
+
+  return {
+    type: "text",
+    text: `\n---\n⚠️  Quota Alert: ${contextName} is at ${percentage_used.toFixed(1)}% capacity (${used.toLocaleString()}/${total.toLocaleString()} queries).\nUpgrade your plan at https://unifiedmemory.ai/pricing to unlock more capacity.\n---`
+  };
+}
+
 /**
  * Handle tools/list request
  * @param {Object} authHeaders - Auth headers
@@ -260,12 +359,18 @@ async function handleCallTool(request, authHeaders, authContext, projectContext)
     const result = await callRemoteMCPTool(name, args, authHeaders, authContext, projectContext);
     console.error(`✓ Tool executed successfully: ${name}`);
 
+    // Check if upgrade prompt needed (cached, 5-min TTL)
+    const upgradePrompt = await checkAndGetUpgradePrompt(authHeaders, authContext);
+
     return {
-      content: result.content || [
-        {
-          type: "text",
-          text: JSON.stringify(result, null, 2),
-        },
+      content: [
+        ...(result.content || [
+          {
+            type: "text",
+            text: JSON.stringify(result, null, 2),
+          },
+        ]),
+        ...(upgradePrompt ? [upgradePrompt] : [])
       ],
     };
   } catch (error) {

package/lib/token-refresh.js

@@ -1,7 +1,6 @@
-import { getToken, saveToken } from './token-storage.js';
+import { saveToken } from './token-storage.js';
 import { config } from './config.js';
 import { parseJWT } from './jwt-utils.js';
-import { getOrgScopedToken } from './clerk-api.js';
 
 /**
  * Check if token has expired
@@ -23,6 +22,7 @@ export function isTokenExpired(tokenData) {
 
 /**
  * Refresh access token using refresh token
+ * Simple refresh - Clerk preserves org context if it was set
  * @param {Object} tokenData - Current token data
  * @returns {Promise<Object>} - New token data
  */
@@ -74,36 +74,22 @@ export async function refreshAccessToken(tokenData) {
 
     // Parse refreshed JWT
     const refreshedToken = newTokenData.id_token || newTokenData.access_token;
-    let finalIdToken = newTokenData.id_token;
-    let decoded = parseJWT(refreshedToken);
+    const decoded = parseJWT(refreshedToken);
 
-    // If we have org context, get org-scoped token to ensure JWT has org claims
-    if (tokenData.selectedOrg?.id && decoded?.sid) {
-      try {
-        const orgToken = await getOrgScopedToken(
-          decoded.sid,
-          tokenData.selectedOrg.id,
-          refreshedToken
-        );
-        finalIdToken = orgToken.jwt;
-        decoded = parseJWT(orgToken.jwt);
-      } catch (error) {
-        // Log warning but continue with base token
-        // The subsequent API call may fail, prompting re-login
-        console.error(`⚠️  Could not refresh org-scoped token: ${error.message}`);
-      }
-    }
-
-    // Build updated token object, preserving selectedOrg
+    // Build updated token object
+    // Note: Clerk preserves org context through refresh automatically
+    // We also preserve selectedOrgId/selectedOrgName from login state
     const updatedToken = {
       accessToken: newTokenData.access_token,
-      idToken: finalIdToken,
+      idToken: newTokenData.id_token,
       tokenType: newTokenData.token_type || 'Bearer',
       expiresIn: newTokenData.expires_in,
       receivedAt: Date.now(),
       decoded: decoded,
-      selectedOrg: tokenData.selectedOrg, // Preserve organization context
       refresh_token: newTokenData.refresh_token || tokenData.refresh_token,
+      sessionId: tokenData.sessionId,  // Preserve sessionId through refresh
+      selectedOrgId: tokenData.selectedOrgId,  // Preserve org selection through refresh
+      selectedOrgName: tokenData.selectedOrgName,  // Preserve org name through refresh
     };
 
     // Save to storage