Add session switching support

Author: thegecko

src/entry-points/browser/extension.ts

@@ -18,7 +18,7 @@ import * as vscode from 'vscode';
 import { AdapterRegistry } from '../../plugin/adapter-registry/adapter-registry';
 import { CAdapter } from '../../plugin/adapter-registry/c-adapter';
 import { ContextTracker } from '../../plugin/context-tracker';
-import { MemoryProvider } from '../../plugin/memory-provider';
+import { MemoryProviderManager } from '../../plugin/memory-provider-manager';
 import { MemoryStorage } from '../../plugin/memory-storage';
 import { MemoryWebview } from '../../plugin/memory-webview-main';
 import { SessionTracker } from '../../plugin/session-tracker';
@@ -27,14 +27,14 @@ export const activate = async (context: vscode.ExtensionContext): Promise<Adapte
     const registry = new AdapterRegistry();
     const sessionTracker = new SessionTracker();
     new ContextTracker(sessionTracker);
-    const memoryProvider = new MemoryProvider(registry, sessionTracker);
-    const memoryView = new MemoryWebview(context.extensionUri, memoryProvider, sessionTracker);
-    const memoryStorage = new MemoryStorage(memoryProvider);
+    const memoryProviderManager = new MemoryProviderManager(registry, sessionTracker);
+    const memoryStorage = new MemoryStorage(sessionTracker, memoryProviderManager);
+    const memoryView = new MemoryWebview(context.extensionUri, memoryProviderManager, sessionTracker, memoryStorage);
     const cAdapter = new CAdapter(registry);
 
     registry.activate(context);
     sessionTracker.activate(context);
-    memoryProvider.activate(context);
+    memoryProviderManager.activate(context);
     memoryView.activate(context);
     memoryStorage.activate(context);
     cAdapter.activate(context);

src/entry-points/desktop/extension.ts

@@ -18,7 +18,7 @@ import * as vscode from 'vscode';
 import { AdapterRegistry } from '../../plugin/adapter-registry/adapter-registry';
 import { CAdapter } from '../../plugin/adapter-registry/c-adapter';
 import { ContextTracker } from '../../plugin/context-tracker';
-import { MemoryProvider } from '../../plugin/memory-provider';
+import { MemoryProviderManager } from '../../plugin/memory-provider-manager';
 import { MemoryStorage } from '../../plugin/memory-storage';
 import { MemoryWebview } from '../../plugin/memory-webview-main';
 import { SessionTracker } from '../../plugin/session-tracker';
@@ -27,14 +27,14 @@ export const activate = async (context: vscode.ExtensionContext): Promise<Adapte
     const registry = new AdapterRegistry();
     const sessionTracker = new SessionTracker();
     new ContextTracker(sessionTracker);
-    const memoryProvider = new MemoryProvider(registry, sessionTracker);
-    const memoryView = new MemoryWebview(context.extensionUri, memoryProvider, sessionTracker);
-    const memoryStorage = new MemoryStorage(memoryProvider);
+    const memoryProviderManager = new MemoryProviderManager(registry, sessionTracker);
+    const memoryStorage = new MemoryStorage(sessionTracker, memoryProviderManager);
+    const memoryView = new MemoryWebview(context.extensionUri, memoryProviderManager, sessionTracker, memoryStorage);
     const cAdapter = new CAdapter(registry);
 
     registry.activate(context);
     sessionTracker.activate(context);
-    memoryProvider.activate(context);
+    memoryProviderManager.activate(context);
     memoryView.activate(context);
     memoryStorage.activate(context);
     cAdapter.activate(context);

src/plugin/memory-provider-manager.ts

@@ -0,0 +1,47 @@
+/********************************************************************************
+ * Copyright (C) 2025 Ericsson, Arm and others.
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Eclipse Public License v. 2.0 which is available at
+ * http://www.eclipse.org/legal/epl-2.0.
+ *
+ * This Source Code may also be made available under the following Secondary
+ * Licenses when the conditions for such availability set forth in the Eclipse
+ * Public License v. 2.0 are satisfied: GNU General Public License, version 2
+ * with the GNU Classpath Exception which is available at
+ * https://www.gnu.org/software/classpath/license.html.
+ *
+ * SPDX-License-Identifier: EPL-2.0 OR GPL-2.0 WITH Classpath-exception-2.0
+ ********************************************************************************/
+
+import * as vscode from 'vscode';
+import { AdapterRegistry } from './adapter-registry/adapter-registry';
+import { MemoryProvider } from './memory-provider';
+import { SessionTracker } from './session-tracker';
+
+export class MemoryProviderManager {
+
+    protected memoryProviders = new Map<vscode.DebugSession, MemoryProvider>();
+
+    constructor(protected adapterRegistry: AdapterRegistry, protected sessionTracker: SessionTracker) {
+    }
+
+    public activate(context: vscode.ExtensionContext): void {
+        const createDebugAdapterTracker = (session: vscode.DebugSession): vscode.ProviderResult<vscode.DebugAdapterTracker> => {
+            const handlerForSession = this.adapterRegistry.getHandlerForSession(session.type);
+            const contributedTracker = handlerForSession?.initializeAdapterTracker?.(session);
+            return contributedTracker;
+        };
+        context.subscriptions.push(vscode.debug.registerDebugAdapterTrackerFactory('*', { createDebugAdapterTracker }));
+    }
+
+    public getProvider(sessionId: string | undefined): MemoryProvider {
+        const session = this.sessionTracker.assertSession(sessionId);
+
+        if (!this.memoryProviders.has(session)) {
+            this.memoryProviders.set(session, new MemoryProvider(session.id, this.adapterRegistry, this.sessionTracker));
+        }
+
+        return this.memoryProviders.get(session)!;
+    }
+}

