From f15ae6ad02376ac17c527ff6215bb15efd3945be Mon Sep 17 00:00:00 2001
From: Cristian Cepeda <43882+pastuxso@users.noreply.github.com>
Date: Thu, 19 Dec 2024 16:50:26 -0500
Subject: [PATCH] Fix features flags loading

---
 src/extension/extension.ts |  1 +
 src/extension/kernel.ts    | 63 ++++++++++++++++++--------------------
 2 files changed, 31 insertions(+), 33 deletions(-)

diff --git a/src/extension/extension.ts b/src/extension/extension.ts
index 87104bbb2..6e41c3d44 100644
--- a/src/extension/extension.ts
+++ b/src/extension/extension.ts
@@ -494,6 +494,7 @@ export class RunmeExtension {
         const session = await StatefulAuthProvider.instance.ensureSession()
         const nunmeIdentity = session ? RunmeIdentity.ALL : getServerLifecycleIdentity()
         await commands.executeCommand('runme.lifecycleIdentitySelection', nunmeIdentity)
+        kernel.updateFeatureContext('statefulAuth', !!session)
       } catch (error) {
         let message
         if (error instanceof Error) {
diff --git a/src/extension/kernel.ts b/src/extension/kernel.ts
index a9db4a819..9f8546451 100644
--- a/src/extension/kernel.ts
+++ b/src/extension/kernel.ts
@@ -116,7 +116,6 @@ import { CommandModeEnum } from './grpc/runner/types'
 import { GrpcReporter } from './reporter'
 import { EnvStoreMonitorWithSession } from './panels/notebook'
 import { SignedIn } from './signedIn'
-import { StatefulAuthProvider } from './provider/statefulAuth'
 
 enum ConfirmationItems {
   Yes = 'Yes',
@@ -203,44 +202,42 @@ export class Kernel implements Disposable {
     this.#onlySignedIn = new SignedIn(this)
     this.#disposables.push(this.#onlySignedIn)
 
-    StatefulAuthProvider.instance.currentSession().then((session) => {
-      const packageJSON = context?.extension?.packageJSON || {}
-      const featContext: FeatureContext = {
-        os: os.platform(),
-        vsCodeVersion: version as string,
-        extensionVersion: packageJSON?.version,
-        githubAuth: false,
-        statefulAuth: !!session,
-        extensionId: context?.extension?.id as ExtensionName,
-      }
+    const packageJSON = this.context?.extension?.packageJSON || {}
+    const featContext: FeatureContext = {
+      os: os.platform(),
+      vsCodeVersion: version as string,
+      extensionVersion: packageJSON?.version,
+      githubAuth: false,
+      statefulAuth: false,
+      extensionId: this.context?.extension?.id as ExtensionName,
+    }
 
-      const runmeFeatureSettings = workspace.getConfiguration('runme.features')
-      const featureNames = Object.keys(FeatureName)
+    const runmeFeatureSettings = workspace.getConfiguration('runme.features')
+    const featureNames = Object.keys(FeatureName)
 
-      featureNames.forEach((feature) => {
-        if (runmeFeatureSettings.has(feature)) {
-          const result = runmeFeatureSettings.get<boolean>(feature, false)
-          this.#featuresSettings.set(feature, result)
-        }
-      })
+    featureNames.forEach((feature) => {
+      if (runmeFeatureSettings.has(feature)) {
+        const result = runmeFeatureSettings.get<boolean>(feature, false)
+        this.#featuresSettings.set(feature, result)
+      }
+    })
 
-      this.featuresState$ = features.loadState(packageJSON, featContext, this.#featuresSettings)
+    this.featuresState$ = features.loadState(packageJSON, featContext, this.#featuresSettings)
 
-      if (this.featuresState$) {
-        const subscription = this.featuresState$
-          .pipe(map((_state) => features.getSnapshot(this.featuresState$)))
-          .subscribe((snapshot) => {
-            ContextState.addKey(FEATURES_CONTEXT_STATE_KEY, snapshot)
-            postClientMessage(this.messaging, ClientMessages.featuresUpdateAction, {
-              snapshot: snapshot,
-            })
+    if (this.featuresState$) {
+      const subscription = this.featuresState$
+        .pipe(map((_state) => features.getSnapshot(this.featuresState$)))
+        .subscribe((snapshot) => {
+          ContextState.addKey(FEATURES_CONTEXT_STATE_KEY, snapshot)
+          postClientMessage(this.messaging, ClientMessages.featuresUpdateAction, {
+            snapshot: snapshot,
           })
-
-        this.#disposables.push({
-          dispose: () => subscription.unsubscribe(),
         })
-      }
-    })
+
+      this.#disposables.push({
+        dispose: () => subscription.unsubscribe(),
+      })
+    }
   }
 
   get envProps() {