feat: Language model provider now updates dynamically

src/services/navieConfigurationService.ts

@@ -39,6 +39,15 @@ export default function navieConfigurationService(context: vscode.ExtensionConte
   );
 }
 
+export async function openAIApiKeyEquals(
+  extensionContext: vscode.ExtensionContext,
+  key: string | undefined
+): Promise<boolean> {
+  const { secrets } = extensionContext;
+  const storedKey = await secrets.get(OPENAI_API_KEY);
+  return key === storedKey;
+}
+
 export async function setOpenAIApiKey(
   extensionContext: vscode.ExtensionContext,
   key: string | undefined

src/services/processWatcher.ts

@@ -76,6 +76,7 @@ export class ProcessWatcher implements vscode.Disposable {
 
   protected _onError: vscode.EventEmitter<Error> = new vscode.EventEmitter<Error>();
   protected _onAbort: vscode.EventEmitter<Error> = new vscode.EventEmitter<Error>();
+  protected _onRestart: vscode.EventEmitter<void> = new vscode.EventEmitter<void>();
 
   protected shouldRun = false;
   protected hasAborted = false;
@@ -101,6 +102,11 @@ export class ProcessWatcher implements vscode.Disposable {
     return this._onAbort.event;
   }
 
+  // Fired when the process is restarted.
+  public get onRestart(): vscode.Event<void> {
+    return this._onRestart.event;
+  }
+
   public get id(): ProcessId {
     return this.options.id;
   }
@@ -169,6 +175,7 @@ export class ProcessWatcher implements vscode.Disposable {
   }
 
   async restart(): Promise<void> {
+    this._onRestart.fire();
     await this.stop();
     await this.start();
   }

src/services/rpcProcessService.ts

@@ -12,6 +12,7 @@ import assert from 'assert';
 import { DEBUG_EXCEPTION, Telemetry } from '../telemetry';
 import ErrorCode from '../telemetry/definitions/errorCodes';
 import AssetService, { AssetIdentifier } from '../assets/assetService';
+import { openAIApiKeyEquals, setOpenAIApiKey } from './navieConfigurationService';
 
 export type RpcConnect = (port: number) => Client;
 
@@ -23,6 +24,14 @@ export interface RpcProcessServiceState {
   killProcess(): void;
 }
 
+interface RpcSettings {
+  useCopilot?: boolean;
+  openAIApiKey?: string;
+
+  // If an env var is set to undefined, it will be removed from the env.
+  env?: Record<string, string | undefined>;
+}
+
 export default class RpcProcessService implements Disposable {
   public readonly _onRpcPortChange = new EventEmitter<number | undefined>();
   public readonly onRpcPortChange = this._onRpcPortChange.event;
@@ -72,6 +81,10 @@ export default class RpcProcessService implements Disposable {
     );
   }
 
+  get onRestart(): vscode.Event<void> {
+    return this.processWatcher.onRestart;
+  }
+
   // Provides some internal state access, primarily for testing purposes.
   public get state(): RpcProcessServiceState {
     return {
@@ -225,4 +238,41 @@ export default class RpcProcessService implements Disposable {
     this._onRpcPortChange.dispose();
     this.diposables.forEach((d) => d.dispose());
   }
+
+  async updateSettings(settings: RpcSettings): Promise<void> {
+    let shouldRestart = false;
+
+    if (settings.useCopilot !== undefined) {
+      const config = vscode.workspace.getConfiguration('appMap.navie');
+      const key = 'useVSCodeLM';
+
+      await config.update(key, settings.useCopilot, true);
+
+      const otherSettings = config.inspect(key);
+      if (otherSettings?.workspaceValue !== undefined) {
+        await config.update(key, undefined, undefined);
+      }
+
+      shouldRestart = true;
+    }
+
+    if (Object.hasOwnProperty.call(settings, 'openAIApiKey')) {
+      const sameKey = await openAIApiKeyEquals(this.context, settings.openAIApiKey);
+      if (!sameKey) {
+        await setOpenAIApiKey(this.context, settings.openAIApiKey);
+        shouldRestart = true;
+      }
+    }
+
+    if (settings.env) {
+      const env = vscode.workspace.getConfiguration('appMap').get('commandLineEnvironment', {});
+      Object.entries(settings.env).forEach(([k, v]) => {
+        env[k] = v;
+      });
+      await vscode.workspace.getConfiguration('appMap').update('commandLineEnvironment', env, true);
+      shouldRestart = true;
+    }
+
+    if (shouldRestart) this.restart();
+  }
 }

src/webviews/chatSearchWebview.ts

@@ -44,7 +44,8 @@ export default class ChatSearchWebview {
   private constructor(
     private readonly context: vscode.ExtensionContext,
     private readonly extensionState: ExtensionState,
-    private readonly dataService: ChatSearchDataService
+    private readonly dataService: ChatSearchDataService,
+    private readonly rpcService: RpcProcessService
   ) {
     this.filterStore = new FilterStore(context);
     this.filterStore.onDidChangeFilters((event) => {
@@ -171,6 +172,16 @@ export default class ChatSearchWebview {
           type: 'update',
           mostRecentAppMaps: appmaps,
         });
+      }),
+      this.rpcService.onRpcPortChange(() => {
+        panel.webview.postMessage({
+          type: 'navie-restarted',
+        });
+      }),
+      this.rpcService.onRestart(() => {
+        panel.webview.postMessage({
+          type: 'navie-restarting',
+        });
       })
     );
 
@@ -248,12 +259,36 @@ export default class ChatSearchWebview {
 
         case 'select-llm-option': {
           const { option } = message;
-          if (option === 'default') {
-            await vscode.commands.executeCommand('appmap.clearNavieAiSettings');
-          } else if (option === 'own-key') {
-            await vscode.commands.executeCommand('appmap.openAIApiKey.set');
-          } else {
-            console.error(`Unknown option: ${option}`);
+          switch (option) {
+            case 'default':
+              this.rpcService.updateSettings({
+                useCopilot: false,
+                openAIApiKey: '',
+                env: {
+                  OPENAI_BASE_URL: undefined,
+                  OPENAI_API_KEY: undefined,
+                  AZURE_OPENAI_API_KEY: undefined,
+                  ANTHROPIC_API_KEY: undefined,
+                },
+              });
+              break;
+
+            case 'copilot':
+              this.rpcService.updateSettings({ useCopilot: true });
+              break;
+
+            case 'own-key':
+              this.rpcService.updateSettings({
+                useCopilot: false,
+                openAIApiKey: await vscode.window.showInputBox({ placeHolder: 'OpenAI API Key' }),
+                env: {
+                  OPENAI_BASE_URL: undefined,
+                },
+              });
+              break;
+
+            default:
+              console.error(`Unknown option: ${option}`);
           }
           break;
         }
@@ -304,6 +339,6 @@ export default class ChatSearchWebview {
   ): ChatSearchWebview {
     const dataService = new ChatSearchDataService(rpcService, appmaps);
 
-    return new ChatSearchWebview(context, extensionState, dataService);
+    return new ChatSearchWebview(context, extensionState, dataService, rpcService);
   }
 }

test/unit/mock/vscode/workspace.ts

@@ -35,9 +35,42 @@ export const EVENTS = {
   onDidChangeConfiguration: listener(),
 };
 
+class Configuration extends Map<string, unknown> {
+  get(key: string, defaultValue?: unknown): unknown {
+    return super.get(key) ?? defaultValue;
+  }
+
+  inspect(key: string): { workspaceValue?: unknown } | undefined {
+    return {};
+  }
+
+  update(key: string, value: unknown, target?: unknown): Promise<void> {
+    if (value === undefined || value === null) {
+      this.delete(key);
+    } else {
+      let filteredValue = value;
+      if (typeof value === 'object') {
+        // Undefined/null values are deleted
+        filteredValue = Object.entries(value).reduce((acc, [k, v]) => {
+          if (v !== undefined && v !== null) acc[k] = v;
+          return acc;
+        }, {});
+      }
+      this.set(key, filteredValue);
+    }
+    return Promise.resolve();
+  }
+}
+
+const configs = new Map<string, Configuration>();
+
 export default {
   fs,
-  getConfiguration: () => new Map<string, unknown>(),
+  getConfiguration: (key: string) => {
+    let config = configs.get(key);
+    if (!config) configs.set(key, (config = new Configuration()));
+    return config;
+  },
   workspaceFolders: [],
   onDidChangeConfiguration: EVENTS.onDidChangeConfiguration,
   onDidChangeWorkspaceFolders: EVENTS.onDidChangeWorkspaceFolders,

test/unit/services/rpcProcessService.test.ts

@@ -12,6 +12,7 @@ import RpcProcessService, { RpcProcessServiceState } from '../../../src/services
 import MockExtensionContext from '../../mocks/mockExtensionContext';
 import EventEmitter from '../mock/vscode/EventEmitter';
 import { waitFor } from '../../waitFor';
+import vscode from '../mock/vscode';
 
 chai.use(chaiAsPromised);
 
@@ -189,4 +190,46 @@ describe('RpcProcessService', () => {
       await waitFor(`Expecting debounced restart`, () => restartSpy.calledOnce);
     });
   });
+
+  describe('updateSettings', () => {
+    it('does not restart if no settings are changed', async () => {
+      const restartSpy = sinon.spy(rpcService, 'restart');
+      await rpcService.updateSettings({});
+      expect(restartSpy.called).to.be.false;
+    });
+
+    it('restarts if the useCopilot setting is changed', async () => {
+      const restartSpy = sinon.spy(rpcService, 'restart');
+      await rpcService.updateSettings({ useCopilot: true });
+      expect(restartSpy.called).to.be.true;
+    });
+
+    it('restarts if the openAIApiKey setting is changed', async () => {
+      const restartSpy = sinon.spy(rpcService, 'restart');
+      await rpcService.updateSettings({ openAIApiKey: 'new-key' });
+      expect(restartSpy.called).to.be.true;
+    });
+
+    it('does not restart if the openAIApiKey setting is the same', async () => {
+      extensionContext.secrets.store('openai.api_key', 'old-key');
+      const restartSpy = sinon.spy(rpcService, 'restart');
+      await rpcService.updateSettings({ openAIApiKey: 'old-key' });
+      expect(restartSpy.called).to.be.false;
+    });
+
+    it('restarts if the env setting is changed', async () => {
+      const restartSpy = sinon.spy(rpcService, 'restart');
+      await rpcService.updateSettings({ env: { foo: 'bar' } });
+      expect(restartSpy.called).to.be.true;
+    });
+
+    it('properly sets env values', async () => {
+      vscode.workspace
+        .getConfiguration('appMap')
+        .update('commandLineEnvironment', { foo: 'bar', bar: 'baz' }, true);
+      await rpcService.updateSettings({ env: { foo: undefined, baz: 'qux' } });
+      const env = vscode.workspace.getConfiguration('appMap').get('commandLineEnvironment');
+      expect(env).to.deep.equal({ bar: 'baz', baz: 'qux' });
+    });
+  });
 });

web/src/chatSearchView.js

@@ -84,6 +84,14 @@ export default function mountChatSearchView() {
         });
     });
 
+    messages
+      .on('navie-restarting', () => {
+        app.$refs.ui.onNavieRestarting();
+      })
+      .on('navie-restarted', () => {
+        app.$refs.ui.loadNavieConfig();
+      });
+
     app.$on('open-install-instructions', () => {
       vscode.postMessage({ command: 'open-install-instructions' });
     });