Repository: stateful/vscode-runme
Pull request #425 (2023-04-18 17:10:28)

CodeLens support
Requires stateful/runme#218 to work properly

Still not done:

 - ~~User setting to toggle CodeLens support~~
 - ~~Only apply to markdown files~~
 - ~~Unit tests~~

sourishkrout: @mxsdev did you publish the protos to the BSR? This does not seem to work for me.
mxsdev: @sourishkrout 

My bad! Just pushed latest protos.

Surprised CI passed, I would've expected it to fail on type errors.
sourishkrout: No sweat. This feature is fantastic!

I suppose the generate types are not used anywhere on the client side? Otherwise, tsc should have blown up. Perhaps something to check into?
mxsdev: Hmm, the new types are definitely used client side (in `codelens.ts`, we access the `textRange` field, which is a new type)

I'll look into this tomorrow 👍 
sourishkrout: Good to merge with typing clean up as separate effort 👍 .

Patch:
diff --git a/__mocks__/vscode.ts b/__mocks__/vscode.ts
index 92675298..7b9cfe97 100644
--- a/__mocks__/vscode.ts
+++ b/__mocks__/vscode.ts
@@ -4,7 +4,7 @@ import { join } from 'node:path'
 
 import { vi } from 'vitest'
 import { URI } from 'vscode-uri'
