marimo

1// marimo: follows the live marimo notebook open under the project.
2//
3// - Status line under the prompt: `marimo: fit.py · running Data loading › Model fit (12s) · 1 error`.
4// - Context: each prompt carries the notebook's current state beside it (prompt.submit's context).
5//   Older copies are left out of every later request (prompt.attachment answers null for them),
6//   so the conversation holds one copy, the latest, instead of one per prompt.
7// - /marimo: status, show, auto, off, or a notebook path to follow.
8// - herdr: when a turn ends with a cell Claude started still running, the pane token `marimo`
9//   says what runs until the kernel goes quiet, then a notification says it finished.
10//
11// The port of pi-marimo; core/ is shared with it unchanged. The host reads on(...) and
12// $.noun.method(...) from source, so every $ call is spelled in full in this file and the
13// core reaches the host only through the Io built in claudeIo.
14import type { EngineInterface, On } from 'claude-code'
15
16import { HERDR_SOURCE, HERDR_TOKEN, type HoldChange, KernelHold } from './core/hold.js'
17import { type Cancellable, type Io, normalize, type StreamEnd } from './core/io.js'
18import { type SnapshotEntry, snapshot, STATE_TAG, statusParts, statusText } from './core/render.js'
19import { pairTargets } from './core/touch.js'
20import { MarimoWatcher, type Mode } from './core/watcher.js'
21
22const COMMAND = 'marimo'
23
24/** What the hooks share across a session. */
25interface State {
26  watcher?: MarimoWatcher
27  /** Browser edits after this change number are flagged as new in the next prompt's state. */
28  /** Per notebook path: browser edits after this change number are flagged as new in the next state. */
29  seen: Map<string, number>
30  /** Revisions of the state this process attached are `<nonce>-<n>`; only the latest is kept. */
31  nonce: string
32  revision: number
33  ticker?: { cancel: () => void }
34  lastStatus?: string
35  /** A cell Claude started that outlives its turn, shown in herdr. */
36  hold: KernelHold
37  /** herdr reports, in order. */
38  herdr: Promise<unknown>
39  /** Whether the `marimo` token is set. */
40  holding: boolean
41}
42
43export function register(on: On): void {
44  const s: State = { seen: new Map(), nonce: Math.random().toString(36).slice(2, 8), revision: 0, hold: new KernelHold(), herdr: Promise.resolve(), holding: false }
45
46  on('session.start', async ($, e, next) => {
47    const started = await next(e)
48    await $.command.register({
49      name: COMMAND,
50      description: 'marimo: status, show, auto, off, or a notebook path to follow',
51      argumentHint: '[status|show|auto|off|<path>]',
52      immediate: true,
53    })
54    const cwd = await $.session.cwd()
55    const token = (await $.env.get('MARIMO_TOKEN'))?.trim() || undefined
56    const pinned = (await $.env.get('MARIMO_NOTEBOOK'))?.trim()
57    startWatcher($, s, cwd, token, pinned)
58    return started
59  })
60
61  on('session.end', async ($, e, next) => {
62    await herdr($, s, s.hold.release())
63    s.ticker?.cancel()
64    s.ticker = undefined
65    await s.watcher?.stop()
66    s.watcher = undefined
67    return next(e)
68  })
69
70  on('prompt.submit', async ($, e, next) => {
71    const watcher = s.watcher
72    if (!watcher) return next(e)
73    watcher.refresh()
74    await watcher.settle(2000)
75    const entries = stateEntries(watcher, s)
76    if (!entries.length) return next(e)
77    s.revision++
78    const text = stateText(entries, `${s.nonce}-${s.revision}`)
79    const result = await next({ ...e, context: [...(e.context ?? []), text] })
80    // Ask every attachment again, so the copies this one supersedes are left out.
81    $.ui.invalidate('prompt.attachment')
82    return result
83  })
84
85  on('prompt.attachment', ($, e, next) => {
86    if (e.origin.kind === 'plugin' && isStaleState(e.text, `${s.nonce}-${s.revision}`)) return { text: null }
87    return next(e)
88  })
89
90  // The notebook this session's agent works in (through marimo-pair) becomes current for it,
91  // whatever other sessions or the browser do; subagents' calls count as this session's.
92  on('tool.call', ($, e, next) => {
93    const targets = pairTargets(JSON.stringify(e))
94    if (targets.length) s.watcher?.touch(targets)
95    return next(e)
96  })
97
98  on('turn.start', ($, e, next) => {
99    void herdr($, s, s.hold.begin(Date.now()))
100    return next(e)
101  })
102
103  on('turn.complete', async ($, e, next) => {
104    if (e.agentId !== undefined) return next(e)
105    const followed = s.watcher?.followed() ?? []
106    for (const f of followed) s.seen.set(f.attachment.path, f.notebook.seq)
107    void herdr($, s, s.hold.end(followed, Date.now()))
108    return next(e)
109  })
110
111  on('command.run', { command: COMMAND }, async ($, e) => {
112    const watcher = s.watcher
113    if (!watcher) return { text: 'marimo: not started' }
114    const arg = e.args.trim()
115    if (arg === 'auto' || arg === 'off') {
116      watcher.setMode(modeFor(arg === 'off' ? 'off' : undefined))
117      return { text: arg === 'off' ? 'marimo: off' : 'marimo: following every notebook open under this directory' }
118    }
119    if (arg === 'show') {
120      const entries = stateEntries(watcher, s)
121      return { text: entries.length ? snapshot(entries, { refresh: 'prompt' }) : statusFor(watcher) ?? 'marimo: no notebook attached' }
122    }
123    if (arg && arg !== 'status') {
124      const cwd = await $.session.cwd()
125      const paths = arg.split(/[\s,]+/).filter(Boolean).map((p) => normalize(p, cwd))
126      watcher.setMode({ kind: 'pinned', paths })
127      return { text: `marimo: following ${paths.join(', ')}` }
128    }
129    const open = await watcher.list()
130    return {
131      text: [
132        statusFor(watcher) ?? 'marimo: no notebook attached',
133        `Mode: ${watcher.mode.kind}${watcher.mode.kind === 'pinned' ? ` (${watcher.mode.paths.join(', ')})` : ''}.`,
134        open.length ? `Open notebooks:\n${open.map((n) => `  ${n.path}  (${n.url}, session ${n.sessionId})`).join('\n')}` : 'No marimo notebooks are open.',
135      ].join('\n'),
136    }
137  })
138
139}
140
141function startWatcher($: EngineInterface, s: State, cwd: string, token: string | undefined, pinned: string | undefined): void {
142  s.watcher = new MarimoWatcher({ io: claudeIo($), cwd, token, onChange: () => showStatus($, s) })
143  s.watcher.mode = modeFor(pinned)
144  s.watcher.start()
145}
146
147/** Pin the status line; tick it every second while a cell runs. */
148function showStatus($: EngineInterface, s: State): void {
149  const watcher = s.watcher
150  if (!watcher) return
151  if (s.hold.active) void herdr($, s, s.hold.update(watcher.followed(), Date.now()))
152  const text = statusFor(watcher)
153  if (text !== s.lastStatus) {
154    s.lastStatus = text
155    // The host labels a plugin's status line with the plugin's name ("marimo: ").
156    $.ui.status(text?.replace(/^marimo: /, ''))
157  }
158  const running = watcher.connection === 'connected' && watcher.notebook.running()
159  if (running && !s.ticker) s.ticker = $.clock.every(1000, () => showStatus($, s))
160  if (!running && s.ticker) {
161    s.ticker.cancel()
162    s.ticker = undefined
163  }
164}
165
166/**
167 * The herdr CLI calls for a hold change. herdr's Claude integration only reports the session and
168 * herdr reads Claude Code's state from the screen; it takes no state from any other reporter once
169 * that session is known, so the pane stays `idle`. While held, the pane token `marimo` says what
170 * runs; when the run ends it is cleared, and a notification follows unless the pane is focused.
171 */
172export function herdrCalls(change: HoldChange, pane: string, reported: boolean): string[][] {
173  const metadata = ['herdr', 'pane', 'report-metadata', pane, '--source', HERDR_SOURCE, '--agent', 'claude']
174  if (change.kind === 'none') return []
175  if (change.kind === 'held') return [[...metadata, '--token', `${HERDR_TOKEN}=${change.value}`, '--ttl-ms', '86400000']]
176  return reported ? [[...metadata, '--clear-token', HERDR_TOKEN]] : []
177}
178
179/** Report a hold change to herdr, inside a herdr pane only; any failure is ignored. */
180function herdr($: EngineInterface, s: State, change: HoldChange): Promise<unknown> {
181  if (change.kind === 'none') return s.herdr
182  s.herdr = s.herdr.then(async () => {
183    if ((await $.env.get('HERDR_ENV')) !== '1') return
184    const pane = await $.env.get('HERDR_PANE_ID')
185    if (!pane) return
186    const calls = herdrCalls(change, pane, s.holding)
187    s.holding = change.kind === 'held'
188    for (const argv of calls) await $.process.run(argv, { timeoutMs: 3_000 }).catch(() => undefined)
189    if (change.kind !== 'finished') return
190    const got = await $.process.run(['herdr', 'pane', 'get', pane], { timeoutMs: 3_000 }).catch(() => undefined)
191    if (got && got.exitCode === 0 && /"focused":true/.test(got.stdout)) return
192    await $.process.run(['herdr', 'notification', 'show', change.title, '--body', change.body, '--sound', 'done'], { timeoutMs: 3_000 }).catch(() => undefined)
193  }).catch(() => undefined)
194  return s.herdr
195}
196
197/**
198 * MARIMO_NOTEBOOK=<path>[,<path>...] follows those notebooks, =off turns the mod off; unset follows
199 * every notebook open under the project.
200 */
201export function modeFor(pinned: string | undefined): Mode {
202  if (pinned === 'off') return { kind: 'off' }
203  const paths = (pinned ?? '').split(',').map((p) => p.trim()).filter(Boolean)
204  return paths.length ? { kind: 'pinned', paths } : { kind: 'auto' }
205}
206
207/** The followed notebooks that are connected, the current one first. */
208export function stateEntries(watcher: MarimoWatcher, s: { seen: Map<string, number> }) {
209  return watcher.followed()
210    .filter((f) => f.connection === 'connected' && f.notebook.ready)
211    .map((f) => ({ notebook: f.notebook, attachment: f.attachment, current: f.current, seenSeq: s.seen.get(f.attachment.path) ?? 0 }))
212}
213
214export function statusFor(watcher: MarimoWatcher): string | undefined {
215  const { connection, attachment, mode } = watcher
216  if (!attachment) return undefined
217  if (connection === 'searching') {
218    return mode.kind === 'pinned' ? statusText({ notebook: mode.paths.map((p) => p.split('/').pop()).join(', '), connection: 'not open', queued: 0, errors: 0 }) : undefined
219  }
220  return statusText(statusParts(watcher.notebook, attachment, connection, Date.now(), watcher.others()))
221}
222
223/** The state block, tagged with its revision so later requests can tell it is superseded. */
224export function stateText(entries: SnapshotEntry[], rev: string): string {
225  return snapshot(entries, { refresh: 'prompt' }).replace(`<${STATE_TAG} `, `<${STATE_TAG} rev="${rev}" `)
226}
227
228/** True for a state block of this plugin that is not the latest revision. */
229export function isStaleState(text: string, latest: string): boolean {
230  const match = new RegExp(`<${STATE_TAG} rev="([^"]+)"`).exec(text)
231  return match !== null && match[1] !== latest
232}
233
234/** The core's host calls, through $: no Node APIs in a mod's sandbox. */
235function claudeIo($: EngineInterface): Io {
236  return {
237    async registryEntries() {
238      const home = (await $.env.get('HOME')) ?? ''
239      const state = (await $.env.get('XDG_STATE_HOME')) || `${home}/.local/state`
240      const dir = `${state}/marimo/servers`
241      const entries: unknown[] = []
242      const files = await $.fs.list(dir).catch(() => [])
243      for (const file of files) {
244        if (file.kind !== 'file' || !file.name.endsWith('.json')) continue
245        try {
246          entries.push(JSON.parse(await $.fs.read(`${dir}/${file.name}`)))
247        } catch {
248          // A file being written or removed; skip it this round.
249        }
250      }
251      return entries
252    },
253
254    async getJson(url, headers, timeoutMs) {
255      const request = $.http.fetch(url, { headers }).then(
256        (response) => (response.ok ? JSON.parse(response.text) : undefined),
257        () => undefined,
258      )
259      const timeout = $.clock.sleep(timeoutMs).then(() => undefined)
260      return Promise.race([request, timeout]).catch(() => undefined)
261    },
262
263    // $.http.fetch buffers a whole body, so the SSE stream is read through curl.
264    stream(url, headers, onText): Cancellable<StreamEnd> {
265      const argv = ['curl', '--silent', '--no-buffer', '--fail']
266      for (const [name, value] of Object.entries(headers)) argv.push('--header', `${name}: ${value}`)
267      argv.push(url)
268      const child = $.process.spawn({ argv })
269      let received = false
270      let cancelled = false
271      const done = (async (): Promise<StreamEnd> => {
272        try {
273          for await (const chunk of child) {
274            if (cancelled) break
275            if (chunk.stream !== 'stdout') continue
276            received = true
277            onText(chunk.text)
278          }
279        } catch {
280          return { ok: cancelled }
281        }
282        return { ok: received || cancelled }
283      })()
284      return {
285        done,
286        cancel: () => {
287          cancelled = true
288          void child.return(undefined as never).catch(() => undefined)
289        },
290      }
291    },
292
293    async realpath(path) {
294      const result = await $.process.run(['realpath', path]).catch(() => undefined)
295      const real = result && result.exitCode === 0 ? result.stdout.trim() : ''
296      return real || path
297    },
298
299    sleep(ms): Cancellable<void> {
300      let cancel = (): void => {}
301      const done = new Promise<void>((resolve) => {
302        const timer = $.clock.after(ms, resolve)
303        cancel = () => {
304          timer.cancel()
305          resolve()
306        }
307      })
308      return { done, cancel }
309    },
310  }
311}
312

