fix: add connectionManager @W-22389160@

src/agent.ts

@@ -17,16 +17,7 @@
 import { inspect } from 'node:util';
 import * as path from 'node:path';
 import { readdir, stat } from 'node:fs/promises';
-import {
-  AuthInfo,
-  Connection,
-  generateApiName,
-  Lifecycle,
-  Logger,
-  Messages,
-  SfError,
-  SfProject,
-} from '@salesforce/core';
+import { Connection, generateApiName, Lifecycle, Logger, Messages, SfError, SfProject } from '@salesforce/core';
 import { ComponentSetBuilder } from '@salesforce/source-deploy-retrieve';
 import { Duration } from '@salesforce/kit';
 import {
@@ -43,9 +34,10 @@ import {
   ScriptAgentOptions,
 } from './types';
 import { MaybeMock } from './maybe-mock';
-import { decodeHtmlEntities, findLocalAgents, useNamedUserJwt } from './utils';
+import { decodeHtmlEntities, findLocalAgents } from './utils';
 import { ScriptAgent } from './agents/scriptAgent';
 import { ProductionAgent } from './agents/productionAgent';
+import { ConnectionManager } from './connectionManager';
 
 /** Instance type returned from Agent.init(); has setSessionId, getHistoryDir, preview, etc. */
 export type AgentInstance = ScriptAgent | ProductionAgent;
@@ -108,24 +100,15 @@ export class Agent {
   public static async init(
     options: ProductionAgentOptions | ScriptAgentOptions
   ): Promise<ScriptAgent | ProductionAgent> {
-    const username = options.connection.getUsername();
-
-    // Create a fresh connection instance for agent operations
-    // This ensures we don't modify the original connection passed in
-    // The original connection remains unchanged and can be used for other operations, mid agent-operation
-    const authInfo = await AuthInfo.create({ username });
-    const isolatedConnection = await Connection.create({ authInfo });
-
-    // Upgrade the isolated connection with JWT
-    const jwtConnection = await useNamedUserJwt(isolatedConnection);
+    // ConnectionManager isolates JWT (for SFAP) and standard (for org) connections so
+    // the caller's connection is never mutated by JWT upgrades or auto-refresh.
+    const connectionManager = await ConnectionManager.create(options.connection);
 
     // Type guard: check if it's ScriptAgentOptions by looking for 'aabName'
     if ('aabName' in options) {
-      // TypeScript now knows this is ScriptAgentOptions
-      return new ScriptAgent({ ...options, connection: jwtConnection });
+      return new ScriptAgent(options, connectionManager);
     } else {
-      // TypeScript now knows this is ProductionAgentOptions
-      const agent = new ProductionAgent({ ...options, connection: jwtConnection });
+      const agent = new ProductionAgent(options, connectionManager);
       await agent.getBotMetadata();
       return agent;
     }
@@ -246,7 +229,11 @@ export class Agent {
     config: AgentCreateConfig
   ): Promise<AgentCreateResponse> {
     const url = '/connect/ai-assist/create-agent';
-    const maybeMock = new MaybeMock(connection);
+
+    // Create ConnectionManager to get JWT connection for SFAP API calls
+    const connectionManager = await ConnectionManager.create(connection);
+    const jwtConnection = connectionManager.getJwtConnection();
+    const maybeMock = new MaybeMock(jwtConnection);
 
     // When previewing agent creation just return the response.
     if (!config.saveAgent) {
@@ -272,19 +259,20 @@ export class Agent {
     if (response.isSuccess) {
       await Lifecycle.getInstance().emit(AgentCreateLifecycleStages.Retrieving, {});
       const defaultPackagePath = project.getDefaultPackage().path ?? 'force-app';
+      const standardConnection = connectionManager.getStandardConnection();
       try {
         const cs = await ComponentSetBuilder.build({
           metadata: {
             metadataEntries: [`Agent:${config.agentSettings.agentApiName}`],
             directoryPaths: [defaultPackagePath],
           },
           org: {
-            username: connection.getUsername() as string,
+            username: standardConnection.getUsername() as string,
             exclude: [],
           },
         });
         const retrieve = await cs.retrieve({
-          usernameOrConnection: connection,
+          usernameOrConnection: standardConnection,
           merge: true,
           format: 'source',
           output: path.resolve(project.getPath(), defaultPackagePath),
@@ -324,7 +312,10 @@ export class Agent {
    * @returns the agent job spec
    */
   public static async createSpec(connection: Connection, config: AgentJobSpecCreateConfig): Promise<AgentJobSpec> {
-    const maybeMock = new MaybeMock(connection);
+    // Create ConnectionManager to get JWT connection for SFAP API calls
+    const connectionManager = await ConnectionManager.create(connection);
+    const jwtConnection = connectionManager.getJwtConnection();
+    const maybeMock = new MaybeMock(jwtConnection);
     verifyAgentSpecConfig(config);
 
     const url = '/connect/ai-assist/draft-agent-topics';

src/agents/agentBase.ts

@@ -18,6 +18,7 @@ import { join } from 'node:path';
 import { Connection, SfError } from '@salesforce/core';
 import { AgentPreviewInterface, type AgentPreviewSendResponse, type PlannerResponse, PreviewMetadata } from '../types';
 import { getHistoryDir, SessionHistoryBuffer, TranscriptEntry } from '../utils';
+import { ConnectionManager } from '../connectionManager';
 
 /**
  * Abstract base class for agent preview functionality.
@@ -28,6 +29,20 @@ export abstract class AgentBase {
    * The display name of the agent (user-friendly name, not API name)
    */
   public name: string | undefined;
+  /**
+   * The standard org connection used for SOQL queries, tooling API calls, and metadata
+   * operations. This is distinct from the JWT-upgraded connection used for SFAP API
+   * calls (held internally by the connection manager).
+   */
+  protected readonly connection: Connection;
+  /**
+   * Holds isolated JWT and standard connections. When a ConnectionManager is supplied
+   * by Agent.init() / Agent.create() / Agent.createSpec(), the agent gets fully
+   * isolated connections (the caller's connection is not mutated). When undefined
+   * (consumers instantiating an agent directly), the supplied connection is used as
+   * both the JWT and standard connection — preserving the pre-isolation behavior.
+   */
+  protected readonly connectionManager: ConnectionManager | undefined;
   protected sessionId: string | undefined;
   protected historyDir: string | undefined;
   protected historyBuffer: SessionHistoryBuffer | undefined;
@@ -36,9 +51,23 @@ export abstract class AgentBase {
   protected planIds = new Set<string>();
   public abstract preview: AgentPreviewInterface;
 
-  protected constructor(protected readonly connection: Connection) {}
+  protected constructor(connection: Connection, connectionManager?: ConnectionManager) {
+    this.connectionManager = connectionManager;
+    this.connection = connectionManager ? connectionManager.getStandardConnection() : connection;
+  }
 
+  /**
+   * Refreshes the access token on the standard connection.
+   *
+   * Retained for backward compatibility. With the connection manager in place the
+   * caller's original Connection object is no longer mutated by agent operations,
+   * so this method only refreshes the internal standard connection.
+   */
   public async restoreConnection(): Promise<void> {
+    if (this.connectionManager) {
+      await this.connectionManager.refreshStandardConnection();
+      return;
+    }
     delete this.connection.accessToken;
     await this.connection.refreshAuth();
   }
@@ -76,6 +105,15 @@ export abstract class AgentBase {
     }
   }
 
+  /**
+   * Returns the connection to use for SFAP API calls (api.salesforce.com/einstein/ai-agent).
+   * When a ConnectionManager is in use this is the JWT-upgraded connection; otherwise it
+   * is the connection supplied at construction time.
+   */
+  protected getJwtConnection(): Connection {
+    return this.connectionManager ? this.connectionManager.getJwtConnection() : this.connection;
+  }
+
   /**
    * Get all traces from the current session
    * Reads traces from the session directory if available, otherwise fetches from API

src/agents/productionAgent.ts

@@ -41,6 +41,7 @@ import {
   SessionHistoryBuffer,
 } from '../utils';
 import { createTraceFlag, findTraceFlag, getDebugLog } from '../apexUtils';
+import { ConnectionManager } from '../connectionManager';
 import { AgentBase } from './agentBase';
 Messages.importMessagesDirectory(__dirname);
 const messages = Messages.loadMessages('@salesforce/agents', 'agents');
@@ -52,8 +53,8 @@ export class ProductionAgent extends AgentBase {
   private apiName: string | undefined;
   private readonly apiBase: string;
 
-  public constructor(private options: ProductionAgentOptions) {
-    super(options.connection);
+  public constructor(private options: ProductionAgentOptions, connectionManager?: ConnectionManager) {
+    super(options.connection, connectionManager);
     this.apiBase = 'https://api.salesforce.com/einstein/ai-agent/v1';
     if (!options.apiNameOrId) {
       throw messages.createError('missingAgentNameOrId');
@@ -264,7 +265,7 @@ export class ProductionAgent extends AgentBase {
       };
       await logTurnToHistory(userEntry, ++this.turnCounter, this.historyDir, this.historyBuffer);
 
-      const response = await requestWithEndpointFallback<AgentPreviewSendResponse>(this.connection, {
+      const response = await requestWithEndpointFallback<AgentPreviewSendResponse>(this.getJwtConnection(), {
         method: 'POST',
         url,
         body: JSON.stringify(body),
@@ -322,7 +323,7 @@ export class ProductionAgent extends AgentBase {
     }
 
     const url = `/connect/bot-versions/${botVersionMetadata.Id}/activation`;
-    const maybeMock = new MaybeMock(this.connection);
+    const maybeMock = new MaybeMock(this.getJwtConnection());
     const response = await maybeMock.request<BotActivationResponse>('POST', url, { status: desiredState });
     if (response.success) {
       const versionToUpdate = this.botMetadata!.BotVersions.records.find(
@@ -360,7 +361,7 @@ export class ProductionAgent extends AgentBase {
     };
 
     try {
-      const response = await requestWithEndpointFallback<AgentPreviewStartResponse>(this.connection, {
+      const response = await requestWithEndpointFallback<AgentPreviewStartResponse>(this.getJwtConnection(), {
         method: 'POST',
         url,
         body: JSON.stringify(body),
@@ -422,7 +423,7 @@ export class ProductionAgent extends AgentBase {
     const url = `${this.apiBase}/sessions/${this.sessionId}`;
     try {
       // https://developer.salesforce.com/docs/einstein/genai/guide/agent-api-examples.html#end-session
-      const response = await requestWithEndpointFallback<AgentPreviewEndResponse>(this.connection, {
+      const response = await requestWithEndpointFallback<AgentPreviewEndResponse>(this.getJwtConnection(), {
         method: 'DELETE',
         url,
         headers: {

src/agents/scriptAgent.ts

@@ -49,6 +49,7 @@ import {
 import { getDebugLog } from '../apexUtils';
 import { generateAgentScript } from '../templates/agentScriptTemplate';
 import { applyStringReplacementsToAgent } from '../stringReplacements';
+import { ConnectionManager } from '../connectionManager';
 import { ScriptAgentPublisher } from './scriptAgentPublisher';
 import { AgentBase } from './agentBase';
 
@@ -63,8 +64,8 @@ export class ScriptAgent extends AgentBase {
   private readonly aabDirectory: string;
   private readonly metaContent: string;
   private readonly agentFilePath: string;
-  public constructor(private options: ScriptAgentOptions) {
-    super(options.connection);
+  public constructor(private options: ScriptAgentOptions, connectionManager?: ConnectionManager) {
+    super(options.connection, connectionManager);
     this.options = options;
     this.apiBase = 'https://api.salesforce.com/einstein/ai-agent';
 
@@ -169,7 +170,7 @@ export class ScriptAgent extends AgentBase {
   }
 
   public async getTrace(planId: string): Promise<PlannerResponse> {
-    return requestWithEndpointFallback<PlannerResponse>(this.connection, {
+    return requestWithEndpointFallback<PlannerResponse>(this.getJwtConnection(), {
       method: 'GET',
       url: `${this.apiBase}/v1.1/preview/sessions/${this.sessionId!}/plans/${planId}`,
       headers: {
@@ -205,7 +206,7 @@ export class ScriptAgent extends AgentBase {
 
     try {
       const response = await requestWithEndpointFallback<CompileAgentScriptResponse>(
-        this.connection,
+        this.getJwtConnection(),
         {
           method: 'POST',
           url,
@@ -271,10 +272,11 @@ export class ScriptAgent extends AgentBase {
     }
 
     const publisher = new ScriptAgentPublisher(
-      this.connection,
+      this.options.connection,
       this.options.project,
       this.agentJson!,
-      skipMetadataRetrieve
+      skipMetadataRetrieve,
+      this.connectionManager
     );
     return publisher.publishAgentJson();
   }
@@ -395,7 +397,7 @@ export class ScriptAgent extends AgentBase {
         this.historyBuffer
       );
 
-      const response = await requestWithEndpointFallback<AgentPreviewSendResponse>(this.connection, {
+      const response = await requestWithEndpointFallback<AgentPreviewSendResponse>(this.getJwtConnection(), {
         method: 'POST',
         url,
         body: JSON.stringify(body),
@@ -472,7 +474,7 @@ export class ScriptAgent extends AgentBase {
     // send bypassUser=false when the compiledAgent.globalConfiguration.defaultAgentUser is INVALID
     let bypassUser =
       (
-        await this.connection.query(
+        await this.connection.query<{ Id: string }>(
           `SELECT Id FROM USER WHERE username='${this.agentJson.globalConfiguration.defaultAgentUser}'`
         )
       ).totalSize === 1;
@@ -509,7 +511,7 @@ export class ScriptAgent extends AgentBase {
       let response: AgentPreviewStartResponse;
       try {
         response = await requestWithEndpointFallback<AgentPreviewStartResponse>(
-          this.connection,
+          this.getJwtConnection(),
           {
             method: 'POST',
             url: `${this.apiBase}/v1.1/preview/sessions`,