Support "Continue On"

package.json

@@ -776,6 +776,12 @@
                 "title": "%DataScience.runInDedicatedExtensionHost%",
                 "enablement": "!jupyter.webExtension",
                 "category": "Jupyter"
+            },
+            {
+                "command": "jupyter.continueEditSessionInCodespace",
+                "title": "Continue On Codespace",
+                "enablement": "true",
+                "category": "Jupyter"
             }
         ],
         "submenus": [
@@ -885,6 +891,11 @@
                     "group": "navigation@2",
                     "when": "notebookType == 'jupyter-notebook' && config.jupyter.showOutlineButtonInNotebookToolbar"
                 },
+                {
+                    "command": "jupyter.continueEditSessionInCodespace",
+                    "group": "navigation@3",
+                    "when": "notebookType == 'jupyter-notebook' && jupyter.kernelSource == 'github-codespaces'"
+                },
                 {
                     "command": "jupyter.notebookeditor.export",
                     "group": "Jupyter",

src/commands.ts

@@ -186,4 +186,5 @@ export interface ICommandNameArgumentTypeMapping {
     [DSCommands.ReplayPylanceLogStep]: [];
     [DSCommands.InstallPythonExtensionViaKernelPicker]: [];
     [DSCommands.InstallPythonViaKernelPicker]: [];
+    [DSCommands.ContinueEditSessionInCodespace]: [];
 }

src/platform/common/constants.ts

@@ -251,6 +251,7 @@ export namespace Commands {
     export const ReplayPylanceLogStep = 'jupyter.replayPylanceLogStep';
     export const InstallPythonExtensionViaKernelPicker = 'jupyter.installPythonExtensionViaKernelPicker';
     export const InstallPythonViaKernelPicker = 'jupyter.installPythonViaKernelPicker';
+    export const ContinueEditSessionInCodespace = 'jupyter.continueEditSessionInCodespace';
 }
 
 export namespace CodeLensCommands {
@@ -286,6 +287,7 @@ export namespace EditorContexts {
     export const HasNativeNotebookOrInteractiveWindowOpen = 'jupyter.hasNativeNotebookOrInteractiveWindowOpen';
     export const ZmqAvailable = 'jupyter.zmqavailable';
     export const ReplayLogLoaded = 'jupyter.replayLogLoaded';
+    export const KernelSource = 'jupyter.kernelSource';
 }
 
 export namespace RegExpValues {

src/standalone/codespace/commandRegistry.ts

@@ -0,0 +1,48 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+import { inject, injectable } from 'inversify';
+import { IExtensionSyncActivationService } from '../../platform/activation/types';
+import { IDisposable, IDisposableRegistry } from '../../platform/common/types';
+import { ICommandManager, IWorkspaceService } from '../../platform/common/application/types';
+import { Commands } from '../../platform/common/constants';
+import { commands } from 'vscode';
+import { ICommandNameArgumentTypeMapping } from '../../commands';
+
+@injectable()
+export class CommandRegistry implements IDisposable, IExtensionSyncActivationService {
+    constructor(
+        @inject(IDisposableRegistry) private readonly disposables: IDisposableRegistry,
+        @inject(IWorkspaceService) private readonly workspace: IWorkspaceService,
+        @inject(ICommandManager) private readonly commandManager: ICommandManager
+    ) {}
+    activate() {
+        this.registerCommandsIfTrusted();
+    }
+    dispose() {
+        this.disposables.forEach((d) => d.dispose());
+    }
+
+    private registerCommandsIfTrusted() {
+        if (!this.workspace.isTrusted) {
+            return;
+        }
+
+        this.registerCommand(Commands.ContinueEditSessionInCodespace, this.continueEditSessionInCodespace);
+    }
+
+    private registerCommand<
+        E extends keyof ICommandNameArgumentTypeMapping,
+        U extends ICommandNameArgumentTypeMapping[E]
+        // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    >(command: E, callback: (...args: U) => any) {
+        const disposable = this.commandManager.registerCommand(command, callback, this);
+        this.disposables.push(disposable);
+    }
+
+    private async continueEditSessionInCodespace() {
+        await commands.executeCommand(
+            '_workbench.editSessions.actions.continueEditSession.github.codespaces.continueEditSessionInCodespaceWithJupyterServer'
+        );
+    }
+}

src/standalone/context/activeEditorContext.ts

@@ -3,7 +3,7 @@
 
 import { inject, injectable, optional } from 'inversify';
 import { NotebookEditor, TextEditor } from 'vscode';
-import { IKernel, IKernelProvider } from '../../kernels/types';
+import { IKernel, IKernelProvider, isRemoteConnection } from '../../kernels/types';
 import { IExtensionSyncActivationService } from '../../platform/activation/types';
 import { ICommandManager, IDocumentManager, IVSCodeNotebook } from '../../platform/common/application/types';
 import { EditorContexts, PYTHON_LANGUAGE } from '../../platform/common/constants';
@@ -15,6 +15,7 @@ import { IInteractiveWindowProvider, IInteractiveWindow } from '../../interactiv
 import { getNotebookMetadata, isJupyterNotebook } from '../../platform/common/utils';
 import { isPythonNotebook } from '../../kernels/helpers';
 import { IControllerRegistration } from '../../notebooks/controllers/types';
+import { IJupyterServerUriStorage, IJupyterUriProviderRegistration } from '../../kernels/jupyter/types';
 
 /**
  * Tracks a lot of the context keys needed in the extension.
@@ -37,6 +38,7 @@ export class ActiveEditorContextService implements IExtensionSyncActivationServi
     private isPythonNotebook: ContextKey;
     private isJupyterKernelSelected: ContextKey;
     private hasNativeNotebookOrInteractiveWindowOpen: ContextKey;
+    private kernelSourceContext: ContextKey<string>;
     constructor(
         @inject(IInteractiveWindowProvider)
         @optional()
@@ -46,7 +48,10 @@ export class ActiveEditorContextService implements IExtensionSyncActivationServi
         @inject(IDisposableRegistry) disposables: IDisposableRegistry,
         @inject(IVSCodeNotebook) private readonly vscNotebook: IVSCodeNotebook,
         @inject(IKernelProvider) private readonly kernelProvider: IKernelProvider,
-        @inject(IControllerRegistration) private readonly controllers: IControllerRegistration
+        @inject(IControllerRegistration) private readonly controllers: IControllerRegistration,
+        @inject(IJupyterServerUriStorage) private readonly serverUriStorage: IJupyterServerUriStorage,
+        @inject(IJupyterUriProviderRegistration)
+        private readonly jupyterUriProviderRegistration: IJupyterUriProviderRegistration
     ) {
         disposables.push(this);
         this.nativeContext = new ContextKey(EditorContexts.IsNativeActive, this.commandManager);
@@ -87,6 +92,7 @@ export class ActiveEditorContextService implements IExtensionSyncActivationServi
             EditorContexts.HasNativeNotebookOrInteractiveWindowOpen,
             this.commandManager
         );
+        this.kernelSourceContext = new ContextKey(EditorContexts.KernelSource, this.commandManager);
     }
     public dispose() {
         this.disposables.forEach((item) => item.dispose());
@@ -181,8 +187,26 @@ export class ActiveEditorContextService implements IExtensionSyncActivationServi
             this.canRestartNotebookKernelContext.set(false).catch(noop);
             this.canInterruptNotebookKernelContext.set(false).catch(noop);
         }
+        this.updateKernelSourceContext(kernel).catch(noop);
         this.updateSelectedKernelContext();
     }
+    private async updateKernelSourceContext(kernel: IKernel | undefined) {
+        if (!kernel || !isRemoteConnection(kernel.kernelConnectionMetadata)) {
+            this.kernelSourceContext.set('').catch(noop);
+            return;
+        }
+
+        const connection = kernel.kernelConnectionMetadata;
+        const uriItem = await this.serverUriStorage.get(connection.serverId);
+        const provider = uriItem && (await this.jupyterUriProviderRegistration.getProvider(uriItem.provider.id));
+
+        if (!provider) {
+            this.kernelSourceContext.set('').catch(noop);
+            return;
+        }
+
+        this.kernelSourceContext.set(provider.id).catch(noop);
+    }
     private updateSelectedKernelContext() {
         const document =
             this.vscNotebook.activeNotebookEditor?.notebook ||

src/standalone/serviceRegistry.node.ts

@@ -28,6 +28,7 @@ import { registerTypes as registerIntellisenseTypes } from './intellisense/servi
 import { PythonExtensionRestartNotification } from './notification/pythonExtensionRestartNotification';
 import { UserJupyterServerUrlProvider } from './userJupyterServer/userServerUrlProvider';
 import { JupyterServerSelectorCommand } from './userJupyterServer/serverSelectorForTests';
+import { CommandRegistry as CodespaceCommandRegistry } from './codespace/commandRegistry';
 import { EagerlyActivateJupyterUriProviders } from './api/activateJupyterProviderExtensions';
 
 export function registerTypes(context: IExtensionContext, serviceManager: IServiceManager, isDevMode: boolean) {
@@ -65,6 +66,11 @@ export function registerTypes(context: IExtensionContext, serviceManager: IServi
         ExportCommandRegistry
     );
 
+    serviceManager.addSingleton<IExtensionSyncActivationService>(
+        IExtensionSyncActivationService,
+        CodespaceCommandRegistry
+    );
+
     serviceManager.addSingleton<ISurveyBanner>(ISurveyBanner, DataScienceSurveyBanner);
     serviceManager.addBinding(ISurveyBanner, IExtensionSyncActivationService);
     // Activation Manager

src/standalone/serviceRegistry.web.ts

@@ -19,6 +19,7 @@ import { PythonExtensionRestartNotification } from './notification/pythonExtensi
 import { ImportTracker } from './import-export/importTracker';
 import { UserJupyterServerUrlProvider } from './userJupyterServer/userServerUrlProvider';
 import { JupyterServerSelectorCommand } from './userJupyterServer/serverSelectorForTests';
+import { CommandRegistry as CodespaceCommandRegistry } from './codespace/commandRegistry';
 import { EagerlyActivateJupyterUriProviders } from './api/activateJupyterProviderExtensions';
 
 export function registerTypes(context: IExtensionContext, serviceManager: IServiceManager, isDevMode: boolean) {
@@ -71,4 +72,9 @@ export function registerTypes(context: IExtensionContext, serviceManager: IServi
         IExtensionSyncActivationService,
         UserJupyterServerUrlProvider
     );
+
+    serviceManager.addSingleton<IExtensionSyncActivationService>(
+        IExtensionSyncActivationService,
+        CodespaceCommandRegistry
+    );
 }