1// The kernel hold: the stretch after the agent's turn ends while a cell it
2// started is still running. It begins when a turn ends with a notebook the agent
3// worked in during that turn busy (a cell running or queued), and ends when
4// every held notebook is quiet, disconnected or no longer followed, or when the
5// next turn begins. The hosts show it in herdr; this file only decides it.
6
7import { basename } from "./io.js";
8import { truncate } from "./notebook.js";
9import { elapsed } from "./render.js";
10import type { Followed } from "./watcher.js";
11
12/** The herdr pane token the ports report while the kernel holds (`$marimo` in a sidebar row). */
13export const HERDR_TOKEN = "marimo";
14/** The herdr metadata source the ports report under. */
15export const HERDR_SOURCE = "marimo";
16
17export type HoldChange =
18  | { kind: "none" }
19  /** The hold began, or what it shows changed. */
20  | { kind: "held"; value: string }
21  /** The hold ended because the kernel went quiet; what to tell the user. */
22  | { kind: "finished"; title: string; body: string }
23  /** The hold ended without the kernel finishing (a new turn, the session ending). */
24  | { kind: "released" };
25
26const NONE: HoldChange = { kind: "none" };
27
28export function busy(f: Followed): boolean {
29  return f.connection === "connected" && (f.notebook.running() !== undefined || f.notebook.queued() > 0);
30}
31
32/** What the pane token says: `fit.py: Model fit`, `fit.py: cell train`, or `fit.py: 3 queued`. */
33export function holdValue(f: Followed): string {
34  const file = basename(f.attachment.path);
35  const run = f.notebook.running();
36  if (!run) return `${file}: ${f.notebook.queued()} queued`;
37  const where = run.section[run.section.length - 1] ?? `cell ${run.cell.name !== "_" ? run.cell.name : run.cell.id}`;
38  return truncate(`${file}: ${where}`, 80);
39}
40
41export class KernelHold {
42  private turnStartedAt: number | undefined;
43  private held = new Set<string>();
44  private since = 0;
45  private last: string | undefined;
46
47  get active(): boolean {
48    return this.held.size > 0;
49  }
50
51  /** The notebooks held, by path. */
52  get paths(): string[] {
53    return [...this.held];
54  }
55
56  /** A turn began: the agent is working again, so any hold ends. */
57  begin(now: number): HoldChange {
58    this.turnStartedAt = now;
59    return this.release();
60  }
61
62  /** The turn ended: hold the busy notebooks the agent worked in during it. */
63  end(followed: Followed[], now: number): HoldChange {
64    const startedAt = this.turnStartedAt;
65    this.turnStartedAt = undefined;
66    if (startedAt === undefined || this.active) return NONE;
67    for (const f of followed) if (f.touchedAt >= startedAt && busy(f)) this.held.add(f.attachment.path);
68    if (!this.active) return NONE;
69    this.since = now;
70    return this.update(followed, now);
71  }
72
73  /** After a notebook changed: what the hold shows now, or that it finished. */
74  update(followed: Followed[], now: number): HoldChange {
75    if (!this.active) return NONE;
76    const held = followed.filter((f) => this.held.has(f.attachment.path));
77    const running = held.filter(busy);
78    if (!running.length) {
79      this.held.clear();
80      this.last = undefined;
81      const files = held.length ? held.map((f) => basename(f.attachment.path)).join(", ") : "notebook";
82      const errors = held.reduce((n, f) => n + (f.connection === "connected" ? f.notebook.errors().length : 0), 0);
83      const gone = held.length < 1 || held.some((f) => f.connection !== "connected");
84      const parts = [`ran ${elapsed(now - this.since)}`];
85      if (errors) parts.push(`${errors} error${errors === 1 ? "" : "s"}`);
86      if (gone) parts.push("disconnected");
87      return { kind: "finished", title: truncate(`${files} ${gone ? "stopped" : "finished"}`, 80), body: parts.join(" · ") };
88    }
89    const value = holdValue(running[0]!);
90    if (value === this.last) return NONE;
91    this.last = value;
92    return { kind: "held", value };
93  }
94
95  /** Stop holding without a notice. */
96  release(): HoldChange {
97    if (!this.active) return NONE;
98    this.held.clear();
99    this.last = undefined;
100    return { kind: "released" };
101  }
102}
103

