Refactor OPA secrets management to use Token Exchange

Author: iamspathan

packages/agent0/src/agent.ts

@@ -231,46 +231,11 @@ export async function initializeLLMConfig(): Promise<void> {
     return;
   }
 
-  console.log('\n Initializing LLM configuration...');
+  console.log('\n Initializing LLM configuration from environment variables...');
 
-  // Try OPA first if configured
-  if (isOPAConfigured()) {
-    console.log(' OPA secret management detected, attempting to fetch credentials...');
-    try {
-      const opaCredentials = await fetchLLMCredentialsFromOPA();
+  // Note: OPA credentials are now fetched per-user-session in getAgentForUserContext()
+  // This function only handles environment variable fallback
 
-      if (opaCredentials) {
-        // Convert OPA credentials to AgentLLMConfig format
-        if (opaCredentials.provider === 'anthropic') {
-          llmConfig = {
-            mcpServerUrl: process.env.MCP_SERVER_URL || '',
-            llmProvider: 'anthropic',
-            anthropicApiKey: opaCredentials.apiKey,
-            anthropicModel: opaCredentials.model,
-          };
-        } else {
-          llmConfig = {
-            mcpServerUrl: process.env.MCP_SERVER_URL || '',
-            llmProvider: 'bedrock',
-            awsRegion: opaCredentials.awsRegion,
-            awsAccessKeyId: opaCredentials.awsAccessKeyId,
-            awsSecretAccessKey: opaCredentials.awsSecretAccessKey,
-            awsSessionToken: opaCredentials.awsSessionToken,
-            bedrockModelId: opaCredentials.bedrockModelId,
-          };
-        }
-        llmConfigSource = 'opa';
-        llmConfigInitialized = true;
-        console.log(' LLM credentials loaded from OPA');
-        return;
-      }
-    } catch (error: any) {
-      console.warn(' Failed to fetch credentials from OPA:', error.message);
-      console.warn(' Falling back to environment variables...');
-    }
-  }
-
-  // Fall back to environment variables
   try {
     llmConfig = validateAgentLLMEnv();
     llmConfigSource = 'env';
@@ -300,17 +265,20 @@ export function isLLMConfigInitialized(): boolean {
 
 // Try synchronous initialization at module load (for backwards compatibility)
 // This will only work if env vars are set and OPA is not configured
+// Try synchronous initialization from env vars at module load
+// OPA credentials are fetched per-user-session, so we don't init them here
 try {
   if (!isOPAConfigured()) {
+    // No OPA configured, must use env vars
     llmConfig = validateAgentLLMEnv();
     llmConfigSource = 'env';
     llmConfigInitialized = true;
   }
+  // If OPA is configured, credentials will be fetched in getAgentForUserContext()
 } catch {
-  // If sync init fails and OPA is configured, that's fine - will init async later
-  // If sync init fails and OPA is not configured, we have a problem
+  // If env var init fails and OPA is configured, that's fine - will use OPA per-session
   if (!isOPAConfigured()) {
-    console.error('LLM configuration failed - no env vars or OPA configured');
+    console.error('LLM configuration failed - no env vars configured and OPA not available');
   }
 }
 
@@ -405,7 +373,55 @@ function buildAgentConfig(): Omit<AgentConfig, 'idToken' | 'userContext'> | null
 }
 
 export async function getAgentForUserContext(idToken: string, userContext: UserContext): Promise<Agent> {
-  // Ensure LLM config is initialized (will be a no-op if already initialized)
+  // Try OPA token exchange for credential retrieval (per-user-session)
+  if (isOPAConfigured()) {
+    console.log('\n🔐 OPA secret management detected, fetching credentials via token exchange...');
+
+    try {
+      const opaCredentials = await fetchLLMCredentialsFromOPA(idToken);
+
+      if (opaCredentials) {
+        // Build config from OPA credentials
+        const baseConfig = {
+          mcpServerUrl: process.env.MCP_SERVER_URL || '',
+          name: 'agent0',
+          version: '1.0.0',
+          tokenExchange: buildTokenExchangeConfig(),
+          enableLLM: true,
+          idToken,
+          userContext,
+        };
+
+        let agentConfig: AgentConfig;
+
+        if (opaCredentials.provider === 'anthropic') {
+          agentConfig = {
+            ...baseConfig,
+            anthropicApiKey: opaCredentials.apiKey,
+            anthropicModel: opaCredentials.model,
+          };
+        } else {
+          agentConfig = {
+            ...baseConfig,
+            awsRegion: opaCredentials.awsRegion,
+            awsAccessKeyId: opaCredentials.awsAccessKeyId,
+            awsSecretAccessKey: opaCredentials.awsSecretAccessKey,
+            awsSessionToken: opaCredentials.awsSessionToken,
+            bedrockModelId: opaCredentials.bedrockModelId,
+          };
+        }
+
+        llmConfigSource = 'opa';
+        console.log('✅ Agent configured with OPA credentials');
+        return new Agent(agentConfig);
+      }
+    } catch (error: any) {
+      console.warn('⚠️  Failed to fetch OPA credentials:', error.message);
+      console.warn('   Falling back to environment variables...');
+    }
+  }
+
+  // Fallback to environment variables
   await initializeLLMConfig();
 
   const agentConfig = buildAgentConfig();

packages/agent0/src/auth/token-exchange.ts

@@ -6,10 +6,6 @@ import { randomUUID } from 'crypto';
 import jwt from 'jsonwebtoken';
 import axios from 'axios';
 
-// Debug mode for verbose auth logging - OFF by default for security
-// Set AUTH_DEBUG=true to enable detailed logging (includes sensitive data)
-const AUTH_DEBUG = process.env.AUTH_DEBUG === 'true';
-
 // ============================================================================
 // Scope Challenge Types and Parser (MCP Authorization Best Practices)
 // ============================================================================
@@ -148,12 +144,6 @@ export class TokenExchangeHandler {
     formData.append('client_assertion_type', 'urn:ietf:params:oauth:client-assertion-type:jwt-bearer');
     formData.append('client_assertion', clientAssertion);
 
-    const requestedScopes = scopes || this.config.agentScopes;
-    console.log(`🔄 Step 1: Exchanging ID token for ID-JAG token...`);
-    console.log(`🎯 Requested scopes: ${requestedScopes}`);
-    console.log(`📍 Audience: ${this.config.authorizationServer}`);
-    console.log(`🆔 Client ID: ${this.config.clientId}`);
-
     const response = await axios.post(
       `https://${this.config.oktaDomain}/oauth2/v1/token`,
       formData,
@@ -164,12 +154,89 @@ export class TokenExchangeHandler {
       }
     );
 
-    console.log(`✅ ID-JAG token obtained`);
-    console.log(`🎯 Issued token type: ${response.data.issued_token_type}`);
-
     return response.data.access_token; // This is actually the ID-JAG token
   }
 
+  // ============================================================================
+  // Exchange ID Token for Vaulted Secret (PAM)
+  // ============================================================================
+
+  /**
+   * Exchange ID token for a vaulted secret from Okta PAM
+   * @param idToken - The user's ID token
+   * @param resourceOrn - The ORN of the secret (e.g., orn:okta:pam:{orgId}:secrets:{secretId})
+   * @param secretName - Optional name for logging purposes
+   * @returns The secret value
+   */
+  async exchangeIdTokenForVaultedSecret(
+    idToken: string,
+    resourceOrn: string,
+    secretName?: string
+  ): Promise<string> {
+    if (!this.privateKey) {
+      throw new Error('Private key not loaded for token exchange');
+    }
+
+    const tokenEndpoint = `https://${this.config.oktaDomain}/oauth2/v1/token`;
+    const clientAssertion = this.createClientAssertion(tokenEndpoint);
+
+    const formData = new URLSearchParams();
+    formData.append('grant_type', 'urn:ietf:params:oauth:grant-type:token-exchange');
+    formData.append('requested_token_type', 'urn:okta:params:oauth:token-type:vaulted-secret');
+    formData.append('subject_token', idToken);
+    formData.append('subject_token_type', 'urn:ietf:params:oauth:token-type:id_token');
+    formData.append('resource', resourceOrn);
+    formData.append('client_assertion_type', 'urn:ietf:params:oauth:client-assertion-type:jwt-bearer');
+    formData.append('client_assertion', clientAssertion);
+
+    try {
+      const response = await axios.post(tokenEndpoint, formData, {
+        headers: {
+          'Content-Type': 'application/x-www-form-urlencoded',
+        },
+      });
+
+      // Extract the actual secret value from the response
+      const vaultedSecret = response.data.vaulted_secret;
+      let secretValue: string;
+
+      if (typeof vaultedSecret === 'string') {
+        secretValue = vaultedSecret;
+      } else if (typeof vaultedSecret === 'object' && vaultedSecret !== null) {
+        // Try common key names for secrets
+        secretValue = vaultedSecret.secret
+          || vaultedSecret.password
+          || vaultedSecret.api_key
+          || vaultedSecret.apiKey
+          || vaultedSecret.token
+          || vaultedSecret.value
+          || Object.values(vaultedSecret)[0] as string;
+      } else {
+        throw new Error(`Invalid secret format received for ${secretName || resourceOrn}`);
+      }
+
+      // Validate non-empty
+      if (!secretValue || (typeof secretValue === 'string' && secretValue.trim() === '')) {
+        throw new Error(`Empty secret value received for ${secretName || resourceOrn}`);
+      }
+
+      return secretValue;
+    } catch (error: any) {
+      const errorData = error.response?.data;
+      console.error(`❌ Failed to retrieve secret${secretName ? ` (${secretName})` : ''}:`, errorData || error.message);
+
+      if (errorData?.error === 'invalid_grant') {
+        throw new Error(`Secret access denied: ${errorData.error_description || 'Invalid grant'}`);
+      } else if (errorData?.error === 'invalid_target') {
+        throw new Error(`Invalid secret resource: ${resourceOrn}`);
+      }
+
+      throw new Error(
+        `Failed to retrieve vaulted secret: ${errorData?.error_description || error.message}`
+      );
+    }
+  }
+
   // ============================================================================
   // Step 2: Exchange ID-JAG for Access Token
   // ============================================================================
@@ -183,9 +250,6 @@ export class TokenExchangeHandler {
     const authorizationServer = this.config.authorizationServer;
     const authorizationServerTokenEndpoint = this.config.authorizationServerTokenEndpoint;
 
-    console.log(`🔄 Step 2: Exchanging ID-JAG for Access Token at Resource Server...`);
-    console.log(`📍 MCP authorization server: ${authorizationServer}`);
-
     const clientAssertion = this.createClientAssertion(authorizationServerTokenEndpoint);
 
     const resourceTokenForm = new URLSearchParams();
@@ -204,10 +268,6 @@ export class TokenExchangeHandler {
       }
     );
 
-    console.log(`✅ Access Token obtained from Resource Server`);
-    console.log(`🎯 Token type: ${response.data.token_type}`);
-    console.log(`⏰ Expires in: ${response.data.expires_in}s`);
-
     return response.data;
   }
 
@@ -235,14 +295,6 @@ export class TokenExchangeHandler {
       throw new Error('Cross-app access not configured properly. Private key not loaded.');
     }
 
-    // Only log tokens in debug mode - contains sensitive credentials
-    if (AUTH_DEBUG) {
-      console.log(`👻 Subject token: ${idToken}`);
-    }
-    if (requestedScopes) {
-      console.log(`🔄 Step-up authorization requested with scopes: ${requestedScopes}`);
-    }
-
     try {
       // Step 1: Exchange ID token for ID-JAG
       const idJag = await this.exchangeIdTokenForIdJag(idToken, requestedScopes);
@@ -253,8 +305,6 @@ export class TokenExchangeHandler {
 
         // Return the access token
         const accessToken = accessTokenResponse.access_token;
-        console.log('✅ Access token obtained successfully');
-        console.log('💡 Token can be used for MCP tool calls');
 
         return {
           success: true,