Implementation for kernel reconnection + tests

package.json

@@ -780,6 +780,30 @@
                 "title": "List Variables",
                 "enablement": "isWorkspaceTrusted && !jupyter.webExtension",
                 "category": "Jupyter"
+            },
+            {
+                "command": "jupyter.startPersistentServer",
+                "title": "Start Persistent Jupyter Server",
+                "category": "Jupyter",
+                "enablement": "isWorkspaceTrusted && !jupyter.webExtension"
+            },
+            {
+                "command": "jupyter.stopPersistentServer",
+                "title": "Stop Persistent Jupyter Server",
+                "category": "Jupyter",
+                "enablement": "isWorkspaceTrusted && !jupyter.webExtension"
+            },
+            {
+                "command": "jupyter.connectToPersistentServer",
+                "title": "Connect to Persistent Jupyter Server",
+                "category": "Jupyter",
+                "enablement": "isWorkspaceTrusted && !jupyter.webExtension"
+            },
+            {
+                "command": "jupyter.cleanupPersistentServers",
+                "title": "Cleanup Persistent Jupyter Servers",
+                "category": "Jupyter",
+                "enablement": "isWorkspaceTrusted && !jupyter.webExtension"
             }
         ],
         "submenus": [

src/commands.ts

@@ -49,6 +49,10 @@ export interface ICommandNameArgumentTypeMapping {
     ['jupyter.selectLocalJupyterServer']: [] | [undefined | string];
     ['workbench.action.openSettings']: ['jupyter.kernels.excludePythonEnvironments'];
     ['jupyter.getUsedAzMLServerHandles']: [];
+    [DSCommands.StartPersistentServer]: [];
+    [DSCommands.StopPersistentServer]: [string | undefined]; // serverId
+    [DSCommands.ConnectToPersistentServer]: [];
+    [DSCommands.CleanupPersistentServers]: [];
     [DSCommands.RunCurrentCell]: [];
     [DSCommands.RunCurrentCellAdvance]: [];
     [DSCommands.CreateNewInteractive]: [];

src/kernels/jupyter/connection/serverUriStorage.ts

@@ -2,14 +2,15 @@
 // Licensed under the MIT License.
 
 import { EventEmitter, Memento, env } from 'vscode';
-import { inject, injectable, named } from 'inversify';
+import { inject, injectable, named, optional } from 'inversify';
 import { IMemento, GLOBAL_MEMENTO, IDisposableRegistry } from '../../../platform/common/types';
 import { logger } from '../../../platform/logging';
 import { generateIdFromRemoteProvider } from '../jupyterUtils';
 import { IJupyterServerUriEntry, IJupyterServerUriStorage, JupyterServerProviderHandle } from '../types';
 import { noop } from '../../../platform/common/utils/misc';
 import { DisposableBase } from '../../../platform/common/utils/lifecycle';
 import { Settings } from '../../../platform/common/constants';
+import { IPersistentServerStorage } from '../launcher/persistentServerStorage';
 
 export type StorageMRUItem = {
     displayName: string;
@@ -61,12 +62,19 @@ export class JupyterServerUriStorage extends DisposableBase implements IJupyterS
     constructor(
         @inject(IMemento) @named(GLOBAL_MEMENTO) globalMemento: Memento,
         @inject(IDisposableRegistry)
-        disposables: IDisposableRegistry
+        disposables: IDisposableRegistry,
+        @inject(IPersistentServerStorage) @optional() private readonly persistentServerStorage?: IPersistentServerStorage
     ) {
         super();
         disposables.push(this);
         // eslint-disable-next-line @typescript-eslint/no-use-before-define
         this.newStorage = this._register(new NewStorage(globalMemento));
+
+        // Listen to persistent server changes
+        if (this.persistentServerStorage) {
+            this._register(this.persistentServerStorage.onDidAdd(() => this._onDidChangeUri.fire()));
+            this._register(this.persistentServerStorage.onDidRemove(() => this._onDidChangeUri.fire()));
+        }
     }
     private hookupStorageEvents() {
         if (this.storageEventsHooked) {
@@ -80,12 +88,48 @@ export class JupyterServerUriStorage extends DisposableBase implements IJupyterS
     private updateStore(): IJupyterServerUriEntry[] {
         this.hookupStorageEvents();
         const previous = this._all;
-        this._all = this.newStorage.getAll();
+        const regularServers = this.newStorage.getAll();
+        const persistentServers = this.getPersistentServers();
+
+        // Combine regular servers with persistent servers, avoiding duplicates
+        const allServers = [...regularServers, ...persistentServers];
+        const uniqueServers = allServers.filter((server, index, arr) => 
+            arr.findIndex(s => generateIdFromRemoteProvider(s.provider) === generateIdFromRemoteProvider(server.provider)) === index
+        );
+
+        this._all = uniqueServers.sort((a, b) => b.time - a.time); // Sort by most recent first
+
         if (previous.length !== this._all.length || JSON.stringify(this._all) !== JSON.stringify(previous)) {
             this._onDidLoad.fire();
         }
         return this._all;
     }
+
+    private getPersistentServers(): IJupyterServerUriEntry[] {
+        if (!this.persistentServerStorage) {
+            return [];
+        }
+
+        const persistentServers = this.persistentServerStorage.all;
+        return persistentServers
+            .filter(server => server.launchedByExtension) // Only show servers launched by the extension
+            .map(server => {
+                const serverHandle: JupyterServerProviderHandle = {
+                    id: 'persistent-server-provider',
+                    handle: server.serverId,
+                    extensionId: 'ms-toolsai.jupyter'
+                };
+
+                const entry: IJupyterServerUriEntry = {
+                    provider: serverHandle,
+                    time: server.time,
+                    displayName: `${server.displayName} (Persistent)`
+                };
+
+                return entry;
+            });
+    }
+
     public async clear(): Promise<void> {
         this.hookupStorageEvents();
         await this.newStorage.clear();
@@ -106,12 +150,30 @@ export class JupyterServerUriStorage extends DisposableBase implements IJupyterS
     }
     public async update(server: JupyterServerProviderHandle) {
         this.hookupStorageEvents();
-        await this.newStorage.update(server);
+
+        // Check if this is a persistent server
+        if (server.id === 'persistent-server-provider' && this.persistentServerStorage) {
+            // Update persistent server time
+            await this.persistentServerStorage.update(server.handle, { time: Date.now() });
+        } else {
+            // Update regular server
+            await this.newStorage.update(server);
+        }
+
         this.updateStore();
     }
     public async remove(server: JupyterServerProviderHandle) {
         this.hookupStorageEvents();
-        await this.newStorage.remove(server);
+
+        // Check if this is a persistent server
+        if (server.id === 'persistent-server-provider' && this.persistentServerStorage) {
+            // Remove from persistent storage
+            await this.persistentServerStorage.remove(server.handle);
+        } else {
+            // Remove from regular storage
+            await this.newStorage.remove(server);
+        }
+
         this.updateStore();
     }
 }

src/kernels/jupyter/connection/serverUriStorage.unit.test.ts

@@ -19,6 +19,7 @@ import { TestEventHandler, createEventHandler } from '../../../test/common';
 import { generateIdFromRemoteProvider } from '../jupyterUtils';
 import { sleep } from '../../../test/core';
 import { mockedVSCodeNamespaces } from '../../../test/vscode-mock';
+import { IPersistentServerStorage, IPersistentServerInfo } from '../launcher/persistentServerStorage';
 
 suite('Server Uri Storage', async () => {
     let serverUriStorage: IJupyterServerUriStorage;
@@ -581,4 +582,154 @@ suite('Server Uri Storage', async () => {
         onDidAddEvent = createEventHandler(serverUriStorage, 'onDidAdd', disposables);
         return itemsInNewStorage;
     }
+
+    suite('Persistent Server Integration', () => {
+        let persistentServerStorage: IPersistentServerStorage;
+        let persistentAddEventEmitter: EventEmitter<IPersistentServerInfo>;
+        let persistentRemoveEventEmitter: EventEmitter<string>;
+
+        const mockPersistentServer: IPersistentServerInfo = {
+            serverId: 'persistent-server-1',
+            displayName: 'Test Persistent Server',
+            url: 'http://localhost:8888/?token=test-token',
+            token: 'test-token',
+            workingDirectory: '/test/workspace',
+            launchedByExtension: true,
+            time: Date.now()
+        };
+
+        setup(() => {
+            persistentServerStorage = mock<IPersistentServerStorage>();
+            persistentAddEventEmitter = new EventEmitter<IPersistentServerInfo>();
+            persistentRemoveEventEmitter = new EventEmitter<string>();
+            disposables.push(persistentAddEventEmitter, persistentRemoveEventEmitter);
+
+            when(persistentServerStorage.all).thenReturn([mockPersistentServer]);
+            when(persistentServerStorage.onDidAdd).thenReturn(persistentAddEventEmitter.event);
+            when(persistentServerStorage.onDidRemove).thenReturn(persistentRemoveEventEmitter.event);
+            when(persistentServerStorage.remove(anything())).thenResolve();
+            when(persistentServerStorage.update(anything(), anything())).thenResolve();
+        });
+
+        test('Should include persistent servers in all list', () => {
+            generateDummyData(1);
+            // Create storage with persistent server storage
+            serverUriStorage = new JupyterServerUriStorage(instance(memento), disposables, instance(persistentServerStorage));
+
+            const all = serverUriStorage.all;
+
+            // Should have regular servers + persistent servers
+            assert.isTrue(all.length >= 2, 'Should include both regular and persistent servers');
+
+            const persistentEntry = all.find(entry => 
+                entry.provider.id === 'persistent-server-provider' && 
+                entry.provider.handle === 'persistent-server-1'
+            );
+
+            assert.isDefined(persistentEntry, 'Should find persistent server entry');
+            assert.isTrue(persistentEntry!.displayName!.includes('(Persistent)'), 'Should mark as persistent');
+        });
+
+        test('Should filter out non-extension persistent servers', () => {
+            const externalServer: IPersistentServerInfo = {
+                ...mockPersistentServer,
+                serverId: 'external-server',
+                launchedByExtension: false
+            };
+
+            when(persistentServerStorage.all).thenReturn([mockPersistentServer, externalServer]);
+            generateDummyData(0);
+            serverUriStorage = new JupyterServerUriStorage(instance(memento), disposables, instance(persistentServerStorage));
+
+            const all = serverUriStorage.all;
+            const persistentEntries = all.filter(entry => entry.provider.id === 'persistent-server-provider');
+
+            assert.equal(persistentEntries.length, 1, 'Should only include extension-launched servers');
+            assert.equal(persistentEntries[0].provider.handle, 'persistent-server-1');
+        });
+
+        test('Should remove persistent server when requested', async () => {
+            generateDummyData(0);
+            serverUriStorage = new JupyterServerUriStorage(instance(memento), disposables, instance(persistentServerStorage));
+
+            const persistentServerHandle: JupyterServerProviderHandle = {
+                id: 'persistent-server-provider',
+                handle: 'persistent-server-1',
+                extensionId: 'ms-toolsai.jupyter'
+            };
+
+            await serverUriStorage.remove(persistentServerHandle);
+
+            verify(persistentServerStorage.remove('persistent-server-1')).once();
+        });
+
+        test('Should update persistent server when requested', async () => {
+            generateDummyData(0);
+            serverUriStorage = new JupyterServerUriStorage(instance(memento), disposables, instance(persistentServerStorage));
+
+            const persistentServerHandle: JupyterServerProviderHandle = {
+                id: 'persistent-server-provider',
+                handle: 'persistent-server-1',
+                extensionId: 'ms-toolsai.jupyter'
+            };
+
+            await serverUriStorage.update(persistentServerHandle);
+
+            verify(persistentServerStorage.update('persistent-server-1', anything())).once();
+        });
+
+        test('Should fire change events when persistent servers change', () => {
+            generateDummyData(0);
+            serverUriStorage = new JupyterServerUriStorage(instance(memento), disposables, instance(persistentServerStorage));
+            onDidChangeEvent = createEventHandler(serverUriStorage, 'onDidChange', disposables);
+
+            const initialCount = onDidChangeEvent.count;
+
+            // Simulate persistent server add event
+            persistentAddEventEmitter.fire(mockPersistentServer);
+
+            assert.equal(onDidChangeEvent.count, initialCount + 1, 'Should fire change event on add');
+
+            // Simulate persistent server remove event
+            persistentRemoveEventEmitter.fire('persistent-server-1');
+
+            assert.equal(onDidChangeEvent.count, initialCount + 2, 'Should fire change event on remove');
+        });
+
+        test('Should work without persistent server storage', () => {
+            generateDummyData(1);
+            // Create storage without persistent server storage
+            serverUriStorage = new JupyterServerUriStorage(instance(memento), disposables);
+
+            const all = serverUriStorage.all;
+
+            // Should only have regular servers
+            const persistentEntries = all.filter(entry => entry.provider.id === 'persistent-server-provider');
+            assert.equal(persistentEntries.length, 0, 'Should not have persistent servers');
+        });
+
+        test('Should avoid duplicates between regular and persistent servers', () => {
+            // Create a regular server with same ID as persistent server
+            const regularServerData = [{
+                displayName: 'Regular Server',
+                serverHandle: {
+                    id: 'persistent-server-provider',
+                    handle: 'persistent-server-1',
+                    extensionId: 'ms-toolsai.jupyter'
+                },
+                time: Date.now() - 1000
+            }];
+
+            when(memento.get(mementoKeyForStoringUsedJupyterProviders, anything())).thenReturn(regularServerData);
+            serverUriStorage = new JupyterServerUriStorage(instance(memento), disposables, instance(persistentServerStorage));
+
+            const all = serverUriStorage.all;
+            const matchingEntries = all.filter(entry => 
+                entry.provider.id === 'persistent-server-provider' && 
+                entry.provider.handle === 'persistent-server-1'
+            );
+
+            assert.equal(matchingEntries.length, 1, 'Should not have duplicates');
+        });
+    });
 });