├── .vscode ├── settings.json ├── extensions.json ├── tasks.json └── launch.json ├── src ├── utils │ ├── TypedString.ts │ ├── types.ts │ ├── wait.ts │ ├── async.ts │ ├── errors.ts │ ├── arrays.ts │ ├── debounce.ts │ ├── invariant.ts │ ├── deepEqual.ts │ ├── versions.ts │ ├── strings.ts │ ├── retry.ts │ ├── deferred.ts │ ├── single-promise.ts │ ├── network.ts │ ├── log.ts │ ├── show.ts │ ├── query.ts │ ├── __tests__ │ │ ├── strings.test.ts │ │ ├── deepEqual.test.ts │ │ ├── url.test.ts │ │ ├── deferred.test.ts │ │ ├── cmd.test.ts │ │ ├── config.test.ts │ │ └── exec.test.ts │ ├── url.ts │ ├── cmd.ts │ ├── DeferredRequestRegistry.ts │ ├── python.ts │ └── exec.ts ├── notebook │ ├── constants.ts │ ├── createMarimoNotebookController.ts │ ├── common │ │ ├── key.ts │ │ └── metadata.ts │ ├── md.ts │ ├── serializer.ts │ ├── marimo │ │ └── types.ts │ ├── extension.ts │ ├── __tests__ │ │ └── md.test.ts │ └── kernel-manager.ts ├── __fixtures__ │ ├── convert │ │ ├── mock.md │ │ └── mock.ipynb │ ├── export │ │ └── app.py │ └── mocks.ts ├── ctx.ts ├── ts-reset.d.ts ├── types.ts ├── services │ ├── context-manager.ts │ ├── health.ts │ └── health.test.ts ├── convert │ ├── __tests__ │ │ └── convert.test.ts │ └── convert.ts ├── export │ ├── export-as-commands.ts │ ├── __tests__ │ │ └── export-as.test.ts │ └── export-as.ts ├── launcher │ ├── start.ts │ ├── new-file.ts │ ├── utils.ts │ ├── __tests__ │ │ └── utils.test.ts │ ├── terminal.ts │ └── controller.ts ├── commands │ ├── tutorial-commands.ts │ ├── __tests__ │ │ └── show-commands.test.ts │ └── show-commands.ts ├── constants.ts ├── logger.ts ├── telemetry.ts ├── browser │ ├── __tests__ │ │ └── panel.test.ts │ └── panel.ts ├── __mocks__ │ └── vscode.ts ├── config.ts ├── __integration__ │ ├── exec.test.ts │ └── server-manager.test.ts ├── ui │ └── status-bar.ts └── extension.ts ├── resources ├── icon.png ├── marimo.png └── icon.svg ├── images └── screenshot.png ├── .vscode-test.js ├── vitest.config.ts ├── .gitignore ├── ISSUE_TEMPLATE ├── config.yml └── bug_report.yaml ├── tsup.config.ts ├── testing ├── app.py └── what.ipynb ├── .github ├── workflows │ ├── release.yml │ └── ci.yml └── renovate.json ├── tsconfig.json ├── biome.json ├── CONTRIBUTING.md ├── LICENSE └── README.md /.vscode/settings.json: -------------------------------------------------------------------------------- 1 | { 2 | "editor.defaultFormatter": "biomejs.biome" 3 | } 4 | -------------------------------------------------------------------------------- /src/utils/TypedString.ts: -------------------------------------------------------------------------------- 1 | export type TypedString = string & { __type__: T }; 2 | -------------------------------------------------------------------------------- /resources/icon.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/marimo-team/vscode-marimo/HEAD/resources/icon.png -------------------------------------------------------------------------------- /images/screenshot.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/marimo-team/vscode-marimo/HEAD/images/screenshot.png -------------------------------------------------------------------------------- /resources/marimo.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/marimo-team/vscode-marimo/HEAD/resources/marimo.png -------------------------------------------------------------------------------- /src/utils/types.ts: -------------------------------------------------------------------------------- 1 | export type DeepPartial = { 2 | [P in keyof T]?: DeepPartial; 3 | }; 4 | -------------------------------------------------------------------------------- /.vscode/extensions.json: -------------------------------------------------------------------------------- 1 | { 2 | "recommendations": ["dbaeumer.vscode-eslint", "marimo-team.vscode-marimo"] 3 | } 4 | -------------------------------------------------------------------------------- /src/utils/wait.ts: -------------------------------------------------------------------------------- 1 | export function wait(ms: number) { 2 | return new Promise((resolve) => setTimeout(resolve, ms)); 3 | } 4 | -------------------------------------------------------------------------------- /src/utils/async.ts: -------------------------------------------------------------------------------- 1 | export const sleep = (ms: number) => { 2 | return new Promise((resolve) => setTimeout(resolve, ms)); 3 | }; 4 | -------------------------------------------------------------------------------- /.vscode-test.js: -------------------------------------------------------------------------------- 1 | const { defineConfig } = require("@vscode/test-cli"); 2 | 3 | module.exports = defineConfig({ files: "tests/**/*.test.js" }); 4 | -------------------------------------------------------------------------------- /vitest.config.ts: -------------------------------------------------------------------------------- 1 | import { defineConfig } from "vitest/config"; 2 | 3 | export default defineConfig({ 4 | test: { 5 | // ... 6 | }, 7 | }); 8 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | .cache 2 | .DS_Store 3 | .idea 4 | *.log 5 | *.tgz 6 | *.vsix 7 | coverage 8 | dist 9 | lib-cov 10 | logs 11 | node_modules 12 | temp 13 | !.vscode 14 | 15 | # python 16 | __pycache__/ 17 | -------------------------------------------------------------------------------- /ISSUE_TEMPLATE/config.yml: -------------------------------------------------------------------------------- 1 | blank_issues_enabled: true 2 | contact_links: 3 | - name: Discord Chat 4 | url: https://marimo.io/discord?ref=issues 5 | about: Ask questions and discuss with other marimo users. 6 | -------------------------------------------------------------------------------- /tsup.config.ts: -------------------------------------------------------------------------------- 1 | import { defineConfig } from "tsup"; 2 | 3 | export default defineConfig({ 4 | entry: ["src/index.ts"], 5 | format: ["cjs"], 6 | shims: false, 7 | dts: false, 8 | external: ["vscode"], 9 | }); 10 | -------------------------------------------------------------------------------- /src/notebook/constants.ts: -------------------------------------------------------------------------------- 1 | export const NOTEBOOK_TYPE = "marimo-notebook"; 2 | export const NOTEBOOK_CONTROLLER_ID = "marimo-kernel"; 3 | export const PYTHON_LANGUAGE_ID = "python"; 4 | export const MARKDOWN_LANGUAGE_ID = "markdown"; 5 | -------------------------------------------------------------------------------- /src/__fixtures__/convert/mock.md: -------------------------------------------------------------------------------- 1 | --- 2 | title: Markdown 3 | marimo-version: 0.4.11 4 | --- 5 | 6 | # marimo in Markdown 7 | 8 | Everything in marimo is pure Python, but sometimes, annotating your notebook 9 | in markdown can be a little cumbersome. 10 | -------------------------------------------------------------------------------- /src/utils/errors.ts: -------------------------------------------------------------------------------- 1 | export function printError(error: unknown) { 2 | if (error instanceof Error) { 3 | return error.message; 4 | } 5 | 6 | if (typeof error === "string") { 7 | return error; 8 | } 9 | 10 | try { 11 | return String(error); 12 | } catch { 13 | return "Unknown error"; 14 | } 15 | } 16 | -------------------------------------------------------------------------------- /src/utils/arrays.ts: -------------------------------------------------------------------------------- 1 | export function toArray(value: T | ReadonlyArray): T[] { 2 | if (Array.isArray(value)) { 3 | return [...value]; 4 | } 5 | if (value == null) { 6 | return []; 7 | } 8 | return [value] as T[]; 9 | } 10 | 11 | export function unique(arr: T[]): T[] { 12 | return Array.from(new Set(arr)); 13 | } 14 | -------------------------------------------------------------------------------- /src/ctx.ts: -------------------------------------------------------------------------------- 1 | import type { ExtensionContext } from "vscode"; 2 | 3 | export let extension: ExtensionContext; 4 | 5 | export function setExtension(ext: ExtensionContext) { 6 | extension = ext; 7 | } 8 | 9 | export function getGlobalState() { 10 | if (!extension) { 11 | throw new Error("Extension not set"); 12 | } 13 | return extension.globalState; 14 | } 15 | -------------------------------------------------------------------------------- /.vscode/tasks.json: -------------------------------------------------------------------------------- 1 | { 2 | "version": "2.0.0", 3 | "tasks": [ 4 | { 5 | "type": "npm", 6 | "script": "build", 7 | "group": { 8 | "kind": "build", 9 | "isDefault": true 10 | }, 11 | "problemMatcher": [], 12 | "label": "npm: build", 13 | "detail": "tsup src/extension.ts --external=vscode" 14 | } 15 | ] 16 | } 17 | -------------------------------------------------------------------------------- /testing/app.py: -------------------------------------------------------------------------------- 1 | import marimo 2 | 3 | __generated_with = "0.1.8" 4 | app = marimo.App() 5 | 6 | 7 | @app.cell 8 | def __(): 9 | import marimo as mo 10 | return mo, 11 | 12 | 13 | @app.cell 14 | def __(): 15 | x = 2 + 2 16 | return x, 17 | 18 | 19 | @app.cell 20 | def __(x): 21 | x 22 | return 23 | 24 | 25 | if __name__ == "__main__": 26 | app.run() 27 | -------------------------------------------------------------------------------- /src/utils/debounce.ts: -------------------------------------------------------------------------------- 1 | export function debounce any>( 2 | func: T, 3 | wait: number, 4 | ): (...args: Parameters) => void { 5 | let timeout: NodeJS.Timeout | null = null; 6 | return (...args: Parameters) => { 7 | if (timeout) { 8 | clearTimeout(timeout); 9 | } 10 | timeout = setTimeout(() => func(...args), wait); 11 | }; 12 | } 13 | -------------------------------------------------------------------------------- /src/notebook/createMarimoNotebookController.ts: -------------------------------------------------------------------------------- 1 | import * as vscode from "vscode"; 2 | import { NOTEBOOK_CONTROLLER_ID, NOTEBOOK_TYPE } from "./constants"; 3 | 4 | export function createNotebookController() { 5 | const controller = vscode.notebooks.createNotebookController( 6 | NOTEBOOK_CONTROLLER_ID, 7 | NOTEBOOK_TYPE, 8 | "marimo kernel", 9 | ); 10 | return controller; 11 | } 12 | -------------------------------------------------------------------------------- /src/ts-reset.d.ts: -------------------------------------------------------------------------------- 1 | interface Body { 2 | json(): Promise; 3 | } 4 | 5 | interface JSON { 6 | parse( 7 | text: string, 8 | // biome-ignore lint/suspicious/noExplicitAny: any is ok 9 | reviver?: (this: any, key: string, value: any) => any, 10 | ): unknown; 11 | } 12 | 13 | interface Array { 14 | filter(predicate: BooleanConstructor): Array>; 15 | } 16 | -------------------------------------------------------------------------------- /src/__fixtures__/export/app.py: -------------------------------------------------------------------------------- 1 | import marimo 2 | 3 | __generated_with = "0.1.8" 4 | app = marimo.App() 5 | 6 | 7 | @app.cell 8 | def __(): 9 | import marimo as mo 10 | 11 | return (mo,) 12 | 13 | 14 | @app.cell 15 | def __(): 16 | x = 2 + 2 17 | return (x,) 18 | 19 | 20 | @app.cell 21 | def __(x): 22 | x 23 | return 24 | 25 | 26 | if __name__ == "__main__": 27 | app.run() 28 | -------------------------------------------------------------------------------- /src/notebook/common/key.ts: -------------------------------------------------------------------------------- 1 | import type { Uri } from "vscode"; 2 | import type { TypedString } from "../../utils/TypedString"; 3 | import { SessionId } from "../marimo/types"; 4 | 5 | export type KernelKey = TypedString<"KernelKey">; 6 | 7 | export function toKernelKey(uriOrNew: Uri | `__new__`): KernelKey { 8 | if (typeof uriOrNew === "string") { 9 | return `__new__${SessionId.create()}` as KernelKey; 10 | } 11 | return uriOrNew.fsPath as KernelKey; 12 | } 13 | -------------------------------------------------------------------------------- /.github/workflows/release.yml: -------------------------------------------------------------------------------- 1 | name: Release 2 | 3 | permissions: 4 | contents: write 5 | 6 | on: 7 | push: 8 | tags: 9 | - "v*" 10 | 11 | jobs: 12 | release: 13 | runs-on: ubuntu-latest 14 | steps: 15 | - uses: actions/checkout@v4 16 | with: 17 | fetch-depth: 0 18 | 19 | - uses: actions/setup-node@v4 20 | with: 21 | node-version: 22.15.0 22 | 23 | - run: npx changelogithub 24 | env: 25 | GITHUB_TOKEN: ${{secrets.GITHUB_TOKEN}} 26 | -------------------------------------------------------------------------------- /src/types.ts: -------------------------------------------------------------------------------- 1 | import type { Disposable } from "vscode"; 2 | import type { MarimoConfig, SkewToken } from "./notebook/marimo/types"; 3 | 4 | export interface ILifecycle extends Disposable { 5 | start?(opts: Options): Promise; 6 | restart?(opts: Options): Promise; 7 | } 8 | 9 | export type ServerStatus = "stopped" | "starting" | "started"; 10 | 11 | export interface StartupResult { 12 | port: number; 13 | skewToken: SkewToken; 14 | version: string; 15 | userConfig: MarimoConfig; 16 | } 17 | -------------------------------------------------------------------------------- /src/utils/invariant.ts: -------------------------------------------------------------------------------- 1 | export function invariant( 2 | condition: boolean, 3 | message: string, 4 | ): asserts condition; 5 | export function invariant( 6 | condition: T, 7 | message: string, 8 | ): asserts condition is NonNullable; 9 | export function invariant( 10 | condition: boolean, 11 | message: string, 12 | ): asserts condition { 13 | if (!condition) { 14 | throw new Error(message); 15 | } 16 | } 17 | 18 | export function logNever(obj: never): never { 19 | throw new Error("Unexpected object", obj); 20 | } 21 | -------------------------------------------------------------------------------- /tsconfig.json: -------------------------------------------------------------------------------- 1 | { 2 | "compilerOptions": { 3 | "target": "es2020", 4 | "module": "esnext", 5 | "outDir": "./dist", 6 | "lib": ["esnext"], 7 | "moduleResolution": "node", 8 | "esModuleInterop": true, 9 | "strict": true, 10 | "strictNullChecks": true, 11 | "resolveJsonModule": true, 12 | "skipLibCheck": true, 13 | "skipDefaultLibCheck": true, 14 | "experimentalDecorators": true, 15 | "noImplicitReturns": true, 16 | "noFallthroughCasesInSwitch": true, 17 | "noUnusedParameters": true 18 | } 19 | } 20 | -------------------------------------------------------------------------------- /src/services/context-manager.ts: -------------------------------------------------------------------------------- 1 | import { commands } from "vscode"; 2 | 3 | export class VscodeContextManager { 4 | private setContext( 5 | context: string, 6 | value: string | boolean | undefined | null, 7 | ) { 8 | void commands.executeCommand("setContext", context, value); 9 | } 10 | 11 | public setMarimoServerRunning(value: true | false | "null") { 12 | this.setContext("marimo.isMarimoServerRunning", value); 13 | } 14 | 15 | public setMarimoApp(value: boolean) { 16 | this.setContext("marimo.isMarimoApp", value); 17 | } 18 | } 19 | -------------------------------------------------------------------------------- /src/utils/deepEqual.ts: -------------------------------------------------------------------------------- 1 | // biome-ignore lint/suspicious/noExplicitAny: any is ok 2 | export function deepEqual(a: any, b: any): boolean { 3 | if (a === b) { 4 | return true; 5 | } 6 | if (typeof a !== "object" || typeof b !== "object") { 7 | return false; 8 | } 9 | if (Object.keys(a).length !== Object.keys(b).length) { 10 | return false; 11 | } 12 | for (const key in a) { 13 | if (!(key in b)) { 14 | return false; 15 | } 16 | if (!deepEqual(a[key], b[key])) { 17 | return false; 18 | } 19 | } 20 | return true; 21 | } 22 | -------------------------------------------------------------------------------- /biome.json: -------------------------------------------------------------------------------- 1 | { 2 | "$schema": "https://biomejs.dev/schemas/1.7.3/schema.json", 3 | "files": { 4 | "ignore": ["dist"] 5 | }, 6 | "organizeImports": { 7 | "enabled": true 8 | }, 9 | "linter": { 10 | "enabled": true, 11 | "rules": { 12 | "recommended": true, 13 | "complexity": { 14 | "noForEach": { 15 | "level": "off" 16 | } 17 | }, 18 | "suspicious": { 19 | "noExplicitAny": "warn" 20 | } 21 | } 22 | }, 23 | "formatter": { 24 | "enabled": true, 25 | "indentStyle": "space", 26 | "indentWidth": 2 27 | } 28 | } 29 | -------------------------------------------------------------------------------- /src/utils/versions.ts: -------------------------------------------------------------------------------- 1 | import { extensions } from "vscode"; 2 | import { version as vscodeVersion } from "vscode"; 3 | import { EXTENSION_PACKAGE } from "../constants"; 4 | import { logger } from "../logger"; 5 | 6 | export function getExtensionVersion(): string { 7 | try { 8 | const extension = extensions.getExtension(EXTENSION_PACKAGE.fullName); 9 | return extension?.packageJSON.version || "unknown"; 10 | } catch (error) { 11 | logger.error("Error getting extension version:", error); 12 | return "unknown"; 13 | } 14 | } 15 | 16 | export function getVscodeVersion(): string { 17 | return vscodeVersion; 18 | } 19 | -------------------------------------------------------------------------------- /CONTRIBUTING.md: -------------------------------------------------------------------------------- 1 | # Contributing 2 | 3 | ## Development 4 | 5 | - Run `pnpm install` in the root directory to install dependencies. 6 | - Press `F5` to open a new window with your extension loaded. 7 | 8 | ### Make changes 9 | 10 | - Relaunch the extension from the debug toolbar after changing code in `src/extension.ts`. 11 | - `Ctrl+R` or `Cmd+R` on Mac will reload the extension with your changes. 12 | 13 | ## Release (only for maintainers) 14 | 15 | - Run `pnpm run release` to create a new release. 16 | - This requires a valid vsce token 17 | - Run `pnpm run openvsx:publish ` to publish the release to OpenVSX. 18 | - This requires a valid OpenVSX token 19 | -------------------------------------------------------------------------------- /src/utils/strings.ts: -------------------------------------------------------------------------------- 1 | export const Strings = { 2 | indent: (str: string, indent: string) => { 3 | if (str.length === 0) { 4 | return str; 5 | } 6 | return str 7 | .split("\n") 8 | .map((line) => `${indent}${line}`) 9 | .join("\n"); 10 | }, 11 | FOUR_SPACES: " ", 12 | dedent: (str: string) => { 13 | const match = str.match(/^[ \t]*(?=\S)/gm); 14 | if (!match) { 15 | return str; // If no indentation, return original string 16 | } 17 | const minIndent = Math.min(...match.map((el) => el.length)); 18 | const re = new RegExp(`^[ \t]{${minIndent}}`, "gm"); 19 | return str.replace(re, ""); 20 | }, 21 | }; 22 | -------------------------------------------------------------------------------- /src/utils/retry.ts: -------------------------------------------------------------------------------- 1 | /** 2 | * Retry a function until it succeeds or the number of retries is reached. 3 | * @param fn - The function to retry. 4 | * @param retries - The number of retries. 5 | * @param backoffMs - The backoff time in milliseconds. 6 | * @returns A promise that resolves when the function succeeds or the number of retries is reached. 7 | */ 8 | export function retry( 9 | fn: () => Promise, 10 | retries = 3, 11 | backoffMs = 100, 12 | ): Promise { 13 | return fn().catch(async (error) => { 14 | if (retries === 0) { 15 | throw error; 16 | } 17 | await new Promise((resolve) => setTimeout(resolve, backoffMs)); 18 | return retry(fn, retries - 1, backoffMs * 2); 19 | }); 20 | } 21 | -------------------------------------------------------------------------------- /.vscode/launch.json: -------------------------------------------------------------------------------- 1 | { 2 | "version": "0.2.0", 3 | "configurations": [ 4 | { 5 | "name": "Run Extension", 6 | "type": "extensionHost", 7 | "request": "launch", 8 | "args": ["--extensionDevelopmentPath=${workspaceFolder}"], 9 | "outFiles": ["${workspaceFolder}/dist/**/*.js"], 10 | "preLaunchTask": "npm: build" 11 | }, 12 | { 13 | "name": "Extension Tests", 14 | "type": "extensionHost", 15 | "request": "launch", 16 | "args": [ 17 | "--extensionDevelopmentPath=${workspaceFolder}", 18 | "--extensionTestsPath=${workspaceFolder}/out/test/suite/index" 19 | ], 20 | "outFiles": ["${workspaceFolder}/out/test/**/*.js"], 21 | "preLaunchTask": "${defaultBuildTask}" 22 | } 23 | ] 24 | } 25 | -------------------------------------------------------------------------------- /src/utils/deferred.ts: -------------------------------------------------------------------------------- 1 | export class Deferred { 2 | promise: Promise; 3 | hasCompleted = false; 4 | hasResolved = false; 5 | hasRejected = false; 6 | resolve!: (value: T | PromiseLike) => void; 7 | reject!: (reason?: unknown) => void; 8 | 9 | constructor( 10 | private opts: { 11 | onRejected?: (reason?: unknown) => void; 12 | onResolved?: (value: T | PromiseLike) => void; 13 | } = {}, 14 | ) { 15 | this.promise = new Promise((resolve, reject) => { 16 | this.reject = (reason) => { 17 | this.hasRejected = true; 18 | this.hasCompleted = true; 19 | this.opts.onRejected?.(reason); 20 | reject(reason); 21 | }; 22 | this.resolve = (value) => { 23 | this.hasResolved = true; 24 | this.hasCompleted = true; 25 | this.opts.onResolved?.(value); 26 | resolve(value); 27 | }; 28 | }); 29 | } 30 | } 31 | -------------------------------------------------------------------------------- /src/utils/single-promise.ts: -------------------------------------------------------------------------------- 1 | import { Deferred } from "./deferred"; 2 | 3 | type Channel = string; 4 | 5 | /** 6 | * Ensures that only one message is being processed at a time for a given channel. 7 | * All queued messages will get the same result. 8 | */ 9 | export class SingleMessage { 10 | public static instance = new SingleMessage(); 11 | private constructor() {} 12 | 13 | private gates: Map> = new Map(); 14 | 15 | async gate(channel: string, promise: () => Promise): Promise { 16 | // If gate exists, wait for it 17 | const deferred = this.gates.get(channel); 18 | if (deferred) { 19 | return deferred.promise; 20 | } 21 | // Create a new gate 22 | const newDeferred = new Deferred(); 23 | this.gates.set(channel, newDeferred); 24 | 25 | try { 26 | await promise(); 27 | } finally { 28 | // Remove the gate 29 | this.gates.delete(channel); 30 | newDeferred.resolve(); 31 | } 32 | } 33 | } 34 | -------------------------------------------------------------------------------- /.github/renovate.json: -------------------------------------------------------------------------------- 1 | { 2 | "$schema": "https://docs.renovatebot.com/renovate-schema.json", 3 | "extends": ["config:base", "schedule:weekly"], 4 | "labels": ["dependencies"], 5 | "pin": false, 6 | "rangeStrategy": "bump", 7 | "packageRules": [ 8 | { 9 | "depTypeList": ["peerDependencies"], 10 | "enabled": false 11 | }, 12 | { 13 | "packageNames": ["node"], 14 | "enabled": false 15 | }, 16 | { 17 | "packageNames": ["vscode", "@types/vscode"], 18 | "enabled": false 19 | }, 20 | { 21 | "matchUpdateTypes": ["minor", "patch"], 22 | "groupName": "all non-major dependencies", 23 | "groupSlug": "all-minor-patch", 24 | "excludePackagePatterns": [".*eslint.*"], 25 | "excludePackageNames": ["typescript", "@types/vscode", "vscode"] 26 | }, 27 | { 28 | "packagePatterns": [".*eslint.*"], 29 | "groupName": "all eslint dependencies", 30 | "groupSlug": "all-eslint", 31 | "matchUpdateTypes": ["minor", "patch", "major"] 32 | } 33 | ] 34 | } 35 | -------------------------------------------------------------------------------- /src/notebook/md.ts: -------------------------------------------------------------------------------- 1 | import { Strings } from "../utils/strings"; 2 | 3 | export function toMarkdown(text: string): string { 4 | // Trim 5 | const value = text.trim(); 6 | 7 | const isMultiline = value.includes("\n"); 8 | if (!isMultiline) { 9 | return `mo.md(r"${value}")`; 10 | } 11 | 12 | return `mo.md(\n${Strings.FOUR_SPACES}r"""\n${Strings.indent(value, Strings.FOUR_SPACES)}\n${Strings.FOUR_SPACES}"""\n)`; 13 | } 14 | 15 | export function maybeMarkdown(text: string): string | null { 16 | // TODO: Python can safely extract the string value with the 17 | // AST, anything done here is a bit of a hack, data should come from server. 18 | const value = text.trim(); 19 | // Regular expression to match the function calls 20 | const regex = /^mo\.md\(\s*r?((["'])(?:\2\2)?)(.*?)\1\s*\)$/gms; // 'g' flag to check all occurrences 21 | const matches = [...value.matchAll(regex)]; 22 | 23 | // Check if there is exactly one match 24 | if (matches.length === 1) { 25 | const extractedString = matches[0][3]; // Extract the string content 26 | return Strings.dedent(extractedString).trim(); 27 | } 28 | return null; 29 | } 30 | -------------------------------------------------------------------------------- /src/utils/network.ts: -------------------------------------------------------------------------------- 1 | import http from "node:http"; 2 | import https from "node:https"; 3 | import { composeUrl } from "../config"; 4 | 5 | /** 6 | * Check if a port is free 7 | */ 8 | async function isPortFree(port: number) { 9 | const healthy = await ping(await composeUrl(port)); 10 | return !healthy; 11 | } 12 | 13 | export async function tryPort(start: number): Promise { 14 | if (await isPortFree(start)) { 15 | return start; 16 | } 17 | return tryPort(start + 1); 18 | } 19 | 20 | /** 21 | * Ping a url to see if it is healthy 22 | */ 23 | export function ping(url: string): Promise { 24 | const promise = new Promise((resolve) => { 25 | const useHttps = url.indexOf("https") === 0; 26 | const module_ = useHttps ? https.request : http.request; 27 | 28 | const pingRequest = module_(url, () => { 29 | resolve(true); 30 | pingRequest.destroy(); 31 | }); 32 | 33 | pingRequest.on("error", () => { 34 | resolve(false); 35 | pingRequest.destroy(); 36 | }); 37 | 38 | pingRequest.write(""); 39 | pingRequest.end(); 40 | }); 41 | return promise; 42 | } 43 | -------------------------------------------------------------------------------- /src/utils/log.ts: -------------------------------------------------------------------------------- 1 | import { logger as defaultLogger } from "../logger"; 2 | import { invariant } from "./invariant"; 3 | 4 | /** 5 | * Decorator that logs method calls. 6 | */ 7 | // biome-ignore lint/suspicious/noExplicitAny: any is ok 8 | export function LogMethodCalls any>() { 9 | return ( 10 | _target: object, 11 | propertyKey: string | symbol, 12 | descriptor: TypedPropertyDescriptor, 13 | ): TypedPropertyDescriptor => { 14 | const originalMethod = descriptor.value; 15 | invariant(originalMethod, "Method not found"); 16 | 17 | // @ts-expect-error ignore 18 | descriptor.value = function (...args: Parameters) { 19 | // biome-ignore lint/suspicious/noExplicitAny: any is ok 20 | const logger = (this as any).logger; 21 | if (logger && "debug" in logger) { 22 | logger.debug(`-> ${String(propertyKey)}()`); 23 | } else { 24 | defaultLogger.debug( 25 | `[${this.constructor.name}] -> ${String(propertyKey)}()`, 26 | ); 27 | } 28 | 29 | return originalMethod.apply(this, args); 30 | }; 31 | 32 | return descriptor; 33 | }; 34 | } 35 | -------------------------------------------------------------------------------- /src/convert/__tests__/convert.test.ts: -------------------------------------------------------------------------------- 1 | import { describe, expect, it, vi } from "vitest"; 2 | 3 | import { createVSCodeMock } from "../../__mocks__/vscode"; 4 | 5 | vi.mock("vscode", () => createVSCodeMock(vi)); 6 | vi.mock("@vscode/python-extension", () => ({})); 7 | 8 | import { unlink } from "node:fs/promises"; 9 | import type { Uri } from "vscode"; 10 | import { ipynbFixtureUri, markdownFixtureUri } from "../../__fixtures__/mocks"; 11 | import { convertIPyNotebook } from "../convert"; 12 | 13 | describe("convert", () => { 14 | it("should convertIPyNotebook", async () => { 15 | const newUri = await convertIPyNotebook(ipynbFixtureUri.fsPath); 16 | expect(newUri).toBeDefined(); 17 | // Clean up 18 | expect((newUri as Uri).fsPath).toMatch(/\.py$/); 19 | // Clean up 20 | await unlink((newUri as Uri).fsPath); 21 | }); 22 | 23 | it("should export as ipynb", async () => { 24 | const newUri = await convertIPyNotebook(markdownFixtureUri.fsPath); 25 | expect(newUri).toBeDefined(); 26 | // Clean up 27 | expect((newUri as Uri).fsPath).toMatch(/\.py$/); 28 | // Clean up 29 | await unlink((newUri as Uri).fsPath); 30 | }); 31 | }); 32 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2023 marimo 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /src/utils/show.ts: -------------------------------------------------------------------------------- 1 | import { 2 | type NotebookDocument, 3 | type NotebookEditor, 4 | commands, 5 | window, 6 | } from "vscode"; 7 | import { logger } from "../logger"; 8 | 9 | export async function showNotebookDocument( 10 | document: NotebookDocument, 11 | ): Promise { 12 | // If existing editor is found, focus it 13 | const existingEditor = window.visibleNotebookEditors.find( 14 | (editor) => editor.notebook === document, 15 | ); 16 | try { 17 | if (existingEditor) { 18 | logger.info("Focusing an existing marimo notebook"); 19 | } else { 20 | logger.info("Showing a new marimo notebook"); 21 | } 22 | return await window.showNotebookDocument(document, { 23 | viewColumn: existingEditor?.viewColumn, 24 | }); 25 | } catch { 26 | // Do nothing 27 | } 28 | 29 | return undefined; 30 | } 31 | 32 | export async function closeNotebookEditor(editor: NotebookEditor) { 33 | // Focus the editor 34 | await window.showNotebookDocument(editor.notebook, { 35 | preview: false, 36 | viewColumn: editor.viewColumn, 37 | }); 38 | // Close the editor 39 | await commands.executeCommand("workbench.action.closeActiveEditor"); 40 | } 41 | -------------------------------------------------------------------------------- /src/utils/query.ts: -------------------------------------------------------------------------------- 1 | import { type TextDocument, window } from "vscode"; 2 | 3 | export function isMarimoApp( 4 | document: TextDocument | undefined, 5 | includeEmpty = true, 6 | ) { 7 | if (!document || !["python", "markdown"].includes(document.languageId)) { 8 | return false; 9 | } 10 | 11 | // If it's empty, return true. 12 | // This is so we can create a new file and start the server. 13 | if (includeEmpty && document.getText().trim() === "") { 14 | return true; 15 | } 16 | 17 | // Cheap way of checking if it's a marimo app 18 | const fileName = document.fileName; 19 | const text = document.getText(); 20 | return ( 21 | (text.includes("app = marimo.App(") && fileName.endsWith(".py")) || 22 | (text.includes("marimo-version") && fileName.endsWith(".md")) 23 | ); 24 | } 25 | 26 | /** 27 | * Get the current text editor if it's a marimo file 28 | */ 29 | export function getFocusedMarimoTextEditor({ 30 | toast = true, 31 | }: { toast?: boolean } = {}) { 32 | const editor = [window.activeTextEditor] 33 | .filter(Boolean) 34 | .find((editor) => isMarimoApp(editor.document, false)); 35 | if (!editor) { 36 | if (toast) { 37 | window.showInformationMessage("No marimo file is open."); 38 | } 39 | return; 40 | } 41 | return editor; 42 | } 43 | -------------------------------------------------------------------------------- /src/export/export-as-commands.ts: -------------------------------------------------------------------------------- 1 | import { type QuickPickItem, type Uri, window } from "vscode"; 2 | import { exportNotebookAs } from "./export-as"; 3 | 4 | interface CommandPickItem extends QuickPickItem { 5 | handler: () => void; 6 | if?: boolean; 7 | } 8 | 9 | export async function exportAsCommands(file: Uri) { 10 | const commands: CommandPickItem[] = [ 11 | { 12 | label: "$(browser) Export as HTML", 13 | async handler() { 14 | await exportNotebookAs(file.fsPath, "html"); 15 | }, 16 | }, 17 | { 18 | label: "$(markdown) Export as Markdown", 19 | async handler() { 20 | await exportNotebookAs(file.fsPath, "md"); 21 | }, 22 | }, 23 | { 24 | label: "$(notebook) Export as Jupyter Notebook", 25 | async handler() { 26 | await exportNotebookAs(file.fsPath, "ipynb"); 27 | }, 28 | }, 29 | { 30 | label: "$(file-code) Export as a flat Python Script", 31 | async handler() { 32 | await exportNotebookAs(file.fsPath, "script"); 33 | }, 34 | }, 35 | ]; 36 | 37 | const filteredCommands = commands.filter((index) => index.if !== false); 38 | const result = await window.showQuickPick(filteredCommands); 39 | 40 | if (result) { 41 | result.handler(); 42 | } 43 | } 44 | -------------------------------------------------------------------------------- /src/utils/__tests__/strings.test.ts: -------------------------------------------------------------------------------- 1 | import { describe, expect, it } from "vitest"; 2 | import { Strings } from "../strings"; 3 | 4 | describe("Strings", () => { 5 | describe("indent", () => { 6 | it("should indent each line with given string", () => { 7 | const input = "line1\nline2\nline3"; 8 | const expected = " line1\n line2\n line3"; 9 | expect(Strings.indent(input, " ")).toBe(expected); 10 | }); 11 | 12 | it("should handle empty string", () => { 13 | expect(Strings.indent("", " ")).toBe(""); 14 | }); 15 | }); 16 | 17 | describe("dedent", () => { 18 | it("should remove common leading whitespace", () => { 19 | const input = " line1\n line2\n line3"; 20 | const expected = "line1\nline2\nline3"; 21 | expect(Strings.dedent(input)).toBe(expected); 22 | }); 23 | 24 | it("should handle mixed indentation", () => { 25 | const input = " line1\n line2\n line3"; 26 | const expected = " line1\nline2\n line3"; 27 | expect(Strings.dedent(input)).toBe(expected); 28 | }); 29 | 30 | it("should handle empty string", () => { 31 | expect(Strings.dedent("")).toBe(""); 32 | }); 33 | }); 34 | 35 | it("should have FOUR_SPACES constant", () => { 36 | expect(Strings.FOUR_SPACES).toBe(" "); 37 | }); 38 | }); 39 | -------------------------------------------------------------------------------- /testing/what.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "cell_type": "code", 5 | "execution_count": 3, 6 | "id": "3be16cc1", 7 | "metadata": {}, 8 | "outputs": [], 9 | "source": [ 10 | "x = 2 + 2" 11 | ] 12 | }, 13 | { 14 | "cell_type": "code", 15 | "execution_count": 4, 16 | "id": "e64015b8", 17 | "metadata": {}, 18 | "outputs": [ 19 | { 20 | "data": { 21 | "text/plain": [ 22 | "4" 23 | ] 24 | }, 25 | "execution_count": 4, 26 | "metadata": {}, 27 | "output_type": "execute_result" 28 | } 29 | ], 30 | "source": [ 31 | "x" 32 | ] 33 | } 34 | ], 35 | "metadata": { 36 | "kernelspec": { 37 | "display_name": "Python 3 (ipykernel)", 38 | "language": "python", 39 | "name": "python3" 40 | }, 41 | "language_info": { 42 | "codemirror_mode": { 43 | "name": "ipython", 44 | "version": 3 45 | }, 46 | "file_extension": ".py", 47 | "mimetype": "text/x-python", 48 | "name": "python", 49 | "nbconvert_exporter": "python", 50 | "pygments_lexer": "ipython3", 51 | "version": "3.11.4" 52 | } 53 | }, 54 | "nbformat": 4, 55 | "nbformat_minor": 5 56 | } 57 | -------------------------------------------------------------------------------- /src/utils/url.ts: -------------------------------------------------------------------------------- 1 | import { logger } from "../logger"; 2 | 3 | export function asURL(url: string): URL { 4 | try { 5 | return new URL(url); 6 | } catch (e) { 7 | logger.error("Failed to parse url", url, e); 8 | throw e; 9 | } 10 | } 11 | 12 | /** 13 | * Similar to path.join, but for URLs. 14 | * We cannot use node:path.join because it messes up the URL scheme on windows. 15 | */ 16 | export function urlJoin(...paths: string[]): string { 17 | if (paths.length === 0) { 18 | return ""; 19 | } 20 | if (paths.length === 1) { 21 | return paths[0]; 22 | } 23 | let normalized = [...paths]; 24 | 25 | // Process the first path to remove its trailing slash 26 | normalized = normalized.map((path, index) => { 27 | // first 28 | if (index === 0) { 29 | return withoutTrailingSlash(path); 30 | } 31 | // last 32 | if (index === normalized.length - 1) { 33 | return withoutLeadingSlash(path); 34 | } 35 | // middle 36 | return withoutLeadingSlash(withoutTrailingSlash(path)); 37 | }); 38 | 39 | return normalized.join("/"); 40 | } 41 | 42 | function withoutLeadingSlash(path: string): string { 43 | return path.startsWith("/") ? path.slice(1) : path; 44 | } 45 | 46 | function withoutTrailingSlash(path: string): string { 47 | return path.endsWith("/") ? path.slice(0, -1) : path; 48 | } 49 | -------------------------------------------------------------------------------- /src/__fixtures__/convert/mock.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "cell_type": "code", 5 | "execution_count": 3, 6 | "id": "3be16cc1", 7 | "metadata": {}, 8 | "outputs": [], 9 | "source": [ 10 | "x = 2 + 2" 11 | ] 12 | }, 13 | { 14 | "cell_type": "code", 15 | "execution_count": 4, 16 | "id": "e64015b8", 17 | "metadata": {}, 18 | "outputs": [ 19 | { 20 | "data": { 21 | "text/plain": [ 22 | "4" 23 | ] 24 | }, 25 | "execution_count": 4, 26 | "metadata": {}, 27 | "output_type": "execute_result" 28 | } 29 | ], 30 | "source": [ 31 | "x" 32 | ] 33 | } 34 | ], 35 | "metadata": { 36 | "kernelspec": { 37 | "display_name": "Python 3 (ipykernel)", 38 | "language": "python", 39 | "name": "python3" 40 | }, 41 | "language_info": { 42 | "codemirror_mode": { 43 | "name": "ipython", 44 | "version": 3 45 | }, 46 | "file_extension": ".py", 47 | "mimetype": "text/x-python", 48 | "name": "python", 49 | "nbconvert_exporter": "python", 50 | "pygments_lexer": "ipython3", 51 | "version": "3.11.4" 52 | } 53 | }, 54 | "nbformat": 4, 55 | "nbformat_minor": 5 56 | } 57 | -------------------------------------------------------------------------------- /ISSUE_TEMPLATE/bug_report.yaml: -------------------------------------------------------------------------------- 1 | # yaml-language-server: $schema=https://raw.githubusercontent.com/SchemaStore/schemastore/refs/heads/master/src/schemas/json/github-issue-forms.json 2 | name: '🐞 Bug report' 3 | description: Report an issue with vscode-marimo 4 | labels: [bug] 5 | type: Bug 6 | body: 7 | - type: textarea 8 | id: bug-description 9 | attributes: 10 | label: Describe the bug 11 | description: | 12 | A clear and concise description of the bug. If you have a workaround or plan to submit a PR for this issue, please let us know. 13 | You may also include a stack trace or screenshots here. 14 | placeholder: What happened? 15 | validations: 16 | required: true 17 | - type: textarea 18 | id: env 19 | attributes: 20 | label: Environment 21 | description: Open the command palette and paste the output of "Show marimo diagnostics" 22 | value: | 23 |
24 | 25 | ``` 26 | Replace this line with the output. Leave the backticks in place. 27 | ``` 28 | 29 |
30 | validations: 31 | required: true 32 | - type: textarea 33 | id: reproduction-steps 34 | attributes: 35 | label: Steps to reproduce 36 | description: | 37 | Help us help you! 38 | If possible, provide steps to reproduce the bug. 39 | If you have a workaround, please let us know. 40 | -------------------------------------------------------------------------------- /src/utils/__tests__/deepEqual.test.ts: -------------------------------------------------------------------------------- 1 | import { describe, expect, it } from "vitest"; 2 | import { deepEqual } from "../deepEqual"; 3 | 4 | describe("deepEqual", () => { 5 | it("returns true for identical primitives", () => { 6 | expect(deepEqual(1, 1)).toBe(true); 7 | expect(deepEqual("test", "test")).toBe(true); 8 | expect(deepEqual(true, true)).toBe(true); 9 | expect(deepEqual(null, null)).toBe(true); 10 | expect(deepEqual(undefined, undefined)).toBe(true); 11 | }); 12 | 13 | it("returns false for different primitives", () => { 14 | expect(deepEqual(1, 2)).toBe(false); 15 | expect(deepEqual("test", "other")).toBe(false); 16 | expect(deepEqual(true, false)).toBe(false); 17 | expect(deepEqual(null, undefined)).toBe(false); 18 | }); 19 | 20 | it("compares objects deeply", () => { 21 | expect(deepEqual({ a: 1, b: 2 }, { a: 1, b: 2 })).toBe(true); 22 | expect(deepEqual({ a: 1, b: 2 }, { b: 2, a: 1 })).toBe(true); 23 | expect(deepEqual({ a: { b: 1 } }, { a: { b: 1 } })).toBe(true); 24 | expect(deepEqual({ a: 1 }, { a: 2 })).toBe(false); 25 | expect(deepEqual({ a: 1 }, { b: 1 })).toBe(false); 26 | }); 27 | 28 | it("compares arrays deeply", () => { 29 | expect(deepEqual([1, 2, 3], [1, 2, 3])).toBe(true); 30 | expect(deepEqual([{ a: 1 }], [{ a: 1 }])).toBe(true); 31 | expect(deepEqual([1, 2], [1, 3])).toBe(false); 32 | expect(deepEqual([1, 2], [1, 2, 3])).toBe(false); 33 | }); 34 | }); 35 | -------------------------------------------------------------------------------- /src/utils/cmd.ts: -------------------------------------------------------------------------------- 1 | export class MarimoCmdBuilder { 2 | private cmd: string[] = ["--yes"]; 3 | 4 | debug(value: boolean) { 5 | if (value) { 6 | this.cmd.push("-d"); 7 | } 8 | return this; 9 | } 10 | 11 | mode(mode: "edit" | "run") { 12 | this.cmd.push(mode); 13 | return this; 14 | } 15 | 16 | fileOrDir(fileOrDir: string) { 17 | if (fileOrDir.includes(" ")) { 18 | this.cmd.push(`"${fileOrDir}"`); 19 | return this; 20 | } 21 | this.cmd.push(fileOrDir); 22 | return this; 23 | } 24 | 25 | host(host: string) { 26 | if (host) { 27 | this.cmd.push(`--host=${host}`); 28 | } 29 | 30 | return this; 31 | } 32 | 33 | port(port: number) { 34 | this.cmd.push(`--port=${port}`); 35 | return this; 36 | } 37 | 38 | headless(value: boolean) { 39 | if (value) { 40 | this.cmd.push("--headless"); 41 | } 42 | return this; 43 | } 44 | 45 | enableToken(value: boolean) { 46 | if (!value) { 47 | this.cmd.push("--no-token"); 48 | } 49 | return this; 50 | } 51 | 52 | tokenPassword(password: string | undefined) { 53 | if (password) { 54 | this.cmd.push(`--token-password=${password}`); 55 | } 56 | 57 | return this; 58 | } 59 | 60 | sandbox(value: boolean) { 61 | if (value) { 62 | this.cmd.push("--sandbox"); 63 | } 64 | return this; 65 | } 66 | 67 | watch(value: boolean) { 68 | if (value) { 69 | this.cmd.push("--watch"); 70 | } 71 | return this; 72 | } 73 | 74 | build() { 75 | return this.cmd.join(" "); 76 | } 77 | } 78 | -------------------------------------------------------------------------------- /src/launcher/start.ts: -------------------------------------------------------------------------------- 1 | import { window } from "vscode"; 2 | import { Config } from "../config"; 3 | import { tryPort } from "../utils/network"; 4 | import type { AppMode, MarimoController } from "./controller"; 5 | 6 | async function start({ 7 | controller, 8 | mode, 9 | }: { 10 | controller: MarimoController; 11 | mode: AppMode; 12 | }) { 13 | // If its already running and the mode is the same, do nothing 14 | if (controller.currentMode === mode && controller.active) { 15 | const button = "Open"; 16 | if (Config.showTerminal) { 17 | controller.terminal.show(); 18 | } 19 | window 20 | .showInformationMessage( 21 | `${controller.appName} is already running in ${mode} mode`, 22 | button, 23 | ) 24 | .then((response) => { 25 | if (response === button) { 26 | controller.open(); 27 | } 28 | }); 29 | return; 30 | } 31 | 32 | // Restart if the mode is different 33 | if (mode !== controller.currentMode && controller.active) { 34 | controller.dispose(); 35 | } 36 | 37 | const defaultPort = mode === "run" ? Config.readPort : Config.port; 38 | 39 | // Start with the current port if available 40 | // Make sure the port is free, otherwise try the next one 41 | const port = await tryPort(controller.port || defaultPort); 42 | window.showInformationMessage( 43 | `Starting ${controller.appName} in ${mode} mode (port: ${port})`, 44 | ); 45 | 46 | await controller.start(mode, port); 47 | } 48 | 49 | function stop(controller: MarimoController) { 50 | controller.dispose(); 51 | } 52 | 53 | export const Launcher = { 54 | start, 55 | stop, 56 | }; 57 | -------------------------------------------------------------------------------- /src/notebook/common/metadata.ts: -------------------------------------------------------------------------------- 1 | import * as vscode from "vscode"; 2 | import type { CellId } from "../marimo/types"; 3 | import type { KernelKey } from "./key"; 4 | 5 | interface CellMetadata { 6 | id?: CellId; 7 | name?: string; 8 | } 9 | 10 | export function getCellMetadata( 11 | cell: vscode.NotebookCell | vscode.NotebookCellData, 12 | ): CellMetadata { 13 | return cell.metadata?.custom || {}; 14 | } 15 | 16 | export async function setCellMetadata( 17 | cell: vscode.NotebookCell, 18 | metadata: CellMetadata, 19 | ): Promise { 20 | // create workspace edit to update metadata 21 | // TODO: could batch these edits 22 | const edit = new vscode.WorkspaceEdit(); 23 | const nbEdit = vscode.NotebookEdit.updateCellMetadata(cell.index, { 24 | ...cell.metadata, 25 | custom: { 26 | ...cell.metadata?.custom, 27 | ...metadata, 28 | }, 29 | }); 30 | edit.set(cell.notebook.uri, [nbEdit]); 31 | await vscode.workspace.applyEdit(edit); 32 | } 33 | 34 | export interface NotebookMetadata { 35 | port?: number; 36 | file?: string; 37 | isNew: boolean; 38 | key?: KernelKey; 39 | loaded?: boolean; 40 | } 41 | 42 | export function getNotebookMetadata( 43 | notebook: vscode.NotebookDocument | vscode.NotebookData, 44 | ): NotebookMetadata { 45 | return (notebook.metadata || {}) as NotebookMetadata; 46 | } 47 | 48 | export function setNotebookMetadata( 49 | notebook: vscode.NotebookDocument | vscode.NotebookData, 50 | metadata: NotebookMetadata, 51 | ): void { 52 | if (!notebook.metadata) { 53 | // @ts-ignore 54 | notebook.metadata = {}; 55 | } 56 | notebook.metadata.port = metadata.port; 57 | notebook.metadata.file = metadata.file; 58 | notebook.metadata.isNew = metadata.isNew; 59 | } 60 | -------------------------------------------------------------------------------- /src/commands/tutorial-commands.ts: -------------------------------------------------------------------------------- 1 | import { type QuickPickItem, ThemeIcon, Uri, env, window } from "vscode"; 2 | 3 | interface CommandPickItem extends QuickPickItem { 4 | handler: () => void; 5 | } 6 | 7 | const TUTORIALS = [ 8 | // Get started with marimo basics 9 | ["Intro", "https://links.marimo.app/tutorial-intro", "book"], 10 | // Learn how cells interact with each other 11 | ["Dataflow", "https://links.marimo.app/tutorial-dataflow", "repo-forked"], 12 | // Create interactive UI components 13 | ["UI Elements", "https://links.marimo.app/tutorial-ui", "layout"], 14 | // Format text with parameterized markdown 15 | ["Markdown", "https://links.marimo.app/tutorial-markdown", "markdown"], 16 | // Create interactive visualizations 17 | ["Plotting", "https://links.marimo.app/tutorial-plotting", "graph"], 18 | // Query databases directly in marimo 19 | ["SQL", "https://links.marimo.app/tutorial-sql", "database"], 20 | // Customize the layout of your cells' output 21 | ["Layout", "https://links.marimo.app/tutorial-layout", "layout-panel-left"], 22 | // Understand marimo's pure-Python file format 23 | ["File Format", "https://links.marimo.app/tutorial-fileformat", "file"], 24 | // Transiting from Jupyter to marimo 25 | ["Coming from Jupyter", "https://links.marimo.app/tutorial-jupyter", "code"], 26 | ]; 27 | 28 | export async function tutorialCommands() { 29 | const commands: CommandPickItem[] = TUTORIALS.map(([label, url, icon]) => ({ 30 | label, 31 | description: url, 32 | iconPath: new ThemeIcon(icon), 33 | handler: () => env.openExternal(Uri.parse(url)), 34 | })); 35 | 36 | const result = await window.showQuickPick(commands); 37 | if (result) { 38 | result.handler(); 39 | } 40 | } 41 | -------------------------------------------------------------------------------- /src/launcher/new-file.ts: -------------------------------------------------------------------------------- 1 | import path from "node:path"; 2 | import { Uri, window, workspace } from "vscode"; 3 | 4 | export async function createNewMarimoFile(): Promise { 5 | const editor = window.activeTextEditor; 6 | // fallback to the first workspace folder 7 | const workspaceFolder = workspace.workspaceFolders?.[0]; 8 | 9 | // prompt for file name 10 | const name = await window.showInputBox({ 11 | prompt: "Enter a name for the new marimo file", 12 | }); 13 | 14 | if (!name) { 15 | window.showErrorMessage("No name provided"); 16 | return; 17 | } 18 | 19 | // ' ' -> '_' 20 | // add py extension 21 | let fileName = name.replaceAll(" ", "_"); 22 | if (!fileName.endsWith(".py")) { 23 | fileName += ".py"; 24 | } 25 | 26 | // Get the directory of the current file 27 | const currentFilePath = editor 28 | ? editor.document.uri.fsPath 29 | : workspaceFolder?.uri.fsPath; 30 | 31 | if (!currentFilePath) { 32 | window.showErrorMessage("No active editor or workspace"); 33 | return; 34 | } 35 | 36 | const directoryPath = path.dirname(currentFilePath); 37 | 38 | // create file 39 | const newFilePath = path.join(directoryPath, fileName); 40 | const newFileUri = Uri.file(newFilePath); 41 | 42 | const encoder = new TextEncoder(); 43 | await workspace.fs.writeFile(newFileUri, encoder.encode(NEW_FILE_CONTENT)); 44 | 45 | const document = await workspace.openTextDocument(newFileUri); 46 | await window.showTextDocument(document); 47 | return newFileUri; 48 | } 49 | 50 | const NEW_FILE_CONTENT = ` 51 | import marimo 52 | 53 | app = marimo.App() 54 | 55 | @app.cell 56 | def __(): 57 | import marimo as mo 58 | return 59 | 60 | if __name__ == "__main__": 61 | app.run() 62 | `.trim(); 63 | -------------------------------------------------------------------------------- /src/constants.ts: -------------------------------------------------------------------------------- 1 | export const DOCUMENTATION_URL = "https://docs.marimo.io"; 2 | export const DISCORD_URL = "https://marimo.io/discord?ref=vscode"; 3 | 4 | export const EXTENSION_PACKAGE = { 5 | publisher: "marimo-team", 6 | name: "vscode-marimo", 7 | get fullName(): string { 8 | return `${this.publisher}.${this.name}`; 9 | }, 10 | }; 11 | 12 | export const EXTENSION_DISPLAY_NAME = "marimo"; 13 | 14 | export const CONFIG_KEY = "marimo"; 15 | 16 | export const CommandsKeys = { 17 | // Start marimo kernel (edit) 18 | edit: "vscode-marimo.edit", 19 | // Start marimo kernel (run) 20 | run: "vscode-marimo.run", 21 | // Refresh marimo explorer 22 | refresh: "vscode-marimo.refresh", 23 | // Restart marimo kernel 24 | restartKernel: "vscode-marimo.restartKernel", 25 | // Stop kernel 26 | stopKernel: "vscode-marimo.stopKernel", 27 | // Show marimo commands 28 | showCommands: "vscode-marimo.showCommands", 29 | // Export notebook as... 30 | exportAsCommands: "vscode-marimo.exportAsCommands", 31 | // Open in system browser 32 | openInBrowser: "vscode-marimo.openInBrowser", 33 | // Show documentation 34 | openDocumentation: "vscode-marimo.openDocumentation", 35 | // Create new marimo file 36 | newMarimoFile: "vscode-marimo.newMarimoFile", 37 | // Reload browser 38 | reloadBrowser: "vscode-marimo.reloadBrowser", 39 | // Convert Jupyter notebook to marimo notebook 40 | convertToMarimoApp: "vscode-marimo.convertToMarimoApp", 41 | 42 | // Start server 43 | startServer: "vscode-marimo.startServer", 44 | // Stop server 45 | stopServer: "vscode-marimo.stopServer", 46 | 47 | // Native vscode notebook commands 48 | openNotebook: "vscode-marimo.openAsVSCodeNotebook", 49 | 50 | // Show marimo diagnostics 51 | showDiagnostics: "vscode-marimo.showDiagnostics", 52 | // Show marimo help 53 | showHelp: "vscode-marimo.showHelp", 54 | }; 55 | -------------------------------------------------------------------------------- /src/utils/__tests__/url.test.ts: -------------------------------------------------------------------------------- 1 | import { createVSCodeMock } from "jest-mock-vscode"; 2 | import { describe, expect, it, vi } from "vitest"; 3 | 4 | vi.mock("vscode", () => createVSCodeMock(vi)); 5 | 6 | import { asURL, urlJoin } from "../url"; 7 | 8 | describe("url utils", () => { 9 | describe("asURL", () => { 10 | it("should parse valid URLs", () => { 11 | const url = asURL("https://example.com"); 12 | expect(url.href).toBe("https://example.com/"); 13 | }); 14 | 15 | it("should throw on invalid URLs", () => { 16 | expect(() => asURL("not-a-url")).toThrow(); 17 | }); 18 | }); 19 | 20 | describe("urlJoin", () => { 21 | it("should handle empty array", () => { 22 | expect(urlJoin()).toBe(""); 23 | }); 24 | 25 | it("should handle single path", () => { 26 | expect(urlJoin("https://example.com")).toBe("https://example.com"); 27 | }); 28 | 29 | it("should join two paths", () => { 30 | expect(urlJoin("https://example.com", "api")).toBe( 31 | "https://example.com/api", 32 | ); 33 | }); 34 | 35 | it("should handle leading/trailing slashes", () => { 36 | expect(urlJoin("https://example.com/", "/api")).toBe( 37 | "https://example.com/api", 38 | ); 39 | expect(urlJoin("https://example.com/", "/api/")).toBe( 40 | "https://example.com/api/", 41 | ); 42 | }); 43 | 44 | it("should join multiple paths", () => { 45 | expect(urlJoin("https://example.com", "api", "v1", "users")).toBe( 46 | "https://example.com/api/v1/users", 47 | ); 48 | }); 49 | 50 | it("should handle paths with query params", () => { 51 | expect(urlJoin("https://example.com", "api?foo=bar")).toBe( 52 | "https://example.com/api?foo=bar", 53 | ); 54 | }); 55 | 56 | it("should handle paths with hash", () => { 57 | expect(urlJoin("https://example.com", "api#section")).toBe( 58 | "https://example.com/api#section", 59 | ); 60 | }); 61 | }); 62 | }); 63 | -------------------------------------------------------------------------------- /src/convert/convert.ts: -------------------------------------------------------------------------------- 1 | import path from "node:path"; 2 | import { Uri, window, workspace } from "vscode"; 3 | import { getUniqueFilename } from "../export/export-as"; 4 | import { logger } from "../logger"; 5 | import { printError } from "../utils/errors"; 6 | import { execMarimoCommand } from "../utils/exec"; 7 | 8 | export async function convertIPyNotebook(filePath: string) { 9 | return convertNotebook(filePath, "ipynb"); 10 | } 11 | 12 | export async function convertMarkdownNotebook(filePath: string) { 13 | return convertNotebook(filePath, "md"); 14 | } 15 | 16 | async function convertNotebook( 17 | filePath: string, 18 | ext: "ipynb" | "md", 19 | ): Promise { 20 | try { 21 | // convert 22 | const directory = path.dirname(filePath); 23 | // Execute marimo via python 24 | const response = await execMarimoCommand(["convert", filePath]); 25 | const appCode = response.toString(); 26 | 27 | try { 28 | // try to save to file system 29 | const currentFilename = path.basename(filePath, `.${ext}`); 30 | const newFilename = getUniqueFilename(directory, currentFilename, "py"); 31 | const newFilePath = Uri.file(path.join(directory, newFilename)); 32 | 33 | await workspace.fs.writeFile(newFilePath, Buffer.from(appCode)); 34 | 35 | // open file 36 | workspace.openTextDocument(newFilePath).then(() => { 37 | // Get relative path if possible 38 | const relativePath = workspace.asRelativePath(newFilePath); 39 | window.showInformationMessage(`Saved to ${relativePath}`); 40 | }); 41 | return newFilePath; 42 | } catch { 43 | // if fails to save to file system, open in new tab 44 | workspace 45 | .openTextDocument({ content: appCode, language: "python" }) 46 | .then((doc) => { 47 | window.showTextDocument(doc); 48 | }); 49 | return true; 50 | } 51 | } catch (error) { 52 | logger.info(error); 53 | window.showErrorMessage( 54 | `Failed to convert notebook: \n${printError(error)}`, 55 | ); 56 | return false; 57 | } 58 | } 59 | -------------------------------------------------------------------------------- /src/logger.ts: -------------------------------------------------------------------------------- 1 | import { type LogOutputChannel, window } from "vscode"; 2 | import { Config } from "./config"; 3 | import { EXTENSION_DISPLAY_NAME } from "./constants"; 4 | 5 | /** 6 | * Logger is a utility class for logging messages with different levels (debug, info, error, warn). 7 | * It supports optional prefixing to provide context for log messages and integrates with the 8 | * VS Code Output Channel for displaying logs. This class is designed to be used throughout 9 | * the extension to ensure consistent and structured logging. 10 | */ 11 | class Logger { 12 | private channel: LogOutputChannel; 13 | 14 | constructor(private prefix: string) { 15 | this.channel = window.createOutputChannel(EXTENSION_DISPLAY_NAME, { 16 | log: true, 17 | }); 18 | } 19 | 20 | debug(...args: unknown[]) { 21 | if (!Config.debug) { 22 | return; 23 | } 24 | if (this.prefix) { 25 | this.channel.debug(`[${this.prefix}]: ${args.map(stringify).join(" ")}`); 26 | } else { 27 | this.channel.debug(args.map(stringify).join(" ")); 28 | } 29 | } 30 | 31 | info(...args: unknown[]) { 32 | if (this.prefix) { 33 | this.channel.info(`[${this.prefix}]: ${args.map(stringify).join(" ")}`); 34 | } else { 35 | this.channel.info(args.map(stringify).join(" ")); 36 | } 37 | } 38 | 39 | error(...args: unknown[]) { 40 | if (this.prefix) { 41 | this.channel.error(`[${this.prefix}] ${args.map(stringify).join(" ")}`); 42 | } else { 43 | this.channel.error(args.map(stringify).join(" ")); 44 | } 45 | } 46 | 47 | warn(...args: unknown[]) { 48 | if (this.prefix) { 49 | this.channel.warn( 50 | `[⚠️ warn] [${this.prefix}] ${args.map(stringify).join(" ")}`, 51 | ); 52 | } else { 53 | this.channel.warn(args.map(stringify).join(" ")); 54 | } 55 | } 56 | 57 | createLogger(prefix: string) { 58 | if (!this.prefix) { 59 | return new Logger(prefix); 60 | } 61 | return new Logger(`${this.prefix} > ${prefix}`); 62 | } 63 | } 64 | 65 | function stringify(value: unknown) { 66 | if (typeof value === "string") { 67 | return value; 68 | } 69 | return JSON.stringify(value); 70 | } 71 | 72 | export const logger = new Logger(""); 73 | -------------------------------------------------------------------------------- /src/utils/DeferredRequestRegistry.ts: -------------------------------------------------------------------------------- 1 | /* eslint-disable @typescript-eslint/no-redeclare */ 2 | 3 | import type { TypedString } from "./TypedString"; 4 | import { Deferred } from "./deferred"; 5 | 6 | export type RequestId = TypedString<"RequestId">; 7 | export const RequestId = { 8 | create(): RequestId { 9 | return Math.random().toString(36).slice(2) as RequestId; 10 | }, 11 | }; 12 | 13 | /** 14 | * Helper class to manage deferred requests. 15 | * We send a request via HTTP and then wait for the response from the kernel 16 | * via a websocket. 17 | */ 18 | export class DeferredRequestRegistry { 19 | public requests = new Map>(); 20 | 21 | constructor( 22 | public operation: string, 23 | private makeRequest: (id: RequestId, req: REQ) => Promise, 24 | private opts: { 25 | /** 26 | * Resolve existing requests with an empty response. 27 | */ 28 | resolveExistingRequests?: () => RES; 29 | } = {}, 30 | ) {} 31 | 32 | async request(opts: REQ): Promise { 33 | if (this.opts.resolveExistingRequests) { 34 | const result = this.opts.resolveExistingRequests(); 35 | this.requests.forEach((deferred) => deferred.resolve(result)); 36 | this.requests.clear(); 37 | } 38 | 39 | const requestId = RequestId.create(); 40 | const deferred = new Deferred(); 41 | 42 | this.requests.set(requestId, deferred); 43 | 44 | await this.makeRequest(requestId, opts).catch((error) => { 45 | deferred.reject(error); 46 | this.requests.delete(requestId); 47 | }); 48 | return deferred.promise; 49 | } 50 | 51 | resolve(requestId: RequestId, response: RES) { 52 | const entry = this.requests.get(requestId); 53 | if (entry === undefined) { 54 | return; 55 | } 56 | 57 | entry.resolve(response); 58 | this.requests.delete(requestId); 59 | } 60 | 61 | rejectAll(error: Error) { 62 | this.requests.forEach((deferred) => deferred.reject(error)); 63 | this.requests.clear(); 64 | } 65 | 66 | reject(requestId: RequestId, error: Error) { 67 | const entry = this.requests.get(requestId); 68 | if (entry === undefined) { 69 | return; 70 | } 71 | 72 | entry.reject(error); 73 | this.requests.delete(requestId); 74 | } 75 | } 76 | -------------------------------------------------------------------------------- /src/notebook/serializer.ts: -------------------------------------------------------------------------------- 1 | import * as vscode from "vscode"; 2 | import { logger } from "../logger"; 3 | import { getNotebookMetadata } from "./common/metadata"; 4 | import type { IKernelManager } from "./kernel-manager"; 5 | 6 | const LOADING_CELL_ID = "loading"; 7 | 8 | export class MarimoNotebookSerializer implements vscode.NotebookSerializer { 9 | constructor(private kernelManager: Pick) {} 10 | 11 | public async deserializeNotebook( 12 | _data: Uint8Array, 13 | _token: vscode.CancellationToken, 14 | ): Promise { 15 | // Add a single markdown cell that says Loading... 16 | const cell = new vscode.NotebookCellData( 17 | vscode.NotebookCellKind.Markup, 18 | "### Loading...", 19 | "markdown", 20 | ); 21 | cell.metadata = { 22 | custom: { 23 | id: LOADING_CELL_ID, 24 | }, 25 | }; 26 | return new vscode.NotebookData([cell]); 27 | } 28 | 29 | public async serializeNotebook( 30 | data: vscode.NotebookData, 31 | _token: vscode.CancellationToken, 32 | ): Promise { 33 | // If there are not cells, throw an error 34 | // This is likely an error, and don't want to save an empty notebook 35 | if (data.cells.length === 0) { 36 | logger.error("No cells found in notebook"); 37 | throw new Error("No cells found in notebook"); 38 | } 39 | 40 | // If the only cell is loading, throw an error 41 | if (data.cells.length === 1) { 42 | const cell = data.cells[0]; 43 | const metadata = cell.metadata?.custom; 44 | if (metadata?.id === LOADING_CELL_ID) { 45 | logger.error("No cells found in notebook"); 46 | throw new Error("No cells found in notebook"); 47 | } 48 | } 49 | 50 | const metadata = getNotebookMetadata(data); 51 | const key = metadata.key; 52 | const kernel = this.kernelManager.getKernel(key); 53 | 54 | if (!kernel) { 55 | logger.error("No kernel found for key", key); 56 | throw new Error(`No kernel found for key: ${key}`); 57 | } 58 | 59 | const code = await kernel.save(data.cells); 60 | if (!code) { 61 | logger.error("No code found for kernel", key); 62 | throw new Error(`No code found for kernel: ${key}`); 63 | } 64 | 65 | return Buffer.from(code); 66 | } 67 | } 68 | -------------------------------------------------------------------------------- /src/__fixtures__/mocks.ts: -------------------------------------------------------------------------------- 1 | import { readTextDocument } from "jest-mock-vscode"; 2 | import { vi } from "vitest"; 3 | import { 4 | type ExtensionContext, 5 | NotebookController, 6 | type NotebookDocument, 7 | Uri, 8 | workspace, 9 | } from "vscode"; 10 | import { Config } from "../config"; 11 | import { setExtension } from "../ctx"; 12 | import { MarimoController } from "../launcher/controller"; 13 | import type { KernelKey } from "../notebook/common/key"; 14 | import { NOTEBOOK_TYPE } from "../notebook/constants"; 15 | import { createNotebookController } from "../notebook/createMarimoNotebookController"; 16 | import { Kernel } from "../notebook/kernel"; 17 | import type { SkewToken } from "../notebook/marimo/types"; 18 | import { ServerManager } from "../services/server-manager"; 19 | 20 | const appFixture = new URL( 21 | "../__fixtures__/export/app.py", 22 | import.meta.url, 23 | ).toString(); 24 | export const appFixtureUri = Uri.parse(appFixture); 25 | export const markdownFixtureUri = Uri.parse( 26 | new URL("../__fixtures__/convert/mock.md", import.meta.url).toString(), 27 | ); 28 | export const ipynbFixtureUri = Uri.parse( 29 | new URL("../__fixtures__/convert/mock.ipynb", import.meta.url).toString(), 30 | ); 31 | 32 | setExtension({ 33 | subscriptions: [], 34 | workspaceState: { 35 | get: vi.fn(), 36 | update: vi.fn(), 37 | }, 38 | globalState: { 39 | get: vi.fn(), 40 | update: vi.fn(), 41 | }, 42 | } as unknown as ExtensionContext); 43 | 44 | export const controller = createNotebookController(); 45 | export const mockNotebookDocument: NotebookDocument = { 46 | uri: Uri.parse(appFixture), 47 | notebookType: NOTEBOOK_TYPE, 48 | version: 1, 49 | isDirty: false, 50 | isUntitled: false, 51 | isClosed: false, 52 | metadata: {}, 53 | cellCount: 0, 54 | cellAt: vi.fn(), 55 | getCells: vi.fn(), 56 | save: vi.fn(), 57 | }; 58 | 59 | export const mockKernel = new Kernel({ 60 | port: 100, 61 | fileUri: Uri.parse(appFixture), 62 | kernelKey: "kernel-key" as KernelKey, 63 | skewToken: "skew-token" as SkewToken, 64 | version: "1.0.0", 65 | userConfig: {}, 66 | controller, 67 | notebookDoc: mockNotebookDocument, 68 | }); 69 | 70 | export async function createMockController(file: Uri = Uri.parse(appFixture)) { 71 | return new MarimoController( 72 | await readTextDocument(file), 73 | ServerManager.getInstance(Config), 74 | ); 75 | } 76 | -------------------------------------------------------------------------------- /.github/workflows/ci.yml: -------------------------------------------------------------------------------- 1 | name: CI 2 | 3 | env: 4 | MARIMO_SKIP_UPDATE_CHECK: true 5 | 6 | on: 7 | push: 8 | branches: [main] 9 | pull_request: 10 | branches: [main] 11 | 12 | jobs: 13 | lint: 14 | runs-on: ubuntu-latest 15 | steps: 16 | - name: 🛑 Cancel Previous Runs 17 | uses: styfle/cancel-workflow-action@0.12.1 18 | 19 | - name: ⬇️ Checkout repo 20 | uses: actions/checkout@v4 21 | 22 | - name: ⎔ Setup pnpm 23 | uses: pnpm/action-setup@v2 24 | with: 25 | version: 8 26 | 27 | - name: ⎔ Setup Node.js 28 | uses: actions/setup-node@v4 29 | with: 30 | node-version: 18 31 | cache: pnpm 32 | 33 | - name: 📥 Install dependencies 34 | run: pnpm install 35 | 36 | - name: 🧹 Lint 37 | run: pnpm lint 38 | 39 | typecheck: 40 | runs-on: ubuntu-latest 41 | steps: 42 | - name: 🛑 Cancel Previous Runs 43 | uses: styfle/cancel-workflow-action@0.12.1 44 | 45 | - name: ⬇️ Checkout repo 46 | uses: actions/checkout@v4 47 | 48 | - name: ⎔ Setup pnpm 49 | uses: pnpm/action-setup@v2 50 | with: 51 | version: 8 52 | 53 | - name: ⎔ Setup Node.js 54 | uses: actions/setup-node@v4 55 | with: 56 | node-version: 18 57 | cache: pnpm 58 | 59 | - name: 📥 Install dependencies 60 | run: pnpm install 61 | 62 | - name: ʦ Typecheck 63 | run: pnpm typecheck 64 | 65 | test: 66 | runs-on: ubuntu-latest 67 | steps: 68 | - name: 🛑 Cancel Previous Runs 69 | uses: styfle/cancel-workflow-action@0.12.1 70 | 71 | - name: ⬇️ Checkout repo 72 | uses: actions/checkout@v4 73 | 74 | - name: ⎔ Setup pnpm 75 | uses: pnpm/action-setup@v2 76 | with: 77 | version: 8 78 | 79 | - name: 🐍 Setup uv 80 | uses: astral-sh/setup-uv@v6 81 | 82 | - name: 🐍 Setup Python 83 | uses: actions/setup-python@v5 84 | with: 85 | python-version: 3.12 86 | 87 | - name: ⎔ Setup Node.js 88 | uses: actions/setup-node@v4 89 | with: 90 | node-version: 20 91 | cache: pnpm 92 | 93 | - name: 📥 Install dependencies 94 | run: | 95 | uv pip install marimo --system 96 | pnpm install 97 | 98 | - name: 🧪 Run tests 99 | run: pnpm test 100 | 101 | 102 | -------------------------------------------------------------------------------- /src/telemetry.ts: -------------------------------------------------------------------------------- 1 | import { workspace } from "vscode"; 2 | import { Config } from "./config"; 3 | import { getGlobalState } from "./ctx"; 4 | import { logger } from "./logger"; 5 | import { getExtensionVersion, getVscodeVersion } from "./utils/versions"; 6 | 7 | export function trackEvent( 8 | event: "vscode-lifecycle", 9 | data: { action: "activate" | "deactivate" }, 10 | ): void; 11 | export function trackEvent( 12 | event: "vscode-command", 13 | data: { command: string }, 14 | ): void; 15 | export function trackEvent( 16 | event: "vscode-configuration", 17 | data: { key: string; value: string }, 18 | ): void; 19 | export function trackEvent(event: string, data: Record): void { 20 | if (!Config.telemetry) { 21 | return; 22 | } 23 | 24 | try { 25 | const metadata = { 26 | anonymous_id: anonymouseId(), 27 | extension_version: getExtensionVersion(), 28 | vscode_version: getVscodeVersion(), 29 | platform: process.platform, 30 | architecture: process.arch, 31 | node_version: process.version, 32 | }; 33 | // Fire and forget 34 | void sendEvent(event, data, metadata).catch((error) => { 35 | logger.info("Error sending telemetry event:", error); 36 | }); 37 | } catch (error) { 38 | logger.info("Error sending telemetry event:", error); 39 | } 40 | } 41 | 42 | async function sendEvent( 43 | event: string, 44 | data: Record, 45 | metadata: Record, 46 | ) { 47 | const res = await fetch("https://metrics.marimo.app/api/v1/telemetry", { 48 | method: "POST", 49 | headers: { 50 | "Content-Type": "application/json", 51 | }, 52 | body: JSON.stringify({ event, data, metadata }), 53 | }); 54 | if (!res.ok) { 55 | logger.info("Error sending telemetry event:", res.statusText); 56 | } 57 | } 58 | 59 | export function anonymouseId(): string { 60 | try { 61 | const globalState = getGlobalState(); 62 | let id = globalState.get("telemetry.anonymousId"); 63 | if (!id) { 64 | id = crypto.randomUUID(); 65 | void globalState.update("telemetry.anonymousId", id); 66 | } 67 | return id; 68 | } catch (error) { 69 | return "unknown"; 70 | } 71 | } 72 | 73 | export function setupConfigTelemetry(): void { 74 | workspace.onDidChangeConfiguration((event) => { 75 | if (event.affectsConfiguration("marimo.telemetry")) { 76 | const enabled = Config.telemetry; 77 | trackEvent("vscode-configuration", { 78 | key: "marimo.telemetry", 79 | value: enabled.toString(), 80 | }); 81 | } 82 | }); 83 | } 84 | -------------------------------------------------------------------------------- /src/utils/__tests__/deferred.test.ts: -------------------------------------------------------------------------------- 1 | import { describe, expect, it, test, vi } from "vitest"; 2 | import { Deferred } from "../deferred"; 3 | 4 | describe("Deferred", () => { 5 | test("should resolve correctly", async () => { 6 | const deferred = new Deferred(); 7 | expect(deferred.hasCompleted).toBe(false); 8 | expect(deferred.hasResolved).toBe(false); 9 | expect(deferred.hasRejected).toBe(false); 10 | 11 | const value = "test"; 12 | deferred.resolve(value); 13 | const result = await deferred.promise; 14 | 15 | expect(result).toBe(value); 16 | expect(deferred.hasCompleted).toBe(true); 17 | expect(deferred.hasResolved).toBe(true); 18 | expect(deferred.hasRejected).toBe(false); 19 | }); 20 | 21 | test("should reject correctly", async () => { 22 | const deferred = new Deferred(); 23 | const error = new Error("test error"); 24 | 25 | deferred.reject(error); 26 | 27 | await expect(deferred.promise).rejects.toThrow(error); 28 | expect(deferred.hasCompleted).toBe(true); 29 | expect(deferred.hasResolved).toBe(false); 30 | expect(deferred.hasRejected).toBe(true); 31 | }); 32 | 33 | test("should work with async resolution", async () => { 34 | const deferred = new Deferred(); 35 | const value = 42; 36 | 37 | setTimeout(() => { 38 | deferred.resolve(value); 39 | }, 10); 40 | 41 | const result = await deferred.promise; 42 | expect(result).toBe(value); 43 | expect(deferred.hasCompleted).toBe(true); 44 | expect(deferred.hasResolved).toBe(true); 45 | }); 46 | 47 | test("should handle promise-like values", async () => { 48 | const deferred = new Deferred(); 49 | const promiseValue = Promise.resolve(42); 50 | 51 | deferred.resolve(promiseValue); 52 | const result = await deferred.promise; 53 | 54 | expect(result).toBe(42); 55 | expect(deferred.hasCompleted).toBe(true); 56 | expect(deferred.hasResolved).toBe(true); 57 | }); 58 | 59 | it("should call onResolved when resolved", async () => { 60 | const onResolved = vi.fn(); 61 | const deferred = new Deferred({ 62 | onResolved, 63 | }); 64 | const value = 42; 65 | 66 | deferred.resolve(value); 67 | expect(onResolved).toHaveBeenCalledWith(value); 68 | }); 69 | 70 | it("should call onRejected when rejected", async () => { 71 | const onRejected = vi.fn(); 72 | const deferred = new Deferred({ 73 | onRejected, 74 | }); 75 | const error = new Error("test error"); 76 | 77 | deferred.reject(error); 78 | await expect(() => deferred.promise).rejects.toThrow(error); 79 | expect(onRejected).toHaveBeenCalledWith(error); 80 | }); 81 | }); 82 | -------------------------------------------------------------------------------- /resources/icon.svg: -------------------------------------------------------------------------------- 1 | 3 | 4 | 5 | 6 | 7 | 8 | 10 | 11 | 12 | -------------------------------------------------------------------------------- /src/notebook/marimo/types.ts: -------------------------------------------------------------------------------- 1 | import type { components } from "../../generated/api"; 2 | import type { TypedString } from "../../utils/TypedString"; 3 | import type { DeepPartial } from "../../utils/types"; 4 | 5 | const lowercase = "abcdefghijklmnopqrstuvwxyz"; 6 | const uppercase = lowercase.toUpperCase(); 7 | const alphabet = lowercase + uppercase; 8 | 9 | /** 10 | * A typed CellId 11 | */ 12 | export type CellId = TypedString<"CellId">; 13 | export const CellId = { 14 | /** 15 | * Create a new CellId, a random 4 letter string. 16 | */ 17 | create(): CellId { 18 | let id = ""; 19 | for (let i = 0; i < 4; i++) { 20 | id += alphabet[Math.floor(Math.random() * alphabet.length)]; 21 | } 22 | return id as CellId; 23 | }, 24 | }; 25 | 26 | export type SessionId = TypedString<"SessionId">; 27 | export const SessionId = { 28 | create(): SessionId { 29 | let id = ""; 30 | for (let i = 0; i < 4; i++) { 31 | id += alphabet[Math.floor(Math.random() * alphabet.length)]; 32 | } 33 | return id as SessionId; 34 | }, 35 | }; 36 | 37 | export type SkewToken = TypedString<"SkewToken">; 38 | 39 | type schemas = components["schemas"]; 40 | 41 | export type MessageOperationType = schemas["MessageOperation"]["name"]; 42 | export type MessageOperation = { 43 | [Type in MessageOperationType]: { 44 | op: Type; 45 | data: Omit, "name">; 46 | }; 47 | }[MessageOperationType]; 48 | export type MessageOperationData = Omit< 49 | Extract, 50 | "name" 51 | >; 52 | 53 | export type CellStatus = schemas["RuntimeState"]; 54 | export type Operation = schemas["MessageOperation"]; 55 | export type InstantiateRequest = schemas["InstantiateRequest"]; 56 | export type InstallMissingPackagesRequest = 57 | schemas["InstallMissingPackagesRequest"]; 58 | export type UpdateCellIdsRequest = schemas["UpdateCellIdsRequest"]; 59 | export type RunRequest = schemas["RunRequest"]; 60 | export type MarimoFile = schemas["MarimoFile"]; 61 | export type WorkspaceFilesResponse = schemas["WorkspaceFilesResponse"]; 62 | export type DeleteCellRequest = schemas["DeleteCellRequest"]; 63 | export type SaveNotebookRequest = schemas["SaveNotebookRequest"]; 64 | export type CellChannel = schemas["CellChannel"]; 65 | export type FunctionCallRequest = schemas["FunctionCallRequest"]; 66 | export type CellConfig = schemas["CellConfig"]; 67 | export type CellOutput = schemas["CellOutput"]; 68 | export type MarimoConfig = DeepPartial; 69 | 70 | export type CellOp = MessageOperationData<"cell-op">; 71 | export type KernelReady = MessageOperationData<"kernel-ready">; 72 | export type FunctionCallResult = MessageOperationData<"function-call-result">; 73 | -------------------------------------------------------------------------------- /src/browser/__tests__/panel.test.ts: -------------------------------------------------------------------------------- 1 | import { describe, expect, it, vi } from "vitest"; 2 | import { createVSCodeMock } from "../../__mocks__/vscode"; 3 | 4 | vi.mock("vscode", async () => { 5 | return createVSCodeMock(vi); 6 | }); 7 | 8 | import { MarimoPanelManager } from "../panel"; 9 | 10 | describe("Panel", () => { 11 | it("should be created", async () => { 12 | const panel = new MarimoPanelManager("app"); 13 | expect(panel).toBeDefined(); 14 | }); 15 | 16 | it("should show panel", async () => { 17 | const panel = new MarimoPanelManager("app"); 18 | await panel.create("https://example.com"); 19 | const nativePanel = panel.nativePanel; 20 | expect(nativePanel).toBeDefined(); 21 | expect(nativePanel?.webview.html).toMatchInlineSnapshot(` 22 | " 23 | 24 | 25 | 26 | 27 | 28 | marimo 29 | 30 | 40 | 49 | 50 | 74 | 75 | " 76 | `); 77 | }); 78 | }); 79 | -------------------------------------------------------------------------------- /src/launcher/utils.ts: -------------------------------------------------------------------------------- 1 | import { parse } from "node-html-parser"; 2 | import type { CancellationToken } from "vscode"; 3 | import { composeUrl } from "../config"; 4 | import { logger } from "../logger"; 5 | import type { MarimoConfig, SkewToken } from "../notebook/marimo/types"; 6 | import { retry } from "../utils/retry"; 7 | import { asURL } from "../utils/url"; 8 | 9 | /** 10 | * Grabs the index.html of the marimo server and extracts 11 | * various startup values. 12 | * - skewToken 13 | * - version 14 | * - userConfig 15 | */ 16 | export async function fetchMarimoStartupValues({ 17 | port, 18 | backoff, 19 | cancellationToken, 20 | }: { 21 | port: number; 22 | backoff?: number; 23 | cancellationToken?: CancellationToken; 24 | }): Promise<{ 25 | skewToken: SkewToken; 26 | version: string; 27 | userConfig: MarimoConfig; 28 | }> { 29 | const url = asURL(await composeUrl(port)); 30 | let response: Response; 31 | try { 32 | response = await retry( 33 | async () => { 34 | if (cancellationToken?.isCancellationRequested) { 35 | throw new Error("Cancelled"); 36 | } 37 | const resp = await fetch(url.toString(), { 38 | headers: { 39 | Accept: "text/html", 40 | }, 41 | }); 42 | if (!resp.ok) { 43 | throw new Error(`HTTP error ${resp.status}`); 44 | } 45 | return resp; 46 | }, 47 | 5, // retries 48 | backoff ?? 1000, // 1s exponential backoff 49 | ); 50 | } catch (e) { 51 | logger.error(`Could not fetch ${url}. Is ${url} healthy?`); 52 | throw new Error(`Could not fetch ${url}. Is ${url} healthy?`); 53 | } 54 | 55 | if (!response.ok) { 56 | throw new Error( 57 | `Could not fetch ${url}. Is ${url} healthy? ${response.status} ${response.statusText}`, 58 | ); 59 | } 60 | 61 | // If was redirected to /auth/login, then show a message that an existing server is running 62 | if (asURL(response.url).pathname.startsWith("/auth/login")) { 63 | const msg = `An existing marimo server created outside of vscode is running at this url: ${url.toString()}`; 64 | logger.warn(msg); 65 | throw new Error(msg); 66 | } 67 | 68 | const html = await response.text(); 69 | const root = parse(html); 70 | const getDomValue = (tagName: string, datasetKey: string) => { 71 | const element = root.querySelector(tagName); 72 | if (!element) { 73 | throw new Error(`Could not find ${tagName}. Is ${url} healthy?`); 74 | } 75 | const value = element.getAttribute(`data-${datasetKey}`); 76 | if (value === undefined) { 77 | throw new Error(`${datasetKey} is undefined`); 78 | } 79 | 80 | return value; 81 | }; 82 | 83 | const skewToken = getDomValue("marimo-server-token", "token") as SkewToken; 84 | const userConfig = JSON.parse( 85 | getDomValue("marimo-user-config", "config"), 86 | ) as MarimoConfig; 87 | const marimoVersion = getDomValue("marimo-version", "version"); 88 | 89 | return { 90 | skewToken, 91 | version: marimoVersion, 92 | userConfig, 93 | }; 94 | } 95 | -------------------------------------------------------------------------------- /src/utils/python.ts: -------------------------------------------------------------------------------- 1 | import { 2 | PythonExtension, 3 | type ResolvedEnvironment, 4 | } from "@vscode/python-extension"; 5 | import { 6 | Disposable, 7 | type Event, 8 | EventEmitter, 9 | type Uri, 10 | commands, 11 | } from "vscode"; 12 | import { extension } from "../ctx"; 13 | import { logger } from "../logger"; 14 | 15 | // Adapted from https://github.com/astral-sh/ruff-vscode/blob/d5d4acf3b63eacbe5a4b2e2f4da56f74cb494889/src/common/python.ts 16 | 17 | export interface IInterpreterDetails { 18 | path?: string[]; 19 | resource?: Uri; 20 | } 21 | 22 | const onDidChangePythonInterpreterEvent = 23 | new EventEmitter(); 24 | export const onDidChangePythonInterpreter: Event = 25 | onDidChangePythonInterpreterEvent.event; 26 | 27 | let _api: PythonExtension | undefined; 28 | async function getPythonExtensionAPI(): Promise { 29 | if (_api) { 30 | return _api; 31 | } 32 | _api = await PythonExtension.api(); 33 | return _api; 34 | } 35 | 36 | export async function initializePython(): Promise { 37 | try { 38 | const api = await getPythonExtensionAPI(); 39 | 40 | if (api) { 41 | extension.subscriptions.push( 42 | api.environments.onDidChangeActiveEnvironmentPath((e) => { 43 | onDidChangePythonInterpreterEvent.fire({ 44 | path: [e.path], 45 | resource: e.resource?.uri, 46 | }); 47 | }), 48 | ); 49 | 50 | logger.info("Waiting for interpreter from python extension."); 51 | onDidChangePythonInterpreterEvent.fire(await getInterpreterDetails()); 52 | } 53 | } catch (error) { 54 | logger.error("Error initializing python: ", error); 55 | } 56 | } 57 | 58 | export async function resolveInterpreter( 59 | interpreter: string[], 60 | ): Promise { 61 | const api = await getPythonExtensionAPI(); 62 | return api?.environments.resolveEnvironment(interpreter[0]); 63 | } 64 | 65 | export async function getInterpreterDetails( 66 | resource?: Uri, 67 | ): Promise { 68 | const api = await getPythonExtensionAPI(); 69 | const environment = await api?.environments.resolveEnvironment( 70 | api?.environments.getActiveEnvironmentPath(resource), 71 | ); 72 | if (environment?.executable.uri && checkVersion(environment)) { 73 | return { path: [environment?.executable.uri.fsPath], resource }; 74 | } 75 | return { path: undefined, resource }; 76 | } 77 | 78 | export async function getDebuggerPath(): Promise { 79 | const api = await getPythonExtensionAPI(); 80 | return api?.debug.getDebuggerPackagePath(); 81 | } 82 | 83 | export async function runPythonExtensionCommand( 84 | command: string, 85 | ...rest: any[] 86 | ) { 87 | await getPythonExtensionAPI(); 88 | return await commands.executeCommand(command, ...rest); 89 | } 90 | 91 | export function checkVersion(resolved: ResolvedEnvironment): boolean { 92 | const version = resolved.version; 93 | if (version?.major === 3 && version?.minor >= 9) { 94 | return true; 95 | } 96 | logger.error( 97 | `Python version ${version?.major}.${version?.minor} is not supported.`, 98 | ); 99 | logger.error(`Selected python path: ${resolved.executable.uri?.fsPath}`); 100 | logger.error("Supported versions are 3.9 and above."); 101 | return false; 102 | } 103 | -------------------------------------------------------------------------------- /src/export/__tests__/export-as.test.ts: -------------------------------------------------------------------------------- 1 | import { describe, expect, it, vi } from "vitest"; 2 | 3 | import { createVSCodeMock } from "../../__mocks__/vscode"; 4 | 5 | vi.mock("vscode", () => createVSCodeMock(vi)); 6 | vi.mock("@vscode/python-extension", () => ({})); 7 | 8 | import { existsSync } from "node:fs"; 9 | import { unlink } from "node:fs/promises"; 10 | import type { Uri } from "vscode"; 11 | import { appFixtureUri } from "../../__fixtures__/mocks"; 12 | import { exportNotebookAs, getUniqueFilename } from "../export-as"; 13 | 14 | vi.mock("node:fs"); 15 | 16 | describe("export-as", () => { 17 | it("should export as html", async () => { 18 | const newUri = await exportNotebookAs(appFixtureUri.fsPath, "html"); 19 | expect(newUri).toBeDefined(); 20 | expect((newUri as Uri).fsPath).toMatch(/\.html$/); 21 | // Clean up 22 | await unlink((newUri as Uri).fsPath); 23 | }); 24 | 25 | it.skip("should export as ipynb", async () => { 26 | const newUri = await exportNotebookAs(appFixtureUri.fsPath, "ipynb"); 27 | expect(newUri).toBeDefined(); 28 | expect((newUri as Uri).fsPath).toMatch(/\.ipynb$/); 29 | // Clean up 30 | await unlink((newUri as Uri).fsPath); 31 | }); 32 | 33 | it("should export as md", async () => { 34 | const newUri = await exportNotebookAs(appFixtureUri.fsPath, "md"); 35 | expect(newUri).toBeDefined(); 36 | expect((newUri as Uri).fsPath).toMatch(/\.md$/); 37 | // Clean up 38 | await unlink((newUri as Uri).fsPath); 39 | }); 40 | 41 | it("should export as script", async () => { 42 | const newUri = await exportNotebookAs(appFixtureUri.fsPath, "script"); 43 | expect(newUri).toBeDefined(); 44 | expect((newUri as Uri).fsPath).toMatch(/\.py$/); 45 | // Clean up 46 | await unlink((newUri as Uri).fsPath); 47 | }); 48 | 49 | it("should export as html without code", async () => { 50 | const newUri = await exportNotebookAs( 51 | appFixtureUri.fsPath, 52 | "html-without-code", 53 | ); 54 | expect(newUri).toBeDefined(); 55 | expect((newUri as Uri).fsPath).toMatch(/\.html$/); 56 | // Clean up 57 | await unlink((newUri as Uri).fsPath); 58 | }); 59 | }); 60 | 61 | describe("getUniqueFilename", () => { 62 | it("should return the original filename if it doesn't exist", () => { 63 | vi.mocked(existsSync).mockReturnValue(false); 64 | 65 | const result = getUniqueFilename("/test/dir", "file", "txt"); 66 | 67 | expect(result).toBe("file.txt"); 68 | expect(existsSync).toHaveBeenCalledWith("/test/dir/file.txt"); 69 | }); 70 | 71 | it("should append a number if the file already exists", () => { 72 | vi.mocked(existsSync).mockImplementation((filePath) => { 73 | return ( 74 | filePath === "/test/dir/file.txt" || filePath === "/test/dir/file_1.txt" 75 | ); 76 | }); 77 | 78 | const result = getUniqueFilename("/test/dir", "file", "txt"); 79 | 80 | expect(result).toBe("file_2.txt"); 81 | expect(existsSync).toHaveBeenCalledWith("/test/dir/file.txt"); 82 | expect(existsSync).toHaveBeenCalledWith("/test/dir/file_1.txt"); 83 | expect(existsSync).toHaveBeenCalledWith("/test/dir/file_2.txt"); 84 | }); 85 | 86 | it("should handle existing postfix", () => { 87 | vi.mocked(existsSync).mockReturnValue(false); 88 | 89 | const result = getUniqueFilename("/test/dir", "file", "txt", 5); 90 | 91 | expect(result).toBe("file_5.txt"); 92 | expect(existsSync).toHaveBeenCalledWith("/test/dir/file_5.txt"); 93 | }); 94 | }); 95 | -------------------------------------------------------------------------------- /src/utils/__tests__/cmd.test.ts: -------------------------------------------------------------------------------- 1 | import { beforeEach, describe, expect, it } from "vitest"; 2 | import { MarimoCmdBuilder } from "../cmd"; 3 | 4 | let b = new MarimoCmdBuilder(); 5 | 6 | describe("MarimoCmdBuilder", () => { 7 | beforeEach(() => { 8 | b = new MarimoCmdBuilder(); 9 | }); 10 | 11 | it("happy path", () => { 12 | const cmd = new MarimoCmdBuilder() 13 | .debug(false) 14 | .mode("edit") 15 | .fileOrDir("path/to/file") 16 | .host("localhost") 17 | .port(2718) 18 | .headless(true) 19 | .enableToken(false) 20 | .tokenPassword("") 21 | .build(); 22 | expect(cmd).toMatchInlineSnapshot( 23 | `"--yes edit path/to/file --host=localhost --port=2718 --headless --no-token"`, 24 | ); 25 | }); 26 | 27 | it("should correctly handle debug mode", () => { 28 | const cmd = new MarimoCmdBuilder() 29 | .debug(true) 30 | .mode("edit") 31 | .fileOrDir("path/to/file") 32 | .host("localhost") 33 | .port(2718) 34 | .headless(true) 35 | .enableToken(true) 36 | .tokenPassword("secret") 37 | .build(); 38 | expect(cmd).toMatchInlineSnapshot( 39 | `"--yes -d edit path/to/file --host=localhost --port=2718 --headless --token-password=secret"`, 40 | ); 41 | }); 42 | 43 | it("it can handle run mode", () => { 44 | const cmd = new MarimoCmdBuilder() 45 | .debug(false) 46 | .mode("run") 47 | .fileOrDir("path/to/file") 48 | .host("localhost") 49 | .port(2718) 50 | .headless(true) 51 | .enableToken(false) 52 | .tokenPassword("") 53 | .build(); 54 | expect(cmd).toMatchInlineSnapshot( 55 | `"--yes run path/to/file --host=localhost --port=2718 --headless --no-token"`, 56 | ); 57 | }); 58 | 59 | it("should correctly handle fileOrDir with and without spaces", () => { 60 | const b = new MarimoCmdBuilder() 61 | .debug(false) 62 | .mode("edit") 63 | .fileOrDir("path/to/some file") 64 | .host("localhost") 65 | .port(2718) 66 | .headless(true) 67 | .enableToken(false) 68 | .tokenPassword("") 69 | .build(); 70 | 71 | expect(b).toMatchInlineSnapshot( 72 | `"--yes edit "path/to/some file" --host=localhost --port=2718 --headless --no-token"`, 73 | ); 74 | }); 75 | 76 | it("should correctly handle host", () => { 77 | const b = new MarimoCmdBuilder() 78 | .debug(false) 79 | .mode("edit") 80 | .fileOrDir("path/to/file") 81 | .host("0.0.0.0") 82 | .port(2718) 83 | .headless(true) 84 | .enableToken(false) 85 | .tokenPassword("") 86 | .build(); 87 | 88 | expect(b).toMatchInlineSnapshot( 89 | `"--yes edit path/to/file --host=0.0.0.0 --port=2718 --headless --no-token"`, 90 | ); 91 | }); 92 | 93 | it("should support sandbox mode", () => { 94 | const cmd = new MarimoCmdBuilder() 95 | .debug(false) 96 | .mode("edit") 97 | .fileOrDir("path/to/file") 98 | .host("localhost") 99 | .port(2718) 100 | .headless(true) 101 | .enableToken(false) 102 | .sandbox(true) 103 | .build(); 104 | expect(cmd).toMatchInlineSnapshot( 105 | `"--yes edit path/to/file --host=localhost --port=2718 --headless --no-token --sandbox"`, 106 | ); 107 | }); 108 | 109 | it("should support watch mode", () => { 110 | const cmd = new MarimoCmdBuilder() 111 | .debug(false) 112 | .mode("edit") 113 | .fileOrDir("path/to/file") 114 | .host("localhost") 115 | .port(2718) 116 | .headless(true) 117 | .enableToken(false) 118 | .watch(true) 119 | .build(); 120 | expect(cmd).toMatchInlineSnapshot( 121 | `"--yes edit path/to/file --host=localhost --port=2718 --headless --no-token --watch"`, 122 | ); 123 | }); 124 | }); 125 | -------------------------------------------------------------------------------- /src/commands/__tests__/show-commands.test.ts: -------------------------------------------------------------------------------- 1 | import { describe, expect, it, vi } from "vitest"; 2 | import { createVSCodeMock } from "../../__mocks__/vscode"; 3 | 4 | vi.mock("vscode", () => createVSCodeMock(vi)); 5 | vi.mock("@vscode/python-extension", () => ({})); 6 | 7 | import { createMockController, mockKernel } from "../../__fixtures__/mocks"; 8 | import { Config } from "../../config"; 9 | import { ServerManager } from "../../services/server-manager"; 10 | import { 11 | miscCommands, 12 | showKernelCommands, 13 | showMarimoControllerCommands, 14 | } from "../show-commands"; 15 | 16 | describe("showCommands", () => { 17 | const serverManager = ServerManager.getInstance(Config); 18 | 19 | it("should show commands for Kernel", async () => { 20 | const commands = await showKernelCommands(mockKernel); 21 | expect(commands.map((c) => c.label)).toMatchInlineSnapshot(` 22 | [ 23 | "$(split-horizontal) Open outputs in embedded browser", 24 | "$(link-external) Open outputs in system browser", 25 | "", 26 | "$(refresh) Restart kernel", 27 | "$(export) Export notebook as...", 28 | "", 29 | ] 30 | `); 31 | }); 32 | 33 | it("should show commands for non active Controller", async () => { 34 | const commands = ( 35 | await showMarimoControllerCommands(await createMockController()) 36 | ).filter((index) => index.if !== false); 37 | expect(commands.map((c) => c.label)).toMatchInlineSnapshot(` 38 | [ 39 | "$(notebook) Start as VSCode notebook", 40 | "$(zap) Start in marimo editor (edit)", 41 | "$(preview) Start in marimo editor (run)", 42 | "", 43 | "$(export) Export notebook as...", 44 | "", 45 | ] 46 | `); 47 | }); 48 | 49 | it("should show commands for active Controller for run", async () => { 50 | const controller = await createMockController(); 51 | controller.active = true; 52 | controller.currentMode = "run"; 53 | const commands = (await showMarimoControllerCommands(controller)).filter( 54 | (index) => index.if !== false, 55 | ); 56 | expect(commands.map((c) => c.label)).toMatchInlineSnapshot(` 57 | [ 58 | "", 59 | "$(split-horizontal) Open in embedded browser", 60 | "$(link-external) Open in system browser", 61 | "$(refresh) Restart marimo kernel", 62 | "$(package) Switch to edit mode", 63 | "$(terminal) Show Terminal", 64 | "$(close) Stop kernel", 65 | "$(export) Export notebook as...", 66 | "", 67 | ] 68 | `); 69 | }); 70 | 71 | it("should show commands for active Controller for edit", async () => { 72 | const controller = await createMockController(); 73 | controller.active = true; 74 | controller.currentMode = "edit"; 75 | const commands = (await showMarimoControllerCommands(controller)).filter( 76 | (index) => index.if !== false, 77 | ); 78 | expect(commands.map((c) => c.label)).toMatchInlineSnapshot(` 79 | [ 80 | "", 81 | "$(split-horizontal) Open in embedded browser", 82 | "$(link-external) Open in system browser", 83 | "$(refresh) Restart marimo kernel", 84 | "$(package) Switch to run mode", 85 | "$(terminal) Show Terminal", 86 | "$(close) Stop kernel", 87 | "$(export) Export notebook as...", 88 | "", 89 | ] 90 | `); 91 | }); 92 | 93 | it("should show commands for active empty", async () => { 94 | const commands = miscCommands(serverManager).filter( 95 | (index) => index.if !== false, 96 | ); 97 | expect(commands.map((c) => c.label)).toMatchInlineSnapshot(` 98 | [ 99 | "$(question) View marimo documentation", 100 | "$(bookmark) View tutorials", 101 | "$(comment-discussion) Join Discord community", 102 | "$(settings) Edit settings", 103 | "$(info) Server status: stopped", 104 | "$(info) View extension logs", 105 | ] 106 | `); 107 | }); 108 | }); 109 | -------------------------------------------------------------------------------- /src/__mocks__/vscode.ts: -------------------------------------------------------------------------------- 1 | import { type ChildProcess, spawn } from "node:child_process"; 2 | import { createVSCodeMock as create } from "jest-mock-vscode"; 3 | import { type VitestUtils, afterAll } from "vitest"; 4 | import type { NotebookController } from "vscode"; 5 | 6 | const openProcesses: ChildProcess[] = []; 7 | 8 | export async function createVSCodeMock(vi: VitestUtils) { 9 | // biome-ignore lint/suspicious/noExplicitAny: any is ok 10 | const vscode = create(vi) as any; 11 | 12 | vscode.workspace = vscode.workspace || {}; 13 | let configMap: Record = {}; 14 | 15 | // Add createTerminal mock 16 | vscode.window.createTerminal = vi.fn().mockImplementation(() => { 17 | let proc: ChildProcess | undefined; 18 | return { 19 | processId: Promise.resolve(1), 20 | dispose: vi.fn().mockImplementation(() => { 21 | proc?.kill(); 22 | }), 23 | sendText: vi.fn().mockImplementation((args: string) => { 24 | proc = spawn(args, { shell: true }); 25 | proc.stdout?.on("data", (data) => { 26 | const line = data.toString(); 27 | if (line) { 28 | console.warn(line); 29 | } 30 | }); 31 | proc.stderr?.on("data", (data) => { 32 | const line = data.toString(); 33 | if (line) { 34 | console.warn(line); 35 | } 36 | }); 37 | proc.on("error", (error) => { 38 | if (error) { 39 | console.warn(error); 40 | } 41 | }); 42 | proc.on("close", (code) => { 43 | console.warn(`Process exited with code ${code}`); 44 | }); 45 | openProcesses.push(proc); 46 | }), 47 | show: vi.fn(), 48 | }; 49 | }); 50 | 51 | vscode.workspace.getConfiguration = vi.fn().mockImplementation(() => { 52 | return { 53 | get: vi.fn().mockImplementation((key) => configMap[key]), 54 | update: vi.fn().mockImplementation((key, value) => { 55 | configMap[key] = value; 56 | }), 57 | set: vi.fn().mockImplementation((key, value) => { 58 | configMap[key] = value; 59 | }), 60 | reset: vi.fn().mockImplementation(() => { 61 | configMap = {}; 62 | }), 63 | }; 64 | }); 65 | 66 | vscode.window.createOutputChannel.mockImplementation(() => { 67 | return { 68 | debug: vi.fn().mockImplementation((...args) => console.log(...args)), 69 | info: vi.fn().mockImplementation((...args) => console.log(...args)), 70 | error: vi.fn().mockImplementation((...args) => console.error(...args)), 71 | warn: vi.fn().mockImplementation((...args) => console.warn(...args)), 72 | createLogger: vi.fn(), 73 | }; 74 | }); 75 | 76 | vscode.env = vscode.env || {}; 77 | vscode.env.asExternalUri = vi.fn().mockImplementation(async (uri) => uri); 78 | enum QuickPickItemKind { 79 | Separator = -1, 80 | Default = 0, 81 | } 82 | vscode.QuickPickItemKind = QuickPickItemKind; 83 | vscode.window.createWebviewPanel = vi.fn().mockImplementation(() => { 84 | return { 85 | webview: { 86 | onDidReceiveMessage: vi.fn(), 87 | html: "", 88 | }, 89 | onDidDispose: vi.fn(), 90 | dispose: vi.fn(), 91 | }; 92 | }); 93 | 94 | vscode.notebooks = vscode.notebooks || {}; 95 | vscode.notebooks.createNotebookController = vi 96 | .fn() 97 | .mockImplementation((id, notebookType, label) => { 98 | const mockNotebookController: NotebookController = { 99 | id, 100 | notebookType, 101 | supportedLanguages: [], 102 | label, 103 | supportsExecutionOrder: false, 104 | createNotebookCellExecution: vi.fn(), 105 | executeHandler: vi.fn(), 106 | interruptHandler: vi.fn(), 107 | onDidChangeSelectedNotebooks: vi.fn(), 108 | updateNotebookAffinity: vi.fn(), 109 | dispose: vi.fn(), 110 | }; 111 | return mockNotebookController; 112 | }); 113 | 114 | return vscode; 115 | } 116 | 117 | afterAll(() => { 118 | openProcesses.forEach((proc) => proc.kill()); 119 | }); 120 | -------------------------------------------------------------------------------- /src/export/export-as.ts: -------------------------------------------------------------------------------- 1 | import { existsSync } from "node:fs"; 2 | import path from "node:path"; 3 | import { Uri, ViewColumn, window, workspace } from "vscode"; 4 | import { logger } from "../logger"; 5 | import { printError } from "../utils/errors"; 6 | import { execMarimoCommand, hasPythonModule } from "../utils/exec"; 7 | 8 | export type ExportType = 9 | | "ipynb" 10 | | "md" 11 | | "html" 12 | | "html-without-code" 13 | | "script"; 14 | 15 | export type ExportExtension = 16 | | "html" 17 | | "ipynb" 18 | | "md" 19 | | "script.py" 20 | | "txt" 21 | | "py"; 22 | 23 | function getExportCommand(type: ExportType): string { 24 | if (type === "html-without-code") { 25 | return "html"; 26 | } 27 | return type; 28 | } 29 | function getExportArgs(type: ExportType): string { 30 | if (type === "html-without-code") { 31 | return "--no-include-code"; 32 | } 33 | return ""; 34 | } 35 | function getExportExtension(type: ExportType): ExportExtension { 36 | if (type === "html-without-code") { 37 | return "html"; 38 | } 39 | if (type === "ipynb") { 40 | return "ipynb"; 41 | } 42 | if (type === "md") { 43 | return "md"; 44 | } 45 | if (type === "html") { 46 | return "html"; 47 | } 48 | if (type === "script") { 49 | return "script.py"; 50 | } 51 | return "txt"; 52 | } 53 | 54 | export async function exportNotebookAs( 55 | filePath: string, 56 | exportType: ExportType, 57 | ): Promise { 58 | try { 59 | // Check the requirements are met 60 | await checkRequirements(exportType); 61 | 62 | // Run export via marimo CLI 63 | const directory = path.dirname(filePath); 64 | const response = await execMarimoCommand([ 65 | "export", 66 | getExportCommand(exportType), 67 | filePath, 68 | getExportArgs(exportType), 69 | ]); 70 | 71 | const appCode = response.toString(); 72 | 73 | try { 74 | // try to save to file system 75 | const ext = getExportExtension(exportType); 76 | const currentFilename = path.basename(filePath, ".py"); 77 | const newFilename = getUniqueFilename(directory, currentFilename, ext); 78 | const newFilePath = Uri.file(path.join(directory, newFilename)); 79 | 80 | await workspace.fs.writeFile(newFilePath, Buffer.from(appCode)); 81 | 82 | // open file 83 | await workspace.openTextDocument(newFilePath).then(() => { 84 | const relativePath = workspace.asRelativePath(newFilePath); 85 | window.showInformationMessage(`Saved to ${relativePath}`); 86 | window.showTextDocument(newFilePath, { viewColumn: ViewColumn.Beside }); 87 | }); 88 | return newFilePath; 89 | } catch { 90 | // if fails to save to file system, open in new tab 91 | await workspace.openTextDocument({ content: appCode }).then((doc) => { 92 | window.showTextDocument(doc); 93 | }); 94 | return false; 95 | } 96 | } catch (error) { 97 | logger.info(error); 98 | window.showErrorMessage( 99 | `Failed to export notebook: \n${printError(error)}`, 100 | ); 101 | return false; 102 | } 103 | } 104 | 105 | async function checkRequirements(format: ExportType) { 106 | if (format === "ipynb") { 107 | // Check that nbformat is installed 108 | try { 109 | await hasPythonModule("nbformat"); 110 | } catch { 111 | throw new Error( 112 | "nbformat is not installed. Please install nbformat, e.g. `pip install nbformat`", 113 | ); 114 | } 115 | } 116 | } 117 | 118 | export function getUniqueFilename( 119 | directory: string, 120 | filename: string, 121 | extension: string, 122 | postfix?: number, 123 | ) { 124 | const uniqueFilename = postfix 125 | ? `${filename}_${postfix}.${extension}` 126 | : `${filename}.${extension}`; 127 | 128 | // If the file already exists, try again with a higher postfix 129 | if (existsSync(path.join(directory, uniqueFilename))) { 130 | return getUniqueFilename( 131 | directory, 132 | filename, 133 | extension, 134 | postfix ? postfix + 1 : 1, 135 | ); 136 | } 137 | 138 | return uniqueFilename; 139 | } 140 | -------------------------------------------------------------------------------- /src/launcher/__tests__/utils.test.ts: -------------------------------------------------------------------------------- 1 | import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; 2 | import { createVSCodeMock } from "../../__mocks__/vscode"; 3 | 4 | vi.mock("vscode", () => createVSCodeMock(vi)); 5 | 6 | import { parse } from "node-html-parser"; 7 | import { composeUrl } from "../../config"; 8 | import { fetchMarimoStartupValues } from "../utils"; 9 | 10 | vi.mock("node-html-parser"); 11 | vi.mock("../../config"); 12 | 13 | describe("fetchMarimoStartupValues", () => { 14 | const mockFetch = vi.fn(); 15 | global.fetch = mockFetch; 16 | 17 | beforeEach(() => { 18 | vi.resetAllMocks(); 19 | }); 20 | 21 | afterEach(() => { 22 | vi.restoreAllMocks(); 23 | }); 24 | 25 | it("should fetch and parse marimo startup values correctly", async () => { 26 | // Mock the composeUrl function 27 | vi.mocked(composeUrl).mockResolvedValue("http://localhost:1234"); 28 | 29 | // Mock the fetch response 30 | const mockHtml = ` 31 | 32 | 33 | 34 | 35 | 36 | `; 37 | mockFetch.mockResolvedValue({ 38 | ok: true, 39 | text: () => Promise.resolve(mockHtml), 40 | url: "http://localhost:1234", 41 | }); 42 | 43 | // Mock the parse function 44 | const mockRoot = { 45 | querySelector: (selector: string) => ({ 46 | getAttribute: (attr: string) => { 47 | if (selector === "marimo-server-token" && attr === "data-token") 48 | return "mock-skew-token"; 49 | if (selector === "marimo-user-config" && attr === "data-config") 50 | return '{"key": "value"}'; 51 | if (selector === "marimo-version" && attr === "data-version") 52 | return "1.0.0"; 53 | return null; 54 | }, 55 | }), 56 | }; 57 | vi.mocked(parse).mockReturnValue(mockRoot as any); 58 | 59 | const result = await fetchMarimoStartupValues({ port: 1234, backoff: 1 }); 60 | 61 | expect(result).toEqual({ 62 | skewToken: "mock-skew-token", 63 | version: "1.0.0", 64 | userConfig: { key: "value" }, 65 | }); 66 | 67 | expect(composeUrl).toHaveBeenCalledWith(1234); 68 | expect(mockFetch).toHaveBeenCalledWith("http://localhost:1234/", { 69 | headers: { 70 | Accept: "text/html", 71 | }, 72 | }); 73 | expect(parse).toHaveBeenCalledWith(mockHtml); 74 | }); 75 | 76 | it("should throw an error if fetch fails", async () => { 77 | vi.mocked(composeUrl).mockResolvedValue("http://localhost:1234"); 78 | mockFetch.mockResolvedValue({ 79 | ok: false, 80 | status: 404, 81 | statusText: "Not Found", 82 | }); 83 | 84 | await expect( 85 | fetchMarimoStartupValues({ port: 1234, backoff: 1 }), 86 | ).rejects.toThrow( 87 | "Could not fetch http://localhost:1234/. Is http://localhost:1234/ healthy?", 88 | ); 89 | }); 90 | 91 | it("should throw an error if redirected to auth/login", async () => { 92 | vi.mocked(composeUrl).mockResolvedValue("http://localhost:1234"); 93 | mockFetch.mockResolvedValue({ 94 | ok: true, 95 | url: "http://localhost:1234/auth/login", 96 | }); 97 | 98 | await expect( 99 | fetchMarimoStartupValues({ port: 1234, backoff: 1 }), 100 | ).rejects.toThrow( 101 | "An existing marimo server created outside of vscode is running at this url: http://localhost:1234/", 102 | ); 103 | }); 104 | 105 | it("should throw an error if required elements are missing", async () => { 106 | vi.mocked(composeUrl).mockResolvedValue("http://localhost:1234"); 107 | mockFetch.mockResolvedValue({ 108 | ok: true, 109 | text: () => Promise.resolve(""), 110 | url: "http://localhost:1234/", 111 | }); 112 | 113 | vi.mocked(parse).mockReturnValue({ 114 | querySelector: () => null, 115 | } as any); 116 | 117 | await expect( 118 | fetchMarimoStartupValues({ port: 1234, backoff: 1 }), 119 | ).rejects.toThrow( 120 | "Could not find marimo-server-token. Is http://localhost:1234/ healthy?", 121 | ); 122 | }); 123 | }); 124 | -------------------------------------------------------------------------------- /src/notebook/extension.ts: -------------------------------------------------------------------------------- 1 | import * as vscode from "vscode"; 2 | import { logger } from "../logger"; 3 | import type { ServerManager } from "../services/server-manager"; 4 | import { isMarimoApp } from "../utils/query"; 5 | import { showNotebookDocument } from "../utils/show"; 6 | import { 7 | type NotebookMetadata, 8 | getNotebookMetadata, 9 | setNotebookMetadata, 10 | } from "./common/metadata"; 11 | import { NOTEBOOK_TYPE, PYTHON_LANGUAGE_ID } from "./constants"; 12 | import { KernelManager } from "./kernel-manager"; 13 | 14 | // export async function createNotebookDocument() { 15 | // const data = new vscode.NotebookData([]); 16 | // setNotebookMetadata(data, { 17 | // isNew: true, 18 | // loaded: false, 19 | // }); 20 | 21 | // // Create NotebookDocument and open it 22 | // const doc = await vscode.workspace.openNotebookDocument(NOTEBOOK_TYPE, data); 23 | 24 | // // Open Notebook 25 | // logger.info("Opening new marimo notebook"); 26 | // await vscode.window.showNotebookDocument(doc); 27 | // } 28 | 29 | export async function getActiveMarimoFile() { 30 | const editor = vscode.window.activeTextEditor; 31 | if (!editor) { 32 | vscode.window.showErrorMessage("No active editor!"); 33 | return; 34 | } 35 | 36 | const document = editor.document; 37 | if (!isMarimoApp(document, false)) { 38 | vscode.window.showErrorMessage("Active editor is not a Marimo file!"); 39 | return; 40 | } 41 | 42 | return document.uri; 43 | } 44 | 45 | export async function openMarimoNotebookDocument(uri: vscode.Uri | undefined) { 46 | if (!uri) { 47 | return; 48 | } 49 | // Create data 50 | const data = new vscode.NotebookData([]); 51 | setNotebookMetadata(data, { 52 | isNew: false, 53 | loaded: false, 54 | file: uri.fsPath, 55 | }); 56 | 57 | // Open notebook 58 | logger.info("Opening existing marimo notebook"); 59 | const doc = await vscode.workspace.openNotebookDocument(uri); 60 | // Show the notebook, if not shown 61 | await showNotebookDocument(doc); 62 | 63 | // Open the panel if the kernel is still active 64 | const kernel = KernelManager.instance.getKernelByUri(uri); 65 | if (kernel) { 66 | await kernel.openKiosk(); 67 | } 68 | } 69 | 70 | export async function handleOnOpenNotebookDocument( 71 | doc: vscode.NotebookDocument, 72 | serverManager: ServerManager, 73 | kernelManager: KernelManager, 74 | ) { 75 | logger.info("Opened notebook document", doc.notebookType); 76 | if (doc.notebookType !== NOTEBOOK_TYPE) { 77 | logger.info("Not a marimo notebook", doc.notebookType); 78 | return; 79 | } 80 | 81 | const metadata = getNotebookMetadata(doc); 82 | 83 | if (metadata.loaded) { 84 | logger.info("Notebook already loaded", metadata); 85 | return; 86 | } 87 | 88 | // Show the notebook, if not shown 89 | await showNotebookDocument(doc); 90 | 91 | await vscode.window.withProgress( 92 | { 93 | location: vscode.ProgressLocation.Notification, 94 | title: "Starting marimo server...", 95 | cancellable: true, 96 | }, 97 | async (_, cancellationToken) => { 98 | // Start Marimo server 99 | logger.info("Checking server..."); 100 | const { port, skewToken, userConfig, version } = 101 | await serverManager.start(cancellationToken); 102 | // If not new, try to hydrate existing notebooks 103 | if (!metadata.isNew) { 104 | await kernelManager.hydrateExistingNotebooks({ 105 | port, 106 | skewToken, 107 | userConfig, 108 | version, 109 | }); 110 | } 111 | 112 | // Create Kernel 113 | const kernel = kernelManager.createKernel({ 114 | port, 115 | uri: metadata.isNew ? "__new__" : doc.uri, 116 | skewToken, 117 | version, 118 | userConfig, 119 | notebookDoc: doc, 120 | }); 121 | 122 | // Edit metadata 123 | const nextMetadata: NotebookMetadata = { 124 | ...metadata, 125 | port: port, 126 | key: kernel.kernelKey, 127 | loaded: true, 128 | }; 129 | const nbEdit = vscode.NotebookEdit.updateNotebookMetadata(nextMetadata); 130 | const edit2 = new vscode.WorkspaceEdit(); 131 | edit2.set(doc.uri, [nbEdit]); 132 | await vscode.workspace.applyEdit(edit2); 133 | 134 | // Start the kernel 135 | await kernel.start(); 136 | await kernel.openKiosk(); 137 | }, 138 | ); 139 | } 140 | -------------------------------------------------------------------------------- /src/utils/exec.ts: -------------------------------------------------------------------------------- 1 | import { execFileSync } from "node:child_process"; 2 | import { workspace } from "vscode"; 3 | import { Config } from "../config"; 4 | import { logger } from "../logger"; 5 | import { getInterpreterDetails } from "./python"; 6 | 7 | function execWithLogger(command: string, args: string[]) { 8 | const cleanedArgs = args.map((arg) => arg.trim()).filter(Boolean); 9 | logger.info(`Executing: ${command} ${JSON.stringify(cleanedArgs)}`); 10 | return execFileSync(command, cleanedArgs, { 11 | shell: false, 12 | encoding: "utf8", 13 | env: { 14 | ...process.env, 15 | // Force dangerous sandbox mode: https://github.com/marimo-team/marimo/pull/5958 16 | MARIMO_DANGEROUS_SANDBOX: "1", 17 | }, 18 | }); 19 | } 20 | 21 | export async function execMarimoCommand(command: string[]): Promise { 22 | // When marimoPath is set, use that directly 23 | const cmd = Config.marimoPath; 24 | if (cmd && cmd !== "marimo") { 25 | if (cmd.startsWith("uv ") || cmd.startsWith("uvx ")) { 26 | const [uvCmd, ...uvArgs] = cmd.split(" "); 27 | return execWithLogger(uvCmd, [...uvArgs, ...command]); 28 | } 29 | return execWithLogger(cmd, command); 30 | } 31 | // Otherwise, use python -m marimo 32 | return execPythonModule(["marimo", ...command]); 33 | } 34 | 35 | /** 36 | * Execute a python module 37 | * e.g. /usr/bin/python -m marimo edit 38 | * or 39 | * e.g. uv run python -m marimo edit 40 | */ 41 | export async function execPythonModule(command: string[]) { 42 | // Otherwise use python interpreter 43 | const interpreter = (await getInterpreter()) || "python"; 44 | logger.info(`Using interpreter: ${interpreter}`); 45 | 46 | // Handle uv/uvx commands specially 47 | if (interpreter.startsWith("uv ") || interpreter.startsWith("uvx ")) { 48 | const [uvCmd, ...uvArgs] = interpreter.split(" "); 49 | return execWithLogger(uvCmd, [...uvArgs, "-m", ...command]); 50 | } 51 | 52 | return execWithLogger(interpreter, ["-m", ...command]); 53 | } 54 | 55 | export async function execPythonFile(command: string[]) { 56 | const interpreter = (await getInterpreter()) || "python"; 57 | logger.info(`Using interpreter: ${interpreter}`); 58 | 59 | // Handle uv/uvx commands specially 60 | if (interpreter.startsWith("uv ") || interpreter.startsWith("uvx ")) { 61 | const [uvCmd, ...uvArgs] = interpreter.split(" "); 62 | return execWithLogger(uvCmd, [...uvArgs, ...command]); 63 | } 64 | 65 | return execWithLogger(interpreter, command); 66 | } 67 | 68 | export async function hasPythonModule(module: string) { 69 | const interpreter = (await getInterpreter()) || "python"; 70 | logger.info(`Using interpreter: ${interpreter}`); 71 | 72 | // Handle uv/uvx commands specially 73 | if (interpreter.startsWith("uv ") || interpreter.startsWith("uvx ")) { 74 | const [uvCmd, ...uvArgs] = interpreter.split(" "); 75 | return execWithLogger(uvCmd, [...uvArgs, "-c", `import ${module}`]); 76 | } 77 | 78 | return execWithLogger(interpreter, ["-c", `import ${module}`]); 79 | } 80 | 81 | // Quote if it has spaces and is not a command like "uv run" 82 | export function maybeQuotes(command: string) { 83 | if ( 84 | command.includes(" ") && 85 | !command.startsWith("uv ") && 86 | !command.startsWith("uvx ") 87 | ) { 88 | return `"${command}"`; 89 | } 90 | return command; 91 | } 92 | 93 | export async function hasExecutable(executable: string): Promise { 94 | try { 95 | await execWithLogger(executable, ["--help"]); 96 | return true; 97 | } catch (error) { 98 | return false; 99 | } 100 | } 101 | 102 | export async function getInterpreter(): Promise { 103 | try { 104 | if (Config.pythonPath) { 105 | return Config.pythonPath; 106 | } 107 | const activeWorkspace = workspace.workspaceFolders?.[0]; 108 | 109 | if (activeWorkspace) { 110 | const interpreters = (await getInterpreterDetails(activeWorkspace.uri)) 111 | .path; 112 | 113 | if (!interpreters) { 114 | logger.error("No interpreters found"); 115 | return undefined; 116 | } 117 | 118 | logger.debug("Found interpreters", interpreters); 119 | 120 | if (interpreters.length > 0) { 121 | return interpreters[0]; 122 | } 123 | } 124 | } catch (error) { 125 | logger.error("Error getting interpreter: ", error); 126 | } 127 | 128 | return undefined; 129 | } 130 | -------------------------------------------------------------------------------- /src/utils/__tests__/config.test.ts: -------------------------------------------------------------------------------- 1 | import { type Mock, beforeEach, describe, expect, test, vi } from "vitest"; 2 | import { createVSCodeMock } from "../../__mocks__/vscode"; 3 | 4 | vi.mock("vscode", () => createVSCodeMock(vi)); 5 | 6 | import * as vscode from "vscode"; 7 | import { Config, composeUrl, composeWsUrl, getConfig } from "../../config"; 8 | 9 | describe("Config", () => { 10 | beforeEach(() => { 11 | vi.clearAllMocks(); 12 | }); 13 | 14 | test("getConfig returns correct values", () => { 15 | const mockGet = vi.fn(); 16 | (vscode.workspace.getConfiguration as Mock).mockReturnValue({ 17 | get: mockGet, 18 | }); 19 | 20 | mockGet.mockReturnValueOnce("test-value"); 21 | expect(getConfig("testKey")).toBe("test-value"); 22 | 23 | mockGet.mockReturnValueOnce(undefined); 24 | expect(getConfig("testKey", "default")).toBe("default"); 25 | }); 26 | 27 | test("Config.root returns correct value", () => { 28 | expect(Config.root).toBe(""); 29 | }); 30 | 31 | test("Config properties return correct values", () => { 32 | const mockGet = vi.fn(); 33 | (vscode.workspace.getConfiguration as Mock).mockReturnValue({ 34 | get: mockGet, 35 | }); 36 | 37 | mockGet.mockReturnValueOnce("system"); 38 | expect(Config.browser).toBe("system"); 39 | 40 | mockGet.mockReturnValueOnce(3000); 41 | expect(Config.port).toBe(3000); 42 | 43 | mockGet.mockReturnValueOnce("127.0.0.1"); 44 | expect(Config.host).toBe("127.0.0.1"); 45 | 46 | mockGet.mockReturnValueOnce(true); 47 | expect(Config.https).toBe(true); 48 | 49 | mockGet.mockReturnValueOnce(true); 50 | expect(Config.enableToken).toBe(true); 51 | 52 | mockGet.mockReturnValueOnce("secret"); 53 | expect(Config.tokenPassword).toBe("secret"); 54 | 55 | mockGet.mockReturnValueOnce(true); 56 | expect(Config.debug).toBe(true); 57 | 58 | mockGet.mockReturnValueOnce("/usr/bin/python"); 59 | expect(Config.pythonPath).toBe("/usr/bin/python"); 60 | 61 | mockGet.mockReturnValueOnce("custom-marimo"); 62 | expect(Config.marimoPath).toBe("custom-marimo"); 63 | mockGet.mockReturnValueOnce("uvx run marimo"); 64 | expect(Config.marimoPath).toBe("uvx run marimo"); 65 | mockGet.mockReturnValueOnce("marimo"); 66 | expect(Config.marimoPath).toBe(undefined); 67 | 68 | mockGet.mockReturnValueOnce(true); 69 | expect(Config.showTerminal).toBe(true); 70 | }); 71 | 72 | test("Config.readPort returns correct value", () => { 73 | const mockGet = vi.fn().mockReturnValue(3000); 74 | (vscode.workspace.getConfiguration as Mock).mockReturnValue({ 75 | get: mockGet, 76 | }); 77 | 78 | expect(Config.readPort).toBe(3010); 79 | }); 80 | }); 81 | 82 | describe("URL composition functions", () => { 83 | beforeEach(() => { 84 | vi.clearAllMocks(); 85 | }); 86 | 87 | test("composeUrl returns correct URL", async () => { 88 | const mockGet = vi.fn().mockReturnValue(false); // https: false 89 | (vscode.workspace.getConfiguration as Mock).mockReturnValue({ 90 | get: mockGet, 91 | }); 92 | (vscode.env.asExternalUri as Mock).mockResolvedValue({ 93 | toString: () => "http://external-host:3000/", 94 | }); 95 | 96 | const url = await composeUrl(3000); 97 | expect(url).toBe("http://external-host:3000/"); 98 | }); 99 | 100 | test("composeWsUrl returns correct WebSocket URL", async () => { 101 | const mockGet = vi.fn().mockReturnValue(true); // https: true 102 | (vscode.workspace.getConfiguration as Mock).mockReturnValue({ 103 | get: mockGet, 104 | }); 105 | (vscode.env.asExternalUri as Mock).mockResolvedValue({ 106 | toString: () => "wss://external-host:3000/", 107 | }); 108 | 109 | const url = await composeWsUrl(3000); 110 | expect(url).toBe("wss://external-host:3000/"); 111 | }); 112 | 113 | test("composeUrl and composeWsUrl handle errors", async () => { 114 | const mockGet = vi.fn().mockReturnValue(false); // https: false 115 | (vscode.workspace.getConfiguration as Mock).mockReturnValue({ 116 | get: mockGet, 117 | }); 118 | (vscode.env.asExternalUri as Mock).mockRejectedValue( 119 | new Error("Test error"), 120 | ); 121 | 122 | const httpUrl = await composeUrl(3000); 123 | expect(httpUrl).toBe("http://localhost:3000/"); 124 | 125 | const wsUrl = await composeWsUrl(3000); 126 | expect(wsUrl).toBe("ws://localhost:3000/"); 127 | }); 128 | }); 129 | -------------------------------------------------------------------------------- /src/notebook/__tests__/md.test.ts: -------------------------------------------------------------------------------- 1 | import { describe, expect, it } from "vitest"; 2 | import { maybeMarkdown, toMarkdown } from "../md"; 3 | 4 | describe("markdown utils", () => { 5 | describe("toMarkdown", () => { 6 | it("should wrap single line in mo.md()", () => { 7 | expect(toMarkdown("hello")).toBe('mo.md(r"hello")'); 8 | }); 9 | 10 | it("should handle empty string", () => { 11 | expect(toMarkdown("")).toBe('mo.md(r"")'); 12 | }); 13 | 14 | it("should handle whitespace", () => { 15 | expect(toMarkdown(" hello ")).toBe('mo.md(r"hello")'); 16 | }); 17 | 18 | it("should indent multiline content with 4 spaces", () => { 19 | const input = `line 1 20 | line 2 21 | line 3`; 22 | expect(toMarkdown(input)).toMatchInlineSnapshot(` 23 | "mo.md( 24 | r""" 25 | line 1 26 | line 2 27 | line 3 28 | """ 29 | )" 30 | `); 31 | }); 32 | 33 | it("should preserve existing indentation in multiline", () => { 34 | const input = `line 1 35 | line 2 36 | line 3`; 37 | expect(toMarkdown(input)).toMatchInlineSnapshot(` 38 | "mo.md( 39 | r""" 40 | line 1 41 | line 2 42 | line 3 43 | """ 44 | )" 45 | `); 46 | }); 47 | it("should handle multiline with empty lines", () => { 48 | const input = `line 1 49 | 50 | line 2 51 | 52 | line 3`; 53 | expect(trimEmptyLines(toMarkdown(input))).toMatchInlineSnapshot(` 54 | "mo.md( 55 | r""" 56 | line 1 57 | 58 | line 2 59 | 60 | line 3 61 | """ 62 | )" 63 | `); 64 | }); 65 | 66 | it("should handle multiline starting with empty line", () => { 67 | const input = ` 68 | line 1 69 | line 2`; 70 | expect(toMarkdown(input)).toMatchInlineSnapshot(` 71 | "mo.md( 72 | r""" 73 | line 1 74 | line 2 75 | """ 76 | )" 77 | `); 78 | }); 79 | 80 | it("should handle multiline ending with empty line", () => { 81 | const input = `line 1 82 | line 2 83 | 84 | `; 85 | expect(toMarkdown(input)).toMatchInlineSnapshot(` 86 | "mo.md( 87 | r""" 88 | line 1 89 | line 2 90 | """ 91 | )" 92 | `); 93 | }); 94 | }); 95 | 96 | describe("maybeMarkdown", () => { 97 | it("should extract content from mo.md() call", () => { 98 | expect(maybeMarkdown('mo.md("hello")')).toBe("hello"); 99 | }); 100 | 101 | it("should handle empty mo.md()", () => { 102 | expect(maybeMarkdown('mo.md("")')).toBe(""); 103 | }); 104 | 105 | it("should handle multiline mo.md()", () => { 106 | const input = `mo.md(""" 107 | line 1 108 | line 2 109 | line 3 110 | """)`; 111 | expect(maybeMarkdown(input)).toBe("line 1\nline 2\nline 3"); 112 | }); 113 | 114 | it("should handle raw strings", () => { 115 | expect(maybeMarkdown('mo.md(r"hello")')).toBe("hello"); 116 | }); 117 | 118 | it("should handle triple quotes", () => { 119 | expect(maybeMarkdown('mo.md("""hello""")')).toBe("hello"); 120 | }); 121 | 122 | it("should return null for invalid markdown", () => { 123 | expect(maybeMarkdown("not markdown")).toBe(null); 124 | expect(maybeMarkdown('mo.md("unclosed')).toBe(null); 125 | expect(maybeMarkdown('mo.md(hello")')).toBe(null); 126 | expect(maybeMarkdown('md("hello")')).toBe(null); 127 | }); 128 | 129 | it("should dedent multiline content", () => { 130 | const input = `mo.md(""" 131 | line 1 132 | line 2 133 | line 3 134 | """)`; 135 | expect(maybeMarkdown(input)).toBe("line 1\n line 2\n line 3"); 136 | }); 137 | 138 | it("should handle mixed quotes", () => { 139 | expect(maybeMarkdown(`mo.md('hello')`)).toBe("hello"); 140 | expect(maybeMarkdown(`mo.md("'hello'")`)).toBe("'hello'"); 141 | }); 142 | 143 | it("should handle whitespace", () => { 144 | expect(maybeMarkdown(' mo.md("hello") ')).toBe("hello"); 145 | expect(maybeMarkdown('mo.md( "hello" )')).toBe("hello"); 146 | }); 147 | }); 148 | }); 149 | 150 | function trimEmptyLines(text: string): string { 151 | return text 152 | .split("\n") 153 | .map((line) => (line.trim() ? line : "")) 154 | .join("\n") 155 | .trim(); 156 | } 157 | -------------------------------------------------------------------------------- /src/launcher/terminal.ts: -------------------------------------------------------------------------------- 1 | import path from "node:path"; 2 | import { 3 | type CancellationToken, 4 | type Disposable, 5 | type Terminal, 6 | type Uri, 7 | window, 8 | } from "vscode"; 9 | import { Config } from "../config"; 10 | import { getGlobalState } from "../ctx"; 11 | import { logger } from "../logger"; 12 | import { LogMethodCalls } from "../utils/log"; 13 | import { wait } from "../utils/wait"; 14 | 15 | export interface IMarimoTerminal extends Disposable { 16 | show(): void; 17 | relativePathFor(file: Uri): string; 18 | is(term: Terminal): boolean; 19 | tryRecoverTerminal(): Promise; 20 | executeCommand(cmd: string, token?: CancellationToken): Promise; 21 | } 22 | 23 | export class MarimoTerminal implements IMarimoTerminal { 24 | private terminal: Terminal | undefined; 25 | private logger = logger.createLogger(this.appName); 26 | 27 | constructor( 28 | private fsPath: string, 29 | private cwd: string | undefined, 30 | private appName: string, 31 | ) {} 32 | 33 | private async ensureTerminal() { 34 | if (this.isTerminalActive()) { 35 | return; 36 | } 37 | 38 | this.terminal = window.createTerminal({ 39 | name: `marimo ${this.appName}`, 40 | cwd: this.cwd, 41 | }); 42 | // Wait 1s for the terminal to 'source' the virtualenv 43 | await wait(1000); 44 | } 45 | 46 | relativePathFor(file: Uri): string { 47 | if (this.cwd && file.fsPath.startsWith(this.cwd)) { 48 | return path.relative(this.cwd, file.fsPath); 49 | } 50 | return file.fsPath; 51 | } 52 | 53 | private isTerminalActive() { 54 | return this.terminal && this.terminal.exitStatus === undefined; 55 | } 56 | 57 | @LogMethodCalls() 58 | dispose() { 59 | this.endProcess(); 60 | this.terminal?.dispose(); 61 | this.terminal = undefined; 62 | this.logger.info("terminal disposed"); 63 | } 64 | 65 | @LogMethodCalls() 66 | async show() { 67 | await this.ensureTerminal(); 68 | this.terminal?.show(true); 69 | } 70 | 71 | private endProcess() { 72 | if (this.isTerminalActive()) { 73 | this.terminal?.sendText("\u0003"); 74 | this.terminal?.sendText("\u0003"); 75 | } 76 | getGlobalState().update(this.keyFor("pid"), undefined); 77 | } 78 | 79 | is(term: Terminal) { 80 | return this.terminal === term; 81 | } 82 | 83 | @LogMethodCalls() 84 | async tryRecoverTerminal(): Promise { 85 | if (this.terminal) { 86 | return false; 87 | } 88 | 89 | const pid = getGlobalState().get(this.keyFor("pid")); 90 | if (!pid) { 91 | return false; 92 | } 93 | this.logger.debug("recovered pid", pid); 94 | 95 | const terminals = await Promise.all( 96 | window.terminals.map(async (index) => 97 | pid === (await index.processId) ? index : undefined, 98 | ), 99 | ); 100 | 101 | const terminal = terminals.find(Boolean); 102 | 103 | if (terminal) { 104 | this.terminal = terminal; 105 | return true; 106 | } 107 | return false; 108 | } 109 | 110 | @LogMethodCalls() 111 | async executeCommand(cmd: string, token?: CancellationToken) { 112 | await this.ensureTerminal(); 113 | if (!this.terminal) { 114 | this.logger.error("terminal not found"); 115 | return; 116 | } 117 | 118 | // Set up abort handler 119 | if (token) { 120 | if (token.isCancellationRequested) { 121 | throw new Error("Command aborted before execution"); 122 | } 123 | token.onCancellationRequested(() => { 124 | this.logger.info("Command aborted"); 125 | this.endProcess(); 126 | }); 127 | } 128 | 129 | try { 130 | // Check if we're in PowerShell and format command accordingly 131 | const isPowerShell = this.terminal.name 132 | ?.toLowerCase() 133 | .includes("powershell"); 134 | 135 | if (isPowerShell) { 136 | this.logger.info( 137 | "Detected PowerShell terminal, running command with &", 138 | ); 139 | } 140 | 141 | const formattedCmd = isPowerShell ? `& ${cmd}` : cmd; 142 | 143 | this.terminal.sendText(formattedCmd); 144 | 145 | if (Config.showTerminal) { 146 | this.terminal.show(true); 147 | } 148 | await wait(2_000); 149 | const pid = await this.terminal.processId; 150 | if (pid) { 151 | getGlobalState().update(this.keyFor("pid"), pid); 152 | } 153 | } catch (error) { 154 | this.terminal.show(true); 155 | window.showErrorMessage(`Command failed: ${error}`); 156 | throw error; 157 | } 158 | } 159 | 160 | private keyFor(key: string) { 161 | return `marimo.${this.fsPath}.${key}`; 162 | } 163 | } 164 | -------------------------------------------------------------------------------- /src/config.ts: -------------------------------------------------------------------------------- 1 | import { Uri, env, workspace } from "vscode"; 2 | import { logger } from "./logger"; 3 | 4 | export function getConfig(key: string): T | undefined; 5 | export function getConfig(key: string, v: T): T; 6 | export function getConfig(key: string, v?: T) { 7 | return workspace.getConfiguration().get(`marimo.${key}`, v) ?? v; 8 | } 9 | 10 | export interface Config { 11 | readonly root: string; 12 | readonly browser: "system" | "embedded"; 13 | readonly showTerminal: boolean; 14 | readonly debug: boolean; 15 | readonly pythonPath: string | undefined; 16 | readonly port: number; 17 | readonly readPort: number; 18 | readonly host: string; 19 | readonly marimoPath: string | undefined; 20 | readonly enableToken: boolean; 21 | readonly tokenPassword: string | undefined; 22 | readonly https: boolean; 23 | readonly sandbox: boolean; 24 | readonly watch: boolean; 25 | readonly telemetry: boolean; 26 | } 27 | 28 | /** 29 | * Configuration options for the marimo extension. 30 | * These options can be set in the user's settings.json file. 31 | */ 32 | export const Config = { 33 | get root() { 34 | return workspace.workspaceFolders?.[0]?.uri?.fsPath || ""; 35 | }, 36 | // Browser settings 37 | /** 38 | * The type of browser to use for opening marimo apps. 39 | * @default "embedded" 40 | */ 41 | get browser(): "embedded" | "system" { 42 | return getConfig("browserType", "embedded"); 43 | }, 44 | 45 | // Server settings 46 | /** 47 | * The port to use for the marimo server. 48 | * @default 2818 49 | */ 50 | get port(): number { 51 | return getConfig("port", 2818); 52 | }, 53 | 54 | get readPort() { 55 | if (typeof Config.port === "string") { 56 | return Number.parseInt(Config.port) + 10; 57 | } 58 | return Config.port + 10; 59 | }, 60 | 61 | /** 62 | * The hostname to use for the marimo server. 63 | * @default "localhost" 64 | */ 65 | get host(): string { 66 | return getConfig("host", "localhost") || "localhost"; 67 | }, 68 | 69 | /** 70 | * Whether to use HTTPS for the marimo server. 71 | * @default false 72 | */ 73 | get https(): boolean { 74 | return getConfig("https", false); 75 | }, 76 | 77 | // Authentication settings 78 | /** 79 | * Whether to enable token authentication for the marimo server. 80 | * @default false 81 | */ 82 | get enableToken(): boolean { 83 | return getConfig("enableToken", false); 84 | }, 85 | 86 | /** 87 | * The token password to use for authentication. 88 | * @default "" 89 | */ 90 | get tokenPassword(): string { 91 | return getConfig("tokenPassword", ""); 92 | }, 93 | 94 | // Debug settings 95 | /** 96 | * Whether to enable debug mode. 97 | * @default false 98 | */ 99 | get debug(): boolean { 100 | return getConfig("debug", false); 101 | }, 102 | 103 | // Python settings 104 | /** 105 | * The path to the Python interpreter to use. 106 | * @default undefined (use the default Python interpreter) 107 | */ 108 | get pythonPath(): string | undefined { 109 | return getConfig("pythonPath"); 110 | }, 111 | 112 | /** 113 | * The path to the marimo package to use. 114 | * @default undefined (use the default marimo package) 115 | */ 116 | get marimoPath(): string | undefined { 117 | const path: string | undefined = getConfig("marimoPath"); 118 | // Ignore just 'marimo' 119 | if (path === "marimo") { 120 | return undefined; 121 | } 122 | return path; 123 | }, 124 | 125 | // UI settings 126 | /** 127 | * Whether to show the terminal when the server starts. 128 | * @default true or when debug is enabled 129 | */ 130 | get showTerminal(): boolean { 131 | return getConfig("showTerminal", true) || Config.debug; 132 | }, 133 | 134 | /** 135 | * Whether to send anonymous usage data to marimo. 136 | * @default true 137 | */ 138 | get telemetry(): boolean { 139 | return getConfig("telemetry", true); 140 | }, 141 | 142 | /** 143 | * Whether to always start marimo in sandbox mode. 144 | * @default false 145 | */ 146 | get sandbox(): boolean { 147 | return getConfig("sandbox", false); 148 | }, 149 | 150 | /** 151 | * Whether to always start marimo with the --watch flag. 152 | * @default true 153 | */ 154 | get watch(): boolean { 155 | return getConfig("watch", true); 156 | }, 157 | }; 158 | 159 | export async function composeUrl(port: number): Promise { 160 | const url = `${Config.https ? "https" : "http"}://${Config.host}:${port}/`; 161 | try { 162 | const externalUri = await env.asExternalUri(Uri.parse(url)); 163 | const externalUrl = externalUri.toString(); 164 | if (externalUrl !== url) { 165 | logger.info("Mapping to external url", externalUrl, "from", url); 166 | } 167 | return externalUrl; 168 | } catch (e) { 169 | logger.error("Failed to create external url", url, e); 170 | return url; 171 | } 172 | } 173 | 174 | export async function composeWsUrl(port: number): Promise { 175 | const url = `${Config.https ? "wss" : "ws"}://${Config.host}:${port}/`; 176 | try { 177 | const externalUri = await env.asExternalUri(Uri.parse(url)); 178 | const externalUrl = externalUri.toString(); 179 | if (externalUrl !== url) { 180 | logger.info("Mapping to external url", externalUrl, "from", url); 181 | } 182 | return externalUrl; 183 | } catch (e) { 184 | logger.error("Failed to create external url", url, e); 185 | return url; 186 | } 187 | } 188 | -------------------------------------------------------------------------------- /src/__integration__/exec.test.ts: -------------------------------------------------------------------------------- 1 | import { beforeEach, describe, expect, it, vi } from "vitest"; 2 | import { createVSCodeMock } from "../__mocks__/vscode"; 3 | 4 | vi.mock("vscode", () => createVSCodeMock(vi)); 5 | vi.mock("@vscode/python-extension", () => ({})); 6 | 7 | import { workspace } from "vscode"; 8 | import { 9 | execMarimoCommand, 10 | execPythonFile, 11 | execPythonModule, 12 | hasExecutable, 13 | hasPythonModule, 14 | } from "../utils/exec"; 15 | 16 | beforeEach(() => { 17 | workspace.getConfiguration().reset(); 18 | }); 19 | 20 | describe("execMarimoCommand integration tests", () => { 21 | it("should work out of the box", async () => { 22 | const output = await execMarimoCommand(["--version"]); 23 | expect(output.toString()).toMatch(/\d+\.\d+\.\d+/); 24 | }); 25 | 26 | it.each(["python", "uv run python"])( 27 | "should run marimo --version with %s", 28 | async (interpreter) => { 29 | workspace.getConfiguration().set("marimo.pythonPath", interpreter); 30 | const output = await execMarimoCommand(["--version"]); 31 | expect(output.toString()).toMatch(/\d+\.\d+\.\d+/); 32 | }, 33 | ); 34 | 35 | it.each(["uv run marimo", "uvx marimo", "uv tool run marimo"])( 36 | "should run marimo --version with %s", 37 | async (interpreter) => { 38 | workspace.getConfiguration().set("marimo.marimoPath", interpreter); 39 | const output = await execMarimoCommand(["--version"]); 40 | expect(output.toString()).toMatch(/\d+\.\d+\.\d+/); 41 | }, 42 | ); 43 | 44 | it("path path gets overridden by marimoPath", async () => { 45 | workspace.getConfiguration().set("marimo.pythonPath", "doesnotexist"); 46 | expect(() => execMarimoCommand(["--version"])).rejects.toThrow( 47 | /doesnotexist/, 48 | ); 49 | 50 | workspace.getConfiguration().set("marimo.marimoPath", "uv run marimo"); 51 | const output = await execMarimoCommand(["--version"]); 52 | expect(output.toString()).toMatch(/\d+\.\d+\.\d+/); 53 | }); 54 | 55 | it("should handle spaces in marimoPath", async () => { 56 | workspace 57 | .getConfiguration() 58 | .set("marimo.marimoPath", "/path with spaces/marimo"); 59 | await expect(execMarimoCommand(["--version"])).rejects.toThrow(/ENOENT/); 60 | }); 61 | 62 | it("should handle invalid marimoPath", async () => { 63 | workspace.getConfiguration().set("marimo.marimoPath", "nonexistent"); 64 | await expect(execMarimoCommand(["--version"])).rejects.toThrow(/ENOENT/); 65 | }); 66 | }); 67 | 68 | describe("execPythonModule integration tests", () => { 69 | it("should work out of the box", async () => { 70 | const output = await execPythonModule(["marimo", "--version"]); 71 | expect(output.toString()).toMatch(/\d+\.\d+\.\d+/); 72 | }); 73 | 74 | it.each(["python", "uv run python"])( 75 | "should run python -m marimo --version with %s", 76 | async (interpreter) => { 77 | workspace.getConfiguration().set("marimo.pythonPath", interpreter); 78 | const output = await execPythonModule(["marimo", "--version"]); 79 | expect(output.toString()).toMatch(/\d+\.\d+\.\d+/); 80 | }, 81 | ); 82 | 83 | it("should handle spaces in pythonPath", async () => { 84 | workspace 85 | .getConfiguration() 86 | .set("marimo.pythonPath", "/path with spaces/python"); 87 | await expect(execPythonModule(["marimo", "--version"])).rejects.toThrow( 88 | /ENOENT/, 89 | ); 90 | }); 91 | 92 | it("should handle invalid pythonPath", async () => { 93 | workspace.getConfiguration().set("marimo.pythonPath", "nonexistent"); 94 | await expect(execPythonModule(["marimo", "--version"])).rejects.toThrow( 95 | /ENOENT/, 96 | ); 97 | }); 98 | }); 99 | 100 | describe("execPythonFile integration tests", () => { 101 | it("should execute a python file", async () => { 102 | const output = await execPythonFile(["-c", "print('hello')"]); 103 | expect(output.toString().trim()).toBe("hello"); 104 | }); 105 | 106 | it("should handle spaces in file path", async () => { 107 | await expect( 108 | execPythonFile(["/path with spaces/script.py"]), 109 | ).rejects.toThrow(/python: can't open file/); 110 | }); 111 | 112 | it("should handle invalid file path", async () => { 113 | await expect(execPythonFile(["nonexistent.py"])).rejects.toThrow( 114 | /python: can't open file/, 115 | ); 116 | }); 117 | }); 118 | 119 | describe("hasPythonModule integration tests", () => { 120 | it("should detect installed module", async () => { 121 | const output = await hasPythonModule("sys"); 122 | expect(output).toBeDefined(); 123 | }); 124 | 125 | it("should handle non-existent module", async () => { 126 | await expect(hasPythonModule("nonexistentmodule")).rejects.toThrow( 127 | /ModuleNotFoundError/, 128 | ); 129 | }); 130 | 131 | it("should handle module with spaces", async () => { 132 | await expect(hasPythonModule("module with spaces")).rejects.toThrow( 133 | /SyntaxError/, 134 | ); 135 | }); 136 | }); 137 | 138 | describe("hasExecutable integration tests", () => { 139 | it("should detect existing executable", async () => { 140 | const exists = await hasExecutable("python"); 141 | expect(exists).toBe(true); 142 | }); 143 | 144 | it("should handle non-existent executable", async () => { 145 | const exists = await hasExecutable("nonexistent"); 146 | expect(exists).toBe(false); 147 | }); 148 | 149 | it("should handle executable with spaces", async () => { 150 | const exists = await hasExecutable("/path with spaces/executable"); 151 | expect(exists).toBe(false); 152 | }); 153 | }); 154 | -------------------------------------------------------------------------------- /src/services/health.ts: -------------------------------------------------------------------------------- 1 | import { window, workspace } from "vscode"; 2 | import { Config } from "../config"; 3 | import { logger } from "../logger"; 4 | import { anonymouseId } from "../telemetry"; 5 | import { execMarimoCommand, getInterpreter } from "../utils/exec"; 6 | import { getExtensionVersion, getVscodeVersion } from "../utils/versions"; 7 | import type { ServerManager } from "./server-manager"; 8 | 9 | export class HealthService { 10 | constructor(private readonly serverManager: ServerManager) {} 11 | 12 | public async isMarimoInstalled(): Promise<{ 13 | isInstalled: boolean; 14 | version: string; 15 | path: string | undefined; 16 | }> { 17 | try { 18 | const bytes = await execMarimoCommand(["--version"]); 19 | const stdout = bytes.toString(); 20 | return { 21 | isInstalled: true, 22 | version: stdout.trim(), 23 | path: Config.marimoPath, 24 | }; 25 | } catch (error) { 26 | logger.error("Error checking marimo installation:", error); 27 | return { 28 | isInstalled: false, 29 | version: "unknown", 30 | path: Config.marimoPath, 31 | }; 32 | } 33 | } 34 | 35 | public async isServerRunning(): Promise<{ 36 | isRunning: boolean; 37 | port: number; 38 | }> { 39 | return { 40 | isRunning: this.serverManager.getStatus() === "started", 41 | port: this.serverManager.getPort() || 0, 42 | }; 43 | } 44 | 45 | /** 46 | * Shows an page with diagnostics the extension and server 47 | */ 48 | public async showDiagnostics() { 49 | let statusText = ""; 50 | try { 51 | statusText = await this.printStatusVerbose(); 52 | } catch (error) { 53 | logger.error("Error showing status:", error); 54 | statusText = `Error showing status: ${error}`; 55 | } 56 | 57 | const document = await workspace.openTextDocument({ 58 | content: statusText, 59 | language: "plaintext", 60 | }); 61 | await window.showTextDocument(document); 62 | return document; 63 | } 64 | 65 | public async printStatusVerbose(): Promise { 66 | const [{ isInstalled, version, path }, pythonInterpreter] = 67 | await Promise.all([this.isMarimoInstalled(), getInterpreter()]); 68 | 69 | if (isInstalled) { 70 | const status = await this.isServerRunning(); 71 | const serverUrl = `${Config.https ? "https" : "http"}://${Config.host}:${status.port}`; 72 | 73 | return [ 74 | "marimo configuration:", 75 | `\tpython interpreter: ${pythonInterpreter}`, 76 | isDefaultMarimoPath(path) ? "" : `\tmarimo executable path: ${path}`, // don't show if default 77 | `\tmarimo version: ${version}`, 78 | `\textension version: ${getExtensionVersion()}`, 79 | "\nserver status:", 80 | status.isRunning 81 | ? [ 82 | `\trunning on port ${status.port}`, 83 | `\turl: ${serverUrl}`, 84 | `\tread port: ${Config.readPort}`, 85 | ].join("\n") 86 | : "\tnot running", 87 | "\nserver configuration:", 88 | `\thost: ${Config.host}`, 89 | `\tdefault port: ${Config.port}`, 90 | `\thttps enabled: ${Config.https}`, 91 | `\ttoken auth enabled: ${Config.enableToken}`, 92 | Config.enableToken 93 | ? `\ttoken password: ${Config.tokenPassword ? "set" : "not set"}` 94 | : "", 95 | "\nenvironment settings:", 96 | `\tsandbox mode: ${Config.sandbox}`, 97 | `\twatch mode: ${Config.watch}`, 98 | "\nUI settings:", 99 | `\tbrowser type: ${Config.browser}`, 100 | `\tshow terminal: ${Config.showTerminal}`, 101 | `\tdebug mode: ${Config.debug}`, 102 | "\nsystem information:", 103 | `\tplatform: ${process.platform}`, 104 | `\tanonymous id: ${anonymouseId()}`, 105 | `\tarchitecture: ${process.arch}`, 106 | `\tnode version: ${process.version}`, 107 | `\tvscode version: ${getVscodeVersion()}`, 108 | ] 109 | .filter(Boolean) 110 | .join("\n"); 111 | } 112 | 113 | return troubleShootingMessage(path, pythonInterpreter); 114 | } 115 | 116 | public async printStatus(): Promise { 117 | const [{ isInstalled, version, path }, pythonInterpreter] = 118 | await Promise.all([this.isMarimoInstalled(), getInterpreter()]); 119 | 120 | if (isInstalled) { 121 | const status = await this.isServerRunning(); 122 | return [ 123 | "marimo is installed", 124 | path === "marimo" ? "" : `\tmarimo executable path: ${path}`, // don't show if default 125 | `\tpython interpreter: ${pythonInterpreter}`, 126 | `\tversion: ${version}`, 127 | status.isRunning 128 | ? `\tserver running: port ${status.port}` 129 | : "\tserver not running", 130 | ] 131 | .filter(Boolean) 132 | .join("\n"); 133 | } 134 | 135 | return troubleShootingMessage(path, pythonInterpreter); 136 | } 137 | } 138 | 139 | function troubleShootingMessage( 140 | marimoPath: string | undefined, 141 | pythonInterpreter: string | undefined, 142 | ) { 143 | return [ 144 | "marimo does not appear to be installed.", 145 | "", 146 | "Current configuration:", 147 | `\tpython interpreter: ${pythonInterpreter || "not set"}`, 148 | isDefaultMarimoPath(marimoPath) 149 | ? "" 150 | : `\tmarimo executable path: ${marimoPath}`, // don't show if default 151 | "", 152 | "Troubleshooting steps:", 153 | `\t1. Verify installation: ${pythonInterpreter} -m marimo`, 154 | "\t2. Install marimo: pip install marimo", 155 | "\t3. Check python.defaultInterpreterPath in VS Code settings", 156 | "\t4. Try creating a new virtual environment", 157 | "\t5. If using a virtual environment, ensure it's activated", 158 | ].join("\n"); 159 | } 160 | 161 | function isDefaultMarimoPath(path: string | undefined) { 162 | return path === "marimo" || path === undefined; 163 | } 164 | -------------------------------------------------------------------------------- /src/ui/status-bar.ts: -------------------------------------------------------------------------------- 1 | import { 2 | type Disposable, 3 | EventEmitter, 4 | StatusBarAlignment, 5 | type StatusBarItem, 6 | type TextDocument, 7 | type TextEditor, 8 | ThemeColor, 9 | commands, 10 | window, 11 | } from "vscode"; 12 | import { CommandsKeys } from "../constants"; 13 | import type { 14 | ControllerManager, 15 | MarimoController, 16 | } from "../launcher/controller"; 17 | import type { Kernel } from "../notebook/kernel"; 18 | import { KernelManager } from "../notebook/kernel-manager"; 19 | import { VscodeContextManager } from "../services/context-manager"; 20 | import type { HealthService } from "../services/health"; 21 | import type { ILifecycle } from "../types"; 22 | import { LogMethodCalls } from "../utils/log"; 23 | import { getFocusedMarimoTextEditor, isMarimoApp } from "../utils/query"; 24 | 25 | enum StatusBarState { 26 | Idle = 0, 27 | Kernel = 1, 28 | ActiveRun = 2, 29 | ActiveEdit = 3, 30 | Inactive = 4, 31 | } 32 | 33 | // Dumb component that handles UI updates 34 | class StatusBarView implements Disposable { 35 | private statusBar: StatusBarItem; 36 | 37 | constructor() { 38 | this.statusBar = window.createStatusBarItem(StatusBarAlignment.Left, 10); 39 | this.statusBar.command = CommandsKeys.showCommands; 40 | } 41 | 42 | update(state: StatusBarState) { 43 | switch (state) { 44 | case StatusBarState.Idle: 45 | this.setIdleAppearance(); 46 | break; 47 | case StatusBarState.Kernel: 48 | case StatusBarState.ActiveRun: 49 | case StatusBarState.ActiveEdit: 50 | this.setActiveAppearance(state); 51 | break; 52 | case StatusBarState.Inactive: 53 | this.setInactiveAppearance(); 54 | break; 55 | } 56 | } 57 | 58 | updateTooltip(tooltip: string) { 59 | this.statusBar.tooltip = tooltip; 60 | } 61 | 62 | private setIdleAppearance() { 63 | this.statusBar.show(); 64 | this.statusBar.text = "$(notebook) marimo"; 65 | this.statusBar.backgroundColor = undefined; 66 | this.statusBar.color = undefined; 67 | } 68 | 69 | private setActiveAppearance(state: StatusBarState) { 70 | this.statusBar.show(); 71 | this.statusBar.text = this.getStatusText(state); 72 | this.statusBar.backgroundColor = new ThemeColor( 73 | "statusBarItem.warningBackground", 74 | ); 75 | this.statusBar.color = new ThemeColor("statusBarItem.warningForeground"); 76 | } 77 | 78 | private setInactiveAppearance() { 79 | this.statusBar.show(); 80 | this.statusBar.text = "$(play) Start marimo"; 81 | this.statusBar.backgroundColor = undefined; 82 | this.statusBar.color = undefined; 83 | } 84 | 85 | private getStatusText(state: StatusBarState): string { 86 | switch (state) { 87 | case StatusBarState.Kernel: 88 | case StatusBarState.ActiveEdit: 89 | return "$(zap) marimo"; 90 | case StatusBarState.ActiveRun: 91 | return "$(zap) marimo (run)"; 92 | default: 93 | return ""; 94 | } 95 | } 96 | 97 | dispose() { 98 | this.statusBar.dispose(); 99 | } 100 | } 101 | 102 | // Smart component that manages state and logic 103 | export class StatusBar implements ILifecycle { 104 | private view: StatusBarView; 105 | private otherDisposables: Disposable[] = []; 106 | private currentState: StatusBarState = StatusBarState.Idle; 107 | private contextManager = new VscodeContextManager(); 108 | 109 | private static _onUpdate = new EventEmitter(); 110 | 111 | @LogMethodCalls() 112 | static update(): void { 113 | StatusBar._onUpdate.fire(undefined); 114 | } 115 | 116 | constructor( 117 | private readonly controllerManager: ControllerManager, 118 | private readonly healthService: HealthService, 119 | ) { 120 | this.view = new StatusBarView(); 121 | this.update(); 122 | this.registerListeners(); 123 | } 124 | 125 | async restart(): Promise { 126 | this.update(); 127 | } 128 | 129 | private registerListeners() { 130 | const updateEvents = [ 131 | window.onDidChangeActiveTextEditor, 132 | window.onDidChangeActiveNotebookEditor, 133 | window.onDidChangeTextEditorViewColumn, 134 | ]; 135 | 136 | this.otherDisposables.push( 137 | ...updateEvents.map((event) => event(() => this.update())), 138 | ); 139 | 140 | this.otherDisposables.push(StatusBar._onUpdate.event(() => this.update())); 141 | } 142 | 143 | dispose() { 144 | this.otherDisposables.forEach((d) => d.dispose()); 145 | this.view.dispose(); 146 | } 147 | 148 | async update() { 149 | const kernel = KernelManager.getFocusedMarimoKernel(); 150 | const editor = getFocusedMarimoTextEditor({ toast: false }); 151 | const activeController = 152 | this.controllerManager.getControllerForActivePanel() || 153 | (editor?.document && this.controllerManager.getOrCreate(editor.document)); 154 | 155 | this.updateState(kernel, editor, activeController); 156 | this.view.update(this.currentState); 157 | this.contextManager.setMarimoApp(isMarimoApp(editor?.document)); 158 | 159 | const status = await this.healthService.printStatus(); 160 | this.view.updateTooltip(status); 161 | } 162 | 163 | private updateState( 164 | kernel: Kernel | undefined, 165 | editor: TextEditor | undefined, 166 | activeController: MarimoController | undefined, 167 | ) { 168 | if (kernel) { 169 | // We are connected from a VSCode notebook to a kernel 170 | this.currentState = StatusBarState.Kernel; 171 | } else if (!editor?.document) { 172 | // Not a marimo file or notebook 173 | this.currentState = StatusBarState.Idle; 174 | } else if (activeController?.active) { 175 | // Marimo file is active: either running or editing 176 | this.currentState = 177 | activeController.currentMode === "run" 178 | ? StatusBarState.ActiveRun 179 | : StatusBarState.ActiveEdit; 180 | } else { 181 | // Marimo file is not active 182 | this.currentState = StatusBarState.Inactive; 183 | } 184 | } 185 | } 186 | -------------------------------------------------------------------------------- /src/browser/panel.ts: -------------------------------------------------------------------------------- 1 | import { Uri, ViewColumn, type WebviewPanel, env, window } from "vscode"; 2 | import { logger } from "../logger"; 3 | import { invariant } from "../utils/invariant"; 4 | import { LogMethodCalls } from "../utils/log"; 5 | 6 | export class MarimoPanelManager { 7 | private static readonly WEBVIEW_TYPE = "marimo"; 8 | private static readonly VSCODE_PARAM = "vscode"; 9 | 10 | public nativePanel: WebviewPanel | undefined; 11 | private url: string | undefined; 12 | private readonly logger = logger.createLogger(this.appName); 13 | 14 | constructor(private readonly appName: string) {} 15 | 16 | public isReady(): boolean { 17 | return !!this.nativePanel; 18 | } 19 | 20 | public isActive(): boolean { 21 | return this.nativePanel?.active ?? false; 22 | } 23 | 24 | private createWebviewPanel(url: string): WebviewPanel { 25 | this.nativePanel = window.createWebviewPanel( 26 | MarimoPanelManager.WEBVIEW_TYPE, 27 | `marimo: ${this.appName}`, 28 | ViewColumn.Beside, 29 | { 30 | enableScripts: true, 31 | enableCommandUris: true, 32 | }, 33 | ); 34 | invariant(this.nativePanel, "Failed to create webview panel"); 35 | 36 | this.nativePanel.webview.html = MarimoPanelManager.getWebviewContent(url); 37 | 38 | this.nativePanel.onDidDispose(() => { 39 | this.nativePanel = undefined; 40 | }); 41 | 42 | return this.nativePanel; 43 | } 44 | 45 | @LogMethodCalls() 46 | public reload(): void { 47 | if (this.nativePanel && this.url) { 48 | this.nativePanel.dispose(); 49 | this.nativePanel = this.createWebviewPanel(this.url); 50 | } else { 51 | this.logger.warn("Cannot reload: panel or URL not set"); 52 | } 53 | } 54 | 55 | public async create(url: string): Promise { 56 | this.logger.info("Creating panel at", url); 57 | this.url = url; 58 | 59 | if (this.nativePanel) { 60 | this.logger.warn("Panel already exists"); 61 | return; 62 | } 63 | 64 | this.nativePanel = this.createWebviewPanel(url); 65 | this.setupMessageHandler(); 66 | } 67 | 68 | @LogMethodCalls() 69 | public show(): void { 70 | if (!this.nativePanel) { 71 | this.logger.warn("Panel not created yet"); 72 | } 73 | this.nativePanel?.reveal(); 74 | } 75 | 76 | @LogMethodCalls() 77 | public dispose(): void { 78 | this.nativePanel?.dispose(); 79 | } 80 | 81 | private setupMessageHandler(): void { 82 | if (!this.nativePanel) { 83 | this.logger.error("Cannot setup message handler: panel not created"); 84 | return; 85 | } 86 | 87 | this.nativePanel.webview.onDidReceiveMessage( 88 | this.handleWebviewMessage, 89 | undefined, 90 | ); 91 | } 92 | 93 | private handleWebviewMessage = async (message: { 94 | command: string; 95 | text?: string; 96 | url?: string; 97 | }): Promise => { 98 | logger.info("Received message from webview", message); 99 | switch (message.command) { 100 | case "copy": 101 | case "cut": 102 | if (!message.text) { 103 | break; 104 | } 105 | await env.clipboard.writeText(message.text).then(() => { 106 | logger.info("Copied text to clipboard", message.text); 107 | }); 108 | break; 109 | case "paste": { 110 | const text = await env.clipboard.readText(); 111 | await this.nativePanel?.webview 112 | .postMessage({ command: "paste", text }) 113 | .then(() => { 114 | logger.info("Pasted text from clipboard", text); 115 | }); 116 | break; 117 | } 118 | case "external_link": 119 | if (!message.url) { 120 | break; 121 | } 122 | await env.openExternal(Uri.parse(message.url)).then(() => { 123 | logger.info("Opened external link", message.url); 124 | }); 125 | break; 126 | case "context_menu": 127 | // Context menu is not supported yet 128 | break; 129 | default: 130 | this.logger.info("Unknown message", message.command); 131 | } 132 | }; 133 | 134 | private static getWebviewContent(urlString: string): string { 135 | const url = new URL(urlString); 136 | url.searchParams.set(MarimoPanelManager.VSCODE_PARAM, "true"); 137 | 138 | const styles = ` 139 | position: absolute; 140 | padding: 0; 141 | margin: 0; 142 | top: 0; 143 | bottom: 0; 144 | left: 0; 145 | right: 0; 146 | display: flex; 147 | `; 148 | 149 | const iframeStyles = ` 150 | flex: 1; 151 | border: none; 152 | `; 153 | 154 | return ` 155 | 156 | 157 | 158 | 159 | 160 | marimo 161 | 162 | 163 | 169 | 170 | 194 | 195 | `; 196 | } 197 | } 198 | -------------------------------------------------------------------------------- /src/__integration__/server-manager.test.ts: -------------------------------------------------------------------------------- 1 | import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; 2 | import { createVSCodeMock } from "../__mocks__/vscode"; 3 | 4 | vi.mock("vscode", () => createVSCodeMock(vi)); 5 | vi.mock("@vscode/python-extension", () => ({})); 6 | 7 | import { tmpdir } from "node:os"; 8 | import { join } from "node:path"; 9 | import { type ExtensionContext, type Memento, window, workspace } from "vscode"; 10 | import type { Config } from "../config"; 11 | import { setExtension } from "../ctx"; 12 | import { ServerManager } from "../services/server-manager"; 13 | import { sleep } from "../utils/async"; 14 | 15 | const TEST_TIMEOUT = 10_000; 16 | 17 | describe("ServerManager integration tests", () => { 18 | let config: Config; 19 | let manager: ServerManager; 20 | 21 | beforeEach(() => { 22 | const tmpDir = tmpdir(); 23 | workspace.getConfiguration().reset(); 24 | config = { 25 | root: join(tmpDir, "nb.py"), 26 | browser: "system", 27 | port: 2918, 28 | readPort: 2919, 29 | host: "localhost", 30 | https: false, 31 | enableToken: false, 32 | tokenPassword: "", 33 | debug: false, 34 | sandbox: false, 35 | watch: false, 36 | } as Config; 37 | manager = ServerManager.getInstance(config); 38 | const globalState = new Map(); 39 | setExtension({ 40 | globalState: { 41 | get: vi.fn().mockImplementation((key) => globalState.get(key)), 42 | update: vi 43 | .fn() 44 | .mockImplementation((key, value) => globalState.set(key, value)), 45 | keys: vi.fn().mockImplementation(() => Array.from(globalState.keys())), 46 | } as unknown as Memento, 47 | } as unknown as ExtensionContext); 48 | manager.init(); 49 | }); 50 | 51 | afterEach(async () => { 52 | await manager.stopServer(); 53 | }); 54 | 55 | it( 56 | "should start and stop server", 57 | async () => { 58 | expect(manager.getStatus()).toBe("stopped"); 59 | 60 | const result = await manager.start(); 61 | expect(result.port).toBeGreaterThan(0); 62 | expect(result.skewToken).toBeDefined(); 63 | expect(result.version).toMatch(/\d+\.\d+\.\d+/); 64 | expect(manager.getStatus()).toBe("started"); 65 | 66 | await manager.stopServer(); 67 | expect(manager.getStatus()).toBe("stopped"); 68 | }, 69 | TEST_TIMEOUT, 70 | ); 71 | 72 | it( 73 | "should start and stop a server with custom python path", 74 | async () => { 75 | workspace.getConfiguration().set("marimo.pythonPath", "uv run python"); 76 | const result = await manager.start(); 77 | expect(result.port).toBeGreaterThan(0); 78 | expect(result.skewToken).toBeDefined(); 79 | expect(result.version).toMatch(/\d+\.\d+\.\d+/); 80 | expect(manager.getStatus()).toBe("started"); 81 | 82 | await manager.stopServer(); 83 | expect(manager.getStatus()).toBe("stopped"); 84 | }, 85 | TEST_TIMEOUT, 86 | ); 87 | 88 | it( 89 | "should start and stop a server with custom marimo path", 90 | async () => { 91 | workspace.getConfiguration().set("marimo.marimoPath", "uv run marimo"); 92 | const result = await manager.start(); 93 | expect(result.port).toBeGreaterThan(0); 94 | expect(result.skewToken).toBeDefined(); 95 | expect(result.version).toMatch(/\d+\.\d+\.\d+/); 96 | expect(manager.getStatus()).toBe("started"); 97 | 98 | await manager.stopServer(); 99 | expect(manager.getStatus()).toBe("stopped"); 100 | }, 101 | TEST_TIMEOUT, 102 | ); 103 | 104 | it( 105 | "should reuse existing server if healthy", 106 | async () => { 107 | console.warn("Starting server"); 108 | const firstStart = await manager.start(); 109 | console.warn("Starting server again"); 110 | const secondStart = await manager.start(); 111 | 112 | expect(secondStart.port).toBe(firstStart.port); 113 | expect(secondStart.skewToken).toBe(firstStart.skewToken); 114 | 115 | // Check if healthy 116 | const isHealthy = await manager.isHealthy(firstStart.port); 117 | expect(isHealthy).toBe(true); 118 | }, 119 | TEST_TIMEOUT, 120 | ); 121 | 122 | it.skip( 123 | "should start a server and cancel it", 124 | async () => { 125 | const subscribers = new Set<() => void>(); 126 | const token = { 127 | isCancellationRequested: false, 128 | onCancellationRequested: vi.fn().mockImplementation((fn) => { 129 | subscribers.add(fn); 130 | }), 131 | cancel: vi.fn().mockImplementation(() => { 132 | subscribers.forEach(async (fn) => { 133 | try { 134 | fn(); 135 | } catch (e) { 136 | // pass 137 | } 138 | }); 139 | }), 140 | }; 141 | 142 | await expect(async () => { 143 | await Promise.all([manager.start(token), token.cancel(), sleep(100)]); 144 | }).rejects.toThrow("Server start was cancelled"); 145 | expect(token.onCancellationRequested).toHaveBeenCalledTimes(1); 146 | expect(manager.getStatus()).toBe("stopped"); 147 | }, 148 | TEST_TIMEOUT, 149 | ); 150 | 151 | it( 152 | "should restart unhealthy server", 153 | async () => { 154 | const firstStart = await manager.start(); 155 | 156 | // Force server into unhealthy state by stopping it directly 157 | await manager.stopServer(); 158 | 159 | // Should detect unhealthy state and restart 160 | const secondStart = await manager.start(); 161 | expect(secondStart.port).not.toBe(firstStart.port); 162 | }, 163 | TEST_TIMEOUT, 164 | ); 165 | 166 | it( 167 | "should handle multiple start requests while starting", 168 | async () => { 169 | const startPromise1 = manager.start(); 170 | const startPromise2 = manager.start(); 171 | 172 | const [result1, result2] = await Promise.all([ 173 | startPromise1, 174 | startPromise2, 175 | ]); 176 | expect(result1.port).toBe(result2.port); 177 | }, 178 | TEST_TIMEOUT, 179 | ); 180 | 181 | it( 182 | "should get active sessions", 183 | async () => { 184 | await manager.start(); 185 | const sessions = await manager.getActiveSessions(); 186 | expect(Array.isArray(sessions)).toBe(true); 187 | }, 188 | TEST_TIMEOUT, 189 | ); 190 | 191 | it.skip("should show warning and handle restart on unhealthy server", async () => { 192 | await manager.start(); 193 | const showWarningMock = vi.spyOn(window, "showWarningMessage"); 194 | 195 | // Simulate health check failure 196 | vi.spyOn(global, "fetch").mockRejectedValueOnce( 197 | new Error("Connection failed"), 198 | ); 199 | 200 | // Trigger health check 201 | await sleep(100); 202 | 203 | expect(showWarningMock).toHaveBeenCalledWith( 204 | "The marimo server is not responding. What would you like to do?", 205 | "Restart Server", 206 | "Ignore", 207 | ); 208 | }); 209 | }); 210 | -------------------------------------------------------------------------------- /src/commands/show-commands.ts: -------------------------------------------------------------------------------- 1 | import { 2 | type QuickPickItem, 3 | QuickPickItemKind, 4 | Uri, 5 | commands, 6 | env, 7 | window, 8 | } from "vscode"; 9 | import { 10 | CommandsKeys, 11 | DISCORD_URL, 12 | DOCUMENTATION_URL, 13 | EXTENSION_DISPLAY_NAME, 14 | EXTENSION_PACKAGE, 15 | } from "../constants"; 16 | import { exportAsCommands } from "../export/export-as-commands"; 17 | import { MarimoController } from "../launcher/controller"; 18 | import { Launcher } from "../launcher/start"; 19 | import { 20 | getActiveMarimoFile, 21 | openMarimoNotebookDocument, 22 | } from "../notebook/extension"; 23 | import { Kernel } from "../notebook/kernel"; 24 | import type { ServerManager } from "../services/server-manager"; 25 | import { tutorialCommands } from "./tutorial-commands"; 26 | 27 | interface CommandPickItem extends QuickPickItem { 28 | handler: () => void; 29 | if?: boolean; 30 | } 31 | 32 | const SEPARATOR = { 33 | label: "", 34 | kind: QuickPickItemKind.Separator, 35 | handler: () => {}, 36 | }; 37 | 38 | export async function showCommands( 39 | controller: MarimoController | Kernel | undefined, 40 | serverManager: ServerManager, 41 | ) { 42 | let commands: CommandPickItem[] = []; 43 | 44 | if (controller instanceof Kernel) { 45 | commands = await showKernelCommands(controller); 46 | } 47 | if (controller instanceof MarimoController) { 48 | commands = await showMarimoControllerCommands(controller); 49 | } 50 | 51 | commands.push(...miscCommands(serverManager)); 52 | 53 | const filteredCommands = commands.filter((index) => index.if !== false); 54 | const result = await window.showQuickPick(filteredCommands); 55 | 56 | if (result) { 57 | result.handler(); 58 | } 59 | } 60 | 61 | export function showKernelCommands(kernel: Kernel): CommandPickItem[] { 62 | return [ 63 | { 64 | label: "$(split-horizontal) Open outputs in embedded browser", 65 | description: kernel.relativePath, 66 | async handler() { 67 | await kernel.openKiosk("embedded"); 68 | }, 69 | }, 70 | { 71 | label: "$(link-external) Open outputs in system browser", 72 | description: kernel.relativePath, 73 | async handler() { 74 | await kernel.openKiosk("system"); 75 | }, 76 | }, 77 | SEPARATOR, 78 | { 79 | label: "$(refresh) Restart kernel", 80 | async handler() { 81 | await kernel.restart(); 82 | await kernel.openKiosk(); 83 | }, 84 | }, 85 | { 86 | label: "$(export) Export notebook as...", 87 | handler() { 88 | exportAsCommands(kernel.fileUri); 89 | }, 90 | }, 91 | SEPARATOR, 92 | ]; 93 | } 94 | 95 | export async function showMarimoControllerCommands( 96 | controller: MarimoController, 97 | ): Promise { 98 | return [ 99 | // Non-active commands 100 | { 101 | label: "$(notebook) Start as VSCode notebook", 102 | async handler() { 103 | await openMarimoNotebookDocument(await getActiveMarimoFile()); 104 | }, 105 | if: !controller.active, 106 | }, 107 | { 108 | label: "$(zap) Start in marimo editor (edit)", 109 | async handler() { 110 | await Launcher.start({ controller, mode: "edit" }); 111 | controller.open(); 112 | }, 113 | if: !controller.active, 114 | }, 115 | { 116 | label: "$(preview) Start in marimo editor (run)", 117 | async handler() { 118 | await Launcher.start({ controller, mode: "run" }); 119 | controller.open(); 120 | }, 121 | if: !controller.active, 122 | }, 123 | SEPARATOR, 124 | // Active commands 125 | { 126 | label: "$(split-horizontal) Open in embedded browser", 127 | description: await controller.url(), 128 | handler() { 129 | controller.open("embedded"); 130 | }, 131 | if: controller.active, 132 | }, 133 | { 134 | label: "$(link-external) Open in system browser", 135 | description: await controller.url(), 136 | handler() { 137 | controller.open("system"); 138 | }, 139 | if: controller.active, 140 | }, 141 | { 142 | label: "$(refresh) Restart marimo kernel", 143 | async handler() { 144 | const mode = controller.currentMode || "edit"; 145 | await Launcher.stop(controller); 146 | await Launcher.start({ mode, controller }); 147 | controller.open(); 148 | }, 149 | if: controller.active, 150 | }, 151 | { 152 | label: 153 | controller.currentMode === "run" 154 | ? "$(package) Switch to edit mode" 155 | : "$(package) Switch to run mode", 156 | async handler() { 157 | const otherMode = controller.currentMode === "run" ? "edit" : "run"; 158 | await Launcher.stop(controller); 159 | await Launcher.start({ mode: otherMode, controller }); 160 | controller.open(); 161 | }, 162 | if: controller.active, 163 | }, 164 | { 165 | label: "$(terminal) Show Terminal", 166 | handler() { 167 | controller.terminal.show(); 168 | }, 169 | if: controller.active, 170 | }, 171 | { 172 | label: "$(close) Stop kernel", 173 | handler() { 174 | Launcher.stop(controller); 175 | }, 176 | if: controller.active, 177 | }, 178 | { 179 | label: "$(export) Export notebook as...", 180 | handler() { 181 | exportAsCommands(controller.file.uri); 182 | }, 183 | }, 184 | 185 | SEPARATOR, 186 | ]; 187 | } 188 | 189 | export function miscCommands(serverManager: ServerManager): CommandPickItem[] { 190 | return [ 191 | { 192 | label: "$(question) View marimo documentation", 193 | handler() { 194 | env.openExternal(Uri.parse(DOCUMENTATION_URL)); 195 | }, 196 | }, 197 | { 198 | label: "$(bookmark) View tutorials", 199 | async handler() { 200 | await tutorialCommands(); 201 | }, 202 | }, 203 | { 204 | label: "$(comment-discussion) Join Discord community", 205 | handler() { 206 | env.openExternal(Uri.parse(DISCORD_URL)); 207 | }, 208 | }, 209 | { 210 | label: "$(settings) Edit settings", 211 | handler() { 212 | void commands.executeCommand("workbench.action.openSettings", "marimo"); 213 | }, 214 | }, 215 | { 216 | label: `$(info) Server status: ${serverManager.getStatus()}`, 217 | handler: () => { 218 | void commands.executeCommand(CommandsKeys.showDiagnostics); 219 | goToLogs(); 220 | }, 221 | }, 222 | { 223 | label: "$(info) View extension logs", 224 | handler: goToLogs, 225 | }, 226 | ]; 227 | } 228 | 229 | async function goToLogs() { 230 | // Open output panel with channel 'marimo' 231 | await commands.executeCommand( 232 | `workbench.action.output.show.${EXTENSION_PACKAGE.fullName}.${EXTENSION_DISPLAY_NAME}`, 233 | ); 234 | await commands.executeCommand("marimo-explorer-running-applications.focus"); 235 | } 236 | -------------------------------------------------------------------------------- /src/services/health.test.ts: -------------------------------------------------------------------------------- 1 | import { beforeEach, describe, expect, it, vi } from "vitest"; 2 | import { extensions, window, workspace } from "vscode"; 3 | import { Config } from "../config"; 4 | import { logger } from "../logger"; 5 | import { execMarimoCommand, getInterpreter } from "../utils/exec"; 6 | import { getExtensionVersion } from "../utils/versions"; 7 | import { HealthService } from "./health"; 8 | import type { ServerManager } from "./server-manager"; 9 | 10 | vi.mock("vscode", () => ({ 11 | extensions: { 12 | getExtension: vi.fn(), 13 | }, 14 | window: { 15 | showTextDocument: vi.fn(), 16 | }, 17 | workspace: { 18 | openTextDocument: vi.fn(), 19 | }, 20 | })); 21 | 22 | vi.mock("../logger", () => ({ 23 | logger: { 24 | error: vi.fn(), 25 | }, 26 | })); 27 | 28 | vi.mock("../utils/exec", () => ({ 29 | execMarimoCommand: vi.fn(), 30 | getInterpreter: vi.fn(), 31 | })); 32 | 33 | describe("HealthService", () => { 34 | let healthService: HealthService; 35 | let mockServerManager: ServerManager; 36 | 37 | beforeEach(() => { 38 | vi.resetAllMocks(); 39 | mockServerManager = { 40 | getStatus: vi.fn(), 41 | getPort: vi.fn(), 42 | } as unknown as ServerManager; 43 | healthService = new HealthService(mockServerManager); 44 | }); 45 | 46 | describe("isMarimoInstalled", () => { 47 | it("should return true when marimo is installed", async () => { 48 | const mockVersion = "1.0.0"; 49 | vi.mocked(execMarimoCommand).mockResolvedValue(mockVersion); 50 | vi.spyOn(Config, "marimoPath", "get").mockReturnValue("/path/to/marimo"); 51 | 52 | const result = await healthService.isMarimoInstalled(); 53 | 54 | expect(result).toEqual({ 55 | isInstalled: true, 56 | version: mockVersion, 57 | path: "/path/to/marimo", 58 | }); 59 | }); 60 | 61 | it("should return false when marimo is not installed", async () => { 62 | vi.mocked(execMarimoCommand).mockRejectedValue( 63 | new Error("Command failed"), 64 | ); 65 | vi.spyOn(Config, "marimoPath", "get").mockReturnValue("/path/to/marimo"); 66 | 67 | const result = await healthService.isMarimoInstalled(); 68 | 69 | expect(result).toEqual({ 70 | isInstalled: false, 71 | version: "unknown", 72 | path: "/path/to/marimo", 73 | }); 74 | expect(logger.error).toHaveBeenCalled(); 75 | }); 76 | }); 77 | 78 | describe("isServerRunning", () => { 79 | it("should return true when server is running", async () => { 80 | vi.mocked(mockServerManager.getStatus).mockReturnValue("started"); 81 | vi.mocked(mockServerManager.getPort).mockReturnValue(8080); 82 | 83 | const result = await healthService.isServerRunning(); 84 | 85 | expect(result).toEqual({ 86 | isRunning: true, 87 | port: 8080, 88 | }); 89 | }); 90 | 91 | it("should return false when server is not running", async () => { 92 | vi.mocked(mockServerManager.getStatus).mockReturnValue("stopped"); 93 | vi.mocked(mockServerManager.getPort).mockReturnValue(0); 94 | 95 | const result = await healthService.isServerRunning(); 96 | 97 | expect(result).toEqual({ 98 | isRunning: false, 99 | port: 0, 100 | }); 101 | }); 102 | }); 103 | 104 | describe("showDiagnostics", () => { 105 | it("should show status in a new document", async () => { 106 | const mockStatusText = "test status"; 107 | const mockDocument = { uri: "test" }; 108 | vi.mocked(workspace.openTextDocument).mockResolvedValue( 109 | mockDocument as any, 110 | ); 111 | vi.spyOn(healthService, "printStatusVerbose").mockResolvedValue( 112 | mockStatusText, 113 | ); 114 | 115 | const result = await healthService.showDiagnostics(); 116 | 117 | expect(workspace.openTextDocument).toHaveBeenCalledWith({ 118 | content: mockStatusText, 119 | language: "plaintext", 120 | }); 121 | expect(window.showTextDocument).toHaveBeenCalledWith(mockDocument); 122 | expect(result).toBe(mockDocument); 123 | }); 124 | 125 | it("should handle errors gracefully", async () => { 126 | const error = new Error("test error"); 127 | vi.spyOn(healthService, "printStatusVerbose").mockRejectedValue(error); 128 | const mockDocument = { uri: "test" }; 129 | vi.mocked(workspace.openTextDocument).mockResolvedValue( 130 | mockDocument as any, 131 | ); 132 | 133 | const result = await healthService.showDiagnostics(); 134 | 135 | expect(workspace.openTextDocument).toHaveBeenCalledWith({ 136 | content: `Error showing status: ${error}`, 137 | language: "plaintext", 138 | }); 139 | expect(result).toBe(mockDocument); 140 | }); 141 | }); 142 | 143 | describe("printExtensionVersion", () => { 144 | it("should return extension version", () => { 145 | const mockVersion = "1.0.0"; 146 | vi.mocked(extensions.getExtension).mockReturnValue({ 147 | packageJSON: { version: mockVersion }, 148 | } as any); 149 | 150 | const result = getExtensionVersion(); 151 | 152 | expect(result).toBe(mockVersion); 153 | }); 154 | 155 | it("should return unknown when extension is not found", () => { 156 | vi.mocked(extensions.getExtension).mockReturnValue(undefined); 157 | 158 | const result = getExtensionVersion(); 159 | 160 | expect(result).toBe("unknown"); 161 | }); 162 | }); 163 | 164 | describe("printStatus", () => { 165 | it("should return status when marimo is installed", async () => { 166 | const mockVersion = "1.0.0"; 167 | const mockPythonPath = "/path/to/python"; 168 | vi.mocked(execMarimoCommand).mockResolvedValue(mockVersion); 169 | vi.mocked(getInterpreter).mockResolvedValue(mockPythonPath); 170 | vi.spyOn(Config, "marimoPath", "get").mockReturnValue( 171 | "/custom/path/to/marimo", 172 | ); 173 | vi.mocked(mockServerManager.getStatus).mockReturnValue("started"); 174 | vi.mocked(mockServerManager.getPort).mockReturnValue(8080); 175 | 176 | const result = await healthService.printStatus(); 177 | 178 | expect(result).toContain("marimo is installed"); 179 | expect(result).toContain( 180 | "marimo executable path: /custom/path/to/marimo", 181 | ); 182 | expect(result).toContain("python interpreter: /path/to/python"); 183 | expect(result).toContain("version: 1.0.0"); 184 | expect(result).toContain("server running: port 8080"); 185 | }); 186 | 187 | it("should return troubleshooting message when marimo is not installed", async () => { 188 | vi.mocked(execMarimoCommand).mockRejectedValue( 189 | new Error("Command failed"), 190 | ); 191 | vi.mocked(getInterpreter).mockResolvedValue("/path/to/python"); 192 | vi.spyOn(Config, "marimoPath", "get").mockReturnValue("marimo"); 193 | 194 | const result = await healthService.printStatus(); 195 | 196 | expect(result).toContain("marimo does not appear to be installed"); 197 | expect(result).toContain("python interpreter: /path/to/python"); 198 | expect(result).toContain("Troubleshooting steps:"); 199 | }); 200 | }); 201 | }); 202 | -------------------------------------------------------------------------------- /src/notebook/kernel-manager.ts: -------------------------------------------------------------------------------- 1 | import * as vscode from "vscode"; 2 | import { logger } from "../logger"; 3 | import { LogMethodCalls } from "../utils/log"; 4 | import { type KernelKey, toKernelKey } from "./common/key"; 5 | import { getNotebookMetadata } from "./common/metadata"; 6 | import { NOTEBOOK_TYPE, PYTHON_LANGUAGE_ID } from "./constants"; 7 | import { createNotebookController } from "./createMarimoNotebookController"; 8 | import { Kernel } from "./kernel"; 9 | import type { MarimoConfig, SkewToken } from "./marimo/types"; 10 | 11 | // Global state 12 | const kernelMap = new Map(); 13 | 14 | interface CreateKernelOptions { 15 | port: number; 16 | uri: vscode.Uri | "__new__"; 17 | skewToken: SkewToken; 18 | version: string; 19 | userConfig: MarimoConfig; 20 | notebookDoc: vscode.NotebookDocument; 21 | } 22 | 23 | export interface IKernelManager extends vscode.Disposable { 24 | createKernel(opts: CreateKernelOptions): Kernel; 25 | getKernel(key: KernelKey | undefined): Kernel | undefined; 26 | getKernelByUri(uri: vscode.Uri): Kernel | undefined; 27 | } 28 | 29 | /** 30 | * Keeps track of running marimo kernels. 31 | * A kernel is associated with a port, notebook, and file. 32 | * 33 | * Multiple kernels could belong to the same server (and thus the same port), 34 | * but they are not required to. 35 | */ 36 | export class KernelManager implements IKernelManager { 37 | private readonly supportedLanguages = [PYTHON_LANGUAGE_ID]; 38 | private logger = logger.createLogger("kernel-manager"); 39 | public static instance: KernelManager = new KernelManager(); 40 | 41 | private controller: vscode.NotebookController; 42 | 43 | private otherDisposables: vscode.Disposable[] = []; 44 | 45 | private constructor() { 46 | this.controller = createNotebookController(); 47 | this.otherDisposables.push(this.controller); 48 | 49 | this.controller.supportedLanguages = this.supportedLanguages; 50 | this.controller.supportsExecutionOrder = false; 51 | this.controller.executeHandler = this.executeAll.bind(this); 52 | this.controller.interruptHandler = async (notebook) => { 53 | const key = getNotebookMetadata(notebook).key; 54 | const kernel = this.getKernel(key); 55 | if (!kernel) { 56 | return this.logger.error("No kernel found for key", key); 57 | } 58 | await kernel.interrupt(); 59 | }; 60 | this.listenForNotebookChanges(); 61 | } 62 | 63 | static getFocusedMarimoKernel(): Kernel | undefined { 64 | const activeNb = vscode.window.activeNotebookEditor?.notebook; 65 | // Directly active notebook 66 | if (activeNb && activeNb.notebookType === NOTEBOOK_TYPE) { 67 | const metadata = getNotebookMetadata(activeNb); 68 | if (metadata.key) { 69 | return kernelMap.get(metadata.key); 70 | } 71 | } 72 | 73 | // Active webview 74 | for (const kernel of kernelMap.values()) { 75 | if (kernel.isWebviewActive()) { 76 | return kernel; 77 | } 78 | } 79 | 80 | return; 81 | } 82 | 83 | @LogMethodCalls() 84 | createKernel(opts: CreateKernelOptions): Kernel { 85 | const { port, uri, skewToken, version, userConfig, notebookDoc } = opts; 86 | 87 | // If already created, return 88 | const key = toKernelKey(uri); 89 | const existing = kernelMap.get(key); 90 | if (existing && existing.notebookDoc === notebookDoc) { 91 | return existing; 92 | } 93 | 94 | const kernel = new Kernel({ 95 | port: port, 96 | kernelKey: key, 97 | skewToken: skewToken, 98 | fileUri: notebookDoc.uri, 99 | version: version, 100 | userConfig: userConfig, 101 | controller: this.controller, 102 | notebookDoc: notebookDoc, 103 | }); 104 | kernelMap.set(key, kernel); 105 | return kernel; 106 | } 107 | 108 | hydrateExistingNotebooks(opts: { 109 | port: number; 110 | skewToken: SkewToken; 111 | version: string; 112 | userConfig: MarimoConfig; 113 | }): void { 114 | // Find all open notebooks 115 | for (const nb of vscode.workspace.notebookDocuments) { 116 | if (nb.notebookType === NOTEBOOK_TYPE) { 117 | this.createKernel({ 118 | uri: nb.uri, 119 | port: opts.port, 120 | skewToken: opts.skewToken, 121 | version: opts.version, 122 | userConfig: opts.userConfig, 123 | notebookDoc: nb, 124 | }); 125 | } 126 | } 127 | } 128 | 129 | getKernel(key: KernelKey | undefined): Kernel | undefined { 130 | if (!key) { 131 | return undefined; 132 | } 133 | return kernelMap.get(key); 134 | } 135 | 136 | getKernelByUri(uri: vscode.Uri): Kernel | undefined { 137 | const key = toKernelKey(uri); 138 | const found = this.getKernel(key); 139 | if (found) { 140 | return found; 141 | } 142 | return; 143 | } 144 | 145 | @LogMethodCalls() 146 | private async unregisterKernel(kernel: Kernel): Promise { 147 | await kernel.dispose(); 148 | kernelMap.delete(kernel.kernelKey); 149 | } 150 | 151 | async clearAllKernels(): Promise { 152 | const disposePromises = Array.from(kernelMap.values()).map((kernel) => 153 | kernel.dispose(), 154 | ); 155 | await Promise.all(disposePromises); 156 | kernelMap.clear(); 157 | } 158 | 159 | async dispose(): Promise { 160 | this.otherDisposables.forEach((d) => d.dispose()); 161 | await this.clearAllKernels(); 162 | // Ensure the controller is disposed 163 | this.controller.dispose(); 164 | } 165 | 166 | private listenForNotebookChanges(): void { 167 | // Listen for added/removed cells 168 | this.otherDisposables.push( 169 | vscode.workspace.onDidChangeNotebookDocument((e) => { 170 | const kernel = this.getKernelForNotebook(e.notebook); 171 | kernel?.handleNotebookChange(e); 172 | }), 173 | ); 174 | 175 | // Listen for closed notebooks 176 | this.otherDisposables.push( 177 | vscode.workspace.onDidCloseNotebookDocument((nb) => { 178 | this.logger.info("notebook closed", nb.uri.toString()); 179 | const kernel = this.getKernelForNotebook(nb); 180 | if (!kernel) { 181 | return; 182 | } 183 | this.unregisterKernel(kernel); 184 | }), 185 | ); 186 | 187 | // Listen for saved notebooks 188 | this.otherDisposables.push( 189 | vscode.workspace.onWillSaveNotebookDocument((e) => { 190 | const kernel = this.getKernelForNotebook(e.notebook); 191 | if (!kernel) { 192 | return; 193 | } 194 | }), 195 | ); 196 | } 197 | 198 | private executeAll( 199 | cells: vscode.NotebookCell[], 200 | notebook: vscode.NotebookDocument, 201 | controller: vscode.NotebookController, 202 | ): void { 203 | this.controller = controller; 204 | const kernel = this.getKernelForNotebook(notebook); 205 | kernel?.executeAll(cells, notebook, controller); 206 | } 207 | 208 | private getKernelForNotebook( 209 | notebook: vscode.NotebookDocument, 210 | ): Kernel | undefined { 211 | if (notebook.notebookType !== NOTEBOOK_TYPE) { 212 | return; 213 | } 214 | 215 | const key = getNotebookMetadata(notebook).key; 216 | if (!key) { 217 | this.logger.error("No key found in notebook metadata"); 218 | return; 219 | } 220 | if (!kernelMap.has(key)) { 221 | this.logger.error("No kernel found for key", key); 222 | return; 223 | } 224 | return this.getKernel(key); 225 | } 226 | } 227 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # vscode marimo 2 | 3 | > [!Important] 4 | > **We are actively rewriting this extension to provide a more native and robust experience**, similar to what you may be used to when using Jupyter in VS Code. **The new extension is being developed in** 5 | > **[https://github.com/marimo-team/marimo-lsp](https://github.com/marimo-team/marimo-lsp), where new issues should be opened.** This repo is deprecated. 6 | 7 | 8 | Visual Studio Marketplace Version 9 | 10 | 11 | Run [marimo](https://github.com/marimo-team/marimo), directly from VS Code. 12 | 13 |

14 | 15 |

16 | 17 | > [!NOTE] 18 | > This extension requires marimo to be installed on your system: `pip install marimo`. 19 | > See the [installation guide](https://docs.marimo.io/getting_started/index.html) for more details. 20 | 21 | Check out the marimo documentation at . 22 | 23 | ## Features 24 | 25 | - 🚀 Launch marimo from VS Code, in both "edit mode" and "run mode". 26 | - 💻 View the terminal output of marimo directly in VS Code. 27 | - 🌐 View the marimo browser window directly in VS Code or in your default browser. 28 | - 📥 Export notebooks as: html, markdown, or scripts. 29 | - 📓 Convert Jupyter notebooks to marimo notebooks. 30 | - 🧪 [experimental] Run marimo in VSCode's native notebook 31 | 32 | ## Known Issues 33 | 34 | VS Code's embedded browser does not support all native browser features. If you encounter any issues, try opening marimo in your default browser instead. 35 | For example, the embedded browser will not support PDF render, audio recording, video recording, and some [copy/paste operations](https://github.com/microsoft/vscode/issues/115935). 36 | 37 | ## Experimental Native Notebook 38 | 39 | This extension includes an experimental feature to run marimo in VSCode's native notebook interface. This feature lets you use VSCode editors and extensions for writing code while the outputs and visualizations are rendered in a view-only marimo editor. This marimo editor displays outputs, console logs, and UI elements to interact with. 40 | 41 | This feature is experimental and may have some limitations. Some known limitations are: 42 | 43 | - VSCode automatically includes "Run above" and "Run below" buttons in the notebook toolbar. While these work, they do not make sense with a reactive notebook. 44 | - Notebooks can still be edited even though there may not be an active marimo server. This can be confusing since saving or running will not work. 45 | - For autocomplete to work when using native VSCode notebooks for many packages (including `marimo`, `numpy`, and more) you may be required to include a `pyproject.toml` file at the root of the workspace. marimo's editor gets around this by default but unfortunately, the VSCode's native notebook does not. 46 | - You cannot access **many** marimo features in the native notebook (and need to use the marimo browser), such as the variable explorer, dependency viewer, grid mode (plus other layouts), and more - so we show the notebook in "Kiosk Mode" which is a read-only view of the outputs and helper panels. 47 | - VSCode's native notebook does not support different string quoting styles (e.g. `r"""`, `"""`, `f"""`, etc.), so we default all markdown cells to use `r"""`. 48 | 49 | ## Python Configuration 50 | 51 | To ensure marimo works correctly with your Python environment, you have several options: 52 | 53 | > [!TIP] 54 | > The extension will use the Python interpreter from the Python extension by default. Make sure you have the [Python extension](https://marketplace.visualstudio.com/items?itemName=ms-python.python) installed and configured. 55 | 56 | 1. **Workspace Settings (Recommended)** 57 | Create or edit `.vscode/settings.json` in your workspace. You can set the default Python interpreter for your entire workspace, or just for marimo. 58 | 59 | For setting the workspace Python interpreter, you can set: 60 | 61 | ```json 62 | { 63 | "python.defaultInterpreterPath": "${workspaceFolder}/.venv/bin/python" 64 | } 65 | ``` 66 | 67 | For setting the Python interpreter only for marimo, you can set: 68 | 69 | ```json 70 | { 71 | "marimo.pythonPath": "${workspaceFolder}/.venv/bin/python" 72 | } 73 | ``` 74 | 75 | If you set `marimo.pythonPath`, the extension will use that interpreter with `-m marimo` to invoke marimo. 76 | 77 | 2. **Global Settings** 78 | You can also configure these settings globally in VS Code's settings: 79 | 80 | - Set `python.defaultInterpreterPath` to your preferred Python interpreter 81 | - Verify that marimo is available in your Python interpreter: `/value/of/defaultInterpreterPath -m marimo` 82 | - (Likely not needed) Set `marimo.marimoPath` to the path of your marimo installation. When set, the extension will use this path directly `/path/to/marimo` instead of `python -m marimo`. 83 | 84 | 3. **Virtual Environments** 85 | If using a virtual environment: 86 | - Create and activate your virtual environment 87 | - Install marimo: `pip install marimo` 88 | - VS Code should automatically detect the Python interpreter 89 | 90 | 4. **uv projects and package environment sandboxes** 91 | If you are using `uv` to manage your Python project (e.g. with a `pyproject.toml` file). You can run `uv add marimo` to install marimo in your project's environment. Then update your settings to use: 92 | 93 | ```json 94 | { 95 | "marimo.marimoPath": "uv run marimo", 96 | "marimo.sandbox": true // optional 97 | } 98 | ``` 99 | 100 | 5. **uvx and package environment sandboxes** 101 | If you are not creating Python projects and don't want to create virtual environments, you can use `uvx` with `marimo edit --sandbox` to run marimo in a sandbox. 102 | 103 | ```json 104 | { 105 | "marimo.marimoPath": "uvx marimo", 106 | "marimo.sandbox": true 107 | } 108 | ``` 109 | 110 | 6. **windows** 111 | If you are on windows, you may have something that looks like this: 112 | 113 | ```json 114 | { 115 | "marimo.marimoPath": ".venv/Scripts/marimo.exe", 116 | "marimo.pythonPath": ".venv/Scripts/python.exe" 117 | } 118 | ``` 119 | 120 | ## Configuration 121 | 122 | - `marimo.browserType`: Browser to open marimo app (`system` or `embedded`, default: `embedded`) 123 | - `marimo.port`: Default port for marimo server (default: `2818`) 124 | - `marimo.sandbox`: Always start marimo in a sandbox, e.g. `marimo edit --sandbox` (default: `false`). Requires [`uv`](https://docs.astral.sh/uv/) to be installed. 125 | - `marimo.watch`: Always start marimo with the `--watch` flag (default: `true`). 126 | - `marimo.host`: Hostname for marimo server (default: `localhost`) 127 | - `marimo.https`: Enable HTTPS for marimo server (default: `false`) 128 | - `marimo.enableToken`: Enable token authentication (default: `false`) 129 | - `marimo.tokenPassword`: Token password (default: _empty_) 130 | - `marimo.showTerminal`: Open the terminal when the server starts (default: `true`) 131 | - `marimo.debug`: Enable debug logging (default: `false`) 132 | - `marimo.pythonPath`: Path to python interpreter (default: the one from python extension). Will be used with `/path/to/python -m marimo` to invoke marimo. 133 | - `marimo.marimoPath`: Path to a marimo executable (default: None). This will override use of the `pythonPath` setting, and instead invoke commands like `/path/to/marimo edit` instead of `python -m marimo edit`. 134 | 135 | ## Troubleshooting 136 | 137 | If you encounter issues, you can open the marimo extension logs by running the `marimo: Show marimo diagnostics` command from the command palette. 138 | 139 | You can also hover over the marimo status bar item in the bottom left of the VSCode window to see the status of the marimo extension. 140 | -------------------------------------------------------------------------------- /src/utils/__tests__/exec.test.ts: -------------------------------------------------------------------------------- 1 | import { createVSCodeMock } from "../../__mocks__/vscode"; 2 | 3 | vi.mock("vscode", () => createVSCodeMock(vi)); 4 | vi.mock("@vscode/python-extension", () => ({})); 5 | 6 | import * as child_process from "node:child_process"; 7 | import { beforeEach, describe, expect, it, vi } from "vitest"; 8 | import { workspace } from "vscode"; 9 | import { 10 | execMarimoCommand, 11 | execPythonFile, 12 | execPythonModule, 13 | hasExecutable, 14 | hasPythonModule, 15 | } from "../exec"; 16 | 17 | vi.mock("node:child_process", () => ({ 18 | execFileSync: vi.fn(), 19 | })); 20 | 21 | beforeEach(() => { 22 | vi.clearAllMocks(); 23 | }); 24 | 25 | describe("exec utilities", () => { 26 | const env = expect.objectContaining({ 27 | MARIMO_DANGEROUS_SANDBOX: "1", 28 | }) 29 | describe("execMarimoCommand", () => { 30 | it("should use marimoPath when set", async () => { 31 | workspace.getConfiguration().set("marimo.marimoPath", "/path/to/marimo"); 32 | await execMarimoCommand(["edit", "--version"]); 33 | expect(child_process.execFileSync).toHaveBeenCalledWith( 34 | "/path/to/marimo", 35 | ["edit", "--version"], 36 | { shell: false, encoding: "utf8", env }, 37 | ); 38 | }); 39 | 40 | it("should handle uv command in marimoPath", async () => { 41 | workspace.getConfiguration().set("marimo.marimoPath", "uv run marimo"); 42 | await execMarimoCommand(["edit", "--version"]); 43 | expect(child_process.execFileSync).toHaveBeenCalledWith( 44 | "uv", 45 | ["run", "marimo", "edit", "--version"], 46 | { shell: false, encoding: "utf8", env }, 47 | ); 48 | }); 49 | 50 | it("should handle uvx command in marimoPath", async () => { 51 | workspace.getConfiguration().set("marimo.marimoPath", "uvx run marimo"); 52 | await execMarimoCommand(["edit", "--version"]); 53 | expect(child_process.execFileSync).toHaveBeenCalledWith( 54 | "uvx", 55 | ["run", "marimo", "edit", "--version"], 56 | { shell: false, encoding: "utf8", env }, 57 | ); 58 | }); 59 | 60 | it("should fallback to python -m marimo when marimoPath is default", async () => { 61 | workspace.getConfiguration().set("marimo.marimoPath", "marimo"); 62 | workspace.getConfiguration().set("marimo.pythonPath", "python"); 63 | await execMarimoCommand(["edit", "--version"]); 64 | expect(child_process.execFileSync).toHaveBeenCalledWith( 65 | "python", 66 | ["-m", "marimo", "edit", "--version"], 67 | { shell: false, encoding: "utf8", env }, 68 | ); 69 | }); 70 | 71 | it("should handle spaces in marimoPath", async () => { 72 | workspace 73 | .getConfiguration() 74 | .set("marimo.marimoPath", "/path with spaces/marimo"); 75 | await execMarimoCommand(["edit", "--version"]); 76 | expect(child_process.execFileSync).toHaveBeenCalledWith( 77 | "/path with spaces/marimo", 78 | ["edit", "--version"], 79 | { shell: false, encoding: "utf8", env }, 80 | ); 81 | }); 82 | }); 83 | 84 | describe("execPythonModule", () => { 85 | it("should handle regular python path", async () => { 86 | workspace.getConfiguration().set("marimo.pythonPath", "python"); 87 | await execPythonModule(["marimo", "--version"]); 88 | expect(child_process.execFileSync).toHaveBeenCalledWith( 89 | "python", 90 | ["-m", "marimo", "--version"], 91 | { shell: false, encoding: "utf8", env }, 92 | ); 93 | }); 94 | 95 | it("should handle python path with spaces", async () => { 96 | workspace 97 | .getConfiguration() 98 | .set("marimo.pythonPath", "/path with spaces/python"); 99 | await execPythonModule(["marimo", "--version"]); 100 | expect(child_process.execFileSync).toHaveBeenCalledWith( 101 | "/path with spaces/python", 102 | ["-m", "marimo", "--version"], 103 | { shell: false, encoding: "utf8", env }, 104 | ); 105 | }); 106 | 107 | it("should handle uv run python", async () => { 108 | workspace.getConfiguration().set("marimo.pythonPath", "uv run python"); 109 | await execPythonModule(["marimo", "--version"]); 110 | expect(child_process.execFileSync).toHaveBeenCalledWith( 111 | "uv", 112 | ["run", "python", "-m", "marimo", "--version"], 113 | { shell: false, encoding: "utf8", env }, 114 | ); 115 | }); 116 | 117 | it("should handle uvx run python", async () => { 118 | workspace.getConfiguration().set("marimo.pythonPath", "uvx run python"); 119 | await execPythonModule(["marimo", "--version"]); 120 | expect(child_process.execFileSync).toHaveBeenCalledWith( 121 | "uvx", 122 | ["run", "python", "-m", "marimo", "--version"], 123 | { shell: false, encoding: "utf8", env }, 124 | ); 125 | }); 126 | }); 127 | 128 | describe("execPythonFile", () => { 129 | it("should handle regular python path", async () => { 130 | workspace.getConfiguration().set("marimo.pythonPath", "python"); 131 | await execPythonFile(["script.py", "--arg"]); 132 | expect(child_process.execFileSync).toHaveBeenCalledWith( 133 | "python", 134 | ["script.py", "--arg"], 135 | { shell: false, encoding: "utf8", env }, 136 | ); 137 | }); 138 | 139 | it("should handle uv run python", async () => { 140 | workspace.getConfiguration().set("marimo.pythonPath", "uv run python"); 141 | await execPythonFile(["script.py", "--arg"]); 142 | expect(child_process.execFileSync).toHaveBeenCalledWith( 143 | "uv", 144 | ["run", "python", "script.py", "--arg"], 145 | { shell: false, encoding: "utf8", env }, 146 | ); 147 | }); 148 | 149 | it("should handle spaces in script path", async () => { 150 | workspace.getConfiguration().set("marimo.pythonPath", "python"); 151 | await execPythonFile(["/path with spaces/script.py", "--arg"]); 152 | expect(child_process.execFileSync).toHaveBeenCalledWith( 153 | "python", 154 | ["/path with spaces/script.py", "--arg"], 155 | { shell: false, encoding: "utf8", env }, 156 | ); 157 | }); 158 | }); 159 | 160 | describe("hasPythonModule", () => { 161 | it("should handle hasPythonModule with uv run python", async () => { 162 | workspace.getConfiguration().set("marimo.pythonPath", "uv run python"); 163 | await hasPythonModule("marimo"); 164 | expect(child_process.execFileSync).toHaveBeenCalledWith( 165 | "uv", 166 | ["run", "python", "-c", "import marimo"], 167 | { shell: false, encoding: "utf8", env }, 168 | ); 169 | }); 170 | 171 | it("should handle module with spaces", async () => { 172 | workspace.getConfiguration().set("marimo.pythonPath", "python"); 173 | await hasPythonModule("module with spaces"); 174 | expect(child_process.execFileSync).toHaveBeenCalledWith( 175 | "python", 176 | ["-c", "import module with spaces"], 177 | { shell: false, encoding: "utf8", env }, 178 | ); 179 | }); 180 | }); 181 | 182 | describe("hasExecutable", () => { 183 | it("should return true when executable exists", async () => { 184 | (child_process.execFileSync as any).mockReturnValue(Buffer.from("")); 185 | const result = await hasExecutable("python"); 186 | expect(result).toBe(true); 187 | expect(child_process.execFileSync).toHaveBeenCalledWith( 188 | "python", 189 | ["--help"], 190 | { shell: false, encoding: "utf8", env }, 191 | ); 192 | }); 193 | 194 | it("should return false when executable doesn't exist", async () => { 195 | (child_process.execFileSync as any).mockImplementation(() => { 196 | throw new Error("Command failed"); 197 | }); 198 | const result = await hasExecutable("nonexistent"); 199 | expect(result).toBe(false); 200 | }); 201 | 202 | it("should handle executable with spaces", async () => { 203 | (child_process.execFileSync as any).mockReturnValue(Buffer.from("")); 204 | const result = await hasExecutable("/path with spaces/executable"); 205 | expect(result).toBe(true); 206 | expect(child_process.execFileSync).toHaveBeenCalledWith( 207 | "/path with spaces/executable", 208 | ["--help"], 209 | { shell: false, encoding: "utf8", env }, 210 | ); 211 | }); 212 | }); 213 | }); 214 | -------------------------------------------------------------------------------- /src/launcher/controller.ts: -------------------------------------------------------------------------------- 1 | import path from "node:path"; 2 | import { 3 | type Disposable, 4 | type Terminal, 5 | type TextDocument, 6 | Uri, 7 | env, 8 | workspace, 9 | } from "vscode"; 10 | import { MarimoPanelManager } from "../browser/panel"; 11 | import { Config, composeUrl } from "../config"; 12 | import { getGlobalState } from "../ctx"; 13 | import { MarimoRunningKernelsProvider } from "../explorer/explorer"; 14 | import { logger } from "../logger"; 15 | import type { ServerManager } from "../services/server-manager"; 16 | import { StatusBar } from "../ui/status-bar"; 17 | import { MarimoCmdBuilder } from "../utils/cmd"; 18 | import { getInterpreter, maybeQuotes } from "../utils/exec"; 19 | import { ping } from "../utils/network"; 20 | import { getFocusedMarimoTextEditor } from "../utils/query"; 21 | import { asURL } from "../utils/url"; 22 | import { type IMarimoTerminal, MarimoTerminal } from "./terminal"; 23 | 24 | export type AppMode = "edit" | "run"; 25 | 26 | export interface IMarimoController { 27 | currentMode: AppMode | undefined; 28 | active: boolean; 29 | port: number | undefined; 30 | terminal: MarimoTerminal; 31 | start(mode: AppMode, port: number): Promise; 32 | open(browser?: "embedded" | "system"): Promise; 33 | reloadPanel(): void; 34 | dispose(): void; 35 | } 36 | 37 | /** 38 | * Controller for a marimo app. 39 | * Manages a running marimo server, the terminal, and the webview panel. 40 | */ 41 | export class MarimoController implements Disposable { 42 | public terminal: IMarimoTerminal; 43 | private panel: MarimoPanelManager; 44 | 45 | public currentMode: AppMode | undefined; 46 | public active = false; 47 | public port: number | undefined; 48 | private logger = logger.createLogger(this.appName); 49 | 50 | constructor( 51 | public file: TextDocument, 52 | public serverManager: ServerManager, 53 | ) { 54 | this.file = file; 55 | const workspaceFolder = workspace.workspaceFolders?.find((folder) => 56 | file.uri.fsPath.startsWith(folder.uri.fsPath), 57 | )?.uri.fsPath; 58 | 59 | this.terminal = new MarimoTerminal( 60 | file.uri.fsPath, 61 | workspaceFolder, 62 | this.appName, 63 | ); 64 | this.panel = new MarimoPanelManager(this.appName); 65 | 66 | // Try to recover state 67 | this.tryRecoverState(); 68 | } 69 | 70 | hasTerminal(term: Terminal) { 71 | return this.terminal.is(term); 72 | } 73 | 74 | async start(mode: AppMode, port: number) { 75 | // If edit mode, use the existing server 76 | if (mode === "edit") { 77 | const { port } = await this.serverManager.start(); 78 | this.active = true; 79 | this.port = port; 80 | this.currentMode = mode; 81 | this.onUpdate(); 82 | return; 83 | } 84 | 85 | getGlobalState().update(this.keyFor("mode"), mode); 86 | this.currentMode = mode; 87 | getGlobalState().update(this.keyFor("port"), port); 88 | this.port = port; 89 | const filePath = this.terminal.relativePathFor(this.file.uri); 90 | 91 | const cmd = new MarimoCmdBuilder() 92 | .debug(Config.debug) 93 | .mode(mode) 94 | .fileOrDir(filePath) 95 | .host(Config.host) 96 | .port(port) 97 | .headless(true) 98 | .enableToken(Config.enableToken) 99 | .tokenPassword(Config.tokenPassword) 100 | .build(); 101 | 102 | const interpreter = await getInterpreter(); 103 | if (Config.marimoPath) { 104 | logger.info(`Using marimo path ${Config.marimoPath}`); 105 | await this.terminal.executeCommand( 106 | `${maybeQuotes(Config.marimoPath)} ${cmd}`, 107 | ); 108 | } else if (interpreter) { 109 | logger.info(`Using interpreter ${interpreter}`); 110 | await this.terminal.executeCommand( 111 | `${maybeQuotes(interpreter)} -m marimo ${cmd}`, 112 | ); 113 | } else { 114 | logger.info("Using default interpreter"); 115 | await this.terminal.executeCommand(cmd); 116 | } 117 | 118 | this.active = true; 119 | this.onUpdate(); 120 | } 121 | 122 | private onUpdate() { 123 | StatusBar.update(); 124 | } 125 | 126 | isWebviewActive() { 127 | return this.panel.isActive(); 128 | } 129 | 130 | async open(browser = Config.browser) { 131 | // If already opened, just focus 132 | if (browser === "embedded" && this.panel.isReady()) { 133 | this.panel.show(); 134 | } 135 | 136 | const url = await this.url(); 137 | if (browser === "system") { 138 | // Close the panel if opened 139 | this.panel.dispose(); 140 | await env.openExternal(Uri.parse(url)); 141 | } else if (browser === "embedded") { 142 | await this.panel.create(url); 143 | this.panel.show(); 144 | } 145 | 146 | // Wait 1s and refresh connections 147 | setTimeout(() => { 148 | MarimoRunningKernelsProvider.refresh(); 149 | }, 1000); 150 | } 151 | 152 | reloadPanel() { 153 | this.panel.reload(); 154 | } 155 | 156 | async tryRecoverState() { 157 | if (!(await this.terminal.tryRecoverTerminal())) { 158 | return; 159 | } 160 | this.logger.info("terminal recovered"); 161 | 162 | const port = +(getGlobalState().get(this.keyFor("port")) || 0); 163 | if (!port) { 164 | return; 165 | } 166 | 167 | const url = await composeUrl(port); 168 | 169 | if (!(await ping(url))) { 170 | return; 171 | } 172 | 173 | this.active = true; 174 | this.port = port; 175 | this.currentMode = getGlobalState().get(this.keyFor("mode")) || "edit"; 176 | this.logger.info("state recovered"); 177 | 178 | this.onUpdate(); 179 | return true; 180 | } 181 | 182 | dispose() { 183 | this.panel.dispose(); 184 | this.terminal.dispose(); 185 | this.active = false; 186 | this.port = undefined; 187 | getGlobalState().update(this.keyFor("mode"), undefined); 188 | getGlobalState().update(this.keyFor("port"), undefined); 189 | this.onUpdate(); 190 | } 191 | 192 | public get appName() { 193 | const filePath = this.file.uri.fsPath; 194 | const fileName = path.basename(filePath) || "app.py"; 195 | const folderName = path.basename(path.dirname(filePath)); 196 | return folderName ? `${folderName}/${fileName}` : fileName; 197 | } 198 | 199 | public async url(): Promise { 200 | if (!this.port) { 201 | return ""; 202 | } 203 | const url = asURL(await composeUrl(this.port)); 204 | if (this.currentMode === "edit") { 205 | url.searchParams.set("file", this.file.uri.fsPath); 206 | } 207 | return url.toString(); 208 | } 209 | 210 | private keyFor(key: string) { 211 | return `marimo.${this.file.uri.fsPath}.${key}`; 212 | } 213 | } 214 | 215 | export class ControllerManager implements Disposable { 216 | private controllers: Map = new Map(); 217 | 218 | constructor(private serverManager: ServerManager) {} 219 | 220 | getControllerForActivePanel(): MarimoController | undefined { 221 | return [...this.controllers.values()].find((c) => c.isWebviewActive()); 222 | } 223 | 224 | getOrCreate(file: TextDocument): MarimoController { 225 | const key = file.uri.fsPath; 226 | let controller = this.controllers.get(key); 227 | if (controller) { 228 | return controller; 229 | } 230 | controller = new MarimoController(file, this.serverManager); 231 | this.controllers.set(key, controller); 232 | return controller; 233 | } 234 | 235 | get(uri: Uri | undefined): MarimoController | undefined { 236 | if (!uri) { 237 | return; 238 | } 239 | return this.controllers.get(uri.fsPath); 240 | } 241 | 242 | findWithTerminal(term: Terminal): MarimoController | undefined { 243 | return [...this.controllers.values()].find((c) => c.terminal.is(term)); 244 | } 245 | 246 | dispose() { 247 | for (const c of this.controllers.values()) c.dispose(); 248 | this.controllers.clear(); 249 | } 250 | 251 | // Run a function an active or new controller 252 | run(fn: (controller: MarimoController) => T) { 253 | return this.runOptional((c) => { 254 | if (c) { 255 | return fn(c); 256 | } 257 | return undefined; 258 | }); 259 | } 260 | 261 | runOptional(fn: (controller: MarimoController | undefined) => T) { 262 | // If we are focused on a panel, use that controller 263 | const activePanelController = this.getControllerForActivePanel(); 264 | if (activePanelController) { 265 | return fn(activePanelController); 266 | } 267 | 268 | // If the active file is a marimo file, use that controller 269 | const file = getFocusedMarimoTextEditor({ toast: false }); 270 | if (!file) { 271 | return fn(undefined); 272 | } 273 | 274 | const controller = this.getOrCreate(file.document); 275 | return fn(controller); 276 | } 277 | } 278 | -------------------------------------------------------------------------------- /src/extension.ts: -------------------------------------------------------------------------------- 1 | import { 2 | type Disposable, 3 | type ExtensionContext, 4 | ProgressLocation, 5 | Uri, 6 | commands, 7 | env, 8 | window, 9 | workspace, 10 | } from "vscode"; 11 | import { showCommands } from "./commands/show-commands"; 12 | import { Config } from "./config"; 13 | import { CommandsKeys, DOCUMENTATION_URL } from "./constants"; 14 | import { convertIPyNotebook, convertMarkdownNotebook } from "./convert/convert"; 15 | import { setExtension } from "./ctx"; 16 | import { 17 | MarimoAppProvider, 18 | MarimoExplorer, 19 | MarimoRunningKernelsProvider, 20 | } from "./explorer/explorer"; 21 | import { exportAsCommands } from "./export/export-as-commands"; 22 | import { ControllerManager } from "./launcher/controller"; 23 | import { createNewMarimoFile } from "./launcher/new-file"; 24 | import { Launcher } from "./launcher/start"; 25 | import { logger } from "./logger"; 26 | import { NOTEBOOK_TYPE } from "./notebook/constants"; 27 | import { 28 | getActiveMarimoFile, 29 | handleOnOpenNotebookDocument, 30 | openMarimoNotebookDocument, 31 | } from "./notebook/extension"; 32 | import { KernelManager } from "./notebook/kernel-manager"; 33 | import { MarimoNotebookSerializer } from "./notebook/serializer"; 34 | import { HealthService } from "./services/health"; 35 | import { ServerManager } from "./services/server-manager"; 36 | import { setupConfigTelemetry, trackEvent } from "./telemetry"; 37 | import { StatusBar } from "./ui/status-bar"; 38 | 39 | class MarimoExtension { 40 | private extension: ExtensionContext; 41 | private kernelManager: KernelManager; 42 | private serverManager: ServerManager; 43 | private statusBar: StatusBar; 44 | private explorer: MarimoExplorer; 45 | private controllerManager: ControllerManager; 46 | private healthService: HealthService; 47 | 48 | constructor(extension: ExtensionContext) { 49 | this.extension = extension; 50 | this.kernelManager = KernelManager.instance; 51 | this.serverManager = ServerManager.getInstance(Config); 52 | this.controllerManager = new ControllerManager(this.serverManager); 53 | this.healthService = new HealthService(this.serverManager); 54 | this.statusBar = new StatusBar(this.controllerManager, this.healthService); 55 | // Sidebar explorer 56 | this.explorer = new MarimoExplorer( 57 | this.serverManager, 58 | this.controllerManager, 59 | ); 60 | } 61 | 62 | public async activate() { 63 | setExtension(this.extension); 64 | logger.info("marimo extension is now active!"); 65 | trackEvent("vscode-lifecycle", { action: "activate" }); 66 | 67 | this.serverManager.init(); 68 | this.addDisposable(this.kernelManager, this.serverManager, this.explorer); 69 | 70 | this.registerCommands(); 71 | this.registerEventListeners(); 72 | 73 | // Serializer 74 | this.registerNotebookSerializer(); 75 | } 76 | 77 | private addDisposable(...disposable: Disposable[]) { 78 | this.extension.subscriptions.push(...disposable); 79 | } 80 | 81 | private async refreshUI() { 82 | MarimoRunningKernelsProvider.refresh(); 83 | await this.statusBar.update(); 84 | } 85 | 86 | private registerCommand(command: string, handler: () => void) { 87 | this.extension.subscriptions.push( 88 | commands.registerCommand(command, () => { 89 | trackEvent("vscode-command", { command }); 90 | handler(); 91 | }), 92 | ); 93 | } 94 | 95 | private registerCommands() { 96 | this.registerCommand(CommandsKeys.startServer, () => this.startServer()); 97 | this.registerCommand(CommandsKeys.stopServer, () => this.stopServer()); 98 | this.registerCommand(CommandsKeys.edit, () => this.edit()); 99 | this.registerCommand(CommandsKeys.run, () => this.run()); 100 | this.registerCommand(CommandsKeys.restartKernel, () => 101 | this.restartKernel(), 102 | ); 103 | this.registerCommand(CommandsKeys.stopKernel, () => this.stopKernel()); 104 | this.registerCommand(CommandsKeys.showCommands, () => this.showCommands()); 105 | this.registerCommand(CommandsKeys.showHelp, () => this.showCommands()); 106 | this.registerCommand(CommandsKeys.exportAsCommands, () => 107 | this.exportAsCommands(), 108 | ); 109 | this.registerCommand(CommandsKeys.openInBrowser, () => 110 | this.openInBrowser(), 111 | ); 112 | this.registerCommand(CommandsKeys.reloadBrowser, () => 113 | this.reloadBrowser(), 114 | ); 115 | this.registerCommand(CommandsKeys.openNotebook, () => this.openNotebook()); 116 | this.registerCommand(CommandsKeys.openDocumentation, () => 117 | this.openDocumentation(), 118 | ); 119 | this.registerCommand(CommandsKeys.newMarimoFile, () => 120 | this.newMarimoFile(), 121 | ); 122 | this.registerCommand(CommandsKeys.convertToMarimoApp, () => 123 | this.convertToMarimoApp(), 124 | ); 125 | this.registerCommand(CommandsKeys.showDiagnostics, () => 126 | this.healthService.showDiagnostics(), 127 | ); 128 | setupConfigTelemetry(); 129 | } 130 | 131 | private async startServer() { 132 | await window.withProgress( 133 | { 134 | location: ProgressLocation.Notification, 135 | title: "Starting marimo server...", 136 | cancellable: true, 137 | }, 138 | async (_, cancellationToken) => { 139 | try { 140 | const response = await this.serverManager.start(cancellationToken); 141 | await this.kernelManager.hydrateExistingNotebooks(response); 142 | } catch (e) { 143 | window.showErrorMessage(`Failed to start marimo server: ${e}`); 144 | } 145 | }, 146 | ); 147 | } 148 | 149 | private async stopServer() { 150 | await window.withProgress( 151 | { 152 | location: ProgressLocation.Notification, 153 | title: "Stopping marimo server...", 154 | cancellable: false, 155 | }, 156 | async () => { 157 | await this.closeAllMarimoNotebookEditors(); 158 | await this.serverManager.stopServer(); 159 | this.kernelManager.clearAllKernels(); 160 | await this.refreshUI(); 161 | }, 162 | ); 163 | } 164 | 165 | private async closeAllMarimoNotebookEditors() { 166 | const editors = window.visibleNotebookEditors.filter( 167 | (editor) => editor.notebook.notebookType === NOTEBOOK_TYPE, 168 | ); 169 | for (const editor of editors) { 170 | await window.showTextDocument(editor.notebook.uri, { preview: false }); 171 | await commands.executeCommand("workbench.action.closeActiveEditor"); 172 | } 173 | } 174 | 175 | private edit() { 176 | this.controllerManager.run(async (controller) => { 177 | await Launcher.start({ controller, mode: "edit" }); 178 | controller.open(); 179 | }); 180 | } 181 | 182 | private run() { 183 | this.controllerManager.run(async (controller) => { 184 | await Launcher.start({ controller, mode: "run" }); 185 | controller.open(); 186 | }); 187 | } 188 | 189 | private async restartKernel() { 190 | const maybeKernel = KernelManager.getFocusedMarimoKernel(); 191 | if (maybeKernel) { 192 | await window.withProgress( 193 | { 194 | location: ProgressLocation.Notification, 195 | title: "Restarting marimo kernel...", 196 | cancellable: true, 197 | }, 198 | async (_, cancellationToken) => { 199 | await this.serverManager.start(cancellationToken); 200 | await maybeKernel.restart(); 201 | await maybeKernel.openKiosk(); 202 | }, 203 | ); 204 | return; 205 | } 206 | 207 | this.controllerManager.run(async (controller) => { 208 | const mode = controller.currentMode || "edit"; 209 | Launcher.stop(controller); 210 | await Launcher.start({ controller, mode }); 211 | controller.open(); 212 | }); 213 | } 214 | 215 | private stopKernel() { 216 | this.controllerManager.run(async (controller) => { 217 | Launcher.stop(controller); 218 | }); 219 | } 220 | 221 | private showCommands() { 222 | const maybeKernel = KernelManager.getFocusedMarimoKernel(); 223 | if (maybeKernel) { 224 | showCommands(maybeKernel, this.serverManager); 225 | return; 226 | } 227 | 228 | this.controllerManager.runOptional(async (controller) => { 229 | showCommands(controller, this.serverManager); 230 | }); 231 | } 232 | 233 | private exportAsCommands() { 234 | const maybeKernel = KernelManager.getFocusedMarimoKernel(); 235 | if (maybeKernel) { 236 | exportAsCommands(maybeKernel.fileUri); 237 | return; 238 | } 239 | 240 | this.controllerManager.run(async (controller) => { 241 | exportAsCommands(controller.file.uri); 242 | }); 243 | } 244 | 245 | private openInBrowser() { 246 | const maybeKernel = KernelManager.getFocusedMarimoKernel(); 247 | if (maybeKernel) { 248 | maybeKernel.openKiosk("system"); 249 | return; 250 | } 251 | 252 | this.controllerManager.run(async (controller) => { 253 | controller.open("system"); 254 | }); 255 | } 256 | 257 | private reloadBrowser() { 258 | const maybeKernel = KernelManager.getFocusedMarimoKernel(); 259 | if (maybeKernel) { 260 | maybeKernel.reloadPanel(); 261 | return; 262 | } 263 | 264 | this.controllerManager.run(async (controller) => { 265 | controller.reloadPanel(); 266 | }); 267 | } 268 | 269 | private async openNotebook() { 270 | await openMarimoNotebookDocument(await getActiveMarimoFile()); 271 | } 272 | 273 | private openDocumentation() { 274 | env.openExternal(Uri.parse(DOCUMENTATION_URL)); 275 | } 276 | 277 | private async newMarimoFile() { 278 | await createNewMarimoFile(); 279 | await openMarimoNotebookDocument(await getActiveMarimoFile()); 280 | } 281 | 282 | private async convertToMarimoApp() { 283 | const editor = window.activeTextEditor; 284 | if (!editor) { 285 | window.showErrorMessage("No active editor"); 286 | return; 287 | } 288 | 289 | const filePath = editor.document.uri.fsPath; 290 | if (filePath.endsWith(".ipynb")) { 291 | await convertIPyNotebook(filePath); 292 | return; 293 | } 294 | if (filePath.endsWith(".md")) { 295 | const content = editor.document.getText(); 296 | if (!content.includes("marimo-version:")) { 297 | await convertMarkdownNotebook(filePath); 298 | return; 299 | } 300 | } 301 | 302 | window.showErrorMessage("Not a notebook file"); 303 | } 304 | 305 | private registerEventListeners() { 306 | window.onDidCloseTerminal(async (terminal) => { 307 | if (this.serverManager.terminal.is(terminal)) { 308 | await this.serverManager.dispose(); 309 | await this.kernelManager.clearAllKernels(); 310 | await this.refreshUI(); 311 | } 312 | 313 | const controller = this.controllerManager.findWithTerminal(terminal); 314 | controller?.dispose(); 315 | }); 316 | 317 | workspace.onDidOpenNotebookDocument(async (doc) => { 318 | await handleOnOpenNotebookDocument( 319 | doc, 320 | this.serverManager, 321 | this.kernelManager, 322 | ); 323 | }); 324 | } 325 | 326 | private registerNotebookSerializer() { 327 | workspace.registerNotebookSerializer( 328 | NOTEBOOK_TYPE, 329 | new MarimoNotebookSerializer(this.kernelManager), 330 | { 331 | transientOutputs: true, 332 | }, 333 | ); 334 | } 335 | } 336 | 337 | export async function activate(extension: ExtensionContext) { 338 | const marimoExtension = new MarimoExtension(extension); 339 | await marimoExtension.activate(); 340 | } 341 | 342 | export async function deactivate() { 343 | logger.info("marimo extension is now deactivated!"); 344 | trackEvent("vscode-lifecycle", { action: "deactivate" }); 345 | 346 | // Make sure to stop any running server on VSCode shutdown 347 | try { 348 | const serverManager = ServerManager.getInstance(Config); 349 | if (serverManager) { 350 | await serverManager.stopServer(); 351 | } 352 | } catch (e) { 353 | logger.error("Error stopping server during deactivation", e); 354 | } 355 | } 356 | --------------------------------------------------------------------------------