Repository: stateful/vscode-runme
Pull request #226 (2023-02-21 20:35:54)

Runner service implementation feature parity
- [x] Run `background=true` as background tasks
- [x] User-prompts (UI) for ENV variables
- [x] Unify cmds/script prep?
- [x] Allow identification of grpc-created terminals

Needs changes in runme service:

- [x] Curl-ing img (last cell in `examples/README.md`) with notebook-native png renderer crash-hangs [#143](https://github.com/stateful/runme/issues/143)
- [x] Cat-ing LICENSE (in `examples/README.md`) does not work ~~[#144](https://github.com/stateful/runme/issues/144)~~ [#155](https://github.com/stateful/runme/issues/155)
 

mxsdev: @sourishkrout

Tested with the main example readme and every script worked for me. 

Could you test against really quick with the latest changes? Thanks!
sourishkrout: > @sourishkrout
> 
> Tested with the main example readme and every script worked for me.
> 
> Could you test against really quick with the latest changes? Thanks!

Confirmed. Great job 👏 

Another thing I noticed is that `closeTerminalOnSuccess` does not seem to have an effect? Just on my end, @mxsdev?
mxsdev: I was able to fix tests, but for the following line in `tests/extension/executors/utils.test.ts`:

```ts
vi.mock('../../../src/extension/utils', () => ({
  replaceOutput: vi.fn(),
}))
```

I was not able to use `importActual` as it would result in the process hanging. I guess this is due to cyclic dependency, but that's [supposed to be fixed](https://github.com/vitest-dev/vitest/issues/2567), so idk.

In the meantime I just mocked the method, but this is less than ideal since there's no need to mock it. Lmk if you think this is worth an issue @sourishkrout.

Anyways, this is still good to merge! 🎉 

Patch:
diff --git a/src/extension/executors/runner.ts b/src/extension/executors/runner.ts
index ddc376cd..ec463367 100644
--- a/src/extension/executors/runner.ts
+++ b/src/extension/executors/runner.ts
@@ -10,7 +10,8 @@ import {
   TaskPanelKind,
   tasks,
   NotebookCellExecution,
-  TextDocument
+  TextDocument,
+  ExtensionContext
 } from 'vscode'
 
 import { OutputType } from '../../constants'
@@ -20,12 +21,14 @@ import { IRunner, IRunnerEnvironment } from '../runner'
 import { getAnnotations, getCmdShellSeq, replaceOutput } from '../utils'
 
 import { closeTerminalByEnvID } from './task'
-import { getShellPath } from './utils'
+import { getShellPath, parseCommandSeq } from './utils'
 
 const LABEL_LIMIT = 15
+const BACKGROUND_TASK_HIDE_TIMEOUT = 2000
 const MIME_TYPES_WITH_CUSTOM_RENDERERS = ['text/plain']
 
 export async function executeRunner(
+  context: ExtensionContext,
   runner: IRunner,
   exec: NotebookCellExecution,
   runningCell: TextDocument,
@@ -36,8 +39,18 @@ export async function executeRunner(
 
   const RUNME_ID = `${runningCell.fileName}:${exec.cell.index}`
 
-  const cellText = exec.cell.document.getText()
-  const script = getCmdShellSeq(cellText, PLATFORM_OS)
+  let cellText = exec.cell.document.getText()
+
+  const envs: Record<string, string> = {
+    RUNME_ID
+  }
+
+  const commands = await parseCommandSeq(cellText)
+  if(!commands) { return false }
+
+  if (commands.length === 0) {
+    commands.push('')
+  }
 
   const annotations = getAnnotations(exec.cell)
   const { interactive, mimeType, background } = annotations
@@ -45,12 +58,10 @@ export async function executeRunner(
   const program = await runner.createProgramSession({
     programName: getShellPath(execKey),
     environment,
-    script: {
-      type: 'script', script
+    exec: {
+      type: 'commands', commands
     },
-    envs: [
-      `RUNME_ID=${RUNME_ID}`
-    ],
+    envs: Object.entries(envs).map(([k, v]) => `${k}=${v}`),
     cwd,
     tty: interactive
   })
@@ -66,6 +77,8 @@ export async function executeRunner(
 
       let item = new NotebookCellOutputItem(Buffer.concat(output), mime)
 
+      const script = getCmdShellSeq(cellText, PLATFORM_OS)
+
       // hacky for now, maybe inheritence is a fitting pattern
       if (script.trim().endsWith('vercel')) {
         // TODO: vercel (see `shellExecutor`)
@@ -89,14 +102,18 @@ export async function executeRunner(
       program.close()
     })
 
+    context.subscriptions.push({
+      dispose: () => program.close()
+    })
+
     await program.run()
   } else {
     const taskExecution = new Task(
       { type: 'shell', name: `Runme Task (${RUNME_ID})` },
       TaskScope.Workspace,
-      cellText.length > LABEL_LIMIT
+      (cellText.length > LABEL_LIMIT
         ? `${cellText.slice(0, LABEL_LIMIT)}...`
-        : cellText,
+        : cellText) + ` (RUNME_ID: ${RUNME_ID})`,
       'exec',
       new CustomExecution(async () => program)
     )
@@ -110,6 +127,10 @@ export async function executeRunner(
 
     const execution = await tasks.executeTask(taskExecution)
 
+    context.subscriptions.push({
+      dispose: () => execution.terminate()
+    })
+
     exec.token.onCancellationRequested(() => {
       try {
         // runs `program.close()` implicitly
@@ -163,5 +184,16 @@ export async function executeRunner(
       // unexpected early return, likely an error
       resolve(false)
     }
+
+    if(background && interactive) {
+      setTimeout(
+        () => {
+          closeTerminalByEnvID(RUNME_ID)
+          resolve(true)
+        },
+        BACKGROUND_TASK_HIDE_TIMEOUT
+      )
+    }
   })
 }
+
diff --git a/src/extension/executors/task.ts b/src/extension/executors/task.ts
index 755c7c72..c909189a 100644
--- a/src/extension/executors/task.ts
+++ b/src/extension/executors/task.ts
@@ -2,12 +2,12 @@ import path from 'node:path'
 
 import {
   Task, TextDocument, NotebookCellExecution, TaskScope, tasks,
-  window, TerminalOptions, TaskRevealKind, TaskPanelKind,
+  window, TaskRevealKind, TaskPanelKind,
   ShellExecution
 } from 'vscode'
 
 // import { ExperimentalTerminal } from "../terminal"
-import { getCmdShellSeq, getAnnotations } from '../utils'
+import { getCmdShellSeq, getAnnotations, getTerminalRunmeId } from '../utils'
 import { PLATFORM_OS, ENV_STORE } from '../constants'
 import type { Kernel } from '../kernel'
 
@@ -18,7 +18,7 @@ const BACKGROUND_TASK_HIDE_TIMEOUT = 2000
 const LABEL_LIMIT = 15
 
 export function closeTerminalByEnvID (id: string) {
-  const terminal = window.terminals.find((t) => (t.creationOptions as TerminalOptions).env?.RUNME_ID === id)
+  const terminal = window.terminals.find(t => getTerminalRunmeId(t) === id)
   if (terminal) {
     terminal.hide()
   }
diff --git a/src/extension/executors/utils.ts b/src/extension/executors/utils.ts
index 10c88790..bc1f291d 100644
--- a/src/extension/executors/utils.ts
+++ b/src/extension/executors/utils.ts
@@ -6,7 +6,7 @@ import { NotebookCellOutput, NotebookCellExecution, NotebookCellOutputItem, wind
 import { ENV_STORE } from '../constants'
 import { OutputType } from '../../constants'
 import type { CellOutputPayload } from '../../types'
-import { replaceOutput } from '../utils'
+import { replaceOutput, getCmdSeq } from '../utils'
 
 const ENV_VAR_REGEXP = /(\$\w+)/g
 /**
@@ -35,6 +35,69 @@ export function populateEnvVar (value: string, env = process.env) {
   return value
 }
 
+export interface CommandExportExtractMatch {
+  type: 'exec'|'prompt'|'direct'
+  key: string
+  value: string
+  match: string
+  regexpMatch: RegExpExecArray
+  hasStringValue: boolean
+}
+
+export async function promptUserForVariable(
+  key: string,
+  placeHolder: string,
+  hasStringValue: boolean,
+): Promise<string | undefined> {
+  return await window.showInputBox({
+    title: `Set Environment Variable "${key}"`,
+    ignoreFocusOut: true,
+    placeHolder,
+    prompt: 'Your shell script wants to set some environment variables, please enter them here.',
+    ...(hasStringValue ? { value: placeHolder } : {})
+  })
+}
+
+export function getCommandExportExtractMatches(
+  rawText: string
+): CommandExportExtractMatch[] {
+  const test = new RegExp(EXPORT_EXTRACT_REGEX)
+
+  const matchStr = (rawText.endsWith('\n') ? rawText : `${rawText}\n`)
+
+  let match: RegExpExecArray | null
+
+  const result: CommandExportExtractMatch[] = []
+
+  while ((match = test.exec(matchStr)) !== null) {
+    const e = match[0]
+
+    const [key, ph] = e.trim().slice('export '.length).split('=')
+    const hasStringValue = ph.startsWith('"') || ph.startsWith('\'')
+    const placeHolder = hasStringValue ? ph.slice(1, -1) : ph
+
+    let matchType: CommandExportExtractMatch['type']
+    let value = placeHolder
+
+    if (placeHolder.startsWith('$(') && placeHolder.endsWith(')')) {
+      matchType = 'exec'
+      value = placeHolder.slice(2, -1)
+    } else if (!placeHolder.includes('\n')) {
+      matchType = 'prompt'
+    } else {
+      matchType = 'direct'
+    }
+
+    result.push({
+      type: matchType,
+      regexpMatch: match,
+      key, value, match: e, hasStringValue,
+    })
+  }
+
+  return result
+}
+
 /**
  * Helper method to parse the shell code and runs the following operations:
  *   - fetches environment variable exports and puts them into ENV_STORE
@@ -47,21 +110,17 @@ export async function retrieveShellCommand (exec: NotebookCellExecution) {
   let cellText = exec.cell.document.getText()
   const cwd = path.dirname(exec.cell.document.uri.fsPath)
   const rawText = exec.cell.document.getText()
-  const exportMatches: string[] = (rawText.endsWith('\n') ? rawText : `${rawText}\n`)
-    .match(EXPORT_EXTRACT_REGEX) || []
+
+  const exportMatches = getCommandExportExtractMatches(rawText)
 
   const stateEnv = Object.fromEntries(ENV_STORE)
-  for (const e of exportMatches) {
-    const [key, ph] = e.trim().slice('export '.length).split('=')
-    const hasStringValue = ph.startsWith('"') || ph.startsWith('\'')
-    const placeHolder = hasStringValue ? ph.slice(1, -1) : ph
 
-    if (placeHolder.startsWith('$(') && placeHolder.endsWith(')')) {
+  for (const { hasStringValue, key, match, type, value } of exportMatches) {
+    if (type === 'exec') {
       /**
        * evaluate expression
        */
-      const expression = placeHolder.slice(2, -1)
-      const expressionProcess = cp.spawn(expression, {
+      const expressionProcess = cp.spawn(value, {
         cwd,
         env: {...process.env, ...stateEnv },
         shell: true
@@ -81,32 +140,29 @@ export async function retrieveShellCommand (exec: NotebookCellExecution) {
       })
 
       if (isError) {
-        window.showErrorMessage(`Failed to evaluate expression "${expression}": ${data}`)
+        window.showErrorMessage(`Failed to evaluate expression "${value}": ${data}`)
         return undefined
       }
 
       stateEnv[key] = data
-    } else if (!placeHolder.includes('\n')) {
+    } else if (type === 'prompt') {
       /**
        * ask user for value only if placeholder has no new line as this would be absorbed by
        * VS Code, see https://github.com/microsoft/vscode/issues/98098
        */
-      stateEnv[key] = populateEnvVar(await window.showInputBox({
-        title: `Set Environment Variable "${key}"`,
-        ignoreFocusOut: true,
-        placeHolder,
-        prompt: 'Your shell script wants to set some environment variables, please enter them here.',
-        ...(hasStringValue ? { value: placeHolder } : {})
-      }) || '', {...process.env, ...stateEnv })
+      stateEnv[key] = populateEnvVar(
+        await promptUserForVariable(key, value, hasStringValue) ?? '',
+        {...process.env, ...stateEnv }
+      )
     } else {
-      stateEnv[key] = populateEnvVar(placeHolder)
+      stateEnv[key] = populateEnvVar(value)
     }
 
     /**
      * we don't want to run these exports anymore as we already stored
      * them in our extension state
      */
-    cellText = cellText.replace(e, '')
+    cellText = cellText.replace(match, '')
 
     /**
      * persist env variable in memory
@@ -126,3 +182,65 @@ export function getShellPath(execKey: 'bash'|'sh'): string
 export function getShellPath(execKey?: string): string|undefined {
   return process.env.SHELL ?? execKey
 }
+
+/**
+ * Parse set of commands, requiring user input for prompted environment
+ * variables, and supporting multiline strings
+ *
+ * Returns `undefined` when a user cancels on prompt
+ */
+export async function parseCommandSeq(
+  cellText: string
+): Promise<string[]|undefined> {
+  const exportMatches = getCommandExportExtractMatches(cellText)
+
+  type CommandBlock =
+    |{
+      type: 'block'
+      content: string
+    }
+    |{
+      type: 'single'
+      content: string
+    }
+
+  const parsedCommandBlocks: CommandBlock[] = []
+
+  let offset = 0
+
+  for (const { hasStringValue, key, match, type, value, regexpMatch } of exportMatches) {
+    let userValue: string
+
+    switch(type) {
+      case 'prompt': {
+        const userInput = await promptUserForVariable(key, value, hasStringValue)
+
+        if(userInput === undefined) {
+          return undefined
+        }
+
+        userValue = userInput
+      } break
+
+      case 'direct': {
+        userValue = value
+      } break
+
+      default: {
+        continue
+      }
+    }
+
+    const prior = cellText.slice(offset, regexpMatch.index)
+    parsedCommandBlocks.push({ type: 'block', content: prior })
+
+    parsedCommandBlocks.push({ type: 'single', content: `export ${key}="${userValue}"` })
+
+    offset = regexpMatch.index + match.length
+  }
+
+  parsedCommandBlocks.push({ type: 'block', content: cellText.slice(offset) })
+
+  return parsedCommandBlocks
+    .flatMap(({ type, content }) => type === 'block' ? getCmdSeq(content) : [content])
+}
diff --git a/src/extension/kernel.ts b/src/extension/kernel.ts
index 25a40168..ec9f053a 100644
--- a/src/extension/kernel.ts
+++ b/src/extension/kernel.ts
@@ -265,6 +265,7 @@ export class Kernel implements Disposable {
       (execKey === 'bash' || execKey === 'sh')
     ) {
       successfulCellExecution = await executeRunner(
+        this.context,
         this.runner,
         exec,
         runningCell,
diff --git a/src/extension/provider/background.ts b/src/extension/provider/background.ts
index 6c33306f..bfbc6c19 100644
--- a/src/extension/provider/background.ts
+++ b/src/extension/provider/background.ts
@@ -29,8 +29,17 @@ export class ShowTerminalProvider implements vscode.NotebookCellStatusBarItemPro
       return
     }
 
+    const terminalButtonParts = [
+      '$(terminal)',
+      'Open Terminal',
+    ]
+
+    if (pid > -1) {
+      terminalButtonParts.push(`(PID: ${pid})`)
+    }
+
     const item = new NotebookCellStatusBarItem(
-      `$(terminal) Open Terminal (PID: ${pid})`,
+      terminalButtonParts.join(' '),
       NotebookCellStatusBarAlignment.Right
     )
     item.command = 'runme.openTerminal'
diff --git a/src/extension/runner.ts b/src/extension/runner.ts
index 6df37f7d..cbe4da2f 100644
--- a/src/extension/runner.ts
+++ b/src/extension/runner.ts
@@ -28,7 +28,7 @@ export interface RunProgramOptions {
   args?: string[]
   cwd?: string
   envs?: string[]
-  script?:
+  exec?:
     |{
       type: 'commands'
       commands: string[]
@@ -369,7 +369,7 @@ export class GrpcRunnerProgramSession implements IRunnerProgramSession {
   }
 
   static runOptionsToExecuteRequest(
-    { programName, args, cwd, environment, script, tty, envs }: RunProgramOptions
+    { programName, args, cwd, environment, exec, tty, envs }: RunProgramOptions
   ): ExecuteRequest {
     if(environment && !(environment instanceof GrpcRunnerEnvironment)) {
       throw new Error('Expected gRPC environment!')
@@ -382,8 +382,8 @@ export class GrpcRunnerProgramSession implements IRunnerProgramSession {
       tty,
       sessionId: environment?.getSessionId(),
       programName,
-      ...script?.type === 'commands' && { commands: script.commands },
-      ...script?.type === 'script' && { script: script.script }
+      ...exec?.type === 'commands' && { commands: exec.commands },
+      ...exec?.type === 'script' && { script: exec.script }
     })
   }
 }
diff --git a/src/extension/utils.ts b/src/extension/utils.ts
index b6fcf341..8e3ab03f 100644
--- a/src/extension/utils.ts
+++ b/src/extension/utils.ts
@@ -87,10 +87,15 @@ export function validateAnnotations(cell: NotebookCell): CellAnnotationsErrorRes
 
 }
 
+export function getTerminalRunmeId(t: vscode.Terminal): string|undefined {
+  return (t.creationOptions as vscode.TerminalOptions).env?.RUNME_ID
+    ?? /\(RUNME_ID: (.*)\)$/.exec(t.name)?.[1]
+    ?? undefined
+}
+
 export function getTerminalByCell(cell: vscode.NotebookCell) {
   return vscode.window.terminals.find((t) => {
-    const taskEnv = (t.creationOptions as vscode.TerminalOptions).env || {}
-    return taskEnv.RUNME_ID === `${cell.document.fileName}:${cell.index}`
+    return getTerminalRunmeId(t) === `${cell.document.fileName}:${cell.index}`
   })
 }
 
@@ -114,8 +119,8 @@ export function getKey(runningCell: vscode.TextDocument): keyof typeof executor
  * treat cells like like a series of individual commands
  * which need to be executed in sequence
  */
-export function getCmdShellSeq(cellText: string, os: string): string {
-  const trimmed = cellText
+export function getCmdSeq(cellText: string): string[] {
+  return cellText
     .trimStart()
     .split('\\\n')
     .map((l) => l.trim())
@@ -138,6 +143,16 @@ export function getCmdShellSeq(cellText: string, os: string): string {
       const hasPrefix = (l.match(HASH_PREFIX_REGEXP) || []).length > 0
       return l !== '' && !hasPrefix
     })
+}
+
+/**
+ * treat cells like like a series of individual commands
+ * which need to be executed in sequence
+ *
+ * packages command sequence into single callable script
+ */
+export function getCmdShellSeq(cellText: string, os: string): string {
+  const trimmed = getCmdSeq(cellText)
     .join('; ')
 
   if (['darwin'].find((entry) => entry === os)) {


Test patch:
diff --git a/tests/extension/executors/utils.test.ts b/tests/extension/executors/utils.test.ts
index 15f54c37..189744d7 100644
--- a/tests/extension/executors/utils.test.ts
+++ b/tests/extension/executors/utils.test.ts
@@ -1,13 +1,15 @@
 import url from 'node:url'
 
 import { window } from 'vscode'
-import { expect, vi, test, beforeEach } from 'vitest'
+import { expect, vi, test, suite, beforeEach } from 'vitest'
 
 import { ENV_STORE } from '../../../src/extension/constants'
-import { retrieveShellCommand } from '../../../src/extension/executors/utils'
+import { retrieveShellCommand, parseCommandSeq } from '../../../src/extension/executors/utils'
 
 const __dirname = url.fileURLToPath(new URL('.', import.meta.url))
 
+vi.mock('vscode-telemetry', () => ({}))
+
 vi.mock('vscode', () => ({
   window: {
     showInputBox: vi.fn(),
@@ -20,6 +22,8 @@ vi.mock('vscode', () => ({
 
 vi.mock('../../../src/extension/utils', () => ({
   replaceOutput: vi.fn(),
+  // TODO: this should use importActual
+  getCmdSeq: vi.fn((cellText: string) => cellText.trim().split('\n').filter(x => x))
 }))
 
 beforeEach(() => {
@@ -190,3 +194,67 @@ test('supports multiline exports', async () => {
   expect(ENV_STORE.get('foobar')).toBe('barfoo')
   expect(window.showInputBox).toBeCalledTimes(2)
 })
+
+suite('parseCommandSeq', () => {
+  test('single-line export', async () => {
+    vi.mocked(window.showInputBox).mockImplementationOnce(async () => 'test value')
+
+    const res = await parseCommandSeq([
+      'export TEST="<placeholder>"'
+    ].join('\n'))
+
+    expect(res).toBeTruthy
+    expect(res).toHaveLength(1)
+    expect(res?.[0]).toBe('export TEST="test value"')
+  })
+
+  test('multiline export', async () => {
+    const exportLine = [
+      'export TEST="I',
+      'am',
+      'doing',
+      'well!"'
+    ].join('\n')
+
+    const res = await parseCommandSeq(exportLine)
+
+    expect(res).toBeTruthy
+    expect(res).toHaveLength(1)
+    expect(res?.[0]).toBe(exportLine)
+  })
+
+  test('exports between normal command sequences', async () => {
+    vi.mocked(window.showInputBox).mockImplementationOnce(async () => 'test value')
+
+    const cmd = [
+      'echo "Hello!"',
+      'echo "Hi!"',
+      'export TEST="<placeholder>"',
+      'echo $TEST',
+      'export TEST_MULTILINE="This',
+      'is',
+      'a',
+      'multiline',
+      'env!"',
+      'echo $TEST_MULTILINE'
+    ].join('\n')
+
+    const res = await parseCommandSeq(cmd)
+
+    expect(res).toBeTruthy()
+    expect(res).toStrictEqual([
+      'echo "Hello!"',
+      'echo "Hi!"',
+      'export TEST="test value"',
+      'echo $TEST',
+      [
+        'export TEST_MULTILINE="This',
+        'is',
+        'a',
+        'multiline',
+        'env!"',
+      ].join('\n'),
+      'echo $TEST_MULTILINE'
+    ])
+  })
+})