From 380f4c52d9a921cefaffff52676a650001c4bc99 Mon Sep 17 00:00:00 2001
From: Don Jayamanne <don.jayamanne@outlook.com>
Date: Mon, 20 Jan 2025 11:53:42 +1100
Subject: [PATCH] Revert "Updated with a few changes"

This reverts commit 797958da16779c8624a85a53392258d1cf96cfee.
---
 src/standalone/api/kernels/backgroundExecution.ts |  2 +-
 src/standalone/api/kernels/index.ts               | 15 ++++++++++-----
 src/standalone/api/kernels/kernel.ts              |  2 +-
 src/standalone/api/kernels/kernel.unit.test.ts    |  2 +-
 4 files changed, 13 insertions(+), 8 deletions(-)

diff --git a/src/standalone/api/kernels/backgroundExecution.ts b/src/standalone/api/kernels/backgroundExecution.ts
index 04e8659e5bb..5732de0b955 100644
--- a/src/standalone/api/kernels/backgroundExecution.ts
+++ b/src/standalone/api/kernels/backgroundExecution.ts
@@ -19,7 +19,7 @@ export async function execCodeInBackgroundThread<T>(
 ) {
     const counter = executionCounters.get(kernel) || 0;
     executionCounters.set(kernel, counter + 1);
-    const api = createKernelApiForExtension(JVSC_EXTENSION_ID, kernel);
+    const { api } = createKernelApiForExtension(JVSC_EXTENSION_ID, kernel);
     const mime = `application/vnd.vscode.bg.execution.${counter}`;
     const mimeFinalResult = `application/vnd.vscode.bg.execution.${counter}.result`;
     const mimeErrorResult = `application/vnd.vscode.bg.execution.${counter}.error`;
diff --git a/src/standalone/api/kernels/index.ts b/src/standalone/api/kernels/index.ts
index 549d7058a3d..47bc4f8bcf1 100644
--- a/src/standalone/api/kernels/index.ts
+++ b/src/standalone/api/kernels/index.ts
@@ -14,6 +14,7 @@ import { createDeferredFromPromise } from '../../../platform/common/utils/async'
 import { logger } from '../../../platform/logging';
 import { StopWatch } from '../../../platform/common/utils/stopWatch';
 import { sendKernelTelemetryEvent } from '../../../kernels/telemetry/sendKernelTelemetryEvent';
+import { KernelExecutionProgressIndicator } from './kernelProgressIndicator';
 
 const extensionAPICache = new Map<
     string,
@@ -26,8 +27,9 @@ const extensionAPICache = new Map<
                   waitUntil(thenable: Thenable<unknown>): void;
               }>
             | undefined;
-        // Kernel cache needs to be scoped per extension
-        kernels: WeakMap<IKernel, Kernel>;
+        // Kernel cache needs to be scoped per extension to make sure that the progress messages
+        // show accurately which extension is actually using it.
+        kernels: WeakMap<IKernel, { api: Kernel; progress: KernelExecutionProgressIndicator }>;
     }
 >();
 
@@ -35,7 +37,7 @@ function getOrCreateExtensionAPI(extensionId: string) {
     if (!extensionAPICache.has(extensionId)) {
         extensionAPICache.set(extensionId, {
             onDidStart: undefined,
-            kernels: new WeakMap<IKernel, Kernel>()
+            kernels: new WeakMap<IKernel, { api: Kernel; progress: KernelExecutionProgressIndicator }>()
         });
     }
     return extensionAPICache.get(extensionId)!;
@@ -72,7 +74,8 @@ export function getKernelsApi(extensionId: string): Kernels {
                     kernel.kernelConnectionMetadata
                 );
             }
-            return getWrappedKernel(kernel, extensionId);
+            const { api } = getWrappedKernel(kernel, extensionId);
+            return api;
         },
         get onDidStart() {
             if (![JVSC_EXTENSION_ID, DATA_WRANGLER_EXTENSION_ID].includes(extensionId)) {
@@ -94,7 +97,7 @@ export function getKernelsApi(extensionId: string): Kernels {
                 disposableRegistry.push(
                     extensionAPI.onDidStart,
                     kernelProvider.onDidPostInitializeKernel(({ kernel, token, waitUntil }) => {
-                        const api = getWrappedKernel(kernel, extensionId);
+                        const { api, progress } = getWrappedKernel(kernel, extensionId);
                         extensionAPI.onDidStart?.fire({
                             uri: kernel.uri,
                             kernel: api,
@@ -106,6 +109,7 @@ export function getKernelsApi(extensionId: string): Kernels {
                                 // it is effectively preventing access to it.
                                 const deferrable = createDeferredFromPromise(Promise.resolve(thenable));
                                 waitUntil(thenable);
+                                const disposable = progress.show();
                                 const stopWatch = new StopWatch();
                                 void deferrable.promise.finally(() => {
                                     logger.trace(
@@ -117,6 +121,7 @@ export function getKernelsApi(extensionId: string): Kernels {
                                         { duration: stopWatch.elapsedTime },
                                         { extensionId }
                                     );
+                                    disposable.dispose();
                                 });
                             }
                         });
diff --git a/src/standalone/api/kernels/kernel.ts b/src/standalone/api/kernels/kernel.ts
index aa373d4a85b..2e74d0dbf02 100644
--- a/src/standalone/api/kernels/kernel.ts
+++ b/src/standalone/api/kernels/kernel.ts
@@ -144,7 +144,7 @@ class WrappedKernelPerExtension extends DisposableBase implements Kernel {
     ) {
         const wrapper = new WrappedKernelPerExtension(extensionId, kernel, execution, controller);
         ServiceContainer.instance.get<IDisposableRegistry>(IDisposableRegistry).push(wrapper);
-        return wrapper._api;
+        return { api: wrapper._api, progress: wrapper.progress };
     }
 
     private async checkAccess() {
diff --git a/src/standalone/api/kernels/kernel.unit.test.ts b/src/standalone/api/kernels/kernel.unit.test.ts
index f44eb100d72..95be3ea2487 100644
--- a/src/standalone/api/kernels/kernel.unit.test.ts
+++ b/src/standalone/api/kernels/kernel.unit.test.ts
@@ -101,7 +101,7 @@ suite('Kernel Api', () => {
 
     test('Verify Access Denied error message has expected value for the property `name`', async () => {
         try {
-            const api = createKernelApiForExtension('xyz', instance(kernel));
+            const { api } = createKernelApiForExtension('xyz', instance(kernel));
             for await (const x of api.executeCode('bogus', token)) {
                 assert.fail(`Should have failed without producing any value such as ${x}`);
             }