feat: AgentManager - foundation for unified execution (#4389) (#4684)

2025-09-24 16:15:26 -04:00
parent ad045bc94c
commit 9d40422e2d
34 changed files with 896 additions and 352 deletions
@@ -17,6 +17,7 @@ import { activateExtension, deleteExtension, toggleExtension, updateExtension }
 import { ExtensionConfig } from '../../../api/types.gen';

 interface ExtensionSectionProps {
+  sessionId: string; // Add required sessionId prop
  deepLinkConfig?: ExtensionConfig;
  showEnvVars?: boolean;
  hideButtons?: boolean;
@@ -26,6 +27,7 @@ interface ExtensionSectionProps {
 }

 export default function ExtensionsSection({
+  sessionId,
  deepLinkConfig,
  showEnvVars,
  hideButtons,
@@ -97,6 +99,7 @@ export default function ExtensionsSection({
        extensionConfig: extensionConfig,
        addToConfig: addExtension,
        toastOptions: { silent: false },
+        sessionId: sessionId,
      });

      await fetchExtensions(); // Refresh the list after successful toggle
@@ -119,7 +122,11 @@ export default function ExtensionsSection({

    const extensionConfig = createExtensionConfig(formData);
    try {
-      await activateExtension({ addToConfig: addExtension, extensionConfig: extensionConfig });
+      await activateExtension({
+        addToConfig: addExtension,
+        extensionConfig: extensionConfig,
+        sessionId: sessionId,
+      });
      // Immediately refresh the extensions list after successful activation
      await fetchExtensions();
    } catch (error) {
@@ -147,6 +154,7 @@ export default function ExtensionsSection({
        addToConfig: addExtension,
        removeFromConfig: removeExtension,
        originalName: originalName,
+        sessionId: sessionId,
      });
    } catch (error) {
      console.error('Failed to update extension:', error);
@@ -162,7 +170,7 @@ export default function ExtensionsSection({
    handleModalClose();

    try {
-      await deleteExtension({ name, removeFromConfig: removeExtension });
+      await deleteExtension({ name, removeFromConfig: removeExtension, sessionId: sessionId });
    } catch (error) {
      console.error('Failed to delete extension:', error);
      // We don't reopen the modal on failure
@@ -61,7 +61,12 @@ describe('Agent API', () => {
      };
      mockFetch.mockResolvedValue(mockResponse);

-      const response = await extensionApiCall('/extensions/add', mockExtensionConfig);
+      const response = await extensionApiCall(
+        '/extensions/add',
+        mockExtensionConfig,
+        {},
+        'test-session'
+      );

      expect(mockFetch).toHaveBeenCalledWith('http://localhost:8080/extensions/add', {
        method: 'POST',
@@ -69,7 +74,7 @@ describe('Agent API', () => {
          'Content-Type': 'application/json',
          'X-Secret-Key': 'secret-key',
        },
-        body: JSON.stringify(mockExtensionConfig),
+        body: JSON.stringify({ ...mockExtensionConfig, session_id: 'test-session' }),
      });

      expect(mockToastService.loading).toHaveBeenCalledWith({
@@ -92,7 +97,12 @@ describe('Agent API', () => {
      };
      mockFetch.mockResolvedValue(mockResponse);

-      const response = await extensionApiCall('/extensions/remove', 'test-extension');
+      const response = await extensionApiCall(
+        '/extensions/remove',
+        'test-extension',
+        {},
+        'test-session'
+      );

      expect(mockFetch).toHaveBeenCalledWith('http://localhost:8080/extensions/remove', {
        method: 'POST',
@@ -100,7 +110,7 @@ describe('Agent API', () => {
          'Content-Type': 'application/json',
          'X-Secret-Key': 'secret-key',
        },
-        body: JSON.stringify('test-extension'),
+        body: JSON.stringify({ name: 'test-extension', session_id: 'test-session' }),
      });

      expect(mockToastService.loading).not.toHaveBeenCalled(); // No loading toast for removal
@@ -120,9 +130,9 @@ describe('Agent API', () => {
      };
      mockFetch.mockResolvedValue(mockResponse);

-      await expect(extensionApiCall('/extensions/add', mockExtensionConfig)).rejects.toThrow(
-        'Server returned 500: Internal Server Error'
-      );
+      await expect(
+        extensionApiCall('/extensions/add', mockExtensionConfig, {}, 'test-session')
+      ).rejects.toThrow('Server returned 500: Internal Server Error');

      expect(mockToastService.error).toHaveBeenCalledWith({
        title: 'test-extension',
@@ -139,9 +149,9 @@ describe('Agent API', () => {
      };
      mockFetch.mockResolvedValue(mockResponse);

-      await expect(extensionApiCall('/extensions/add', mockExtensionConfig)).rejects.toThrow(
-        'Agent is not initialized. Please initialize the agent first.'
-      );
+      await expect(
+        extensionApiCall('/extensions/add', mockExtensionConfig, {}, 'test-session')
+      ).rejects.toThrow('Agent is not initialized. Please initialize the agent first.');

      expect(mockToastService.error).toHaveBeenCalledWith({
        title: 'test-extension',
@@ -157,9 +167,9 @@ describe('Agent API', () => {
      };
      mockFetch.mockResolvedValue(mockResponse);

-      await expect(extensionApiCall('/extensions/remove', 'test-extension')).rejects.toThrow(
-        'Error deactivating extension: Extension not found'
-      );
+      await expect(
+        extensionApiCall('/extensions/remove', 'test-extension', {}, 'test-session')
+      ).rejects.toThrow('Error deactivating extension: Extension not found');

      expect(mockToastService.error).toHaveBeenCalledWith({
        title: 'test-extension',
@@ -175,7 +185,12 @@ describe('Agent API', () => {
      };
      mockFetch.mockResolvedValue(mockResponse);

-      const response = await extensionApiCall('/extensions/add', mockExtensionConfig);
+      const response = await extensionApiCall(
+        '/extensions/add',
+        mockExtensionConfig,
+        {},
+        'test-session'
+      );

      expect(mockToastService.success).toHaveBeenCalledWith({
        title: 'test-extension',
@@ -189,9 +204,9 @@ describe('Agent API', () => {
      const networkError = new Error('Network error');
      mockFetch.mockRejectedValue(networkError);

-      await expect(extensionApiCall('/extensions/add', mockExtensionConfig)).rejects.toThrow(
-        'Network error'
-      );
+      await expect(
+        extensionApiCall('/extensions/add', mockExtensionConfig, {}, 'test-session')
+      ).rejects.toThrow('Network error');

      expect(mockToastService.error).toHaveBeenCalledWith({
        title: 'test-extension',
@@ -207,7 +222,12 @@ describe('Agent API', () => {
      };
      mockFetch.mockResolvedValue(mockResponse);

-      await extensionApiCall('/extensions/add', mockExtensionConfig, { silent: true });
+      await extensionApiCall(
+        '/extensions/add',
+        mockExtensionConfig,
+        { silent: true },
+        'test-session'
+      );

      expect(mockToastService.configure).toHaveBeenCalledWith({ silent: true });
    });
@@ -232,7 +252,7 @@ describe('Agent API', () => {
      const { replaceWithShims } = await import('./utils');
      vi.mocked(replaceWithShims).mockResolvedValue('/path/to/python');

-      await addToAgent(mockExtensionConfig);
+      await addToAgent(mockExtensionConfig, {}, 'test-session');

      expect(mockFetch).toHaveBeenCalledWith('http://localhost:8080/extensions/add', {
        method: 'POST',
@@ -244,6 +264,7 @@ describe('Agent API', () => {
          ...mockExtensionConfig,
          name: 'testextension',
          cmd: '/path/to/python',
+          session_id: 'test-session',
        }),
      });
    });
@@ -256,7 +277,7 @@ describe('Agent API', () => {
      };
      mockFetch.mockResolvedValue(mockResponse);

-      await expect(addToAgent(mockExtensionConfig)).rejects.toThrow(
+      await expect(addToAgent(mockExtensionConfig, {}, 'test-session')).rejects.toThrow(
        'Agent is not initialized. Please initialize the agent first.'
      );
    });
@@ -274,7 +295,7 @@ describe('Agent API', () => {
      };
      mockFetch.mockResolvedValue(mockResponse);

-      await addToAgent(sseConfig);
+      await addToAgent(sseConfig, {}, 'test-session');

      expect(mockFetch).toHaveBeenCalledWith('http://localhost:8080/extensions/add', {
        method: 'POST',
@@ -285,6 +306,7 @@ describe('Agent API', () => {
        body: JSON.stringify({
          ...sseConfig,
          name: 'sseextension',
+          session_id: 'test-session',
        }),
      });
    });
@@ -298,7 +320,7 @@ describe('Agent API', () => {
      };
      mockFetch.mockResolvedValue(mockResponse);

-      await removeFromAgent('Test Extension');
+      await removeFromAgent('Test Extension', {}, 'test-session');

      expect(mockFetch).toHaveBeenCalledWith('http://localhost:8080/extensions/remove', {
        method: 'POST',
@@ -306,7 +328,7 @@ describe('Agent API', () => {
          'Content-Type': 'application/json',
          'X-Secret-Key': 'secret-key',
        },
-        body: JSON.stringify('testextension'),
+        body: JSON.stringify({ name: 'testextension', session_id: 'test-session' }),
      });
    });

@@ -318,7 +340,7 @@ describe('Agent API', () => {
      };
      mockFetch.mockResolvedValue(mockResponse);

-      await expect(removeFromAgent('Test Extension')).rejects.toThrow();
+      await expect(removeFromAgent('Test Extension', {}, 'test-session')).rejects.toThrow();

      expect(mockToastService.error).toHaveBeenCalled();
    });
@@ -14,7 +14,8 @@ interface ApiResponse {
 export async function extensionApiCall(
  endpoint: string,
  payload: ExtensionConfig | string,
-  options: ToastServiceOptions & { isDelete?: boolean } = {}
+  options: ToastServiceOptions & { isDelete?: boolean } = {},
+  sessionId: string
 ): Promise<Response> {
  // Configure toast notifications
  toastService.configure(options);
@@ -43,6 +44,16 @@ export async function extensionApiCall(
  }

  try {
+    // Build the request body
+    let requestBody: ExtensionConfig | { name: string; session_id: string };
+    if (typeof payload === 'object') {
+      // For adding extensions (ExtensionConfig)
+      requestBody = { ...payload, session_id: sessionId };
+    } else {
+      // For removing extensions (just the name string)
+      requestBody = { name: payload, session_id: sessionId };
+    }
+
    // Step 2: Make the API call
    const response = await fetch(getApiUrl(endpoint), {
      method: 'POST',
@@ -50,7 +61,7 @@ export async function extensionApiCall(
        'Content-Type': 'application/json',
        'X-Secret-Key': await window.electron.getSecretKey(),
      },
-      body: JSON.stringify(payload),
+      body: JSON.stringify(requestBody),
    });

    // Step 3: Handle non-successful responses
@@ -142,7 +153,8 @@ async function parseResponseData(response: Response): Promise<ApiResponse> {
 */
 export async function addToAgent(
  extension: ExtensionConfig,
-  options: ToastServiceOptions = {}
+  options: ToastServiceOptions = {},
+  sessionId: string
 ): Promise<Response> {
  try {
    if (extension.type === 'stdio') {
@@ -151,7 +163,7 @@ export async function addToAgent(

    extension.name = sanitizeName(extension.name);

-    return await extensionApiCall('/extensions/add', extension, options);
+    return await extensionApiCall('/extensions/add', extension, options, sessionId);
  } catch (error) {
    // Check if this is a 428 error and make the message more descriptive
    if (error instanceof Error && error.message && error.message.includes('428')) {
@@ -170,10 +182,11 @@ export async function addToAgent(
 */
 export async function removeFromAgent(
  name: string,
-  options: ToastServiceOptions & { isDelete?: boolean } = {}
+  options: ToastServiceOptions & { isDelete?: boolean } = {},
+  sessionId: string
 ): Promise<Response> {
  try {
-    return await extensionApiCall('/extensions/remove', sanitizeName(name), options);
+    return await extensionApiCall('/extensions/remove', sanitizeName(name), options, sessionId);
  } catch (error) {
    const action = options.isDelete ? 'remove' : 'deactivate';
    console.error(`Failed to ${action} extension ${name} from agent:`, error);
@@ -1,6 +1,5 @@
 import type { ExtensionConfig } from '../../../api';
 import { toastService } from '../../../toasts';
-import { activateExtension } from './extension-manager';
 import { DEFAULT_EXTENSION_TIMEOUT } from './utils';

 /**
@@ -180,7 +179,11 @@ export async function addExtensionFromDeepLink(

  try {
    console.log('No env vars required, activating extension directly');
-    await activateExtension({ extensionConfig: config, addToConfig: addExtensionFn });
+    // Note: deeplink activation doesn't have access to sessionId
+    // The extension will be added to config but not activated in the current session
+    // It will be activated when the next session starts
+    console.warn('Extension will be added to config but requires a session to activate');
+    await addExtensionFn(config.name, config, true);
  } catch (error) {
    console.error('Failed to activate extension from deeplink:', error);
    throw error;
@@ -43,10 +43,15 @@ describe('Extension Manager', () => {

      await activateExtension({
        addToConfig: mockAddToConfig,
+        sessionId: 'test-session',
        extensionConfig: mockExtensionConfig,
      });

-      expect(mockAddToAgent).toHaveBeenCalledWith(mockExtensionConfig, { silent: false });
+      expect(mockAddToAgent).toHaveBeenCalledWith(
+        mockExtensionConfig,
+        { silent: false },
+        'test-session'
+      );
      expect(mockAddToConfig).toHaveBeenCalledWith('test-extension', mockExtensionConfig, true);
    });

@@ -57,11 +62,16 @@ describe('Extension Manager', () => {
      await expect(
        activateExtension({
          addToConfig: mockAddToConfig,
+          sessionId: 'test-session',
          extensionConfig: mockExtensionConfig,
        })
      ).rejects.toThrow('Agent failed');

-      expect(mockAddToAgent).toHaveBeenCalledWith(mockExtensionConfig, { silent: false });
+      expect(mockAddToAgent).toHaveBeenCalledWith(
+        mockExtensionConfig,
+        { silent: false },
+        'test-session'
+      );
      expect(mockAddToConfig).toHaveBeenCalledWith('test-extension', mockExtensionConfig, false);
    });

@@ -73,13 +83,18 @@ describe('Extension Manager', () => {
      await expect(
        activateExtension({
          addToConfig: mockAddToConfig,
+          sessionId: 'test-session',
          extensionConfig: mockExtensionConfig,
        })
      ).rejects.toThrow('Config failed');

-      expect(mockAddToAgent).toHaveBeenCalledWith(mockExtensionConfig, { silent: false });
+      expect(mockAddToAgent).toHaveBeenCalledWith(
+        mockExtensionConfig,
+        { silent: false },
+        'test-session'
+      );
      expect(mockAddToConfig).toHaveBeenCalledWith('test-extension', mockExtensionConfig, true);
-      expect(mockRemoveFromAgent).toHaveBeenCalledWith('test-extension');
+      expect(mockRemoveFromAgent).toHaveBeenCalledWith('test-extension', {}, 'test-session');
    });
  });

@@ -89,10 +104,15 @@ describe('Extension Manager', () => {

      await addToAgentOnStartup({
        addToConfig: mockAddToConfig,
+        sessionId: 'test-session',
        extensionConfig: mockExtensionConfig,
      });

-      expect(mockAddToAgent).toHaveBeenCalledWith(mockExtensionConfig, { silent: true });
+      expect(mockAddToAgent).toHaveBeenCalledWith(
+        mockExtensionConfig,
+        { silent: true },
+        'test-session'
+      );
      expect(mockAddToConfig).not.toHaveBeenCalled();
    });

@@ -101,11 +121,16 @@ describe('Extension Manager', () => {

      await addToAgentOnStartup({
        addToConfig: mockAddToConfig,
+        sessionId: 'test-session',
        extensionConfig: mockExtensionConfig,
        toastOptions: { silent: false },
      });

-      expect(mockAddToAgent).toHaveBeenCalledWith(mockExtensionConfig, { silent: false });
+      expect(mockAddToAgent).toHaveBeenCalledWith(
+        mockExtensionConfig,
+        { silent: false },
+        'test-session'
+      );
      expect(mockAddToConfig).not.toHaveBeenCalled();
    });

@@ -118,6 +143,7 @@ describe('Extension Manager', () => {

      await addToAgentOnStartup({
        addToConfig: mockAddToConfig,
+        sessionId: 'test-session',
        extensionConfig: mockExtensionConfig,
      });

@@ -132,6 +158,7 @@ describe('Extension Manager', () => {

      await addToAgentOnStartup({
        addToConfig: mockAddToConfig,
+        sessionId: 'test-session',
        extensionConfig: mockExtensionConfig,
      });

@@ -153,6 +180,7 @@ describe('Extension Manager', () => {
      await updateExtension({
        enabled: true,
        addToConfig: mockAddToConfig,
+        sessionId: 'test-session',
        removeFromConfig: mockRemoveFromConfig,
        extensionConfig: mockExtensionConfig,
        originalName: 'test-extension',
@@ -160,7 +188,8 @@ describe('Extension Manager', () => {

      expect(mockAddToAgent).toHaveBeenCalledWith(
        { ...mockExtensionConfig, name: 'test-extension' },
-        { silent: true }
+        { silent: true },
+        'test-session'
      );
      expect(mockAddToConfig).toHaveBeenCalledWith(
        'test-extension',
@@ -183,16 +212,22 @@ describe('Extension Manager', () => {
      await updateExtension({
        enabled: true,
        addToConfig: mockAddToConfig,
+        sessionId: 'test-session',
        removeFromConfig: mockRemoveFromConfig,
        extensionConfig: { ...mockExtensionConfig, name: 'new-extension' },
        originalName: 'old-extension',
      });

-      expect(mockRemoveFromAgent).toHaveBeenCalledWith('old-extension', { silent: true });
+      expect(mockRemoveFromAgent).toHaveBeenCalledWith(
+        'old-extension',
+        { silent: true },
+        'test-session'
+      );
      expect(mockRemoveFromConfig).toHaveBeenCalledWith('old-extension');
      expect(mockAddToAgent).toHaveBeenCalledWith(
        { ...mockExtensionConfig, name: 'new-extension' },
-        { silent: true }
+        { silent: true },
+        'test-session'
      );
      expect(mockAddToConfig).toHaveBeenCalledWith(
        'new-extension',
@@ -208,6 +243,7 @@ describe('Extension Manager', () => {
      await updateExtension({
        enabled: false,
        addToConfig: mockAddToConfig,
+        sessionId: 'test-session',
        removeFromConfig: mockRemoveFromConfig,
        extensionConfig: mockExtensionConfig,
        originalName: 'test-extension',
@@ -235,9 +271,10 @@ describe('Extension Manager', () => {
        toggle: 'toggleOn',
        extensionConfig: mockExtensionConfig,
        addToConfig: mockAddToConfig,
+        sessionId: 'test-session',
      });

-      expect(mockAddToAgent).toHaveBeenCalledWith(mockExtensionConfig, {});
+      expect(mockAddToAgent).toHaveBeenCalledWith(mockExtensionConfig, {}, 'test-session');
      expect(mockAddToConfig).toHaveBeenCalledWith('test-extension', mockExtensionConfig, true);
    });

@@ -249,9 +286,10 @@ describe('Extension Manager', () => {
        toggle: 'toggleOff',
        extensionConfig: mockExtensionConfig,
        addToConfig: mockAddToConfig,
+        sessionId: 'test-session',
      });

-      expect(mockRemoveFromAgent).toHaveBeenCalledWith('test-extension', {});
+      expect(mockRemoveFromAgent).toHaveBeenCalledWith('test-extension', {}, 'test-session');
      expect(mockAddToConfig).toHaveBeenCalledWith('test-extension', mockExtensionConfig, false);
    });

@@ -265,10 +303,11 @@ describe('Extension Manager', () => {
          toggle: 'toggleOn',
          extensionConfig: mockExtensionConfig,
          addToConfig: mockAddToConfig,
+          sessionId: 'test-session',
        })
      ).rejects.toThrow('Agent failed');

-      expect(mockAddToAgent).toHaveBeenCalledWith(mockExtensionConfig, {});
+      expect(mockAddToAgent).toHaveBeenCalledWith(mockExtensionConfig, {}, 'test-session');
      // addToConfig is called during the rollback (toggleOff)
      expect(mockAddToConfig).toHaveBeenCalledWith('test-extension', mockExtensionConfig, false);
    });
@@ -283,12 +322,13 @@ describe('Extension Manager', () => {
          toggle: 'toggleOn',
          extensionConfig: mockExtensionConfig,
          addToConfig: mockAddToConfig,
+          sessionId: 'test-session',
        })
      ).rejects.toThrow('Config failed');

-      expect(mockAddToAgent).toHaveBeenCalledWith(mockExtensionConfig, {});
+      expect(mockAddToAgent).toHaveBeenCalledWith(mockExtensionConfig, {}, 'test-session');
      expect(mockAddToConfig).toHaveBeenCalledWith('test-extension', mockExtensionConfig, true);
-      expect(mockRemoveFromAgent).toHaveBeenCalledWith('test-extension', {});
+      expect(mockRemoveFromAgent).toHaveBeenCalledWith('test-extension', {}, 'test-session');
    });

    it('should update config even if agent removal fails when toggling off', async () => {
@@ -301,10 +341,11 @@ describe('Extension Manager', () => {
          toggle: 'toggleOff',
          extensionConfig: mockExtensionConfig,
          addToConfig: mockAddToConfig,
+          sessionId: 'test-session',
        })
      ).rejects.toThrow('Agent removal failed');

-      expect(mockRemoveFromAgent).toHaveBeenCalledWith('test-extension', {});
+      expect(mockRemoveFromAgent).toHaveBeenCalledWith('test-extension', {}, 'test-session');
      expect(mockAddToConfig).toHaveBeenCalledWith('test-extension', mockExtensionConfig, false);
    });
  });
@@ -317,9 +358,14 @@ describe('Extension Manager', () => {
      await deleteExtension({
        name: 'test-extension',
        removeFromConfig: mockRemoveFromConfig,
+        sessionId: 'test-session',
      });

-      expect(mockRemoveFromAgent).toHaveBeenCalledWith('test-extension', { isDelete: true });
+      expect(mockRemoveFromAgent).toHaveBeenCalledWith(
+        'test-extension',
+        { isDelete: true },
+        'test-session'
+      );
      expect(mockRemoveFromConfig).toHaveBeenCalledWith('test-extension');
    });

@@ -332,10 +378,15 @@ describe('Extension Manager', () => {
        deleteExtension({
          name: 'test-extension',
          removeFromConfig: mockRemoveFromConfig,
+          sessionId: 'test-session',
        })
      ).rejects.toThrow('Agent removal failed');

-      expect(mockRemoveFromAgent).toHaveBeenCalledWith('test-extension', { isDelete: true });
+      expect(mockRemoveFromAgent).toHaveBeenCalledWith(
+        'test-extension',
+        { isDelete: true },
+        'test-session'
+      );
      expect(mockRemoveFromConfig).toHaveBeenCalledWith('test-extension');
    });

@@ -349,10 +400,15 @@ describe('Extension Manager', () => {
        deleteExtension({
          name: 'test-extension',
          removeFromConfig: mockRemoveFromConfig,
+          sessionId: 'test-session',
        })
      ).rejects.toThrow('Config removal failed');

-      expect(mockRemoveFromAgent).toHaveBeenCalledWith('test-extension', { isDelete: true });
+      expect(mockRemoveFromAgent).toHaveBeenCalledWith(
+        'test-extension',
+        { isDelete: true },
+        'test-session'
+      );
      expect(mockRemoveFromConfig).toHaveBeenCalledWith('test-extension');
    });
  });
@@ -5,6 +5,7 @@ import { addToAgent, removeFromAgent, sanitizeName } from './agent-api';
 interface ActivateExtensionProps {
  addToConfig: (name: string, extensionConfig: ExtensionConfig, enabled: boolean) => Promise<void>;
  extensionConfig: ExtensionConfig;
+  sessionId: string;
 }

 type ExtensionError = {
@@ -55,10 +56,11 @@ async function retryWithBackoff<T>(fn: () => Promise<T>, options: RetryOptions =
 export async function activateExtension({
  addToConfig,
  extensionConfig,
+  sessionId,
 }: ActivateExtensionProps): Promise<void> {
  try {
    // AddToAgent
-    await addToAgent(extensionConfig, { silent: false });
+    await addToAgent(extensionConfig, { silent: false }, sessionId);
  } catch (error) {
    console.error('Failed to add extension to agent:', error);
    // add to config with enabled = false
@@ -74,7 +76,7 @@ export async function activateExtension({
    console.error('Failed to add extension to config:', error);
    // remove from Agent
    try {
-      await removeFromAgent(extensionConfig.name);
+      await removeFromAgent(extensionConfig.name, {}, sessionId);
    } catch (removeError) {
      console.error('Failed to remove extension from agent after config failure:', removeError);
    }
@@ -87,6 +89,7 @@ interface AddToAgentOnStartupProps {
  addToConfig: (name: string, extensionConfig: ExtensionConfig, enabled: boolean) => Promise<void>;
  extensionConfig: ExtensionConfig;
  toastOptions?: ToastServiceOptions;
+  sessionId: string;
 }

 /**
@@ -96,9 +99,10 @@ export async function addToAgentOnStartup({
  addToConfig,
  extensionConfig,
  toastOptions = { silent: true },
+  sessionId,
 }: AddToAgentOnStartupProps): Promise<void> {
  try {
-    await retryWithBackoff(() => addToAgent(extensionConfig, toastOptions), {
+    await retryWithBackoff(() => addToAgent(extensionConfig, toastOptions, sessionId), {
      retries: 3,
      delayMs: 1000,
      shouldRetry: (error: ExtensionError) =>
@@ -121,6 +125,7 @@ export async function addToAgentOnStartup({
        extensionConfig,
        addToConfig,
        toastOptions: { silent: true },
+        sessionId,
      });
    } catch (toggleErr) {
      console.error('Failed to toggle off after error:', toggleErr);
@@ -134,6 +139,7 @@ interface UpdateExtensionProps {
  removeFromConfig: (name: string) => Promise<void>;
  extensionConfig: ExtensionConfig;
  originalName?: string;
+  sessionId: string;
 }

 /**
@@ -145,6 +151,7 @@ export async function updateExtension({
  removeFromConfig,
  extensionConfig,
  originalName,
+  sessionId,
 }: UpdateExtensionProps) {
  // Sanitize the new name to match the behavior when adding extensions
  const sanitizedNewName = sanitizeName(extensionConfig.name);
@@ -158,7 +165,7 @@ export async function updateExtension({

    // First remove the old extension from agent (using original name)
    try {
-      await removeFromAgent(originalName!, { silent: true }); // Suppress removal toast since we'll show update toast
+      await removeFromAgent(originalName!, { silent: true }, sessionId); // Suppress removal toast since we'll show update toast
    } catch (error) {
      console.error('Failed to remove old extension from agent during rename:', error);
      // Continue with the process even if agent removal fails
@@ -182,7 +189,7 @@ export async function updateExtension({
    if (enabled) {
      try {
        // AddToAgent with silent option to avoid duplicate toasts
-        await addToAgent(sanitizedExtensionConfig, { silent: true });
+        await addToAgent(sanitizedExtensionConfig, { silent: true }, sessionId);
      } catch (error) {
        console.error('[updateExtension]: Failed to add renamed extension to agent:', error);
        throw error;
@@ -212,7 +219,7 @@ export async function updateExtension({
    if (enabled) {
      try {
        // AddToAgent with silent option to avoid duplicate toasts
-        await addToAgent(sanitizedExtensionConfig, { silent: true });
+        await addToAgent(sanitizedExtensionConfig, { silent: true }, sessionId);
      } catch (error) {
        console.error('[updateExtension]: Failed to add extension to agent during update:', error);
        // Failed to add to agent -- show that error to user and do not update the config file
@@ -254,6 +261,7 @@ interface ToggleExtensionProps {
  extensionConfig: ExtensionConfig;
  addToConfig: (name: string, extensionConfig: ExtensionConfig, enabled: boolean) => Promise<void>;
  toastOptions?: ToastServiceOptions;
+  sessionId: string;
 }

 /**
@@ -264,14 +272,19 @@ export async function toggleExtension({
  extensionConfig,
  addToConfig,
  toastOptions = {},
+  sessionId,
 }: ToggleExtensionProps) {
  // disabled to enabled
  if (toggle == 'toggleOn') {
    try {
      // add to agent with toast options
-      await addToAgent(extensionConfig, {
-        ...toastOptions,
-      });
+      await addToAgent(
+        extensionConfig,
+        {
+          ...toastOptions,
+        },
+        sessionId
+      );
    } catch (error) {
      console.error('Error adding extension to agent. Will try to toggle back off.');
      try {
@@ -280,6 +293,7 @@ export async function toggleExtension({
          extensionConfig,
          addToConfig,
          toastOptions: { silent: true }, // otherwise we will see a toast for removing something that was never added
+          sessionId,
        });
      } catch (toggleError) {
        console.error('Failed to toggle extension off after agent error:', toggleError);
@@ -294,7 +308,7 @@ export async function toggleExtension({
      console.error('Failed to update config after enabling extension:', error);
      // remove from agent
      try {
-        await removeFromAgent(extensionConfig.name, toastOptions);
+        await removeFromAgent(extensionConfig.name, toastOptions, sessionId);
      } catch (removeError) {
        console.error('Failed to remove extension from agent after config failure:', removeError);
      }
@@ -304,7 +318,7 @@ export async function toggleExtension({
    // enabled to disabled
    let agentRemoveError = null;
    try {
-      await removeFromAgent(extensionConfig.name, toastOptions);
+      await removeFromAgent(extensionConfig.name, toastOptions, sessionId);
    } catch (error) {
      // note there was an error, but attempt to remove from config anyway
      console.error('Error removing extension from agent', extensionConfig.name, error);
@@ -329,16 +343,17 @@ export async function toggleExtension({
 interface DeleteExtensionProps {
  name: string;
  removeFromConfig: (name: string) => Promise<void>;
+  sessionId: string;
 }

 /**
 * Deletes an extension completely from both agent and config
 */
-export async function deleteExtension({ name, removeFromConfig }: DeleteExtensionProps) {
+export async function deleteExtension({ name, removeFromConfig, sessionId }: DeleteExtensionProps) {
  // remove from agent
  let agentRemoveError = null;
  try {
-    await removeFromAgent(name, { isDelete: true });
+    await removeFromAgent(name, { isDelete: true }, sessionId);
  } catch (error) {
    console.error('Failed to remove extension from agent during deletion:', error);
    agentRemoveError = error;