Force Session Outputs in Platform Auth (#1656)

* Enable session outputs for Platform * Refresh terminal * Making session outputs back with Stateful auth for Runme * Passing tests * Improves feature flag encapsulation * Revert condition * Disable Share for runme to avoid conflicts * Restore validation * Enhance debuggin message * Fix escalate when autoSave is off --------- Co-authored-by: Cristian Cepeda <[email protected]>
stateful · Sep 18, 2024 · 708a281 · 708a281
1 parent 1a48894
commit 708a281
Show file tree

Hide file tree

Showing 14 changed files with 240 additions and 170 deletions.
diff --git a/package.json b/package.json
@@ -58,6 +58,18 @@
   "main": "./out/extension.js",
   "runme": {
     "features": {
+      "SignedIn": {
+        "enabled": true,
+        "conditions": {
+          "os": "All",
+          "vsCodeVersion": ">=1.58.0",
+          "statefulAuthRequired": true,
+          "enabledForExtensions": {
+            "stateful.platform": true,
+            "stateful.runme": true
+          }
+        }
+      },
       "ForceLogin": {
         "enabled": true,
         "conditions": {
@@ -88,7 +100,7 @@
           "vsCodeVersion": ">=1.58.0",
           "enabledForExtensions": {
             "stateful.platform": true,
-            "stateful.runme": true
+            "stateful.runme": false
           }
         }
       },

diff --git 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,15 +852,15 @@ 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>`
           },
           () => {},
         )}
         ${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
@@ -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
@@ -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/extension.ts b/src/extension/extension.ts
@@ -14,15 +14,14 @@ 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,
   getRunmeAppUrl,
   getSessionOutputs,
 } from '../utils/configuration'
 import { AuthenticationProviders, WebViews } from '../constants'
-import { FeatureName } from '../features'
 
 import { Kernel } from './kernel'
 import KernelServer from './server/kernelServer'
@@ -385,10 +384,10 @@ export class RunmeExtension {
     context.subscriptions.push(new StatefulAuthProvider(context, uriHandler))
 
     getGithubAuthSession(false).then((session) => {
-      kernel.updateFeatureState('githubAuth', !!session)
+      kernel.updateFeatureContext('githubAuth', !!session)
     })
 
-    const createIfNone = kernel.isFeatureActive(FeatureName.ForceLogin)
+    const createIfNone = kernel.isFeatureOn(FeatureName.ForceLogin)
     const silent = createIfNone ? undefined : true
 
     getPlatformAuthSession(createIfNone, silent)
@@ -425,12 +424,12 @@ export class RunmeExtension {
     authentication.onDidChangeSessions((e) => {
       if (e.provider.id === AuthenticationProviders.Stateful) {
         getPlatformAuthSession(false, true).then((session) => {
-          kernel.updateFeatureState('statefulAuth', !!session)
+          kernel.updateFeatureContext('statefulAuth', !!session)
         })
       }
       if (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
@@ -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/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, {
@@ -204,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, {
@@ -232,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) {
@@ -631,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/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(