Refactor StatefulAuthProvider for static initialization and improved session management

.github/scripts/overwrites/stateful.sh

@@ -5,7 +5,7 @@ npm pkg set name="$EXTENSION_NAME"
 npm pkg set displayName="Stateful Notebooks for DevOps"
 npm pkg set description="DevOps Notebooks built on Runme, connected for collaboration."
 npm pkg set homepage="https://stateful.com"
-npm pkg set contributes.configuration[0].properties[runme.app.baseDomain].default="platform.stateful.com"
+npm pkg set contributes.configuration[0].properties[runme.app.baseDomain].default="cloud.stateful.com"
 npm pkg set contributes.configuration[0].properties[runme.app.platformAuth].default=true --json
 npm pkg set contributes.configuration[0].properties[runme.server.lifecycleIdentity].default=1 --json
 npm pkg set contributes.configuration[0].properties[runme.app.notebookAutoSave].default="yes"

README-platform.md

@@ -45,7 +45,7 @@ The Stateful cloud offers a suite of tools and services designed to enhance your
 
 ## Getting Started
 
-1. **Sign Up**: Create your free account at [Stateful](https://platform.stateful.com/).
+1. **Sign Up**: Create your free account at [Stateful](https://cloud.stateful.com/).
 2. **Install VS Code Extension**: Download the Stateful extension from the [VS Code Marketplace](https://marketplace.visualstudio.com/items?itemName=stateful.platform). The extension is fully compatible with Runme, but adds authentication, collaboration, and security features, making it secure for teams and inside companies.
 3. **Explore**: Start creating, running, sharing, and discussing your first DevOps Notebook and its commands using the Stateful cloud.
 

package.json

@@ -992,7 +992,7 @@
           },
           "runme.app.baseDomain": {
             "type": "string",
-            "default": "platform.stateful.com",
+            "default": "cloud.stateful.com",
             "scope": "window",
             "markdownDescription": "Base domain to be use for Runme app"
           },

src/extension/api/client.ts

@@ -5,21 +5,26 @@ import { Uri } from 'vscode'
 
 import { getRunmeAppUrl } from '../../utils/configuration'
 import { getFeaturesContext } from '../features'
+import { StatefulAuthProvider } from '../provider/statefulAuth'
 
-export function InitializeClient({
-  uri,
-  runmeToken,
-}: {
-  uri?: string | undefined
-  runmeToken: string
-}) {
+export async function InitializeCloudClient(uri?: string) {
+  const session = await StatefulAuthProvider.instance.currentSession()
+
+  if (!session) {
+    throw new Error('You must authenticate with your Stateful account')
+  }
+
+  return InitializeClient({ uri, token: session.accessToken })
+}
+
+function InitializeClient({ uri, token }: { uri?: string | undefined; token: string }) {
   const authLink = setContext((_, { headers }) => {
     const context = getFeaturesContext()
     return {
       headers: {
         ...headers,
         'Auth-Provider': 'platform',
-        authorization: runmeToken ? `Bearer ${runmeToken}` : '',
+        authorization: token ? `Bearer ${token}` : '',
         'X-Extension-Id': context?.extensionId,
         'X-Extension-Os': context?.os,
         'X-Extension-Version': context?.extensionVersion,

src/extension/authSessionChangeHandler.ts

@@ -0,0 +1,74 @@
+import { Subject, Subscription } from 'rxjs'
+import { debounceTime, distinctUntilChanged } from 'rxjs/operators'
+import { authentication, AuthenticationSessionsChangeEvent, Disposable } from 'vscode'
+
+export default class AuthSessionChangeHandler implements Disposable {
+  static #instance: AuthSessionChangeHandler | null = null
+
+  #disposables: Disposable[] = []
+  #eventSubject: Subject<AuthenticationSessionsChangeEvent>
+  #subscriptions: Subscription[] = []
+  #listeners: ((event: AuthenticationSessionsChangeEvent) => void)[] = []
+
+  private constructor(private debounceTimeMs: number = 100) {
+    this.#eventSubject = new Subject<AuthenticationSessionsChangeEvent>()
+    this.#subscriptions.push(
+      this.#eventSubject
+        .pipe(distinctUntilChanged(this.eventComparer), debounceTime(this.debounceTimeMs))
+        .subscribe((event) => {
+          this.notifyListeners(event)
+        }),
+    )
+
+    this.#disposables.push(
+      authentication.onDidChangeSessions((e) => {
+        this.#eventSubject.next(e)
+      }),
+    )
+  }
+
+  public static get instance(): AuthSessionChangeHandler {
+    if (!this.#instance) {
+      this.#instance = new AuthSessionChangeHandler()
+    }
+
+    return this.#instance
+  }
+
+  public addListener(listener: (event: AuthenticationSessionsChangeEvent) => void): void {
+    this.#listeners.push(listener)
+  }
+
+  public removeListener(listener: (event: AuthenticationSessionsChangeEvent) => void): void {
+    this.#listeners = this.#listeners.filter((l) => l !== listener)
+  }
+
+  private notifyListeners(event: AuthenticationSessionsChangeEvent): void {
+    for (const listener of this.#listeners) {
+      try {
+        listener(event)
+      } catch (err) {
+        console.error('Error in listener:', err)
+      }
+    }
+  }
+
+  private eventComparer(
+    previous: AuthenticationSessionsChangeEvent,
+    current: AuthenticationSessionsChangeEvent,
+  ): boolean {
+    return (
+      previous.provider.id === current.provider.id &&
+      JSON.stringify(previous) === JSON.stringify(current)
+    )
+  }
+
+  public async dispose() {
+    this.#disposables.forEach((d) => d.dispose())
+    this.#subscriptions = []
+    this.#eventSubject.complete()
+    this.#listeners = []
+
+    AuthSessionChangeHandler.#instance = null
+  }
+}

src/extension/commands/index.ts

@@ -32,7 +32,6 @@ import { Kernel } from '../kernel'
 import {
   getAnnotations,
   getNotebookCategories,
-  getPlatformAuthSession,
   getTerminalByCell,
   openFileAsRunmeNotebook,
   promptUserSession,
@@ -52,7 +51,7 @@ import {
 } from '../../constants'
 import ContextState from '../contextState'
 import { createGist } from '../services/github/gist'
-import { InitializeClient } from '../api/client'
+import { InitializeCloudClient } from '../api/client'
 import { GetUserEnvironmentsDocument } from '../__generated-platform__/graphql'
 import { EnvironmentManager } from '../environment/manager'
 import features from '../features'
@@ -562,8 +561,7 @@ export async function createCellGistCommand(cell: NotebookCell, context: Extensi
 }
 
 export async function selectEnvironment(manager: EnvironmentManager) {
-  const session = await getPlatformAuthSession()
-  const graphClient = InitializeClient({ runmeToken: session?.accessToken! })
+  const graphClient = await InitializeCloudClient()
 
   const result = await graphClient.query({
     query: GetUserEnvironmentsDocument,

src/extension/extension.ts

@@ -9,7 +9,6 @@ import {
   env,
   Uri,
   NotebookCell,
-  authentication,
 } from 'vscode'
 import { TelemetryReporter } from 'vscode-telemetry'
 import Channel from 'tangle/webviews'
@@ -41,8 +40,6 @@ import {
   getDefaultWorkspace,
   bootFile,
   resetNotebookSettings,
-  getPlatformAuthSession,
-  getGithubAuthSession,
   openFileAsRunmeNotebook,
 } from './utils'
 import { RunmeTaskProvider } from './provider/runmeTask'
@@ -95,6 +92,7 @@ import { EnvironmentManager } from './environment/manager'
 import ContextState from './contextState'
 import { RunmeIdentity } from './grpc/serializerTypes'
 import * as features from './features'
+import AuthSessionChangeHandler from './authSessionChangeHandler'
 
 export class RunmeExtension {
   protected serializer?: SerializerBase
@@ -218,7 +216,6 @@ export class RunmeExtension {
     // extension is deactivated.
     context.subscriptions.push(aiManager)
 
-    const uriHandler = new RunmeUriHandler(context, kernel, getForceNewWindowConfig())
     const winCodeLensRunSurvey = new survey.SurveyWinCodeLensRun(context)
     const surveys: Disposable[] = [
       winCodeLensRunSurvey,
@@ -260,6 +257,8 @@ export class RunmeExtension {
       serializer,
       server,
       treeViewer,
+      StatefulAuthProvider.instance,
+      AuthSessionChangeHandler.instance,
       ...this.registerPanels(kernel, context),
       ...surveys,
       workspace.registerNotebookSerializer(Kernel.type, serializer, {
@@ -338,7 +337,7 @@ export class RunmeExtension {
       /**
        * Uri handler
        */
-      window.registerUriHandler(uriHandler),
+      window.registerUriHandler(new RunmeUriHandler(context, kernel, getForceNewWindowConfig())),
 
       /**
        * Runme Message Display commands
@@ -398,7 +397,7 @@ export class RunmeExtension {
         commands.executeCommand('runme.lifecycleIdentitySelection', RunmeIdentity.CELL),
       ),
 
-      RunmeExtension.registerCommand(
+      commands.registerCommand(
         'runme.lifecycleIdentitySelection',
         async (identity?: RunmeIdentity) => {
           if (identity === undefined) {
@@ -412,6 +411,10 @@ export class RunmeExtension {
             return
           }
 
+          TelemetryReporter.sendTelemetryEvent('extension.command', {
+            command: 'runme.lifecycleIdentitySelection',
+          })
+
           await ContextState.addKey(NOTEBOOK_LIFECYCLE_ID, identity)
 
           await Promise.all(
@@ -486,71 +489,45 @@ export class RunmeExtension {
     }
 
     if (kernel.isFeatureOn(FeatureName.RequireStatefulAuth)) {
-      const statefulAuthProvider = new StatefulAuthProvider(context, uriHandler)
-      context.subscriptions.push(statefulAuthProvider)
-
-      const session = await getPlatformAuthSession(false, true)
-      let sessionFromToken = false
-      if (!session) {
-        sessionFromToken = await statefulAuthProvider.bootstrapFromToken()
-      }
-
-      const forceLogin = kernel.isFeatureOn(FeatureName.ForceLogin) || sessionFromToken
-      const silent = forceLogin ? undefined : true
-
-      getPlatformAuthSession(forceLogin, silent)
-        .then((session) => {
-          if (session) {
-            statefulAuthProvider.showLoginNotification()
-          }
-        })
-        .catch((error) => {
-          let message
-          if (error instanceof Error) {
-            message = error.message
-          } else {
-            message = JSON.stringify(error)
-          }
-
-          // https://github.com/microsoft/vscode/blob/main/src/vs/workbench/api/browser/mainThreadAuthentication.ts#L238
-          // throw new Error('User did not consent to login.')
-          // Calling again to ensure User Menu Badge
-          if (forceLogin && message === 'User did not consent to login.') {
-            getPlatformAuthSession(false)
-          }
-        })
+      await StatefulAuthProvider.instance.ensureSession()
+      StatefulAuthProvider.instance.currentSession().then(async (session) => {
+        if (session) {
+          await commands.executeCommand('runme.lifecycleIdentitySelection', RunmeIdentity.ALL)
+        } else {
+          const settingsDefault = getServerLifecycleIdentity()
+          await commands.executeCommand('runme.lifecycleIdentitySelection', settingsDefault)
+        }
+      })
     }
 
     if (kernel.isFeatureOn(FeatureName.Gist)) {
-      context.subscriptions.push(new GithubAuthProvider(context))
-      getGithubAuthSession(false).then((session) => {
-        kernel.updateFeatureContext('githubAuth', !!session)
-      })
+      context.subscriptions.push(new GithubAuthProvider(context, kernel))
     }
 
-    authentication.onDidChangeSessions((e) => {
+    AuthSessionChangeHandler.instance.addListener((e) => {
       if (
+        StatefulAuthProvider.instance &&
         kernel.isFeatureOn(FeatureName.RequireStatefulAuth) &&
         e.provider.id === AuthenticationProviders.Stateful
       ) {
-        getPlatformAuthSession(false, true).then(async (session) => {
-          if (!!session) {
+        StatefulAuthProvider.instance.currentSession().then(async (session) => {
+          if (session) {
             await commands.executeCommand('runme.lifecycleIdentitySelection', RunmeIdentity.ALL)
+            kernel.emitPanelEvent('runme.cloud', 'onCommand', {
+              name: 'signIn',
+              panelId: 'runme.cloud',
+            })
           } else {
             const settingsDefault = getServerLifecycleIdentity()
             await commands.executeCommand('runme.lifecycleIdentitySelection', settingsDefault)
+            kernel.emitPanelEvent('runme.cloud', 'onCommand', {
+              name: 'signOut',
+              panelId: 'runme.cloud',
+            })
           }
           kernel.updateFeatureContext('statefulAuth', !!session)
         })
       }
-      if (
-        kernel.isFeatureOn(FeatureName.Gist) &&
-        e.provider.id === AuthenticationProviders.GitHub
-      ) {
-        getGithubAuthSession(false).then((session) => {
-          kernel.updateFeatureContext('githubAuth', !!session)
-        })
-      }
     })
 
     // only ever enabled in hosted playground

src/extension/handler/uri.ts

@@ -13,7 +13,6 @@ import {
   TaskScope,
   ShellExecution,
   tasks,
-  EventEmitter,
   Disposable,
 } from 'vscode'
 import got from 'got'
@@ -23,6 +22,7 @@ import { TelemetryReporter } from 'vscode-telemetry'
 import getLogger from '../logger'
 import { Kernel } from '../kernel'
 import { AuthenticationProviders } from '../../constants'
+import { StatefulAuthProvider } from '../provider/statefulAuth'
 
 import {
   getProjectDir,
@@ -45,8 +45,6 @@ const extensionNames: { [key: string]: string } = {
 
 export class RunmeUriHandler implements UriHandler, Disposable {
   #disposables: Disposable[] = []
-  readonly #onAuth = this.register(new EventEmitter<Uri>())
-  readonly onAuthEvent = this.#onAuth.event
 
   constructor(
     private context: ExtensionContext,
@@ -70,7 +68,7 @@ export class RunmeUriHandler implements UriHandler, Disposable {
         command,
         type: AuthenticationProviders.Stateful,
       })
-      this.#onAuth.fire(uri)
+      StatefulAuthProvider.instance.fireOnAuthEvent(uri)
       return
     } else if (command === 'setup') {
       const { fileToOpen, repository } = parseParams(params)

src/extension/index.ts

@@ -4,6 +4,7 @@ import { TelemetryReporter } from 'vscode-telemetry'
 import { RunmeExtension } from './extension'
 import getLogger from './logger'
 import { isTelemetryEnabled } from './utils'
+import { StatefulAuthProvider } from './provider/statefulAuth'
 
 declare const CONNECTION_STR: string
 
@@ -33,6 +34,7 @@ export async function activate(context: ExtensionContext) {
 
   log.info('Activating Extension')
   try {
+    StatefulAuthProvider.initialize(context)
     await ext.initialize(context)
     log.info('Extension successfully activated')
   } catch (err: any) {