@unifiedmemory/cli 1.3.12 → 1.3.14

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

package/lib/token-storage.js

@@ -14,12 +14,6 @@ export function saveToken(tokenData) {
   // Ensure directory permissions are correct (in case it was created with wrong permissions)
   fs.chmodSync(TOKEN_DIR, 0o700);
 
-  // Preserve existing organization context if not explicitly overwritten
-  const existingData = getToken();
-  if (existingData && existingData.selectedOrg && !tokenData.selectedOrg) {
-    tokenData.selectedOrg = existingData.selectedOrg;
-  }
-
   // Write token file with owner-only read/write permissions (0600)
   fs.writeFileSync(TOKEN_FILE, JSON.stringify(tokenData, null, 2));
   // Explicitly set file permissions (writeFileSync mode option doesn't always work with extended attributes)
@@ -46,31 +40,43 @@ export function clearToken() {
 }
 
 /**
- * Update the selected organization context
- * @param {Object|null} orgData - Organization data or null for personal context
+ * Get organization context from JWT claims
+ * JWT is the single source of truth for org context
+ * @returns {Object|null} Organization context or null if personal context
  */
-export function updateSelectedOrg(orgData) {
-  const tokenData = getToken();
-  if (!tokenData) {
-    throw new Error('No token found. Please login first.');
+export function getOrgContext() {
+  const token = getToken();
+
+  // Try nested 'o' object first (original OAuth token format)
+  if (token?.decoded?.o?.o_id) {
+    return {
+      id: token.decoded.o.o_id,
+      name: token.decoded.o.o_name,
+      slug: token.decoded.o.o_slug,
+      role: token.decoded.o.o_role
+    };
   }
 
-  // Ensure directory permissions are correct
-  if (fs.existsSync(TOKEN_DIR)) {
-    fs.chmodSync(TOKEN_DIR, 0o700);
+  // Fall back to flat org_* claims (org-scoped token from Clerk)
+  if (token?.decoded?.org_id) {
+    return {
+      id: token.decoded.org_id,
+      name: token.decoded.org_name,
+      slug: token.decoded.org_slug,
+      role: token.decoded.org_role
+    };
   }
 
-  tokenData.selectedOrg = orgData;
-  // Write with owner-only read/write permissions (0600)
-  fs.writeFileSync(TOKEN_FILE, JSON.stringify(tokenData, null, 2));
-  fs.chmodSync(TOKEN_FILE, 0o600);
-}
+  // Fall back to selectedOrgId from login state parameter
+  // This is used when JWT doesn't have org claims (Clerk setActive limitation)
+  if (token?.selectedOrgId) {
+    return {
+      id: token.selectedOrgId,
+      name: token.selectedOrgName || null,
+      slug: null,
+      role: null
+    };
+  }
 
-/**
- * Get the currently selected organization context
- * @returns {Object|null} Organization data or null if personal context
- */
-export function getSelectedOrg() {
-  const tokenData = getToken();
-  return tokenData?.selectedOrg || null;
+  return null;
 }

package/lib/token-validation.js

@@ -1,11 +1,9 @@
-import { getToken, saveToken } from './token-storage.js';
+import { getToken } from './token-storage.js';
 import { isTokenExpired, refreshAccessToken } from './token-refresh.js';
-import { getOrgScopedToken } from './clerk-api.js';
-import { parseJWT } from './jwt-utils.js';
-import { config } from './config.js';
 
 /**
  * Load token and refresh if expired
+ * JWT is the single source of truth - no recovery logic needed
  * @param {boolean} requireAuth - Whether to throw if not authenticated (default true)
  * @returns {Promise<Object|null>} - Token data or null
  * @throws {Error} - If not authenticated and requireAuth is true, or if refresh fails
@@ -46,86 +44,5 @@ export async function loadAndRefreshToken(requireAuth = true) {
     }
   }
 
-  // Check if token has org context but lacks org claims in JWT
-  // This can happen if getOrgScopedToken() failed during login
-  if (tokenData.selectedOrg?.id && !tokenData.decoded?.o) {
-    console.error('⚠️  Token missing org claims, attempting to fix...');
-
-    // Check for sid in decoded token or preserved originalSid from login
-    let sessionId = tokenData.decoded?.sid || tokenData.originalSid;
-    let currentToken = tokenData.idToken || tokenData.accessToken;
-
-    // If no sid, try to refresh first to get a new base token with sid
-    if (!sessionId && tokenData.refresh_token) {
-      console.error('  Refreshing to obtain session ID...');
-      try {
-        const tokenParams = {
-          client_id: config.clerkClientId,
-          refresh_token: tokenData.refresh_token,
-          grant_type: 'refresh_token',
-        };
-
-        if (config.clerkClientSecret) {
-          tokenParams.client_secret = config.clerkClientSecret;
-        }
-
-        const response = await fetch(`https://${config.clerkDomain}/oauth/token`, {
-          method: 'POST',
-          headers: { 'Content-Type': 'application/x-www-form-urlencoded' },
-          body: new URLSearchParams(tokenParams),
-        });
-
-        if (response.ok) {
-          const newTokenData = await response.json();
-          const refreshedToken = newTokenData.id_token || newTokenData.access_token;
-          const refreshedDecoded = parseJWT(refreshedToken);
-
-          sessionId = refreshedDecoded?.sid;
-          currentToken = refreshedToken;
-
-          // Update tokenData with refreshed values
-          tokenData.accessToken = newTokenData.access_token;
-          tokenData.idToken = newTokenData.id_token;
-          tokenData.refresh_token = newTokenData.refresh_token || tokenData.refresh_token;
-          tokenData.decoded = refreshedDecoded;
-
-          if (sessionId) {
-            console.error('  ✓ Obtained session ID');
-          }
-        }
-      } catch (error) {
-        console.error(`  Could not refresh token: ${error.message}`);
-      }
-    }
-
-    // Now try to get org-scoped token if we have sessionId
-    if (sessionId) {
-      try {
-        const orgToken = await getOrgScopedToken(
-          sessionId,
-          tokenData.selectedOrg.id,
-          currentToken
-        );
-
-        const decoded = parseJWT(orgToken.jwt);
-        const updatedToken = {
-          ...tokenData,
-          idToken: orgToken.jwt,
-          decoded: decoded,
-        };
-
-        saveToken(updatedToken);
-        console.error('✓ Token updated with org claims');
-        return updatedToken;
-      } catch (error) {
-        console.error(`⚠️  Could not get org-scoped token: ${error.message}`);
-        // Continue with existing token - API calls may fail
-      }
-    } else {
-      console.error('⚠️  Could not obtain session ID for org-scoped token');
-      console.error('   Please run: um login');
-    }
-  }
-
   return tokenData;
 }

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@unifiedmemory/cli",
-  "version": "1.3.12",
+  "version": "1.3.14",
   "description": "UnifiedMemory CLI - AI code assistant integration",
   "main": "index.js",
   "type": "module",

package/tests/unit/mcp-proxy.test.js

@@ -221,7 +221,7 @@ describe('mcp-proxy', () => {
 
       const authContext = {
         decoded: { sub: 'user_123' },
-        selectedOrg: { id: 'org_456' },
+        orgContext: { id: 'org_456' },
       };
 
       const projectContext = {
@@ -245,7 +245,7 @@ describe('mcp-proxy', () => {
       expect(capturedBody.params.arguments.headers['X-Org-Id']).toBe('org_456');
     });
 
-    it('should fallback to user_id for org when no selectedOrg', async () => {
+    it('should not inject org when no orgContext (personal context)', async () => {
       let capturedBody = null;
 
       server.use(
@@ -263,7 +263,7 @@ describe('mcp-proxy', () => {
 
       const authContext = {
         decoded: { sub: 'user_123' },
-        // No selectedOrg
+        // No orgContext - personal context
       };
 
       await callRemoteMCPTool(
@@ -274,9 +274,12 @@ describe('mcp-proxy', () => {
         null
       );
 
-      // Org should fallback to user_id
-      expect(capturedBody.params.arguments.pathParams.org).toBe('user_123');
-      expect(capturedBody.params.arguments.headers['X-Org-Id']).toBe('user_123');
+      // No org injection for personal context - gateway handles this
+      expect(capturedBody.params.arguments.pathParams.org).toBeUndefined();
+      expect(capturedBody.params.arguments.headers['X-Org-Id']).toBeUndefined();
+      // User should still be injected
+      expect(capturedBody.params.arguments.pathParams.user).toBe('user_123');
+      expect(capturedBody.params.arguments.headers['X-User-Id']).toBe('user_123');
     });
 
     it('should return tool execution result', async () => {

package/tests/unit/token-refresh.test.js

@@ -44,17 +44,9 @@ vi.mock('../../lib/jwt-utils.js', () => ({
   }),
 }));
 
-vi.mock('../../lib/clerk-api.js', () => ({
-  getOrgScopedToken: vi.fn().mockResolvedValue({
-    jwt: 'org_scoped_jwt',
-    object: 'token',
-  }),
-}));
-
 // Import after mocking
 import { isTokenExpired, refreshAccessToken } from '../../lib/token-refresh.js';
 import { saveToken } from '../../lib/token-storage.js';
-import { getOrgScopedToken } from '../../lib/clerk-api.js';
 
 describe('isTokenExpired', () => {
   it('returns true for null tokenData', () => {
@@ -181,7 +173,7 @@ describe('refreshAccessToken', () => {
     );
   });
 
-  it('preserves selectedOrg in refreshed token', async () => {
+  it('preserves selectedOrgId and selectedOrgName in refreshed token', async () => {
     global.fetch = vi.fn().mockResolvedValue({
       ok: true,
       json: () => Promise.resolve({
@@ -193,60 +185,42 @@ describe('refreshAccessToken', () => {
       }),
     });
 
-    const selectedOrg = { id: 'org_456', name: 'Test Org', role: 'admin' };
-    await refreshAccessToken({ refresh_token: 'rt_123', selectedOrg });
+    await refreshAccessToken({
+      refresh_token: 'rt_123',
+      selectedOrgId: 'org_456',
+      selectedOrgName: 'Test Org',
+    });
 
     expect(saveToken).toHaveBeenCalledWith(
       expect.objectContaining({
-        selectedOrg: selectedOrg,
+        selectedOrgId: 'org_456',
+        selectedOrgName: 'Test Org',
       })
     );
   });
 
-  it('gets org-scoped token when selectedOrg exists and refreshed token has sid', async () => {
+  it('preserves sessionId through refresh', async () => {
     global.fetch = vi.fn().mockResolvedValue({
       ok: true,
       json: () => Promise.resolve({
         access_token: 'new_access_token',
-        id_token: 'new_id_token_with_sid',
+        id_token: 'new_id_token',
         refresh_token: 'new_refresh_token',
         token_type: 'Bearer',
         expires_in: 3600,
       }),
     });
 
-    const selectedOrg = { id: 'org_456', name: 'Test Org', role: 'admin' };
-    await refreshAccessToken({ refresh_token: 'rt_123', selectedOrg });
-
-    expect(getOrgScopedToken).toHaveBeenCalledWith(
-      'sess_new_123',
-      'org_456',
-      'new_id_token_with_sid'
-    );
-  });
-
-  it('continues with base token when getOrgScopedToken fails', async () => {
-    global.fetch = vi.fn().mockResolvedValue({
-      ok: true,
-      json: () => Promise.resolve({
-        access_token: 'new_access_token',
-        id_token: 'new_id_token_with_sid',
-        refresh_token: 'new_refresh_token',
-        token_type: 'Bearer',
-        expires_in: 3600,
-      }),
+    await refreshAccessToken({
+      refresh_token: 'rt_123',
+      sessionId: 'sess_123',
     });
 
-    // Make getOrgScopedToken fail
-    vi.mocked(getOrgScopedToken).mockRejectedValueOnce(new Error('API error'));
-
-    const selectedOrg = { id: 'org_456', name: 'Test Org', role: 'admin' };
-
-    // Should not throw - continues with base token
-    const result = await refreshAccessToken({ refresh_token: 'rt_123', selectedOrg });
-
-    expect(result).toBeDefined();
-    expect(result.idToken).toBe('new_id_token_with_sid'); // Falls back to base token
+    expect(saveToken).toHaveBeenCalledWith(
+      expect.objectContaining({
+        sessionId: 'sess_123',
+      })
+    );
   });
 
   it('throws error when OAuth refresh fails with 400', async () => {

package/tests/unit/token-storage.test.js

@@ -132,7 +132,6 @@ describe('token-storage module behavior', () => {
     expect(typeof ts.saveToken).toBe('function');
     expect(typeof ts.getToken).toBe('function');
     expect(typeof ts.clearToken).toBe('function');
-    expect(typeof ts.updateSelectedOrg).toBe('function');
-    expect(typeof ts.getSelectedOrg).toBe('function');
+    expect(typeof ts.getOrgContext).toBe('function');
   });
 });