1// Everything the core needs from its host, so it runs unchanged under Node and
2// Bun (pi, opencode: node-io.ts) and in sandboxes that offer only their own
3// file, HTTP, process and clock calls (Claude Code mods). The rest of src/core
4// is plain ECMAScript.
5
6export interface Cancellable<T> {
7  done: Promise<T>;
8  cancel(): void;
9}
10
11export interface StreamEnd {
12  /** The server answered and the body ended (or was cancelled). */
13  ok: boolean;
14  status?: number;
15}
16
17export interface Io {
18  /** The parsed JSON of every entry in marimo's server registry; servers known to be dead may be left out. */
19  registryEntries(): Promise<unknown[]>;
20  /** GET a JSON document; undefined on any failure. */
21  getJson(url: string, headers: Record<string, string>, timeoutMs: number): Promise<unknown>;
22  /** GET a streaming body, handing over its text as it arrives. */
23  stream(url: string, headers: Record<string, string>, onText: (text: string) => void): Cancellable<StreamEnd>;
24  /** The path with symlinks resolved (macOS /tmp is /private/tmp); the path itself when that fails. */
25  realpath(path: string): Promise<string>;
26  sleep(ms: number): Cancellable<void>;
27}
28
29/** Posix path helpers, enough for comparing absolute notebook paths. */
30export function normalize(path: string, cwd = "/"): string {
31  const parts: string[] = [];
32  for (const part of (path.startsWith("/") ? path : `${cwd}/${path}`).split("/")) {
33    if (!part || part === ".") continue;
34    if (part === "..") parts.pop();
35    else parts.push(part);
36  }
37  return `/${parts.join("/")}`;
38}
39
40export function basename(path: string): string {
41  return path.replace(/\/+$/, "").split("/").pop() ?? path;
42}
43
44/** Whether path is dir or inside it (both absolute and normalized). */
45export function within(path: string, dir: string): boolean {
46  const d = normalize(dir);
47  const p = normalize(path);
48  return p === d || p.startsWith(d === "/" ? "/" : `${d}/`);
49}
50

1// Text views of a NotebookState: a one-line status for footers and status lines,
2// and the compact state block injected into the agent's context.
3
4import { basename } from "./io.js";
5import { type Cell, NotebookState, truncate } from "./notebook.js";
6
7export interface Attachment {
8  url: string;
9  sessionId: string;
10  path: string;
11}
12
13export const STATE_TAG = "marimo_notebook_state";
14
15export function elapsed(ms: number): string {
16  const s = Math.max(0, Math.floor(ms / 1000));
17  if (s < 60) return `${s}s`;
18  const m = Math.floor(s / 60);
19  if (m < 60) return `${m}m${String(s % 60).padStart(2, "0")}s`;
20  return `${Math.floor(m / 60)}h${String(m % 60).padStart(2, "0")}m`;
21}
22
23export interface StatusParts {
24  notebook: string;
25  /** "connected" | "connecting" | "disconnected" */
26  connection: string;
27  running?: { section: string; cell: string; elapsed?: string };
28  queued: number;
29  errors: number;
30  /** File names of the other notebooks followed, most recently used first. */
31  others?: string[];
32}
33
34export function statusParts(nb: NotebookState, attachment: Attachment, connection: string, now = Date.now(), others: Attachment[] = []): StatusParts {
35  const run = connection === "connected" ? nb.running() : undefined;
36  return {
37    notebook: basename(attachment.path),
38    connection,
39    running: run
40      ? {
41        section: run.section.join(" › "),
42        cell: run.cell.name !== "_" ? run.cell.name : run.cell.id,
43        elapsed: run.since !== undefined ? elapsed(now - run.since) : undefined,
44      }
45      : undefined,
46    queued: connection === "connected" ? nb.queued() : 0,
47    errors: connection === "connected" ? nb.errors().length : 0,
48    others: others.map((o) => basename(o.path)),
49  };
50}
51
52/**
53 * One line, e.g. `marimo: fit.py · running Data loading › Model fit (12s) · 2 queued · 1 error · also prep.py, plots.py`.
54 * The shape is stable so footers can parse it: `marimo: <current file>`, then ` · ` separated parts
55 * about it, then `also <file>, <file>` naming the other notebooks followed.
56 */
57export function statusText(parts: StatusParts): string {
58  const out = [`marimo: ${parts.notebook}`];
59  if (parts.connection !== "connected") out.push(parts.connection);
60  if (parts.running) {
61    const where = parts.running.section || `cell ${parts.running.cell}`;
62    out.push(`running ${where}${parts.running.elapsed ? ` (${parts.running.elapsed})` : ""}`);
63  }
64  if (parts.queued) out.push(`${parts.queued} queued`);
65  if (parts.errors) out.push(`${parts.errors} error${parts.errors === 1 ? "" : "s"}`);
66  if (parts.others?.length) out.push(`also ${parts.others.join(", ")}`);
67  return out.join(" · ");
68}
69
70function firstLine(code: string): string {
71  const line = code.split("\n").map((l) => l.trim()).find((l) => l && !l.startsWith("#")) ?? "";
72  return truncate(line, 60);
73}
74
75export interface SnapshotOptions {
76  /** When the block is refreshed: before every model request, or once per user prompt. */
77  refresh?: "request" | "prompt";
78  /** Above this many cells, quiet code cells of the current notebook are folded into counts. */
79  maxCells?: number;
80  now?: number;
81}
82
83export interface SnapshotEntry {
84  notebook: NotebookState;
85  attachment: Attachment;
86  current?: boolean;
87  /** Edits from the browser after this change number are flagged as new. */
88  seenSeq?: number;
89}
90
91interface SectionOptions {
92  seen: number;
93  now: number;
94  maxCells: number;
95  /** Deeper headings are left out (for the notebooks that are not current). */
96  maxLevel: number;
97}
98
99/**
100 * One notebook as an outline: headings, then one line per code cell with its
101 * id, what it defines, and anything that needs attention.
102 */
103function section(nb: NotebookState, attachment: Attachment, current: boolean, o: SectionOptions): string[] {
104  const { now, seen } = o;
105  const cells = nb.ordered();
106  const fold = cells.length > o.maxCells;
107  // When most cells are stale (a lazy notebook after a restart), say so once instead of on every line.
108  const staleCount = cells.filter((c) => c.stale).length;
109  const staleCommon = staleCount > cells.length / 2;
110  // Likewise a Run all: the queue is counted on the Kernel line, so list queued cells only when few.
111  const queuedCount = cells.filter((c) => c.status === "queued").length;
112  const queuedCommon = queuedCount > 10 || queuedCount > cells.length / 2;
113
114  const notes = (cell: Cell): string[] => {
115    const out: string[] = [];
116    if (cell.status === "running") out.push(`RUNNING${cell.runningSince ? ` ${elapsed(now - cell.runningSince)}` : ""}`);
117    else if (cell.status === "queued" && !queuedCommon) out.push("queued");
118    if (cell.error) out.push(`ERROR ${cell.error}`);
119    if (cell.blocked) out.push("not run: an ancestor failed");
120    if (cell.disabled) out.push("disabled");
121    else if (cell.status === "disabled-transitively") out.push("disabled by an ancestor");
122    if (cell.stale && !staleCommon) out.push("stale");
123    if (nb.edited(cell)) out.push("edited, not rerun");
124    if (cell.editedBy === "frontend" && cell.editSeq > seen) out.push("changed by the user in the browser");
125    return out;
126  };
127
128  const lines: string[] = [];
129  let folded = 0;
130  const flushFolded = (): void => {
131    if (folded) lines.push(`  … ${folded} more cell${folded === 1 ? "" : "s"}`);
132    folded = 0;
133  };
134  const heading = (level: number, text: string, tail: string): void => {
135    if (level <= o.maxLevel) { flushFolded(); lines.push(`${"#".repeat(level)} ${text}${tail}`); }
136  };
137  for (const cell of cells) {
138    const headings = nb.headings(cell);
139    const cellNotes = notes(cell);
140    if (headings.length && nb.isMarkdown(cell)) {
141      headings.forEach((h, i) => heading(h.level, h.text, i === 0 ? `  [${cell.id}]${cellNotes.length ? ` ${cellNotes.join("; ")}` : ""}` : ""));
142      if (cellNotes.length && !headings.some((h) => h.level <= o.maxLevel)) { flushFolded(); lines.push(`  ${cell.id} (markdown) ${cellNotes.join("; ")}`); }
143      continue;
144    }
145    if (nb.isMarkdown(cell)) {
146      if (cellNotes.length) { flushFolded(); lines.push(`  ${cell.id} (markdown) ${cellNotes.join("; ")}`); }
147      else if (fold) folded++;
148      else lines.push(`  ${cell.id} (markdown)`);
149      continue;
150    }
151    // Folded or not, the outline keeps every heading (down to maxLevel).
152    for (const h of headings) heading(h.level, h.text, `  [${cell.id}]`);
153    if (fold && !cellNotes.length) { folded++; continue; }
154    flushFolded();
155    const defs = nb.defs.get(cell.id) ?? [];
156    const label = cell.name !== "_" ? `${cell.id} "${cell.name}"` : cell.id;
157    const what = defs.length ? `defines ${truncate(defs.join(", "), 80)}` : firstLine(cell.code) || "(empty)";
158    lines.push(`  ${label}: ${what}${cellNotes.length ? ` -- ${cellNotes.join("; ")}` : ""}`);
159  }
160  flushFolded();
161
162  const run = nb.running();
163  const summary: string[] = [];
164  if (run) summary.push(`running ${run.cell.id}${run.section.length ? ` under "${run.section.join(" › ")}"` : ""}`);
165  const queued = nb.queued();
166  if (queued) summary.push(`${queued} queued`);
167  const errors = nb.errors().length;
168  if (errors) summary.push(`${errors} cell${errors === 1 ? "" : "s"} with errors`);
169  if (staleCommon) summary.push(`${staleCount} of ${cells.length} cells stale (inputs changed, not rerun)`);
170
171  return [
172    `=== ${attachment.path} (${current ? "current" : "also followed"}; ${attachment.url}, session ${attachment.sessionId}) ===`,
173    `Kernel: ${summary.length ? summary.join(", ") : "idle"}.`,
174    ...lines,
175  ];
176}
177
178/**
179 * The followed notebooks in one block, the current one (used most recently)
180 * first and in full; the others as a short outline (headings down to ##) with
181 * the cells that need attention.
182 */
183export function snapshot(entries: SnapshotEntry[], options: SnapshotOptions = {}): string {
184  const now = options.now ?? Date.now();
185  const sections = entries.map((e, i) => {
186    const current = e.current ?? i === 0;
187    return section(e.notebook, e.attachment, current, {
188      seen: e.seenSeq ?? Number.POSITIVE_INFINITY,
189      now,
190      maxCells: current ? options.maxCells ?? 60 : 0,
191      maxLevel: current ? 6 : 2,
192    });
193  });
194  const many = entries.length > 1;
195  return [
196    `<${STATE_TAG} notebooks="${entries.length}">`,
197    `Live state of the marimo notebook${many ? "s" : ""} open in the user's browser, added automatically by pi-marimo (not`,
198    ...(options.refresh === "prompt"
199      ? ["written by the user). Taken when the user sent their latest message and not updated during the turn;",
200        "earlier copies are removed. Check current values through marimo-pair before relying on them."]
201      : ["written by the user). It is replaced with a fresh copy on every request, so trust this over older reads."]),
202    ...(many ? ["The current notebook (the one used most recently) comes first and in full; the others are outlines",
203      "with only the cells that need attention."] : []),
204    "Cells are listed in notebook order under their markdown headings, by cell id. Inspect or change cells",
205    "through the marimo-pair skill, not by editing the .py file.",
206    ...sections.flatMap((lines) => ["", ...lines]),
207    `</${STATE_TAG}>`,
208  ].join("\n");
209}
210

