Repository: stateful/vscode-runme
Pull request #1660 (2024-09-13 23:57:48)

Pass languageId into resolve program API
This PR is pending releasing the latest Runme kernel (binary). While the extension implementation is somewhat agnostic, the CLI interface is coupled with the runme binary bundled inside the extension.

sonarcloud[bot]: ## [![Quality Gate Passed](https://sonarsource.github.io/sonarcloud-github-static-resources/v2/checks/QualityGateBadge/qg-passed-20px.png 'Quality Gate Passed')](https://sonarcloud.io/dashboard?id=stateful_vscode-runme&pullRequest=1660) **Quality Gate passed**  
Issues  
![](https://sonarsource.github.io/sonarcloud-github-static-resources/v2/common/passed-16px.png '') [0 New issues](https://sonarcloud.io/project/issues?id=stateful_vscode-runme&pullRequest=1660&issueStatuses=OPEN,CONFIRMED&sinceLeakPeriod=true)  
![](https://sonarsource.github.io/sonarcloud-github-static-resources/v2/common/accepted-16px.png '') [0 Accepted issues](https://sonarcloud.io/project/issues?id=stateful_vscode-runme&pullRequest=1660&issueStatuses=ACCEPTED)

Measures  
![](https://sonarsource.github.io/sonarcloud-github-static-resources/v2/common/passed-16px.png '') [0 Security Hotspots](https://sonarcloud.io/project/security_hotspots?id=stateful_vscode-runme&pullRequest=1660&issueStatuses=OPEN,CONFIRMED&sinceLeakPeriod=true)  
![](https://sonarsource.github.io/sonarcloud-github-static-resources/v2/common/passed-16px.png '') [50.0% Coverage on New Code](https://sonarcloud.io/component_measures?id=stateful_vscode-runme&pullRequest=1660&metric=new_coverage&view=list)  
![](https://sonarsource.github.io/sonarcloud-github-static-resources/v2/common/passed-16px.png '') [0.0% Duplication on New Code](https://sonarcloud.io/component_measures?id=stateful_vscode-runme&pullRequest=1660&metric=new_duplicated_lines_density&view=list)  
  
[See analysis details on SonarCloud](https://sonarcloud.io/dashboard?id=stateful_vscode-runme&pullRequest=1660)



Patch:
diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml
index c1c096b4..d8da3aba 100644
--- a/.github/workflows/release.yml
+++ b/.github/workflows/release.yml
@@ -411,7 +411,7 @@ jobs:
         if: env.EXTENSION_NAME == 'runme'
         uses: ncipollo/release-action@v1
         with:
-          artifacts: "./runme-*"
+          artifacts: "./*.vsix"
           bodyFile: ${{ github.workspace }}-CHANGELOG.txt
           tag: ${{ env.GIT_TAG }}
           prerelease: ${{ github.event.inputs.releaseChannel == 'edge' }}
diff --git a/package-lock.json b/package-lock.json
index 7f688678..3bb7e82c 100644
--- a/package-lock.json
+++ b/package-lock.json
@@ -1,12 +1,12 @@
 {
   "name": "runme",
-  "version": "3.7.6-edge.0",
+  "version": "3.7.6-edge.2",
   "lockfileVersion": 3,
   "requires": true,
   "packages": {
     "": {
       "name": "runme",
-      "version": "3.7.6-edge.0",
+      "version": "3.7.6-edge.2",
       "hasInstallScript": true,
       "license": "Apache-2.0",
       "dependencies": {
@@ -14,10 +14,10 @@
         "@aws-sdk/client-ec2": "^3.649.0",
         "@aws-sdk/client-eks": "^3.649.0",
         "@aws-sdk/credential-providers": "^3.649.0",
-        "@buf/grpc_grpc.community_timostamm-protobuf-ts": "^2.9.4-20240809200651-8507e5a24938.4",
+        "@buf/grpc_grpc.community_timostamm-protobuf-ts": "^2.9.4-20240904201038-e0425cfebb28.4",
         "@buf/jlewi_foyle.bufbuild_es": "^1.10.0-00000000000000-d14934cb2733.1",
         "@buf/jlewi_foyle.connectrpc_es": "^1.4.0-00000000000000-d14934cb2733.3",
-        "@buf/stateful_runme.community_timostamm-protobuf-ts": "^2.9.4-20240903224300-611a892c2067.4",
+        "@buf/stateful_runme.community_timostamm-protobuf-ts": "^2.9.4-20240913234806-45813f39881a.4",
         "@connectrpc/connect": "^1.4.0",
         "@connectrpc/connect-node": "^1.5.0",
         "@google-cloud/compute": "^4.1.0",
@@ -2383,8 +2383,8 @@
       }
     },
     "node_modules/@buf/googleapis_googleapis.community_timostamm-protobuf-ts": {
-      "version": "2.9.4-20240729201550-8bc2c51e08c4.4",
-      "resolved": "https://buf.build/gen/npm/v1/@buf/googleapis_googleapis.community_timostamm-protobuf-ts/-/googleapis_googleapis.community_timostamm-protobuf-ts-2.9.4-20240729201550-8bc2c51e08c4.4.tgz",
+      "version": "2.9.4-20240827201746-e7f8d366f526.4",
+      "resolved": "https://buf.build/gen/npm/v1/@buf/googleapis_googleapis.community_timostamm-protobuf-ts/-/googleapis_googleapis.community_timostamm-protobuf-ts-2.9.4-20240827201746-e7f8d366f526.4.tgz",
       "peerDependencies": {
         "@protobuf-ts/runtime": "^2.9.4",
         "@protobuf-ts/runtime-rpc": "^2.9.4"
@@ -2408,10 +2408,10 @@
       }
     },
     "node_modules/@buf/grpc_grpc.community_timostamm-protobuf-ts": {
-      "version": "2.9.4-20240809200651-8507e5a24938.4",
-      "resolved": "https://buf.build/gen/npm/v1/@buf/grpc_grpc.community_timostamm-protobuf-ts/-/grpc_grpc.community_timostamm-protobuf-ts-2.9.4-20240809200651-8507e5a24938.4.tgz",
+      "version": "2.9.4-20240904201038-e0425cfebb28.4",
+      "resolved": "https://buf.build/gen/npm/v1/@buf/grpc_grpc.community_timostamm-protobuf-ts/-/grpc_grpc.community_timostamm-protobuf-ts-2.9.4-20240904201038-e0425cfebb28.4.tgz",
       "dependencies": {
-        "@buf/googleapis_googleapis.community_timostamm-protobuf-ts": "2.9.4-20240729201550-8bc2c51e08c4.4"
+        "@buf/googleapis_googleapis.community_timostamm-protobuf-ts": "2.9.4-20240827201746-e7f8d366f526.4"
       },
       "peerDependencies": {
         "@protobuf-ts/runtime": "^2.9.4",
@@ -2490,8 +2490,8 @@
       }
     },
     "node_modules/@buf/stateful_runme.community_timostamm-protobuf-ts": {
-      "version": "2.9.4-20240903224300-611a892c2067.4",
-      "resolved": "https://buf.build/gen/npm/v1/@buf/stateful_runme.community_timostamm-protobuf-ts/-/stateful_runme.community_timostamm-protobuf-ts-2.9.4-20240903224300-611a892c2067.4.tgz",
+      "version": "2.9.4-20240913234806-45813f39881a.4",
+      "resolved": "https://buf.build/gen/npm/v1/@buf/stateful_runme.community_timostamm-protobuf-ts/-/stateful_runme.community_timostamm-protobuf-ts-2.9.4-20240913234806-45813f39881a.4.tgz",
       "dependencies": {
         "@buf/bufbuild_protovalidate.community_timostamm-protobuf-ts": "2.9.4-20240212200630-3014d81c3a48.4"
       },
diff --git a/package.json b/package.json
index c9d5461b..9f2f8c6b 100644
--- a/package.json
+++ b/package.json
@@ -1,6 +1,6 @@
 {
   "name": "runme",
-  "version": "3.7.6-edge.0",
+  "version": "3.7.6-edge.2",
   "publisher": "stateful",
   "displayName": "Runme Notebooks for DevOps",
   "description": "Notebooks for DevOps written in Markdown. Runnable, configurable commands/code and embedded terminals. Integrated with cloud consoles, containers and pipelines.",
@@ -58,11 +58,28 @@
   "main": "./out/extension.js",
   "runme": {
     "features": {
+      "RequireStatefulAuth": {
+        "enabled": true,
+        "conditions": {
+          "enabledForExtensions": {
+            "stateful.platform": true,
+            "stateful.runme": false
+          }
+        }
+      },
+      "SignedIn": {
+        "enabled": true,
+        "conditions": {
+          "statefulAuthRequired": true,
+          "enabledForExtensions": {
+            "stateful.platform": true,
+            "stateful.runme": true
+          }
+        }
+      },
       "ForceLogin": {
         "enabled": true,
         "conditions": {
-          "os": "All",
-          "vsCodeVersion": ">=1.58.0",
           "enabledForExtensions": {
             "stateful.platform": true,
             "stateful.runme": false
@@ -72,8 +89,6 @@
       "Escalate": {
         "enabled": true,
         "conditions": {
-          "os": "All",
-          "vsCodeVersion": ">=1.58.0",
           "statefulAuthRequired": true,
           "enabledForExtensions": {
             "stateful.platform": true,
@@ -84,19 +99,15 @@
       "Share": {
         "enabled": true,
         "conditions": {
-          "os": "All",
-          "vsCodeVersion": ">=1.58.0",
           "enabledForExtensions": {
             "stateful.platform": true,
-            "stateful.runme": true
+            "stateful.runme": false
           }
         }
       },
       "Gist": {
         "enabled": true,
         "conditions": {
-          "os": "All",
-          "vsCodeVersion": ">=1.58.0",
           "enabledForExtensions": {
             "stateful.platform": false,
             "stateful.runme": true
@@ -599,12 +610,6 @@
             "scope": "window",
             "markdownDescription": "Relative path to a file within the workspace to open within a Runme notebook view, when VS Code starts (e.g. `./CONTRIBUTING.md`)."
           },
-          "runme.experiments.aiLogs": {
-            "type": "boolean",
-            "scope": "window",
-            "default": false,
-            "markdownDescription": "If set to `true`, enable logging server logs used for training AIs."
-          },
           "runme.experiments.aiAutoCell": {
             "type": "boolean",
             "scope": "window",
@@ -1209,10 +1214,10 @@
     "@aws-sdk/client-ec2": "^3.649.0",
     "@aws-sdk/client-eks": "^3.649.0",
     "@aws-sdk/credential-providers": "^3.649.0",
-    "@buf/grpc_grpc.community_timostamm-protobuf-ts": "^2.9.4-20240809200651-8507e5a24938.4",
+    "@buf/grpc_grpc.community_timostamm-protobuf-ts": "^2.9.4-20240904201038-e0425cfebb28.4",
     "@buf/jlewi_foyle.bufbuild_es": "^1.10.0-00000000000000-d14934cb2733.1",
     "@buf/jlewi_foyle.connectrpc_es": "^1.4.0-00000000000000-d14934cb2733.3",
-    "@buf/stateful_runme.community_timostamm-protobuf-ts": "^2.9.4-20240903224300-611a892c2067.4",
+    "@buf/stateful_runme.community_timostamm-protobuf-ts": "^2.9.4-20240913234806-45813f39881a.4",
     "@connectrpc/connect": "^1.4.0",
     "@connectrpc/connect-node": "^1.5.0",
     "@google-cloud/compute": "^4.1.0",
diff --git a/src/client/components/terminal/index.ts b/src/client/components/terminal/index.ts
index c3a8c908..4bc2ce58 100644
--- a/src/client/components/terminal/index.ts
+++ b/src/client/components/terminal/index.ts
@@ -12,7 +12,7 @@ import { ClientMessages, RENDERERS, OutputType, WebViews } from '../../../consta
 import { closeOutput, getContext } from '../../utils'
 import { onClientMessage, postClientMessage } from '../../../utils/messaging'
 import { stripANSI } from '../../../utils/ansi'
-import { APIMethod } from '../../../types'
+import { APIMethod, FeatureObserver, FeatureName } from '../../../types'
 import type { TerminalConfiguration } from '../../../utils/configuration'
 import '../closeCellButton'
 import '../copyButton'
@@ -24,12 +24,7 @@ import {
   CreateExtensionCellOutputMutation,
   UpdateCellOutputMutation,
 } from '../../../extension/__generated-platform__/graphql'
-import {
-  isFeatureActive,
-  loadFeatureSnapshot,
-  FeatureObserver,
-  FeatureName,
-} from '../../../features'
+import features from '../../../features'
 
 interface IWindowSize {
   width: number
@@ -453,7 +448,7 @@ export class TerminalView extends LitElement {
         switch (e.type) {
           case ClientMessages.featuresResponse:
           case ClientMessages.featuresUpdateAction:
-            this.featureState$ = loadFeatureSnapshot(e.output.snapshot)
+            this.featureState$ = features.loadSnapshot(e.output.snapshot)
             break
           case ClientMessages.activeThemeChanged:
             this.#updateTerminalTheme()
@@ -857,7 +852,7 @@ export class TerminalView extends LitElement {
           }}"
         ></copy-button>
         ${when(
-          this.isSessionOutputsEnabled && isFeatureActive(FeatureName.Gist, this.featureState$),
+          features.isOn(FeatureName.Gist, this.featureState$),
           () => {
             return html`<gist-cell @onGist="${this.#openSessionOutput}"></gist-cell>`
           },
@@ -865,7 +860,7 @@ export class TerminalView extends LitElement {
         )}
         ${when(
           (this.exitCode === undefined || this.exitCode === 0 || !this.platformId) &&
-            isFeatureActive(FeatureName.Share, this.featureState$),
+            features.isOn(FeatureName.Share, this.featureState$),
           () => {
             return html` <action-button
               ?loading=${this.isLoading}
@@ -881,7 +876,7 @@ export class TerminalView extends LitElement {
         ${when(
           this.exitCode !== 0 &&
             this.platformId &&
-            isFeatureActive(FeatureName.Escalate, this.featureState$),
+            features.isOn(FeatureName.Escalate, this.featureState$),
           () => {
             return html` <action-button
               ?disabled=${!this.isSlackReady}
diff --git a/src/extension/cell.ts b/src/extension/cell.ts
index 821a9b5a..08a1cba2 100644
--- a/src/extension/cell.ts
+++ b/src/extension/cell.ts
@@ -25,6 +25,7 @@ import {
   AWSState,
   CellOutputPayload,
   DenoState,
+  FeatureName,
   GCPState,
   GitHubState,
   Serializer,
@@ -38,6 +39,7 @@ import {
   isPlatformAuthEnabled,
 } from '../utils/configuration'
 
+import features from './features'
 import { RUNME_TRANSIENT_REVISION } from './constants'
 import { getAnnotations, replaceOutput, validateAnnotations } from './utils'
 import {
@@ -48,7 +50,6 @@ import {
 } from './terminal/terminalState'
 import ContextState from './contextState'
 import { IRunnerEnvironment } from './runner/environment'
-import { GrpcSerializer } from './serializer'
 
 const NOTEBOOK_SELECTION_COMMAND = '_notebook.selectKernel'
 
@@ -524,9 +525,17 @@ export class NotebookCellOutputManager {
    *
    */
   async refreshTerminal(terminalState: ITerminalState | undefined): Promise<void> {
-    if (!ContextState.getKey(NOTEBOOK_AUTOSAVE_ON)) {
+    const isSignedIn = features.isOnInContextState(FeatureName.SignedIn)
+    const isForceLogin = features.isOnInContextState(FeatureName.ForceLogin)
+
+    const isAutoSaveOn = ContextState.getKey(NOTEBOOK_AUTOSAVE_ON)
+
+    if (!isSignedIn && !isAutoSaveOn) {
+      return Promise.resolve()
+    } else if (!isSignedIn && isForceLogin) {
       return Promise.resolve()
     }
+
     await this.withLock(async () => {
       // eslint-disable-next-line @typescript-eslint/no-unused-vars
       await this.getExecutionUnsafe(async (exec) => {
@@ -550,7 +559,7 @@ export class NotebookCellOutputManager {
           }
         }
 
-        if (!GrpcSerializer.sessionOutputsEnabled() || !terminalOutput || !terminalOutputItem) {
+        if (!getSessionOutputs() || !terminalOutput || !terminalOutputItem) {
           return
         }
 
diff --git a/src/extension/commands/index.ts b/src/extension/commands/index.ts
index 63ea7b91..23d40c55 100644
--- a/src/extension/commands/index.ts
+++ b/src/extension/commands/index.ts
@@ -37,7 +37,7 @@ import {
   openFileAsRunmeNotebook,
   promptUserSession,
 } from '../utils'
-import { NotebookToolbarCommand, NotebookUiEvent } from '../../types'
+import { NotebookToolbarCommand, NotebookUiEvent, FeatureName } from '../../types'
 import getLogger from '../logger'
 import { RecommendExtensionMessage } from '../messaging'
 import {
@@ -53,6 +53,7 @@ import { createGist } from '../services/github/gist'
 import { InitializeClient } from '../api/client'
 import { GetUserEnvironmentsDocument } from '../__generated-platform__/graphql'
 import { EnvironmentManager } from '../environment/manager'
+import features from '../features'
 
 const log = getLogger('Commands')
 
@@ -388,7 +389,9 @@ export async function addToRecommendedExtensions(context: ExtensionContext) {
 }
 
 export async function toggleAutosave(autoSaveIsOn: boolean) {
-  if (autoSaveIsOn) {
+  const createIfNone = features.isOnInContextState(FeatureName.ForceLogin)
+
+  if (autoSaveIsOn && createIfNone) {
     await promptUserSession()
   }
   return ContextState.addKey(NOTEBOOK_AUTOSAVE_ON, autoSaveIsOn)
diff --git a/src/extension/executors/runner/index.ts b/src/extension/executors/runner/index.ts
index 4f603ef2..8143ac03 100644
--- a/src/extension/executors/runner/index.ts
+++ b/src/extension/executors/runner/index.ts
@@ -504,7 +504,7 @@ export async function resolveRunProgramExecution(
   const resolver = await runner.createProgramResolver(promptMode, envs)
   // todo(sebastian): removing $-prompts from shell scripts should move kernel-side
   const rawCommands = prepareCommandSeq(script, languageId)
-  const result = await resolver.resolveProgram(rawCommands, runnerEnv?.getSessionId())
+  const result = await resolver.resolveProgram(rawCommands, languageId, runnerEnv?.getSessionId())
   const vars = result.response.vars
 
   // todo(sebastian): once normalization is all kernel-side, it should return commands
diff --git a/src/extension/extension.ts b/src/extension/extension.ts
index a779d12f..81b8aeba 100644
--- a/src/extension/extension.ts
+++ b/src/extension/extension.ts
@@ -14,7 +14,7 @@ import {
 import { TelemetryReporter } from 'vscode-telemetry'
 import Channel from 'tangle/webviews'
 
-import { NotebookUiEvent, Serializer, SyncSchema } from '../types'
+import { NotebookUiEvent, Serializer, SyncSchema, FeatureName } from '../types'
 import {
   getDocsUrlFor,
   getForceNewWindowConfig,
@@ -22,7 +22,6 @@ import {
   getSessionOutputs,
 } from '../utils/configuration'
 import { AuthenticationProviders, WebViews } from '../constants'
-import { FeatureName } from '../features'
 
 import { Kernel } from './kernel'
 import KernelServer from './server/kernelServer'
@@ -95,14 +94,12 @@ export class RunmeExtension {
     const grpcSerializer = kernel.hasExperimentEnabled('grpcSerializer')
     const grpcServer = kernel.hasExperimentEnabled('grpcServer')
     const grpcRunner = kernel.hasExperimentEnabled('grpcRunner')
-    const aiLogs = kernel.hasExperimentEnabled('aiLogs')
 
     const server = new KernelServer(
       context.extensionUri,
       {
         retryOnFailure: true,
         maxNumberOfIntents: 10,
-        aiLogs: aiLogs,
       },
       !grpcServer,
       grpcRunner,
@@ -383,56 +380,65 @@ export class RunmeExtension {
         })
     }
 
-    context.subscriptions.push(new GithubAuthProvider(context))
-    context.subscriptions.push(new StatefulAuthProvider(context, uriHandler))
-
-    getGithubAuthSession(false).then((session) => {
-      kernel.updateFeatureState('githubAuth', !!session)
-    })
+    if (kernel.isFeatureOn(FeatureName.RequireStatefulAuth)) {
+      const forceLogin = kernel.isFeatureOn(FeatureName.ForceLogin)
+      context.subscriptions.push(new StatefulAuthProvider(context, uriHandler))
+      const silent = forceLogin ? undefined : true
+
+      getPlatformAuthSession(forceLogin, silent)
+        .then((session) => {
+          if (session) {
+            const openDashboardStr = 'Open Dashboard'
+            window
+              .showInformationMessage('Logged into the Stateful Platform', openDashboardStr)
+              .then((answer) => {
+                if (answer === openDashboardStr) {
+                  const dashboardUri = getRunmeAppUrl(['app'])
+                  const uri = Uri.parse(dashboardUri)
+                  env.openExternal(uri)
+                }
+              })
+          }
+        })
+        .catch((error) => {
+          let message
+          if (error instanceof Error) {
+            message = error.message
+          } else {
+            message = String(error)
+          }
 
-    const createIfNone = kernel.isFeatureActive(FeatureName.ForceLogin)
-    const silent = createIfNone ? undefined : true
-
-    getPlatformAuthSession(createIfNone, silent)
-      .then((session) => {
-        if (session) {
-          const openDashboardStr = 'Open Dashboard'
-          window
-            .showInformationMessage('Logged into the Stateful Platform', openDashboardStr)
-            .then((answer) => {
-              if (answer === openDashboardStr) {
-                const dashboardUri = getRunmeAppUrl(['app'])
-                const uri = Uri.parse(dashboardUri)
-                env.openExternal(uri)
-              }
-            })
-        }
-      })
-      .catch((error) => {
-        let message
-        if (error instanceof Error) {
-          message = error.message
-        } else {
-          message = String(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)
+          }
+        })
+    }
 
-        // 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 (createIfNone && message === 'User did not consent to login.') {
-          getPlatformAuthSession(false)
-        }
+    if (kernel.isFeatureOn(FeatureName.Gist)) {
+      context.subscriptions.push(new GithubAuthProvider(context))
+      getGithubAuthSession(false).then((session) => {
+        kernel.updateFeatureContext('githubAuth', !!session)
       })
+    }
 
     authentication.onDidChangeSessions((e) => {
-      if (e.provider.id === AuthenticationProviders.Stateful) {
+      if (
+        kernel.isFeatureOn(FeatureName.RequireStatefulAuth) &&
+        e.provider.id === AuthenticationProviders.Stateful
+      ) {
         getPlatformAuthSession(false, true).then((session) => {
-          kernel.updateFeatureState('statefulAuth', !!session)
+          kernel.updateFeatureContext('statefulAuth', !!session)
         })
       }
-      if (e.provider.id === AuthenticationProviders.GitHub) {
+      if (
+        kernel.isFeatureOn(FeatureName.Gist) &&
+        e.provider.id === AuthenticationProviders.GitHub
+      ) {
         getGithubAuthSession(false).then((session) => {
-          kernel.updateFeatureState('githubAuth', !!session)
+          kernel.updateFeatureContext('githubAuth', !!session)
         })
       }
     })
diff --git a/src/extension/features.ts b/src/extension/features.ts
new file mode 100644
index 00000000..edb304ff
--- /dev/null
+++ b/src/extension/features.ts
@@ -0,0 +1,19 @@
+import features, { FEATURES_CONTEXT_STATE_KEY } from '../features'
+import { FeatureName } from '../types'
+
+import ContextState from './contextState'
+
+export function isOnInContextState(featureName: FeatureName): boolean {
+  const snapshot = ContextState.getKey<string>(FEATURES_CONTEXT_STATE_KEY)
+  if (!snapshot) {
+    return false
+  }
+
+  const featureState$ = features.loadSnapshot(snapshot)
+
+  return features.isOn(featureName, featureState$)
+}
+
+export default {
+  isOnInContextState,
+}
diff --git a/src/extension/handler/uri.ts b/src/extension/handler/uri.ts
index cb09ec8a..94d9d27a 100644
--- a/src/extension/handler/uri.ts
+++ b/src/extension/handler/uri.ts
@@ -57,7 +57,6 @@ export class RunmeUriHandler implements UriHandler, Disposable {
   async handleUri(uri: Uri) {
     log.info(`triggered RunmeUriHandler with ${uri}`)
     const params = new URLSearchParams(uri.query)
-    const windowId = params.get('windowId')
     const state = params.get('state')
     const code = params.get('code')
     const command = params.get('command')
@@ -66,7 +65,7 @@ export class RunmeUriHandler implements UriHandler, Disposable {
       window.showErrorMessage('No query parameter "command" provided')
       return
     }
-    if (command === 'auth' && windowId && state && code) {
+    if (command === 'auth' && state && code) {
       TelemetryReporter.sendTelemetryEvent('extension.uriHandler', {
         command,
         type: AuthenticationProviders.Stateful,
@@ -133,7 +132,7 @@ export class RunmeUriHandler implements UriHandler, Disposable {
       }
     }
 
-    window.showErrorMessage(`Couldn't recognise command "${command}"`)
+    window.showErrorMessage(`Couldn't recognize command "${command}"`)
   }
 
   private async _setupProject(fileToOpen: string, repository?: string | null) {
diff --git a/src/extension/kernel.ts b/src/extension/kernel.ts
index 4c20e68d..5278f2bb 100644
--- a/src/extension/kernel.ts
+++ b/src/extension/kernel.ts
@@ -34,6 +34,9 @@ import {
   type ClientMessage,
   type RunmeTerminal,
   type Serializer,
+  type ExtensionName,
+  type FeatureContext,
+  FeatureName,
 } from '../types'
 import {
   ClientMessages,
@@ -48,16 +51,7 @@ import {
 import { API } from '../utils/deno/api'
 import { postClientMessage } from '../utils/messaging'
 import { getNotebookExecutionOrder, registerExtensionEnvVarsMutation } from '../utils/configuration'
-import {
-  isFeatureActive,
-  FeatureContext,
-  getFeatureSnapshot,
-  loadFeaturesState,
-  updateFeatureContext,
-  FEATURES_CONTEXT_STATE_KEY,
-  FeatureName,
-  ExtensionName,
-} from '../features'
+import features, { FEATURES_CONTEXT_STATE_KEY } from '../features'
 
 import getLogger from './logger'
 import executor, {
@@ -158,7 +152,6 @@ export class Kernel implements Disposable {
     this.#experiments.set('grpcRunner', config.get<boolean>('grpcRunner', true))
     this.#experiments.set('grpcServer', config.get<boolean>('grpcServer', true))
     this.#experiments.set('smartEnvStore', config.get<boolean>('smartEnvStore', false))
-    this.#experiments.set('aiLogs', config.get<boolean>('aiLogs', false))
     this.#experiments.set('shellWarning', config.get<boolean>('shellWarning', false))
     this.#experiments.set('reporter', config.get<boolean>('reporter', false))
 
@@ -205,21 +198,21 @@ export class Kernel implements Disposable {
       extensionId: context?.extension?.id as ExtensionName,
     }
 
-    this.featuresState$ = loadFeaturesState(packageJSON, featContext, this.#featuresSettings)
+    this.featuresState$ = features.loadState(packageJSON, featContext, this.#featuresSettings)
 
     if (this.featuresState$) {
-      const features = workspace.getConfiguration('runme.features')
+      const runmeFeatures = workspace.getConfiguration('runme.features')
       const featureNames = Object.keys(FeatureName).map((f) => f.toLowerCase())
       if (features) {
         featureNames.forEach((feature) => {
-          if (features.has(feature)) {
-            this.#featuresSettings.set(feature, features.get<boolean>(feature, false))
+          if (runmeFeatures.has(feature)) {
+            this.#featuresSettings.set(feature, runmeFeatures.get<boolean>(feature, false))
           }
         })
       }
 
       const subscription = this.featuresState$
-        .pipe(map((_state) => getFeatureSnapshot(this.featuresState$)))
+        .pipe(map((_state) => features.getSnapshot(this.featuresState$)))
         .subscribe((snapshot) => {
           ContextState.addKey(FEATURES_CONTEXT_STATE_KEY, snapshot)
           postClientMessage(this.messaging, ClientMessages.featuresUpdateAction, {
@@ -233,16 +226,16 @@ export class Kernel implements Disposable {
     }
   }
 
-  isFeatureActive(featureName: FeatureName): boolean {
+  isFeatureOn(featureName: FeatureName): boolean {
     if (!this.featuresState$) {
       return false
     }
 
-    return isFeatureActive(featureName, this.featuresState$)
+    return features.isOn(featureName, this.featuresState$)
   }
 
-  updateFeatureState<K extends keyof FeatureContext>(key: K, value: FeatureContext[K]) {
-    updateFeatureContext(this.featuresState$, key, value, this.#featuresSettings)
+  updateFeatureContext<K extends keyof FeatureContext>(key: K, value: FeatureContext[K]) {
+    features.updateContext(this.featuresState$, key, value, this.#featuresSettings)
   }
 
   registerNotebookCell(cell: NotebookCell) {
@@ -632,7 +625,7 @@ export class Kernel implements Disposable {
     } else if (message.type.startsWith('terminal:')) {
       return
     } else if (message.type === ClientMessages.featuresRequest) {
-      const snapshot = getFeatureSnapshot(this.featuresState$)
+      const snapshot = features.getSnapshot(this.featuresState$)
       postClientMessage(this.messaging, ClientMessages.featuresResponse, {
         snapshot: snapshot,
       })
diff --git a/src/extension/runner/program.ts b/src/extension/runner/program.ts
index b53834c7..950d5892 100644
--- a/src/extension/runner/program.ts
+++ b/src/extension/runner/program.ts
@@ -14,12 +14,13 @@ export class GrpcRunnerProgramResolver implements IRunnerChild {
     private readonly envs: Record<string, string>,
   ) {}
 
-  async resolveProgram(commands: string[], sessionId: string | undefined) {
+  async resolveProgram(commands: string[], languageId: string, sessionId: string | undefined) {
     const mode = this.mode
     const env = Object.entries(this.envs).map(([key, value]: [string, string]) => `${key}=${value}`)
 
     const req = ResolveProgramRequestImpl().create({
       source: { oneofKind: 'commands', commands: { lines: commands } },
+      languageId,
       mode,
       sessionId,
       env,
diff --git a/src/extension/serializer.ts b/src/extension/serializer.ts
index cccbf8a0..25e6f0a0 100644
--- a/src/extension/serializer.ts
+++ b/src/extension/serializer.ts
@@ -24,7 +24,7 @@ import { ulid } from 'ulidx'
 import { maskString } from 'data-guardian'
 import YAML from 'yaml'
 
-import { Serializer } from '../types'
+import { FeatureName, Serializer } from '../types'
 import {
   NOTEBOOK_AUTOSAVE_ON,
   NOTEBOOK_HAS_OUTPUTS,
@@ -37,7 +37,6 @@ import {
   ServerLifecycleIdentity,
   getServerConfigurationValue,
   getSessionOutputs,
-  isPlatformAuthEnabled,
 } from '../utils/configuration'
 
 import {
@@ -61,6 +60,7 @@ import { getCellById } from './cell'
 import { IProcessInfoState } from './terminal/terminalState'
 import ContextState from './contextState'
 import * as ghost from './ai/ghost'
+import * as features from './features'
 
 declare var globalThis: any
 const DEFAULT_LANG_ID = 'text'
@@ -551,7 +551,6 @@ export class GrpcSerializer extends SerializerBase {
   }
 
   protected async saveNotebookOutputsByCacheId(cacheId: string): Promise<number> {
-    // if session outputs are disabled, we don't write anything
     if (!GrpcSerializer.sessionOutputsEnabled()) {
       this.togglePreviewButton(false)
       return -1
@@ -579,18 +578,20 @@ export class GrpcSerializer extends SerializerBase {
       return -1
     }
 
-    const sessionFile = GrpcSerializer.getOutputsUri(srcDocUri, sessionId)
-    if (!sessionFile) {
+    // Don't write to disk if authenticated and share are disabled
+    if (
+      features.isOnInContextState(FeatureName.SignedIn) &&
+      features.isOnInContextState(FeatureName.Share)
+    ) {
       this.togglePreviewButton(false)
-      return -1
+      // But still return a valid bytes length so the cache keeps working
+      return bytes.length
     }
 
-    // Don't write to disk if platform auth is enabled
-    const isPlatform = isPlatformAuthEnabled()
-    if (isPlatform) {
+    const sessionFile = GrpcSerializer.getOutputsUri(srcDocUri, sessionId)
+    if (!sessionFile) {
       this.togglePreviewButton(false)
-      // But still return a valid bytes length so the cache keeps working
-      return bytes.length
+      return -1
     }
 
     await workspace.fs.writeFile(sessionFile, bytes)
@@ -719,7 +720,10 @@ export class GrpcSerializer extends SerializerBase {
   }
 
   static sessionOutputsEnabled() {
-    return getSessionOutputs() && ContextState.getKey<boolean>(NOTEBOOK_AUTOSAVE_ON)
+    const isAutoSaveOn = ContextState.getKey<boolean>(NOTEBOOK_AUTOSAVE_ON)
+    const isSessionOutputs = getSessionOutputs()
+
+    return isSessionOutputs && isAutoSaveOn
   }
 
   private async cacheNotebookOutputs(
diff --git a/src/extension/server/kernelServer.ts b/src/extension/server/kernelServer.ts
index 6c307dd0..f7e823e8 100644
--- a/src/extension/server/kernelServer.ts
+++ b/src/extension/server/kernelServer.ts
@@ -33,7 +33,6 @@ export interface IServerConfig {
     intents: number
     interval: number
   }
-  aiLogs?: boolean
 }
 
 const log = getLogger('KernelServer')
@@ -65,7 +64,6 @@ class KernelServer implements IServer {
   #transport?: GrpcTransport
   #serverDisposables: Disposable[] = []
   #forceExternalServer: boolean
-  #aiLogs: boolean
 
   readonly #onClose = this.register(new EventEmitter<{ code: number | null }>())
   readonly #onTransportReady = this.register(
@@ -96,7 +94,6 @@ class KernelServer implements IServer {
     this.#acceptsIntents = options.acceptsConnection?.intents || 50
     this.#acceptsInterval = options.acceptsConnection?.interval || 200
     this.#forceExternalServer = externalServer
-    this.#aiLogs = options.aiLogs || false
   }
 
   dispose() {
@@ -235,11 +232,6 @@ class KernelServer implements IServer {
       args.push('--runner')
     }
 
-    if (this.#aiLogs) {
-      log.info('AI logs enabled')
-      args.push('--ai-logs=true')
-    }
-
     if (getTLSEnabled()) {
       args.push('--tls', this.getTLSDir())
     } else {
diff --git a/src/extension/utils.ts b/src/extension/utils.ts
index b4862015..a16cd2fc 100644
--- a/src/extension/utils.ts
+++ b/src/extension/utils.ts
@@ -33,6 +33,7 @@ import {
   NotebookAutoSaveSetting,
   RunmeTerminal,
   Serializer,
+  FeatureName,
 } from '../types'
 import { SafeCellAnnotationsSchema, CellAnnotationsSchema } from '../schema'
 import {
@@ -54,15 +55,9 @@ import {
   getPortNumber,
   getTLSDir,
   getTLSEnabled,
-  isPlatformAuthEnabled,
 } from '../utils/configuration'
-import {
-  FeatureName,
-  FEATURES_CONTEXT_STATE_KEY,
-  isFeatureActive,
-  loadFeatureSnapshot,
-} from '../features'
 
+import features from './features'
 import CategoryQuickPickItem from './quickPickItems/category'
 import getLogger from './logger'
 import { Kernel } from './kernel'
@@ -586,11 +581,15 @@ export async function resolveAuthToken(createIfNone: boolean = true) {
 }
 
 export async function resolveAppToken(createIfNone: boolean = true) {
-  const session = await getPlatformAuthSession(createIfNone)
-  if (!session) {
-    return null
+  if (features.isOnInContextState(FeatureName.RequireStatefulAuth)) {
+    const session = await getPlatformAuthSession(createIfNone)
+    if (!session) {
+      return null
+    }
+    return { token: session.accessToken }
   }
-  return { token: session.accessToken }
+
+  return null
 }
 
 export function fetchStaticHtml(appUrl: string) {
@@ -712,31 +711,25 @@ export function asWorkspaceRelativePath(documentPath: string): {
   return { relativePath, outside: false }
 }
 
-export async function resolveUserSession(
-  createIfNone: boolean,
-): Promise<AuthenticationSession | undefined> {
-  return isPlatformAuthEnabled()
-    ? await getPlatformAuthSession(createIfNone)
-    : await getGithubAuthSession(createIfNone)
-}
-
 /**
  * Handles the first time experience for saving a cell.
  * It informs the user that a Login with a GitHub account is required before prompting the user.
  * This only happens once. Subsequent saves will not display the prompt.
  * @returns AuthenticationSession
  */
-export async function promptUserSession(): Promise<AuthenticationSession | undefined> {
-  let session = await resolveUserSession(false)
+export async function promptUserSession() {
+  const createIfNone = features.isOnInContextState(FeatureName.ForceLogin)
+  const silent = createIfNone ? undefined : true
 
-  const provider = isPlatformAuthEnabled() ? 'Stateful' : 'GitHub'
-  const featureState$ = loadFeatureSnapshot(ContextState.getKey(FEATURES_CONTEXT_STATE_KEY))
-  const displayLoginPrompt = getLoginPrompt() && isFeatureActive(FeatureName.Share, featureState$)
+  const session = await getPlatformAuthSession(false, silent)
+
+  const displayLoginPrompt =
+    getLoginPrompt() && createIfNone && features.isOnInContextState(FeatureName.Share)
 
   if (!session && displayLoginPrompt !== false) {
     const option = await window.showInformationMessage(
       `Securely store your cell outputs.
-      Sign in with ${provider} is required, do you want to proceed?`,
+      Sign in with Stateful is required, do you want to proceed?`,
       'Yes',
       'No',
       'Open Settings',
@@ -749,13 +742,28 @@ export async function promptUserSession(): Promise<AuthenticationSession | undef
       return commands.executeCommand('runme.openSettings', 'runme.app.loginPrompt')
     }
 
-    session = await resolveUserSession(true)
-    if (!session) {
-      throw new Error(`You must authenticate with your ${provider} account`)
-    }
-  }
+    getPlatformAuthSession(createIfNone)
+      .then((session) => {
+        if (!session) {
+          throw new Error('You must authenticate with your Stateful account')
+        }
+      })
+      .catch((error) => {
+        let message
+        if (error instanceof Error) {
+          message = error.message
+        } else {
+          message = String(error)
+        }
 
-  return session
+        // 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 (createIfNone && message === 'User did not consent to login.') {
+          getPlatformAuthSession(false)
+        }
+      })
+  }
 }
 
 export async function checkSession(context: ExtensionContext) {
diff --git a/src/features.ts b/src/features.ts
index e730892f..7a95eddc 100644
--- a/src/features.ts
+++ b/src/features.ts
@@ -1,73 +1,35 @@
 import { BehaviorSubject } from 'rxjs'
 import { satisfies } from 'semver'
 
-export const FEATURES_CONTEXT_STATE_KEY = 'features'
-
-export type FeatureContext = {
-  os?: string
-  vsCodeVersion?: string
-  runmeVersion?: string
-  extensionVersion?: string
-  githubAuth?: boolean // `true`, `false`, or `undefined`
-  statefulAuth?: boolean // `true`, `false`, or `undefined`
-  extensionId?: ExtensionName
-}
-
-export enum ExtensionName {
-  StatefulRunme = 'stateful.runme',
-  StatefulPlatform = 'stateful.platform',
-}
-
-type EnabledForExtensions = Partial<Record<ExtensionName, boolean>>
-
-export type FeatureCondition = {
-  os?: string
-  vsCodeVersion?: string
-  runmeVersion?: string
-  extensionVersion?: string
-  githubAuthRequired?: boolean
-  statefulAuthRequired?: boolean
-  enabledForExtensions?: EnabledForExtensions
-}
-
-export enum FeatureName {
-  Gist = 'Gist',
-  Share = 'Share',
-  Escalate = 'Escalate',
-  ForceLogin = 'ForceLogin',
-}
+import {
+  EnabledForExtensions,
+  ExtensionName,
+  Feature,
+  FeatureContext,
+  FeatureName,
+  FeatureObserver,
+  Features,
+  FeatureState,
+} from './types'
 
-export type Feature = {
-  enabled: boolean
-  activated: boolean
-  conditions: FeatureCondition
-}
-
-export type Features = Partial<Record<keyof typeof FeatureName, Feature>>
-
-export type FeatureState = {
-  features: Features
-  context?: FeatureContext
-}
-
-export type FeatureObserver = BehaviorSubject<FeatureState>
+export const FEATURES_CONTEXT_STATE_KEY = 'features'
 
-function loadFeaturesFromPackageJson(packageJSON: any): Features {
+function loadFromPackageJson(packageJSON: any): Features {
   const features = (packageJSON?.runme?.features || {}) as Features
   return features
 }
 
-export function loadFeaturesState(
+function loadState(
   packageJSON: any,
   context?: FeatureContext,
   overrides: Map<string, boolean> = new Map(),
 ): FeatureObserver {
-  const initialFeatures = loadFeaturesFromPackageJson(packageJSON)
+  const initialFeatures = loadFromPackageJson(packageJSON)
   const state = new BehaviorSubject<FeatureState>({
     features: initialFeatures,
     context,
   })
-  updateFeatureState(state, context, overrides)
+  updateState(state, context, overrides)
   return state
 }
 
@@ -187,7 +149,7 @@ function isActive(
 
   if (!checkExtensionId(enabledForExtensions, context?.extensionId)) {
     console.log(
-      `Feature "${featureName}" is inactive due to checkExtensionId. Expected: ${enabledForExtensions}, actual: ${context?.extensionId}`,
+      `Feature "${featureName}" is inactive due to checkExtensionId. Expected: ${JSON.stringify(enabledForExtensions)}, actual: ${context?.extensionId}`,
     )
     return false
   }
@@ -196,7 +158,7 @@ function isActive(
   return true
 }
 
-export function updateFeatureState(
+function updateState(
   featureState$: FeatureObserver,
   context?: FeatureContext,
   overrides: Map<string, boolean> = new Map(),
@@ -208,7 +170,7 @@ export function updateFeatureState(
       ...acc,
       [key]: {
         ...value,
-        activated: isActive(key as FeatureName, value, context, overrides),
+        on: isActive(key as FeatureName, value, context, overrides),
       },
     }),
     {} as Features,
@@ -218,7 +180,7 @@ export function updateFeatureState(
   return featureState$
 }
 
-export function updateFeatureContext<K extends keyof FeatureContext>(
+function updateContext<K extends keyof FeatureContext>(
   featureState$: FeatureObserver | undefined,
   key: K,
   value: FeatureContext[K],
@@ -231,11 +193,11 @@ export function updateFeatureContext<K extends keyof FeatureContext>(
   const newContext = { ...(currentState.context || {}), [key]: value }
 
   if (newContext[key] !== currentState?.context?.[key]) {
-    updateFeatureState(featureState$, newContext, overrides)
+    updateState(featureState$, newContext, overrides)
   }
 }
 
-export function getFeatureSnapshot(featureState$: FeatureObserver | undefined): string {
+function getSnapshot(featureState$: FeatureObserver | undefined): string {
   if (!featureState$) {
     return ''
   }
@@ -243,7 +205,7 @@ export function getFeatureSnapshot(featureState$: FeatureObserver | undefined):
   return JSON.stringify(featureState$.getValue())
 }
 
-export function loadFeatureSnapshot(snapshot: string): FeatureObserver {
+function loadSnapshot(snapshot: string): FeatureObserver {
   const { features, context } = JSON.parse(snapshot)
   const featureState$ = new BehaviorSubject<FeatureState>({
     features,
@@ -254,14 +216,20 @@ export function loadFeatureSnapshot(snapshot: string): FeatureObserver {
   return featureState$
 }
 
-export function isFeatureActive(
-  featureName: FeatureName,
-  featureState$?: FeatureObserver,
-): boolean {
+function isOn(featureName: FeatureName, featureState$?: FeatureObserver): boolean {
   if (!featureState$) {
     return false
   }
 
   const feature = featureState$.getValue().features[featureName]
-  return feature ? feature.activated : false
+  return feature?.on ?? false
+}
+
+export default {
+  loadState,
+  updateState,
+  updateContext,
+  getSnapshot,
+  loadSnapshot,
+  isOn,
 }
diff --git a/src/types.ts b/src/types.ts
index f2f07c80..7104193f 100644
--- a/src/types.ts
+++ b/src/types.ts
@@ -19,6 +19,7 @@ import {
 } from '@aws-sdk/client-ec2'
 import { google } from '@google-cloud/run/build/protos/protos'
 import { protos } from '@google-cloud/compute'
+import { type BehaviorSubject } from 'rxjs'
 
 import { OutputType, ClientMessages, RUNME_FRONTMATTER_PARSED } from './constants'
 import { SafeCellAnnotationsSchema, SafeNotebookAnnotationsSchema } from './schema'
@@ -732,3 +733,54 @@ export type NotebookUiEvent = {
 export type Settings = {
   docsUrl?: string
 }
+
+export type FeatureContext = {
+  os?: string
+  vsCodeVersion?: string
+  runmeVersion?: string
+  extensionVersion?: string
+  githubAuth?: boolean // `true`, `false`, or `undefined`
+  statefulAuth?: boolean // `true`, `false`, or `undefined`
+  extensionId?: ExtensionName
+}
+
+export enum ExtensionName {
+  StatefulRunme = 'stateful.runme',
+  StatefulPlatform = 'stateful.platform',
+}
+
+export type EnabledForExtensions = Partial<Record<ExtensionName, boolean>>
+
+export type FeatureCondition = {
+  os?: string
+  vsCodeVersion?: string
+  runmeVersion?: string
+  extensionVersion?: string
+  githubAuthRequired?: boolean
+  statefulAuthRequired?: boolean
+  enabledForExtensions?: EnabledForExtensions
+}
+
+export enum FeatureName {
+  Gist = 'Gist',
+  Share = 'Share',
+  Escalate = 'Escalate',
+  ForceLogin = 'ForceLogin',
+  SignedIn = 'SignedIn',
+  RequireStatefulAuth = 'RequireStatefulAuth',
+}
+
+export type Feature = {
+  enabled: boolean
+  on: boolean
+  conditions: FeatureCondition
+}
+
+export type Features = Partial<Record<keyof typeof FeatureName, Feature>>
+
+export type FeatureState = {
+  features: Features
+  context?: FeatureContext
+}
+
+export type FeatureObserver = BehaviorSubject<FeatureState>


Test patch:
diff --git a/tests/extension/cell.test.ts b/tests/extension/cell.test.ts
index a91888f7..4ab07fa5 100644
--- a/tests/extension/cell.test.ts
+++ b/tests/extension/cell.test.ts
@@ -26,10 +26,12 @@ vi.mock('vscode', async () => {
 vi.mock('vscode-telemetry')
 
 vi.mock('../../src/extension/grpc/client', () => ({}))
-vi.mock('../../../src/extension/grpc/runner/v1', () => ({
+vi.mock('../../src/extension/grpc/runner/v1', () => ({
   ResolveProgramRequest_Mode: vi.fn(),
 }))
 
+vi.mock('../../src/extension/features')
+
 describe('NotebookCellManager', () => {
   it('can register cells', async () => {
     const manager = new NotebookCellManager({} as any)
diff --git a/tests/extension/features.test.ts b/tests/extension/features.test.ts
index a80d7029..9d5c80bf 100644
--- a/tests/extension/features.test.ts
+++ b/tests/extension/features.test.ts
@@ -1,17 +1,13 @@
-import { describe, it, expect, beforeEach } from 'vitest'
+import { describe, it, expect, beforeEach, vi } from 'vitest'
 
+import features from '../../src/features'
 import {
-  updateFeatureState,
-  getFeatureSnapshot,
-  loadFeaturesState,
-  isFeatureActive,
-  loadFeatureSnapshot,
+  ExtensionName,
   FeatureContext,
-  FeatureState,
-  FeatureObserver,
   FeatureName,
-  ExtensionName,
-} from '../../src/features'
+  FeatureObserver,
+  FeatureState,
+} from '../../src/types'
 
 const packageJSON = {
   runme: {
@@ -50,11 +46,14 @@ const packageJSON = {
   },
 }
 
+vi.mock('vscode')
+vi.mock('../../src/extension/contextState')
+
 describe('Feature Store', () => {
   let featureState$: FeatureObserver
 
   beforeEach(() => {
-    featureState$ = loadFeaturesState(packageJSON, {
+    featureState$ = features.loadState(packageJSON, {
       os: 'linux',
     })
   })
@@ -70,12 +69,12 @@ describe('Feature Store', () => {
       extensionId: ExtensionName.StatefulRunme,
     }
 
-    updateFeatureState(featureState$, initialContext)
+    features.updateState(featureState$, initialContext)
 
     const currentFeatures = (featureState$.getValue() as FeatureState).features
 
-    expect(currentFeatures.Escalate?.activated).toBe(false)
-    expect(currentFeatures.Gist?.activated).toBe(true)
+    expect(currentFeatures.Escalate?.on).toBe(false)
+    expect(currentFeatures.Gist?.on).toBe(true)
   })
 
   it('should take a snapshot of the current state', () => {
@@ -86,11 +85,11 @@ describe('Feature Store', () => {
       runmeVersion: '1.3.0',
       githubAuth: true,
       statefulAuth: true,
-      extensionId: 'stateful.runme',
+      extensionId: ExtensionName.StatefulRunme,
     }
 
-    updateFeatureState(featureState$, initialContext)
-    const snapshot = getFeatureSnapshot(featureState$)
+    features.updateState(featureState$, initialContext)
+    const snapshot = features.getSnapshot(featureState$)
 
     expect(snapshot).toBe(JSON.stringify(featureState$.getValue()))
   })
@@ -106,15 +105,15 @@ describe('Feature Store', () => {
       extensionId: ExtensionName.StatefulRunme,
     }
 
-    updateFeatureState(featureState$, initialContext)
-    const snapshot = getFeatureSnapshot(featureState$)
+    features.updateState(featureState$, initialContext)
+    const snapshot = features.getSnapshot(featureState$)
 
     featureState$.next({
       context: initialContext,
       features: {
         Escalate: {
           enabled: true,
-          activated: false,
+          on: false,
           conditions: {
             os: 'All',
             vsCodeVersion: '>=1.58.0',
@@ -125,7 +124,7 @@ describe('Feature Store', () => {
         },
         Gist: {
           enabled: true,
-          activated: false,
+          on: false,
           conditions: {
             os: 'win32',
             vsCodeVersion: '>=1.60.0',
@@ -136,7 +135,7 @@ describe('Feature Store', () => {
       },
     })
 
-    const featureStateCopy$ = loadFeatureSnapshot(snapshot)
+    const featureStateCopy$ = features.loadSnapshot(snapshot)
 
     const restoredFeatures = featureStateCopy$.getValue()
     expect(restoredFeatures).toEqual(JSON.parse(snapshot))
@@ -152,11 +151,11 @@ describe('Feature Store', () => {
       statefulAuth: false,
       extensionId: ExtensionName.StatefulRunme,
     }
-    updateFeatureState(featureState$, newContext)
+    features.updateState(featureState$, newContext)
 
     const currentFeatures = (featureState$.getValue() as FeatureState).features
-    expect(currentFeatures.Escalate?.activated).toBe(false)
-    expect(currentFeatures.Gist?.activated).toBe(false)
+    expect(currentFeatures.Escalate?.on).toBe(false)
+    expect(currentFeatures.Gist?.on).toBe(false)
   })
 
   it('should correctly identify if a feature is enabled by name', () => {
@@ -169,10 +168,10 @@ describe('Feature Store', () => {
       statefulAuth: true,
       extensionId: ExtensionName.StatefulPlatform,
     }
-    updateFeatureState(featureState$, ctx)
+    features.updateState(featureState$, ctx)
 
-    expect(isFeatureActive(FeatureName.Escalate, featureState$)).toBe(true)
-    expect(isFeatureActive(FeatureName.Gist, featureState$)).toBe(false)
+    expect(features.isOn(FeatureName.Escalate, featureState$)).toBe(true)
+    expect(features.isOn(FeatureName.Gist, featureState$)).toBe(false)
   })
 
   it('should correctly identify if a feature is enabled by extensionId', () => {
@@ -185,9 +184,9 @@ describe('Feature Store', () => {
       statefulAuth: true,
       extensionId: ExtensionName.StatefulRunme,
     }
-    updateFeatureState(featureState$, ctx)
+    features.updateState(featureState$, ctx)
 
-    expect(isFeatureActive(FeatureName.Escalate, featureState$)).toBe(false)
-    expect(isFeatureActive(FeatureName.Gist, featureState$)).toBe(true)
+    expect(features.isOn(FeatureName.Escalate, featureState$)).toBe(false)
+    expect(features.isOn(FeatureName.Gist, featureState$)).toBe(true)
   })
 })
diff --git a/tests/extension/handler/uri.test.ts b/tests/extension/handler/uri.test.ts
index 53c4caf7..b63529ca 100644
--- a/tests/extension/handler/uri.test.ts
+++ b/tests/extension/handler/uri.test.ts
@@ -63,7 +63,7 @@ describe('RunmeUriHandler', () => {
     it('should fail if no command was recognised', async () => {
       vi.mocked(Uri.parse).mockReturnValue({ query: { command: 'foobar' } } as any)
       expect(await handler.handleUri(Uri.parse('vscode://stateful.runme?foo=bar'))).toBe(undefined)
-      expect(window.showErrorMessage).toBeCalledWith('Couldn\'t recognise command "foobar"')
+      expect(window.showErrorMessage).toBeCalledWith('Couldn\'t recognize command "foobar"')
     })
 
     it('runs _setupProject if command was "setup"', async () => {
diff --git a/tests/extension/serializer.test.ts b/tests/extension/serializer.test.ts
index fb1dade6..5394c3ec 100644
--- a/tests/extension/serializer.test.ts
+++ b/tests/extension/serializer.test.ts
@@ -67,6 +67,8 @@ vi.mock('../../src/extension/utils', () => ({
   initWasm: vi.fn(),
 }))
 
+vi.mock('../../src/extension/features')
+
 function newKernel(): Kernel {
   return {} as unknown as Kernel
 }