From fa76d001a5dd516fb7d981a614477a9f36424e8a Mon Sep 17 00:00:00 2001
From: Anthony Kim <anthonykim@microsoft.com>
Date: Thu, 9 May 2024 13:35:41 -0700
Subject: [PATCH] wrap around experiment, context key

---
 package.json                            |  2 +-
 src/client/common/experiments/groups.ts |  5 +++++
 src/client/extensionActivation.ts       | 17 +++++++++++++++--
 3 files changed, 21 insertions(+), 3 deletions(-)

diff --git a/package.json b/package.json
index 100756802e90..f27f1dc24d00 100644
--- a/package.json
+++ b/package.json
@@ -1439,7 +1439,7 @@
                 {
                     "command": "python.execInREPL",
                     "group": "Python",
-                    "when": "editorFocus && editorLangId == python && !virtualWorkspace && shellExecutionSupported"
+                    "when": "editorFocus && editorLangId == python && !virtualWorkspace && shellExecutionSupported && pythonRunREPL"
                 }
             ],
             "editor/title": [
diff --git a/src/client/common/experiments/groups.ts b/src/client/common/experiments/groups.ts
index d43f376ddc87..3402605cd1ca 100644
--- a/src/client/common/experiments/groups.ts
+++ b/src/client/common/experiments/groups.ts
@@ -24,3 +24,8 @@ export enum EnableTestAdapterRewrite {
 export enum RecommendTensobardExtension {
     experiment = 'pythonRecommendTensorboardExt',
 }
+
+// Experiment to enable running Python REPL using IW.
+export enum EnableRunREPL {
+    experiment = 'pythonRunREPL',
+}
diff --git a/src/client/extensionActivation.ts b/src/client/extensionActivation.ts
index 359a5fea9719..1fed0ffc802b 100644
--- a/src/client/extensionActivation.ts
+++ b/src/client/extensionActivation.ts
@@ -3,7 +3,7 @@
 
 'use strict';
 
-import { DebugConfigurationProvider, debug, languages, window } from 'vscode';
+import { DebugConfigurationProvider, debug, languages, window, commands } from 'vscode';
 
 import { registerTypes as activationRegisterTypes } from './activation/serviceRegistry';
 import { IExtensionActivationManager } from './activation/types';
@@ -16,6 +16,7 @@ import { IFileSystem } from './common/platform/types';
 import {
     IConfigurationService,
     IDisposableRegistry,
+    IExperimentService,
     IExtensions,
     IInterpreterPathService,
     ILogOutputChannel,
@@ -54,6 +55,7 @@ import { logAndNotifyOnLegacySettings } from './logging/settingLogs';
 import { DebuggerTypeName } from './debugger/constants';
 import { StopWatch } from './common/utils/stopWatch';
 import { registerReplCommands } from './repl/replCommands';
+import { EnableRunREPL } from './common/experiments/groups';
 
 export async function activateComponents(
     // `ext` is passed to any extra activation funcs.
@@ -115,7 +117,18 @@ export function activateFeatures(ext: ExtensionState, _components: Components):
     //     registerReplCommands(ext.disposables, interpreterService);
     // }
     // uncomment
-    registerReplCommands(ext.disposables, interpreterService);
+
+    // Only register if they are in experiment for pythonRunREPL.
+    const experimentService = ext.legacyIOC.serviceContainer.get<IExperimentService>(IExperimentService);
+    commands.executeCommand('setContext', 'pythonRunREPL', false);
+    if (experimentService) {
+        const replExperimentValue = experimentService.inExperimentSync(EnableRunREPL.experiment);
+        if (replExperimentValue) {
+            registerReplCommands(ext.disposables, interpreterService);
+            commands.executeCommand('setContext', 'pythonRunREPL', true);
+        }
+    }
+    // registerReplCommands(ext.disposables, interpreterService); // Register regardless
 }
 
 /// //////////////////////////