1// Which notebooks an agent's tool call works in, read from its marimo-pair
2// commands, so each agent session can keep the notebook it works in current
3// even while another session (or the user) works in a different one.
4//
5// marimo-pair runs `execute-code.sh --url URL [--file PATH | --session ID]` or
6// `marimo pair execute --url URL [--file FILE] [--session ID]`. A tool call may
7// hold several (one bash command, a codemode script), in any quoting.
8
9export interface PairTarget {
10  url?: string;
11  file?: string;
12  session?: string;
13}
14
15const CALL = /execute-code\.sh|marimo\s+pair\s+execute/g;
16
17function flag(segment: string, name: string): string | undefined {
18  const match = new RegExp(`--${name}(?:=|\\s+)(?:"([^"]+)"|'([^']+)'|([^\\s"'\\\\;|&)]+))`).exec(segment);
19  const value = match?.[1] ?? match?.[2] ?? match?.[3];
20  // A shell variable or substitution cannot be resolved here.
21  return value && !value.includes("$") ? value : undefined;
22}
23
24/** The notebooks the marimo-pair calls in a tool call's text target, in order. */
25export function pairTargets(raw: string): PairTarget[] {
26  // Tool input arrives as JSON too, with its quotes escaped.
27  const text = raw.replace(/\\"/g, '"');
28  const starts: number[] = [];
29  CALL.lastIndex = 0;
30  for (let m = CALL.exec(text); m; m = CALL.exec(text)) starts.push(m.index);
31  return starts.map((start, i) => {
32    const segment = text.slice(start, starts[i + 1] ?? text.length).split(/\n|\\n/)[0] ?? "";
33    const target: PairTarget = {};
34    const url = flag(segment, "url");
35    const file = flag(segment, "file");
36    const session = flag(segment, "session");
37    if (url) target.url = url;
38    if (file) target.file = file;
39    if (session) target.session = session;
40    return target;
41  });
42}
43
44/** One server, however it is written (localhost or 127.0.0.1, a trailing slash). */
45export function sameServer(a: string, b: string): boolean {
46  const norm = (u: string) => u.replace(/\/+$/, "").replace("://localhost", "://127.0.0.1").replace("://[::1]", "://127.0.0.1");
47  return norm(a) === norm(b);
48}
49

1// Keeps the live state of the notebooks you work in current. It finds them,
2// subscribes to each one's session as a kiosk consumer, follows them across
3// page reloads (which change the session id) and server restarts, and reports
4// every change. The followed notebooks are the pinned ones, or in auto mode
5// every one open under the cwd. The current one is the one this session's agent
6// worked in last (through marimo-pair), so two sessions in two notebooks each
7// keep their own; before the agent touches any, the one used most recently by
8// anyone (a running cell, then the latest cell run or edit).
9
10import { type Cancellable, type Io, normalize, within } from "./io.js";
11import { NotebookState } from "./notebook.js";
12import { findNotebooks, type NotebookSession } from "./registry.js";
13import type { Attachment } from "./render.js";
14import { streamSession } from "./sse.js";
15import { type PairTarget, sameServer } from "./touch.js";
16
17export type Mode =
18  | { kind: "auto" }
19  | { kind: "pinned"; paths: string[] }
20  | { kind: "off" };
21
22export type Connection = "idle" | "searching" | "connecting" | "connected" | "disconnected";
23
24export interface WatcherOptions {
25  io: Io;
26  cwd: string;
27  token?: string;
28  extraUrls?: string[];
29  pollMs?: number;
30  /** Most notebooks followed at once in auto mode. */
31  maxNotebooks?: number;
32  onChange: () => void;
33}
34
35type Target = NotebookSession & { real: string; shared: boolean };
36
37const consumerId = (): string => `pi-marimo-${Math.random().toString(36).slice(2, 10)}`;
38
39/** A followed notebook, as the hosts see it. */
40export interface Followed {
41  attachment: Attachment;
42  notebook: NotebookState;
43  connection: Connection;
44  current: boolean;
45  /** When this session's agent last worked in it (0: never). */
46  touchedAt: number;
47}
48
49/** Read a saved or configured mode, accepting the single-notebook form `{ kind: "pinned", path }`. */
50export function parseMode(raw: unknown): Mode | undefined {
51  const mode = raw as { kind?: unknown; path?: unknown; paths?: unknown } | undefined;
52  if (mode?.kind === "auto" || mode?.kind === "off") return { kind: mode.kind };
53  if (mode?.kind !== "pinned") return undefined;
54  const paths = Array.isArray(mode.paths) ? mode.paths.filter((p): p is string => typeof p === "string") : typeof mode.path === "string" ? [mode.path] : [];
55  return paths.length ? { kind: "pinned", paths } : undefined;
56}
57
58/** One notebook's subscription, reconnecting until stopped. */
59class Follower {
60  readonly notebook = new NotebookState();
61  connection: Connection = "connecting";
62  /** Whether browser edits reach us (false when connected by the browser's own session id). */
63  seesBrowserEdits = false;
64  /** When this session's agent last worked in the notebook (0: never). */
65  touchedAt = 0;
66  private stopped = false;
67  private current?: Cancellable<unknown>;
68  private sleep?: Cancellable<void>;
69  private readonly id = consumerId();
70  readonly loop: Promise<void>;
71
72  constructor(public target: Target, private readonly options: WatcherOptions, private readonly changed: () => void) {
73    this.loop = this.run();
74  }
75
76  get attachment(): Attachment {
77    return { url: this.target.url, sessionId: this.target.sessionId, path: this.target.path };
78  }
79
80  /** Point at a new session for the same file (a page reload) and reconnect. */
81  retarget(target: Target): void {
82    if (target.sessionId === this.target.sessionId && target.shared === this.target.shared) return;
83    this.target = target;
84    this.current?.cancel();
85    this.sleep?.cancel();
86  }
87
88  stop(): void {
89    this.stopped = true;
90    this.current?.cancel();
91    this.sleep?.cancel();
92  }
93
94  private set(connection: Connection): void {
95    if (connection === this.connection) return;
96    this.connection = connection;
97    this.changed();
98  }
99
100  private async run(): Promise<void> {
101    const pollMs = this.options.pollMs ?? 5000;
102    while (!this.stopped) {
103      const target = this.target;
104      this.seesBrowserEdits = !target.shared;
105      this.notebook.reset();
106      this.set("connecting");
107      const stream = streamSession({
108        io: this.options.io,
109        url: target.url,
110        // marimo names a consumer after the session id it connected with and
111        // does not echo a browser's own edits back to that id. Connecting
112        // under our own id (marimo then finds the session by file) receives
113        // them. That lookup takes the oldest session for the file, so when
114        // several hold it, connect by the exact id and miss browser edits.
115        sessionId: target.shared ? target.sessionId : this.id,
116        file: target.path,
117        token: this.options.token,
118        onMessage: (op, data) => {
119          const changed = this.notebook.apply(op, data);
120          if (op === "kernel-ready") this.set("connected");
121          if (changed) this.changed();
122        },
123      });
124      this.current = stream;
125      const result = await stream.done.catch(() => ({ kind: "http-error" as const }));
126      if (this.stopped) break;
127      this.set("disconnected");
128      if (target !== this.target) continue;
129      // A dropped stream reconnects quickly; a refused one waits for the next poll.
130      this.sleep = this.options.io.sleep(result.kind === "ended" ? 1000 : pollMs);
131      await this.sleep.done;
132    }
133  }
134}
135
136export class MarimoWatcher {
137  mode: Mode = { kind: "auto" };
138  private followers = new Map<string, Follower>();
139  private stopped = true;
140  private scanned = false;
141  private wake?: () => void;
142  private loop?: Promise<void>;
143  private realCwd?: string;
144  private readonly realPaths = new Map<string, string>();
145  private readonly empty = new NotebookState();
146
147  constructor(private readonly options: WatcherOptions) {}
148
149  /** Followers, the current one first, then by most recent use. */
150  private ranked(): Follower[] {
151    const rank = (x: Follower): number[] => [x.connection === "connected" ? 1 : 0, x.touchedAt, x.notebook.running() ? 1 : 0, x.notebook.lastActivity];
152    const cmp = (a: Follower, b: Follower): number => {
153      const [x, y] = [rank(a), rank(b)];
154      for (let i = 0; i < x.length; i++) if (x[i] !== y[i]) return y[i]! - x[i]!;
155      return 0;
156    };
157    return [...this.followers.values()].sort(cmp);
158  }
159
160  /** The current notebook: the one used most recently. */
161  private get active(): Follower | undefined {
162    return this.ranked()[0];
163  }
164
165  /** Every followed notebook, the current one first, then by most recent use. */
166  followed(): Followed[] {
167    return this.ranked().map((f, i) => ({ attachment: f.attachment, notebook: f.notebook, connection: f.connection, current: i === 0, touchedAt: f.touchedAt }));
168  }
169
170  get notebook(): NotebookState {
171    return this.active?.notebook ?? this.empty;
172  }
173
174  get attachment(): Attachment | undefined {
175    const active = this.active;
176    if (active) return active.attachment;
177    return this.mode.kind === "pinned" ? { url: "", sessionId: "", path: this.mode.paths[0]! } : undefined;
178  }
179
180  get connection(): Connection {
181    const active = this.active;
182    if (active) return active.connection;
183    return this.mode.kind === "off" ? "idle" : "searching";
184  }
185
186  get seesBrowserEdits(): boolean {
187    return this.active?.seesBrowserEdits ?? false;
188  }
189
190  /**
191   * The one followed notebook a marimo-pair call targets, if it is unambiguous: by session id, by
192   * file (as marimo-pair takes it: absolute, or relative to the server's directory), or by server
193   * alone when only one followed notebook is on it.
194   */
195  private match(t: PairTarget) {
196    let matches = [...this.followers.values()];
197    if (t.url) matches = matches.filter((f) => sameServer(f.target.url, t.url!));
198    if (t.session) matches = matches.filter((f) => f.target.sessionId === t.session);
199    else if (t.file) {
200      const file = t.file.replace(/^\.\//, "");
201      matches = matches.filter((f) => f.target.path === file || f.target.real === file || f.target.path.endsWith(`/${file}`) || f.target.real.endsWith(`/${file}`));
202    } else if (!t.url) matches = [];
203    return matches.length === 1 ? matches[0] : undefined;
204  }
205
206  /** The paths of the followed notebooks these marimo-pair calls target. */
207  targeted(targets: PairTarget[]): string[] {
208    return targets.map((t) => this.match(t)?.attachment.path).filter((p): p is string => p !== undefined);
209  }
210
211  /**
212   * Note that this session's agent works in these notebooks (from its marimo-pair calls), making
213   * the last one current for this session. Returns whether any followed notebook matched.
214   */
215  touch(targets: PairTarget[], now = Date.now()): boolean {
216    let touched = false;
217    targets.forEach((t, i) => {
218      const match = this.match(t);
219      if (!match) return;
220      // Later calls in one tool call win.
221      match.touchedAt = now + i;
222      touched = true;
223    });
224    if (touched) this.options.onChange();
225    return touched;
226  }
227
228  /** The other notebooks followed besides the current one, most recently used first. */
229  others(): Attachment[] {
230    return this.ranked().slice(1).map((f) => f.attachment);
231  }
232
233  start(): void {
234    if (!this.stopped) return;
235    this.stopped = false;
236    this.loop = this.run();
237  }
238
239  /** Stop, waiting at most a second for streams to close (a host may only end one at its next chunk). */
240  async stop(): Promise<void> {
241    this.stopped = true;
242    this.wake?.();
243    const loops = [...this.followers.values()].map((f) => { f.stop(); return f.loop; });
244    this.followers.clear();
245    await Promise.race([Promise.all([this.loop, ...loops]).catch(() => undefined), this.options.io.sleep(1000).done]);
246  }
247
248  setMode(mode: Mode): void {
249    this.mode = mode;
250    this.scanned = false;
251    this.wake?.();
252  }
253
254  /** Wait (up to ms) for the first scan and any connection in progress, so a first request has the state. */
255  async settle(ms: number): Promise<void> {
256    const busy = () => !this.scanned || [...this.followers.values()].some((f) => f.connection === "connecting");
257    for (let waited = 0; !this.stopped && busy() && waited < ms; waited += 50) {
258      await this.options.io.sleep(50).done;
259    }
260  }
261
262  /** Rescan now instead of at the next poll. */
263  refresh(): void {
264    this.wake?.();
265  }
266
267  async list(): Promise<NotebookSession[]> {
268    return findNotebooks(this.options.io, { token: this.options.token, extraUrls: this.options.extraUrls });
269  }
270
271  private async real(path: string): Promise<string> {
272    const absolute = normalize(path, this.options.cwd);
273    let real = this.realPaths.get(absolute);
274    if (real === undefined) {
275      real = normalize(await this.options.io.realpath(absolute));
276      this.realPaths.set(absolute, real);
277    }
278    return real;
279  }
280
281  /** The notebooks to follow now, keyed by server and file. */
282  private async targets(): Promise<Map<string, Target>> {
283    const mode = this.mode;
284    const out = new Map<string, Target>();
285    if (mode.kind === "off") return out;
286    const sessions = await this.list();
287    const real = new Map<NotebookSession, string>();
288    for (const s of sessions) real.set(s, await this.real(s.path));
289    // Several sessions can hold one file (a closed tab's session lingers);
290    // marimo lists them oldest first, so the newest wins.
291    const add = (s: NotebookSession): void => {
292      const path = real.get(s)!;
293      const shared = sessions.filter((o) => o.url === s.url && real.get(o) === path).length > 1;
294      out.set(`${s.url} ${path}`, { ...s, real: path, shared });
295    };
296    if (mode.kind === "pinned") {
297      for (const pinned of mode.paths) {
298        const path = await this.real(pinned);
299        const chosen = sessions.filter((s) => real.get(s) === path).pop();
300        if (chosen) add(chosen);
301      }
302      return out;
303    }
304    this.realCwd ??= await this.real(this.options.cwd);
305    const cwd = this.realCwd;
306    // Under the cwd, but not in a hidden directory there (.worktrees, .claude/worktrees: other checkouts).
307    const local = sessions.filter((s) => {
308      const path = real.get(s)!;
309      return within(path, cwd) && !path.slice(cwd.length).split("/").slice(0, -1).some((part) => part.startsWith("."));
310    });
311    for (const s of local) add(s);
312    const max = this.options.maxNotebooks ?? 8;
313    return new Map([...out].slice(0, max));
314  }
315
316  private reconcile(targets: Map<string, Target>): void {
317    let changed = false;
318    for (const [key, follower] of this.followers) {
319      if (targets.has(key)) continue;
320      follower.stop();
321      this.followers.delete(key);
322      changed = true;
323    }
324    for (const [key, target] of targets) {
325      const follower = this.followers.get(key);
326      if (follower) follower.retarget(target);
327      else {
328        this.followers.set(key, new Follower(target, this.options, () => this.options.onChange()));
329        changed = true;
330      }
331    }
332    if (changed) this.options.onChange();
333  }
334
335  private async run(): Promise<void> {
336    const pollMs = this.options.pollMs ?? 5000;
337    while (!this.stopped) {
338      const targets = await this.targets().catch(() => undefined);
339      if (this.stopped) break;
340      if (targets) this.reconcile(targets);
341      if (!this.scanned) { this.scanned = true; this.options.onChange(); }
342      const sleep = this.options.io.sleep(this.mode.kind === "off" ? 60_000 : pollMs);
343      this.wake = sleep.cancel;
344      await sleep.done;
345    }
346  }
347}
348

1// A notebook's live state, rebuilt from the messages marimo streams to its
2// frontends: the cell document (kernel-ready, notebook-document-transaction),
3// each cell's run state and errors (cell-op), and the dataflow graph
4// (variables). Only what an agent needs is kept; outputs are dropped.
5
6import { type Heading, headingsFromCode, isMarkdownCell, sameCode } from "./markdown.js";
7
8export type EditSource = "frontend" | "code-mode" | "kernel" | string;
9
10export interface Cell {
11  id: string;
12  code: string;
13  name: string;
14  disabled: boolean;
15  /** "idle" | "queued" | "running" | "disabled-transitively" */
16  status: string;
17  /** When the current run started (ms since epoch), while running. */
18  runningSince?: number;
19  /** Inputs changed and the cell has not rerun (lazy runtime). */
20  stale: boolean;
21  /** One-line description of the error the cell's last run raised. */
22  error?: string;
23  /** The cell did not run because an ancestor failed. */
24  blocked?: string;
25  /** Code the kernel last ran for this cell, when known. */
26  lastRunCode?: string;
27  /** Who last changed the code, and at which change number. */
28  editedBy?: EditSource;
29  editSeq: number;
30}
31
32export interface RunningCell {
33  cell: Cell;
34  section: string[];
35  since?: number;
36}
37
38const record = (value: unknown): Record<string, unknown> =>
39  value !== null && typeof value === "object" && !Array.isArray(value) ? value as Record<string, unknown> : {};
40const strings = (value: unknown): string[] => Array.isArray(value) ? value.filter((v): v is string => typeof v === "string") : [];
41const isInternal = (id: string): boolean => id.startsWith("__");
42
43export function describeError(output: Record<string, unknown>): { error?: string; blocked?: string } {
44  const errors = Array.isArray(output.data) ? output.data.map(record) : [];
45  const parts: string[] = [];
46  let blocked: string | undefined;
47  for (const err of errors) {
48    const type = typeof err.type === "string" ? err.type : "error";
49    const msg = typeof err.msg === "string" ? err.msg.trim().split("\n")[0] : "";
50    if (type === "ancestor-prevented" || type === "ancestor-stopped") {
51      blocked = msg || "an ancestor cell failed";
52      continue;
53    }
54    if (type === "exception") {
55      const name = typeof err.exception_type === "string" ? err.exception_type : "Exception";
56      parts.push(msg ? `${name}: ${msg}` : name);
57    } else if (type === "multiple-defs") {
58      const name = typeof err.name === "string" ? err.name : "?";
59      parts.push(`multiple definitions of ${name}`);
60    } else if (type === "interruption") {
61      parts.push("interrupted");
62    } else {
63      parts.push(msg ? `${type}: ${msg}` : type);
64    }
65  }
66  const error = parts.length ? truncate(parts.join("; "), 200) : undefined;
67  return { error, blocked: error ? undefined : blocked };
68}
69
70export function truncate(text: string, max: number): string {
71  return text.length <= max ? text : `${text.slice(0, max - 1)}…`;
72}
73
74export class NotebookState {
75  order: string[] = [];
76  cells = new Map<string, Cell>();
77  defs = new Map<string, string[]>();
78  ready = false;
79  /** Increments on every document change. */
80  seq = 0;
81  /**
82   * When the notebook was last used (ms since epoch): the latest cell run the
83   * kernel reports (replayed on connect, so a notebook's history counts), or an
84   * edit seen live.
85   */
86  lastActivity = 0;
87  private readyAt = 0;
88  private headingCache = new Map<string, Heading[]>();
89
90  reset(): void {
91    this.order = [];
92    this.cells.clear();
93    this.defs.clear();
94    this.ready = false;
95    this.lastActivity = 0;
96    this.readyAt = 0;
97    this.headingCache.clear();
98  }
99
100  private ensure(id: string): Cell {
101    let cell = this.cells.get(id);
102    if (!cell) {
103      cell = { id, code: "", name: "_", disabled: false, status: "idle", stale: false, editSeq: 0 };
104      this.cells.set(id, cell);
105    }
106    return cell;
107  }
108
109  headings(cell: Cell): Heading[] {
110    let headings = this.headingCache.get(cell.code);
111    if (!headings) {
112      headings = headingsFromCode(cell.code);
113      if (this.headingCache.size > 2000) this.headingCache.clear();
114      this.headingCache.set(cell.code, headings);
115    }
116    return headings;
117  }
118
119  isMarkdown(cell: Cell): boolean {
120    return isMarkdownCell(cell.code);
121  }
122
123  /** Apply one streamed message. Returns true when agent-visible state changed. */
124  apply(op: string, raw: unknown): boolean {
125    const data = record(raw);
126    switch (op) {
127      case "kernel-ready":
128        return this.kernelReady(data);
129      case "notebook-document-transaction":
130        return this.transaction(record(data.transaction));
131      case "cell-op":
132        return this.cellOp(data);
133      case "variables":
134        return this.variables(data);
135      default:
136        return false;
137    }
138  }
139
140  private kernelReady(data: Record<string, unknown>): boolean {
141    this.reset();
142    const ids = strings(data.cell_ids);
143    const codes = strings(data.codes);
144    const names = strings(data.names);
145    const configs = Array.isArray(data.configs) ? data.configs.map(record) : [];
146    const lastRun = record(data.last_executed_code);
147    ids.forEach((id, i) => {
148      const cell = this.ensure(id);
149      cell.code = codes[i] ?? "";
150      cell.name = names[i] || "_";
151      cell.disabled = configs[i]?.disabled === true;
152      if (typeof lastRun[id] === "string") cell.lastRunCode = lastRun[id] as string;
153    });
154    this.order = ids;
155    this.ready = true;
156    this.readyAt = Date.now();
157    return true;
158  }
159
160  private place(id: string, before: unknown, after: unknown): void {
161    this.order = this.order.filter((other) => other !== id);
162    if (typeof before === "string" && this.order.includes(before)) {
163      this.order.splice(this.order.indexOf(before), 0, id);
164    } else if (typeof after === "string" && this.order.includes(after)) {
165      this.order.splice(this.order.indexOf(after) + 1, 0, id);
166    } else {
167      this.order.push(id);
168    }
169  }
170
171  private transaction(tx: Record<string, unknown>): boolean {
172    const source = typeof tx.source === "string" ? tx.source : "unknown";
173    const changes = Array.isArray(tx.changes) ? tx.changes.map(record) : [];
174    let changed = false;
175    for (const change of changes) {
176      const id = typeof change.cellId === "string" ? change.cellId : "";
177      switch (change.type) {
178        case "create-cell": {
179          const cell = this.ensure(id);
180          cell.code = typeof change.code === "string" ? change.code : "";
181          cell.name = typeof change.name === "string" && change.name ? change.name : "_";
182          cell.disabled = record(change.config).disabled === true;
183          cell.editedBy = source;
184          cell.editSeq = ++this.seq;
185          this.place(id, change.before, change.after);
186          changed = true;
187          break;
188        }
189        case "delete-cell":
190          this.cells.delete(id);
191          this.order = this.order.filter((other) => other !== id);
192          this.defs.delete(id);
193          this.seq++;
194          changed = true;
195          break;
196        case "move-cell":
197          this.place(id, change.before, change.after);
198          this.seq++;
199          changed = true;
200          break;
201        case "reorder-cells": {
202          const ids = strings(change.cellIds).filter((cid) => this.cells.has(cid));
203          const rest = this.order.filter((cid) => !ids.includes(cid));
204          const next = [...ids, ...rest];
205          if (next.join() !== this.order.join()) { this.order = next; this.seq++; changed = true; }
206          break;
207        }
208        case "set-code": {
209          const cell = this.ensure(id);
210          const code = typeof change.code === "string" ? change.code : "";
211          // The browser respelling a markdown cell is not an edit.
212          if (cell.code !== code && sameCode(cell.code, code)) cell.code = code;
213          else if (cell.code !== code) {
214            cell.code = code;
215            cell.editedBy = source;
216            cell.editSeq = ++this.seq;
217            changed = true;
218          }
219          break;
220        }
221        case "set-name": {
222          const cell = this.ensure(id);
223          cell.name = typeof change.name === "string" && change.name ? change.name : "_";
224          changed = true;
225          break;
226        }
227        case "set-config": {
228          const cell = this.ensure(id);
229          const disabled = change.disabled === true;
230          if (cell.disabled !== disabled) { cell.disabled = disabled; changed = true; }
231          break;
232        }
233      }
234    }
235    // An edit counts as use when it arrives live, not in the replay on connect.
236    if (source !== "kernel" && changed && this.readyAt && Date.now() - this.readyAt > 1500) this.lastActivity = Date.now();
237    return changed;
238  }
239
240  private cellOp(data: Record<string, unknown>): boolean {
241    const id = typeof data.cell_id === "string" ? data.cell_id : "";
242    if (!id || isInternal(id) || !this.cells.has(id)) return false;
243    const cell = this.ensure(id);
244    let changed = false;
245    if (typeof data.status === "string" && typeof data.timestamp === "number") {
246      this.lastActivity = Math.max(this.lastActivity, data.timestamp * 1000);
247    }
248    if (typeof data.status === "string" && data.status !== cell.status) {
249      cell.status = data.status;
250      if (data.status === "queued") {
251        // The previous run's error is about to be replaced.
252        cell.error = undefined;
253        cell.blocked = undefined;
254      }
255      if (data.status === "running") {
256        const at = typeof data.timestamp === "number" ? data.timestamp * 1000 : Date.now();
257        cell.runningSince = at;
258        cell.lastRunCode = cell.code;
259        cell.error = undefined;
260        cell.blocked = undefined;
261        cell.stale = false;
262      } else {
263        cell.runningSince = undefined;
264      }
265      changed = true;
266    }
267    if (typeof data.stale_inputs === "boolean" && data.stale_inputs !== cell.stale) {
268      cell.stale = data.stale_inputs;
269      changed = true;
270    }
271    if (data.output !== null && typeof data.output === "object") {
272      const output = record(data.output);
273      const { error, blocked } = output.channel === "marimo-error" ? describeError(output) : {};
274      if (error !== cell.error || blocked !== cell.blocked) {
275        cell.error = error;
276        cell.blocked = blocked;
277        changed = true;
278      }
279    }
280    return changed;
281  }
282
283  private variables(data: Record<string, unknown>): boolean {
284    const vars = Array.isArray(data.variables) ? data.variables.map(record) : [];
285    const defs = new Map<string, string[]>();
286    for (const v of vars) {
287      if (typeof v.name !== "string") continue;
288      for (const id of strings(v.declared_by)) {
289        if (!defs.has(id)) defs.set(id, []);
290        defs.get(id)!.push(v.name);
291      }
292    }
293    for (const names of defs.values()) names.sort();
294    const before = JSON.stringify([...this.defs.entries()].sort());
295    this.defs = defs;
296    return before !== JSON.stringify([...defs.entries()].sort());
297  }
298
299  /** Cells in notebook order. */
300  ordered(): Cell[] {
301    return this.order.map((id) => this.cells.get(id)).filter((cell): cell is Cell => cell !== undefined);
302  }
303
304  /** The heading path ("Section", "Subsection") a cell sits under, including its own headings. */
305  sectionPath(cellId: string): string[] {
306    const stack: Heading[] = [];
307    for (const cell of this.ordered()) {
308      for (const heading of this.headings(cell)) {
309        while (stack.length && stack[stack.length - 1]!.level >= heading.level) stack.pop();
310        stack.push(heading);
311      }
312      if (cell.id === cellId) break;
313    }
314    return stack.map((h) => h.text);
315  }
316
317  running(): RunningCell | undefined {
318    const cell = this.ordered().find((c) => c.status === "running");
319    return cell ? { cell, section: this.sectionPath(cell.id), since: cell.runningSince } : undefined;
320  }
321
322  queued(): number {
323    return this.ordered().filter((c) => c.status === "queued").length;
324  }
325
326  errors(): Cell[] {
327    return this.ordered().filter((c) => c.error);
328  }
329
330  /** Code changed since the kernel last ran it (only known once a run was seen). */
331  edited(cell: Cell): boolean {
332    return cell.lastRunCode !== undefined && !sameCode(cell.lastRunCode, cell.code);
333  }
334}
335

1// Find running marimo servers and the notebooks open on them.
2//
3// Servers started with --no-token register a JSON file under
4// $XDG_STATE_HOME/marimo/servers (marimo/_server/server_registry.py). A
5// token-protected server is reached with MARIMO_TOKEN.
6
7import type { Io } from "./io.js";
8
9export interface NotebookSession {
10  url: string;
11  sessionId: string;
12  path: string;
13}
14
15export function authHeaders(token?: string): Record<string, string> {
16  return token ? { Authorization: `Bearer ${token}` } : {};
17}
18
19function hostFor(host: unknown): string {
20  if (typeof host !== "string" || host === "" || host === "0.0.0.0") return "127.0.0.1";
21  if (host === "::") return "[::1]";
22  return host.includes(":") && !host.startsWith("[") ? `[${host}]` : host;
23}
24
25/** Base URLs of the registered servers. */
26export async function serverUrls(io: Io): Promise<string[]> {
27  const urls: string[] = [];
28  for (const raw of await io.registryEntries()) {
29    const entry = (raw ?? {}) as Record<string, unknown>;
30    if (typeof entry.port !== "number") continue;
31    const base = typeof entry.base_url === "string" ? entry.base_url.replace(/\/$/, "") : "";
32    urls.push(`http://${hostFor(entry.host)}:${entry.port}${base}`);
33  }
34  return urls;
35}
36
37export async function listSessions(io: Io, url: string, token?: string): Promise<NotebookSession[]> {
38  const body = await io.getJson(`${url}/api/sessions`, authHeaders(token), 2000);
39  if (body === null || typeof body !== "object") return [];
40  return Object.entries(body as Record<string, Record<string, unknown> | null>).flatMap(([sessionId, info]) => {
41    const path = typeof info?.path === "string" ? info.path : typeof info?.filename === "string" ? info.filename : "";
42    return path ? [{ url, sessionId, path }] : [];
43  });
44}
45
46/** Every notebook session open on every registered server, plus any extra server URLs. */
47export async function findNotebooks(io: Io, options: { token?: string; extraUrls?: string[] } = {}): Promise<NotebookSession[]> {
48  const urls = new Set(await serverUrls(io));
49  for (const url of options.extraUrls ?? []) urls.add(url.replace(/\/$/, ""));
50  const lists = await Promise.all([...urls].map((url) => listSessions(io, url, options.token)));
51  return lists.flat();
52}
53

1// A read-only (kiosk) subscription to a marimo session's message stream.
2//
3// marimo serves the same kernel messages over /ws and /sse. A kiosk consumer
4// is never the session's main consumer: it cannot edit or run code and does
5// not take the notebook over from the browser. Connecting replays the
6// session's current state (kernel-ready, then every cell's last messages), so
7// a fresh NotebookState is complete after the replay.
8
9import type { Cancellable, Io } from "./io.js";
10import { authHeaders } from "./registry.js";
11
12export interface SseEvent {
13  event?: string;
14  data: string;
15}
16
17/** Incremental parser for a text/event-stream body. */
18export class SseParser {
19  private buffer = "";
20
21  push(chunk: string): SseEvent[] {
22    this.buffer += chunk.replace(/\r\n?/g, "\n");
23    const events: SseEvent[] = [];
24    let end = this.buffer.indexOf("\n\n");
25    while (end >= 0) {
26      const block = this.buffer.slice(0, end);
27      this.buffer = this.buffer.slice(end + 2);
28      let event: string | undefined;
29      const data: string[] = [];
30      for (const line of block.split("\n")) {
31        if (!line || line.startsWith(":")) continue;
32        const colon = line.indexOf(":");
33        const field = colon < 0 ? line : line.slice(0, colon);
34        const value = colon < 0 ? "" : line.slice(colon + 1).replace(/^ /, "");
35        if (field === "event") event = value;
36        else if (field === "data") data.push(value);
37      }
38      if (data.length) events.push({ event, data: data.join("\n") });
39      end = this.buffer.indexOf("\n\n");
40    }
41    return events;
42  }
43}
44
45export interface StreamResult {
46  /** "closed": the server ended the stream with a close event. "ended": the stream dropped. */
47  kind: "closed" | "ended" | "http-error";
48  code?: number;
49  reason?: string;
50}
51
52export interface StreamOptions {
53  io: Io;
54  url: string;
55  sessionId: string;
56  file?: string;
57  token?: string;
58  onMessage: (op: string, data: unknown) => void;
59}
60
61export function streamSession(options: StreamOptions): Cancellable<StreamResult> {
62  const query: Record<string, string> = { session_id: options.sessionId, kiosk: "true" };
63  if (options.file) query.file = options.file;
64  if (options.token) query.access_token = options.token;
65  const params = Object.entries(query).map(([k, v]) => `${k}=${encodeURIComponent(v)}`).join("&");
66  const parser = new SseParser();
67  let closed: StreamResult | undefined;
68  const stream = options.io.stream(
69    `${options.url}/sse?${params}`,
70    { Accept: "text/event-stream", ...authHeaders(options.token) },
71    (text) => {
72      if (closed) return;
73      for (const event of parser.push(text)) {
74        if (event.event === "close") {
75          try {
76            const body = JSON.parse(event.data) as { code?: number; reason?: string };
77            closed = { kind: "closed", code: body.code, reason: body.reason };
78          } catch {
79            closed = { kind: "closed" };
80          }
81          stream.cancel();
82          return;
83        }
84        if (event.event && event.event !== "message") continue;
85        let message: { op?: unknown; data?: unknown };
86        try {
87          message = JSON.parse(event.data);
88        } catch {
89          continue;
90        }
91        if (typeof message.op === "string") options.onMessage(message.op, message.data);
92      }
93    },
94  );
95  const done = stream.done.then((end): StreamResult => closed ?? (end.ok ? { kind: "ended" } : { kind: "http-error", code: end.status }));
96  return { done, cancel: stream.cancel };
97}
98

1// Headings a cell renders through mo.md(...), read from its source code.
2//
3// marimo has no "markdown cell" type: a markdown cell is a Python cell whose
4// body is `mo.md("""...""")`. Headings are taken from the string literals passed
5// to mo.md, never from Python comments, and fenced code blocks are skipped.
6
7export interface Heading {
8  level: number;
9  text: string;
10}
11
12const CALL = /\bmo\s*\.\s*md\s*\(\s*([rRfFuUbB]{0,2})("""|'''|"|')/g;
13
14/** The string literals passed to mo.md(...) in a cell's code. */
15export function mdLiterals(code: string): string[] {
16  const out: string[] = [];
17  CALL.lastIndex = 0;
18  for (let match = CALL.exec(code); match; match = CALL.exec(code)) {
19    const prefix = (match[1] ?? "").toLowerCase();
20    const quote = match[2] ?? '"';
21    const start = match.index + match[0].length;
22    const raw = prefix.includes("r");
23    let i = start;
24    let end = -1;
25    while (i < code.length) {
26      const ch = code[i];
27      if (ch === "\\") { i += 2; continue; }
28      if (quote.length === 1 && ch === "\n") break;
29      if (code.startsWith(quote, i)) { end = i; break; }
30      i++;
31    }
32    if (end < 0) continue;
33    const body = code.slice(start, end);
34    out.push(raw ? body : body.replace(/\\(["'\\])/g, "$1").replace(/\\n/g, "\n"));
35    CALL.lastIndex = end + quote.length;
36  }
37  return out;
38}
39
40function dedent(text: string): string {
41  const lines = text.split("\n");
42  let min = Number.POSITIVE_INFINITY;
43  for (const line of lines.slice(1)) {
44    if (!line.trim()) continue;
45    min = Math.min(min, line.length - line.trimStart().length);
46  }
47  if (!Number.isFinite(min) || min === 0) return text;
48  return [lines[0], ...lines.slice(1).map((line) => line.slice(Math.min(min, line.length - line.trimStart().length)))].join("\n");
49}
50
51/** Strip inline markdown so a heading reads as plain text. */
52export function plainHeading(text: string): string {
53  return text
54    .replace(/!?\[([^\]]*)\]\([^)]*\)/g, "$1")
55    .replace(/<[^>]+>/g, "")
56    .replace(/(\*\*|__|\*|_|`|~~)(.+?)\1/g, "$2")
57    .replace(/\s+#+\s*$/, "")
58    .replace(/\s+/g, " ")
59    .trim();
60}
61
62export function headingsFromMarkdown(markdown: string): Heading[] {
63  const headings: Heading[] = [];
64  let fence: string | null = null;
65  for (const line of dedent(markdown).split("\n")) {
66    const fenceMatch = /^\s{0,3}(`{3,}|~{3,})/.exec(line);
67    if (fenceMatch) {
68      const marker = (fenceMatch[1] ?? "`").charAt(0);
69      if (fence === null) fence = marker;
70      else if (marker === fence) fence = null;
71      continue;
72    }
73    if (fence !== null) continue;
74    const heading = /^\s{0,3}(#{1,6})\s+(.+)$/.exec(line);
75    if (!heading) continue;
76    const text = plainHeading(heading[2] ?? "");
77    if (text) headings.push({ level: (heading[1] ?? "#").length, text });
78  }
79  return headings;
80}
81
82export function headingsFromCode(code: string): Heading[] {
83  return mdLiterals(code).flatMap(headingsFromMarkdown);
84}
85
86/** True when the cell does nothing but render markdown. */
87export function isMarkdownCell(code: string): boolean {
88  const trimmed = code.trim();
89  if (!trimmed.startsWith("mo.md(")) return false;
90  const literals = mdLiterals(trimmed);
91  return literals.length === 1 && /\)\s*$/.test(trimmed);
92}
93
94/** The markdown a markdown cell renders, whatever quoting and indentation spell it; undefined for other cells. */
95export function markdownText(code: string): string | undefined {
96  if (!isMarkdownCell(code)) return undefined;
97  return dedent(mdLiterals(code.trim())[0] ?? "").split("\n").map((line) => line.trimEnd()).join("\n").trim();
98}
99
100/**
101 * The same code, or two spellings of the same markdown. marimo's editor rewrites a markdown cell
102 * into its own form (`mo.md("""` on its own line, dedented) as soon as a browser shows it, and
103 * posts that as a browser edit, though no one typed anything.
104 */
105export function sameCode(a: string, b: string): boolean {
106  if (a === b) return true;
107  const markdown = markdownText(a);
108  return markdown !== undefined && markdown === markdownText(b);
109}
110