-import type { Disposable } from 'vscode'
+import type { Command, Disposable } from 'vscode'
 
 
 export const notebooks = {
@@ -31,6 +31,7 @@ export class Uri extends URI {
 
 export const workspace = {
   getConfiguration: vi.fn().mockReturnValue(new Map()),
+  onDidChangeConfiguration: vi.fn(),
   openNotebookDocument: vi.fn().mockReturnValue({ uri: 'new notebook uri' }),
   openTextDocument: vi.fn().mockReturnValue({
     getText: vi.fn().mockReturnValue(readFileSync(join(__dirname, 'gitignore.mock'), 'utf8'))
@@ -97,7 +98,8 @@ export const window = {
 export const tasks = {
   registerTaskProvider: vi.fn(),
   onDidEndTaskProcess: vi.fn(),
-  onDidStartTaskProcess: vi.fn()
+  onDidStartTaskProcess: vi.fn(),
+  executeTask: vi.fn()
 }
 
 export const commands = {
@@ -224,3 +226,35 @@ export class EventEmitter<T> {
 
   dispose() { }
 }
+
+export const languages = {
+  registerCodeLensProvider: vi.fn(),
+}
+
+export class Position {
+  constructor(
+    public line: number,
+    public character: number
+  ) { }
+
+  with(line?: number, character?: number): Position {
+    return new Position(
+      line ?? this.line,
+      character ?? this.character,
+    )
+  }
+}
+
+export class Range {
+  constructor(
+    public start: Position,
+    public end: Position,
+  ) { }
+}
+
+export class CodeLens {
+  constructor(
+    public range: Range,
+    public command?: Command
+  ) { }
+}
diff --git a/package.json b/package.json
index 54553a6d..978ce040 100644
--- a/package.json
+++ b/package.json
@@ -135,6 +135,10 @@
         "title": "Try Runme",
         "category": "Runme",
         "description": "Walk through some interactive notebook examples."
+      },
+      {
+        "command": "runme.codelens.action",
+        "title": "Runme CodeLens Action"
       }
     ],
     "menus": {
@@ -166,6 +170,10 @@
         {
           "command": "runme.expandTreeView",
           "when": "false"
+        },
+        {
+          "command": "runme.codelens.action",
+          "when": "false"
         }
       ],
       "editor/title": [
@@ -351,6 +359,11 @@
             "type": "integer",
             "default": 10,
             "markdownDescription": "Default number of rows in notebook terminal"
+          },
+          "runme.codelens.enable": {
+            "type": "boolean",
+            "default": true,
+            "markdownDescription": "Enable CodeLens support in markdown documents"
           }
         }
       }
diff --git a/src/extension/commands/index.ts b/src/extension/commands/index.ts
index e16f3e8a..1a84fbd0 100644
--- a/src/extension/commands/index.ts
+++ b/src/extension/commands/index.ts
@@ -77,7 +77,7 @@ export function runCLICommand(
   server: RunmeServer,
   kernel: Kernel
 ) {
- return async function(cell: NotebookCell) {
+ return async function(cell: { metadata?: any, document: TextDocument }) {
     const cwd = path.dirname(cell.document.uri.fsPath)
 
     const args = [
@@ -104,7 +104,7 @@ export function runCLICommand(
       }
     }
 
-    const annotations = getAnnotations(cell)
+    const annotations = getAnnotations(cell.metadata)
     const term = window.createTerminal({
       name: `CLI: ${annotations.name}`,
       cwd,
diff --git a/src/extension/extension.ts b/src/extension/extension.ts
index 392db476..34809cd1 100644
--- a/src/extension/extension.ts
+++ b/src/extension/extension.ts
@@ -28,6 +28,7 @@ import { BOOTFILE } from './constants'
 import { GrpcRunner, IRunner } from './runner'
 import { CliProvider } from './provider/cli'
 import * as survey from './survey'
+import { RunmeCodeLensProvider } from './provider/codelens'
 
 export class RunmeExtension {
   async initialize(context: ExtensionContext) {
@@ -68,6 +69,10 @@ export class RunmeExtension {
     const winSurvey = new survey.WinDefaultShell(context)
     const stopBackgroundTaskProvider = new StopBackgroundTaskProvider()
 
+    const runCLI = runCLICommand(context.extensionUri, !!grpcRunner, server, kernel)
+
+    const codeLensProvider = new RunmeCodeLensProvider(serializer, runCLI, runner, kernel)
+
     context.subscriptions.push(
       kernel,
       serializer,
@@ -91,12 +96,14 @@ export class RunmeExtension {
 
       stopBackgroundTaskProvider,
 
+      codeLensProvider,
+
       commands.registerCommand('runme.resetEnv', resetEnv),
       RunmeExtension.registerCommand('runme.openTerminal', openTerminal(kernel, !!grpcRunner)),
       RunmeExtension.registerCommand('runme.openIntegratedTerminal', openTerminal(kernel, false)),
       RunmeExtension.registerCommand(
         'runme.runCliCommand',
-        runCLICommand(context.extensionUri, !!grpcRunner, server, kernel)
+        runCLI
       ),
       RunmeExtension.registerCommand('runme.copyCellToClipboard', copyCellToClipboard),
       RunmeExtension.registerCommand('runme.stopBackgroundTask', stopBackgroundTask),
diff --git a/src/extension/grpc/genericTypes.ts b/src/extension/grpc/genericTypes.ts
new file mode 100644
index 00000000..93931718
--- /dev/null
+++ b/src/extension/grpc/genericTypes.ts
@@ -0,0 +1,1 @@
+export { UInt32Value } from '@buf/stateful_runme.community_timostamm-protobuf-ts/google/protobuf/wrappers_pb'
diff --git a/src/extension/provider/codelens.ts b/src/extension/provider/codelens.ts
new file mode 100644
index 00000000..dc4841ac
--- /dev/null
+++ b/src/extension/provider/codelens.ts
@@ -0,0 +1,184 @@
+import {
+  CancellationToken,
+  CodeLens,
+  CodeLensProvider,
+  Disposable,
+  Event,
+  EventEmitter,
+  NotebookCellKind,
+  TextDocument,
+  Range,
+  languages,
+  commands,
+  workspace,
+  tasks
+} from 'vscode'
+
+import { SerializerBase } from '../serializer'
+import type { runCLICommand } from '../commands'
+import { IRunner } from '../runner'
+import { Kernel } from '../kernel'
+import { getAnnotations } from '../utils'
+import { Serializer } from '../../types'
+import { getCodeLensEnabled } from '../../utils/configuration'
+
+import { RunmeTaskProvider } from './runmeTask'
+
+export const ActionCommand = 'runme.codelens.action' as const
+
+const ActionTypes = [
+  'run',
+  'open',
+] as const satisfies readonly string[]
+
+type ActionType = (typeof ActionTypes)[number]
+
+type ActionArguments = [
+  document: TextDocument,
+  token: CancellationToken,
+  cell: Serializer.Cell,
+  index: number,
+  action: ActionType
+]
+
+type ActionCallback = (...arg: ActionArguments) => void
+
+export class RunmeCodeLensProvider implements CodeLensProvider, Disposable {
+  private disposables: Disposable[] = []
+
+  private _onDidChangeCodeLenses: EventEmitter<void> = this.register(new EventEmitter<void>())
+	public readonly onDidChangeCodeLenses: Event<void> = this._onDidChangeCodeLenses.event
+
+  constructor(
+    protected serializer: SerializerBase,
+    protected runCLI: ReturnType<typeof runCLICommand>,
+    protected runner?: IRunner,
+    protected kernel?: Kernel
+  ) {
+    this.register(
+      languages.registerCodeLensProvider(
+        [
+          { language: 'markdown', scheme: 'file' },
+          { language: 'mdx', scheme: 'file' },
+        ]
+      , this)
+    )
+
+    const cmd: ActionCallback = this.codeLensActionCallback.bind(this)
+
+    this.register(
+      commands.registerCommand('runme.codelens.action', cmd),
+    )
+
+    this.register(
+      workspace.onDidChangeConfiguration(() =>
+        this._onDidChangeCodeLenses.fire()
+      )
+    )
+  }
+
+  // eslint-disable-next-line @typescript-eslint/no-unused-vars
+  async provideCodeLenses(document: TextDocument, token: CancellationToken): Promise<CodeLens[]> {
+    if (!this.runner || !getCodeLensEnabled()) {
+      return []
+    }
+
+    const contentBytes = Buffer.from(document.getText())
+    const { cells } = await this.serializer['reviveNotebook'](contentBytes, token)
+
+    return cells.flatMap((cell, i) => {
+      if (cell.kind !== NotebookCellKind.Code || !cell.textRange) { return [] }
+
+      let start = document.positionAt(cell.textRange.start)
+      let end = document.positionAt(cell.textRange.end)
+
+      start = start.with(start.line - 1)
+      end = end.with(end.line - 1)
+
+      const range = new Range(
+        start, end
+      )
+
+      return ActionTypes.map((v) => {
+        const args: ActionArguments = [document, token, cell, i, v]
+
+        /* c8 ignore start */
+        switch (v) {
+          case 'run': {
+            return new CodeLens(range, {
+                title: '$(play) Run Block',
+                command: ActionCommand,
+                arguments: args,
+              })
+          } break
+
+          case 'open': {
+            return new CodeLens(range, {
+                title: '$(notebook) Open in Notebook',
+                command: ActionCommand,
+                arguments: args,
+              })
+          } break
+        }
+        /* c8 ignore stop */
+      })
+    })
+  }
+
+  // eslint-disable-next-line @typescript-eslint/no-unused-vars
+  async resolveCodeLens(codeLens: CodeLens, token: CancellationToken): Promise<CodeLens> {
+    return codeLens
+  }
+
+  protected register<T extends Disposable>(d: T) {
+    this.disposables.push(d)
+    return d
+  }
+
+  dispose() {
+    this.disposables.forEach(({ dispose }) => dispose())
+  }
+
+  protected async codeLensActionCallback(
+    document: TextDocument,
+    token: CancellationToken,
+    cell: Serializer.Cell,
+    index: number,
+    action: ActionType
+  ) {
+    switch (action) {
+      case 'open': {
+        await commands.executeCommand('vscode.openWith', document.uri, Kernel.type)
+
+        // TODO(mxs): surely there's a better way to do this
+        // probably we need to bring this logic to `workspace.onDidOpenNotebookDocument`
+        await new Promise(cb => setTimeout(cb, 200))
+
+        await commands.executeCommand('notebook.focusTop')
+
+        await Promise.all(
+          Array.from(
+            {length: index}, () => commands.executeCommand('notebook.focusNextEditor')
+          )
+        )
+
+        // to execute the command:
+        // await commands.executeCommand('notebook.cell.execute')
+        // await commands.executeCommand('notebook.cell.focusInOutput')
+      } break
+
+      case 'run': {
+        const task = await RunmeTaskProvider.getRunmeTask(
+          document.uri.fsPath,
+          getAnnotations(cell.metadata).name,
+          cell,
+          {},
+          this.runner!,
+          this.kernel?.getRunnerEnvironment()
+        )
+
+        await tasks.executeTask(task)
+      } break
+    }
+  }
+}
diff --git a/src/extension/provider/runmeTask.ts b/src/extension/provider/runmeTask.ts
index 27c7055c..9b5e6853 100644
--- a/src/extension/provider/runmeTask.ts
+++ b/src/extension/provider/runmeTask.ts
@@ -18,11 +18,11 @@ import {
   NotebookCellData,
 } from 'vscode'
 
-import { getAnnotations, processEnviron } from '../utils'
+import { getAnnotations, prepareCmdSeq, processEnviron } from '../utils'
 import { Serializer, RunmeTaskDefinition } from '../../types'
 import { SerializerBase } from '../serializer'
 import type { IRunner, IRunnerEnvironment, RunProgramOptions } from '../runner'
-import { getShellPath } from '../executors/utils'
+import { getShellPath, parseCommandSeq } from '../executors/utils'
 import { Kernel } from '../kernel'
 
 type TaskOptions = Pick<RunmeTaskDefinition, 'closeTerminalOnSuccess' | 'isBackground' | 'cwd'>
@@ -92,7 +92,7 @@ export class RunmeTaskProvider implements TaskProvider {
   static async getRunmeTask (
     filePath: string,
     command: string,
-    cell: NotebookCell|NotebookCellData,
+    cell: NotebookCell|NotebookCellData|Serializer.Cell,
     options: TaskOptions = {},
     runner: IRunner,
     environment?: IRunnerEnvironment,
@@ -106,11 +106,14 @@ export class RunmeTaskProvider implements TaskProvider {
     const cwd = options.cwd || path.dirname(filePath)
     const isBackground = options.isBackground || background
 
+    const cellContent = 'value' in cell ? cell.value : cell.document.getText()
+    const commands = await parseCommandSeq(cellContent, prepareCmdSeq)
+
     const runOpts: RunProgramOptions = {
       programName: getShellPath() ?? 'sh',
       exec: {
-        type: 'script',
-        script: 'value' in cell ? cell.value : cell.document.getText(),
+        type: 'commands',
+        commands: commands ?? [''],
       },
       cwd,
       environment,
@@ -130,6 +133,10 @@ export class RunmeTaskProvider implements TaskProvider {
       source,
       new CustomExecution(async () => {
         const program = await runner.createProgramSession(runOpts)
+
+        program.registerTerminalWindow('vscode')
+        program.setActiveTerminalWindow('vscode')
+
         return program
       })
     )
diff --git a/src/types.ts b/src/types.ts
index a3040e37..3aa35907 100644
--- a/src/types.ts
+++ b/src/types.ts
@@ -11,16 +11,19 @@ export namespace Serializer {
     metadata?: Metadata
   }
 
+  type TextPosition = number
+
+  export type TextRange = {
+    start: TextPosition
+    end: TextPosition
+  }
+
   export type Cell = {
     metadata?: Metadata
     languageId?: string
     value: string
-    kind: NotebookCellKind.Markup
-  } | {
-    metadata?: Metadata
-    languageId?: string
-    value: string
-    kind: NotebookCellKind.Code
+    textRange?: Serializer.TextRange
+    kind: NotebookCellKind.Markup | NotebookCellKind.Code
   }
 
   export interface Wasm {
diff --git a/src/utils/configuration.ts b/src/utils/configuration.ts
index 785cbf27..da32fe44 100644
--- a/src/utils/configuration.ts
+++ b/src/utils/configuration.ts
@@ -9,6 +9,7 @@ import { getAnnotations, isWindows } from '../extension/utils'
 
 const SERVER_SECTION_NAME = 'runme.server'
 const TERMINAL_SECTION_NAME= 'runme.terminal'
+const CODELENS_SECTION_NAME= 'runme.codelens'
 export const DEFAULT_TLS_DIR = path.join(os.tmpdir(), 'runme', uuidv4(), 'tls')
 
 type NotebookTerminalValue = keyof typeof configurationSchema.notebookTerminal
@@ -42,6 +43,9 @@ const configurationSchema = {
         fontSize: z.number().optional(),
         fontFamily: z.string().optional(),
         rows: z.number().int()
+    },
+    codelens: {
+      enable: z.boolean().default(true)
     }
 }
 
@@ -65,6 +69,16 @@ const getRunmeTerminalConfigurationValue = <T>(configName: NotebookTerminalValue
 return defaultValue
 }
 
+const getCodeLensConfigurationValue = <T>(configName: keyof typeof configurationSchema.codelens, defaultValue: T) => {
+  const configurationSection = workspace.getConfiguration(CODELENS_SECTION_NAME)
+  const configurationValue = configurationSection.get<T>(configName)!
+  const parseResult = configurationSchema.codelens[configName].safeParse(configurationValue)
+  if (parseResult.success) {
+      return parseResult.data as T
+  }
+  return defaultValue
+}
+
 const getPortNumber = (): number => {
     return getServerConfigurationValue<number>('port', 7863)
 }
@@ -130,6 +144,10 @@ const getNotebookTerminalRows = (): number => {
   return getRunmeTerminalConfigurationValue<number>('rows', 10)
 }
 
+const getCodeLensEnabled = (): boolean => {
+  return getCodeLensConfigurationValue<boolean>('enable', true)
+}
+
 export {
     getPortNumber,
     getBinaryPath,
@@ -142,4 +160,5 @@ export {
     getNotebookTerminalFontFamily,
     getNotebookTerminalFontSize,
     getNotebookTerminalRows,
+    getCodeLensEnabled,
 }


Test patch:
diff --git a/tests/extension/configuration.test.ts b/tests/extension/configuration.test.ts
index ed85e6a8..6e8ec074 100644
--- a/tests/extension/configuration.test.ts
+++ b/tests/extension/configuration.test.ts
@@ -11,7 +11,8 @@ import {
   getTLSDir,
   DEFAULT_TLS_DIR,
   getNotebookTerminalFontFamily,
-  getNotebookTerminalFontSize
+  getNotebookTerminalFontSize,
+  getCodeLensEnabled
 } from '../../src/utils/configuration'
 
 vi.mock('../../src/extension/grpc/client', () => ({}))
@@ -142,6 +143,12 @@ suite('Configuration', () => {
       ).toStrictEqual('/binary/path')
     })
 
+    test('getCodeLensEnabled should return true by default', () => {
+      expect(
+        getCodeLensEnabled()
+      ).toStrictEqual(true)
+    })
+
     suite('posix', () => {
       beforeEach(() => {
         platformPathMocks(path.posix)
diff --git a/tests/extension/extension.test.ts b/tests/extension/extension.test.ts
index c3950442..a7cee88b 100644
--- a/tests/extension/extension.test.ts
+++ b/tests/extension/extension.test.ts
@@ -54,7 +54,7 @@ test('initializes all providers', async () => {
   await ext.initialize(context)
   expect(notebooks.registerNotebookCellStatusBarItemProvider).toBeCalledTimes(6)
   expect(workspace.registerNotebookSerializer).toBeCalledTimes(1)
-  expect(commands.registerCommand).toBeCalledTimes(18)
+  expect(commands.registerCommand).toBeCalledTimes(19)
   expect(window.registerTreeDataProvider).toBeCalledTimes(1)
   expect(window.registerUriHandler).toBeCalledTimes(1)
 
diff --git a/tests/extension/runner.test.ts b/tests/extension/runner.test.ts
index ffc811d2..fc6b9311 100644
--- a/tests/extension/runner.test.ts
+++ b/tests/extension/runner.test.ts
@@ -2,7 +2,7 @@ import path from 'node:path'
 
 import { vi, suite, test, expect, beforeEach } from 'vitest'
 import { type GrpcTransport } from '@protobuf-ts/grpc-transport'
-import { EventEmitter } from 'vscode'
+import { EventEmitter, NotebookCellKind, Position, tasks, commands } from 'vscode'
 
 import {
   GrpcRunner,
@@ -11,16 +11,20 @@ import {
   RunProgramOptions
 } from '../../src/extension/runner'
 import type { ExecuteResponse } from '../../src/extension/grpc/runnerTypes'
-
+import { ActionCommand, RunmeCodeLensProvider } from '../../src/extension/provider/codelens'
+import { RunmeTaskProvider } from '../../src/extension/provider/runmeTask'
 
 vi.mock('../../src/extension/utils', () => ({
-  getGrpcHost: vi.fn().mockReturnValue('127.0.0.1:7863')
+  getGrpcHost: vi.fn().mockReturnValue('127.0.0.1:7863'),
+  getAnnotations: vi.fn().mockReturnValue({ })
 }))
 
 vi.mock('vscode', async () => ({
   ...await import(path.join(process.cwd(), '__mocks__', 'vscode')) ,
 }))
 
+vi.mock('vscode-telemetry')
+
 vi.mock('@protobuf-ts/grpc-transport', () => ({
   GrpcTransport: class {
     constructor() {}
@@ -105,6 +109,14 @@ class MockedRunmeServer {
   onClose = this._onClose.event
 }
 
+vi.mock('../../src/extension/provider/runmeTask', async () => {
+  return {
+    RunmeTaskProvider: {
+      getRunmeTask: vi.fn().mockResolvedValue({}),
+    }
+  }
+})
+
 beforeEach(() => {
   resetId()
   deleteSession.mockClear()
@@ -458,6 +470,89 @@ suite('grpc Runner', () => {
   })
 })
 
+suite('RunmeCodeLensProvider', () => {
+  test('returns nothing without runner', async () => {
+    const provider = new RunmeCodeLensProvider({} as any, vi.fn())
+
+    const lenses = await provider.provideCodeLenses({} as any, {} as any)
+
+    expect(lenses).toStrictEqual([])
+  })
+
+  test('returns serializer result normally', async () => {
+    const serializer = {
+      reviveNotebook: vi.fn(async () => ({
+        cells: [
+          {
+            kind: NotebookCellKind.Code,
+            textRange: {
+              start: 2,
+              end: 3,
+            }
+          }
+        ]
+      }))
+    } as any
+
+    const provider = new RunmeCodeLensProvider(
+      serializer,
+      vi.fn(),
+      {} as any
+    )
+
+    const lenses = await provider.provideCodeLenses(
+      {
+        getText: vi.fn().mockReturnValue(''),
+        positionAt: vi.fn((c) => new Position(1, c)),
+      } as any,
+      {} as any
+    )
+
+    expect(serializer.reviveNotebook).toBeCalledTimes(1)
+
+    expect(lenses).toHaveLength(2)
+
+    for (const lens of lenses) {
+      expect(lens.range.start.character).toStrictEqual(2)
+      expect(lens.range.end.character).toStrictEqual(3)
+
+      expect(lens.command).toBeTruthy()
+
+      expect(lens.command!.command).toStrictEqual(ActionCommand)
+
+      expect(lens.command!.tooltip).toBeUndefined()
+    }
+  })
+
+  test('action callback for run command', async () => {
+    const provider = new RunmeCodeLensProvider(
+      {} as any,
+      vi.fn(),
+      {} as any
+    )
+
+    await provider['codeLensActionCallback']({uri: { fsPath: '' }} as any, {} as any, {} as any, 0, 'run')
+
+    expect(RunmeTaskProvider.getRunmeTask).toBeCalledTimes(1)
+
+    expect(tasks.executeTask).toBeCalledTimes(1)
+    expect(tasks.executeTask).toBeCalledWith({})
+  })
+
+  test('action callback for open command', async () => {
+    const provider = new RunmeCodeLensProvider(
+      {} as any,
+      vi.fn(),
+      {} as any
+    )
+
+    const document = { uri: { fsPath: '' } }
+    await provider['codeLensActionCallback'](document as any, {} as any, {} as any, 0, 'open')
+
+    expect(commands.executeCommand).toBeCalledWith('vscode.openWith', document.uri, 'runme')
+  })
+})
+
 function getMockedDuplex(session: GrpcRunnerProgramSession): MockedDuplexClientStream {
   return session['session'] as unknown as MockedDuplexClientStream
 }