src/plugin/memory-provider.ts

@@ -15,7 +15,6 @@
  ********************************************************************************/
 
 import { DebugProtocol } from '@vscode/debugprotocol';
-import * as vscode from 'vscode';
 import { sendRequest } from '../common/debug-requests';
 import { stringToBytesMemory } from '../common/memory';
 import { VariableRange } from '../common/memory-range';
@@ -24,37 +23,17 @@ import { MemoryDisplaySettingsContribution } from '../common/webview-configurati
 import { AdapterRegistry } from './adapter-registry/adapter-registry';
 import { isSessionEvent, SessionTracker } from './session-tracker';
 
-export interface LabeledUint8Array extends Uint8Array {
-    label?: string;
-}
-
 export class MemoryProvider {
     protected scheduledOnDidMemoryWriteEvents: { [sessionidmemoryReference: string]: ((response: WriteMemoryResult) => void) | undefined } = {};
 
-    protected _sessionId: string | undefined;
-    public get sessionId(): string | undefined { return this._sessionId; }
-
-    constructor(protected adapterRegistry: AdapterRegistry, protected sessionTracker: SessionTracker) {
+    constructor(protected sessionId: string, protected adapterRegistry: AdapterRegistry, protected sessionTracker: SessionTracker) {
         this.sessionTracker.onSessionEvent(event => {
-            if (isSessionEvent('memory-written', event)) {
+            if (isSessionEvent('memory-written', event) && event.session.raw.id === this.sessionId) {
                 delete this.scheduledOnDidMemoryWriteEvents[event.session.raw.id + '_' + event.data.memoryReference];
             }
         });
     }
 
-    public activate(context: vscode.ExtensionContext): void {
-        const createDebugAdapterTracker = (session: vscode.DebugSession): vscode.ProviderResult<vscode.DebugAdapterTracker> => {
-            const handlerForSession = this.adapterRegistry.getHandlerForSession(session.type);
-            const contributedTracker = handlerForSession?.initializeAdapterTracker?.(session);
-            return contributedTracker;
-        };
-        context.subscriptions.push(vscode.debug.registerDebugAdapterTrackerFactory('*', { createDebugAdapterTracker }));
-    }
-
-    public setSessionId(value: string | undefined): void {
-        this._sessionId = value;
-    }
-
     public async readMemory(args: DebugProtocol.ReadMemoryArguments): Promise<ReadMemoryResult> {
         const session = this.sessionTracker.assertSession(this.sessionId);
         return sendRequest(this.sessionTracker.assertDebugCapability(session, 'supportsReadMemoryRequest', 'read memory'), 'readMemory', args);

src/plugin/memory-storage.ts

@@ -28,9 +28,11 @@ import { toHexStringWithRadixMarker } from '../common/memory-range';
 import { ApplyMemoryArguments, ApplyMemoryResult, MemoryOptions, StoreMemoryArguments } from '../common/messaging';
 import { isWebviewContext } from '../common/webview-context';
 import { MemoryProvider } from './memory-provider';
+import { MemoryProviderManager } from './memory-provider-manager';
+import { SessionTracker } from './session-tracker';
 
-export const StoreCommandType = `${manifest.PACKAGE_NAME}.store-file`;
-export const ApplyCommandType = `${manifest.PACKAGE_NAME}.apply-file`;
+const StoreCommandType = `${manifest.PACKAGE_NAME}.store-file`;
+const ApplyCommandType = `${manifest.PACKAGE_NAME}.apply-file`;
 
 const VALID_FILE_NAME_CHARS = /[^a-zA-Z0-9 _-]/g;
 
@@ -49,19 +51,26 @@ interface ApplyMemoryOptions {
     uri: vscode.Uri;
 }
 
+const getActiveSession = () => vscode.debug.activeDebugSession;
+
 export class MemoryStorage {
-    constructor(protected memoryProvider: MemoryProvider) {
+    constructor(protected sessionTracker: SessionTracker, protected memoryProviderManager: MemoryProviderManager) {
     }
 
     public activate(context: vscode.ExtensionContext): void {
         context.subscriptions.push(
-            vscode.commands.registerCommand(StoreCommandType, args => this.storeMemory(args)),
-            vscode.commands.registerCommand(ApplyCommandType, args => this.applyMemory(args))
+            vscode.commands.registerCommand(StoreCommandType, args => this.storeMemory(getActiveSession()?.id, args)),
+            vscode.commands.registerCommand(ApplyCommandType, args => this.applyMemory(getActiveSession()?.id, args))
         );
     }
 
-    public async storeMemory(args?: StoreMemoryArguments): Promise<void> {
-        const providedDefaultOptions = await this.storeArgsToOptions(args);
+    public async storeMemory(sessionId: string | undefined, args?: StoreMemoryArguments): Promise<void> {
+        // Even if we disable the command in VS Code through enablement or when condition, programmatic execution is still possible.
+        // However, we want to fail early in case the user tries to execute a disabled command
+        this.sessionTracker.assertDebugCapability(this.sessionTracker.assertSession(sessionId), 'supportsReadMemoryRequest', 'store memory');
+
+        const memoryProvider = this.memoryProviderManager.getProvider(sessionId);
+        const providedDefaultOptions = await this.storeArgsToOptions(memoryProvider, args);
         const options = await this.getStoreMemoryOptions(providedDefaultOptions);
         if (!options) {
             // user aborted process
@@ -70,7 +79,7 @@ export class MemoryStorage {
 
         const { outputFile, ...readArgs } = options;
         try {
-            const memoryResponse = await this.memoryProvider.readMemory(readArgs);
+            const memoryResponse = await memoryProvider.readMemory(readArgs);
             const memory = createMemoryFromRead(memoryResponse);
             const memoryMap = new MemoryMap({ [Number(memory.address)]: memory.bytes });
             await vscode.workspace.fs.writeFile(outputFile, new TextEncoder().encode(memoryMap.asHexString()));
@@ -89,7 +98,7 @@ export class MemoryStorage {
         }
     }
 
-    protected async storeArgsToOptions(args?: StoreMemoryArguments): Promise<Partial<StoreMemoryOptions>> {
+    protected async storeArgsToOptions(memoryProvider: MemoryProvider, args?: StoreMemoryArguments): Promise<Partial<StoreMemoryOptions>> {
         if (!args) {
             return {};
         }
@@ -99,8 +108,8 @@ export class MemoryStorage {
         if (isVariablesContext(args)) {
             try {
                 const variableName = args.variable.evaluateName ?? args.variable.name;
-                const count = await this.memoryProvider.getSizeOfVariable(variableName);
-                const memoryReference = args.variable.memoryReference ?? await this.memoryProvider.getAddressOfVariable(variableName);
+                const count = await memoryProvider.getSizeOfVariable(variableName);
+                const memoryReference = args.variable.memoryReference ?? await memoryProvider.getAddressOfVariable(variableName);
                 return { count: Number(count), memoryReference, offset: 0, proposedOutputName: variableName };
             } catch (error) {
                 // ignore, we are just using them as default values
@@ -153,7 +162,12 @@ export class MemoryStorage {
         return { memoryReference, offset: Number(offset), count: Number(count), outputFile };
     }
 
-    public async applyMemory(args?: ApplyMemoryArguments): Promise<ApplyMemoryResult> {
+    public async applyMemory(sessionId: string | undefined, args?: ApplyMemoryArguments): Promise<ApplyMemoryResult> {
+        // Even if we disable the command in VS Code through enablement or when condition, programmatic execution is still possible.
+        // However, we want to fail early in case the user tries to execute a disabled command
+        this.sessionTracker.assertDebugCapability(this.sessionTracker.assertSession(sessionId), 'supportsWriteMemoryRequest', 'apply memory');
+
+        const memoryProvider = this.memoryProviderManager.getProvider(sessionId);
         const providedDefaultOptions = await this.applyArgsToOptions(args);
         const options = await this.getApplyMemoryOptions(providedDefaultOptions);
         if (!options) {
@@ -169,7 +183,7 @@ export class MemoryStorage {
                 memoryReference = toHexStringWithRadixMarker(address);
                 count = memory.length;
                 const data = bytesToStringMemory(memory);
-                await this.memoryProvider.writeMemory({ memoryReference, data });
+                await memoryProvider.writeMemory({ memoryReference, data });
             }
             await vscode.window.showInformationMessage(`Memory from '${vscode.workspace.asRelativePath(options.uri)}' applied.`);
             return { memoryReference, count, offset: 0 };