Collapse hide_code cells on notebook open

extension/src/commands.ts

@@ -122,6 +122,8 @@ export type VscodeBuiltinCommand =
   | "workbench.action.quickOpen"
   // Toggle Outputs
   | "notebook.cell.toggleOutputs"
+  // Collapse a cell's input (code) editor
+  | "notebook.cell.collapseCellInput"
   // Fold Cell
   | "notebook.fold"
   // Unfold Cell

extension/src/features/HideCodeSync.ts

@@ -0,0 +1,103 @@
+import { Effect, HashSet, Layer, Option, Ref, Stream } from "effect";
+
+import { VsCode } from "../platform/VsCode.ts";
+import {
+  type MarimoNotebookCell,
+  MarimoNotebookDocument,
+  type NotebookId,
+} from "../schemas/MarimoNotebookDocument.ts";
+
+/**
+ * A contiguous, end-exclusive range of cells, addressed by index.
+ *
+ * This matches VS Code's `ICellRange`, the shape `notebook.cell.collapseCellInput`
+ * accepts to pick which cells to act on.
+ */
+export interface CellRange {
+  readonly start: number;
+  readonly end: number;
+}
+
+/**
+ * Returns one single-cell range per `hide_code` cell, in document order.
+ *
+ * It is a plain function rather than part of the layer below so the selection
+ * logic can be tested without standing up a notebook editor.
+ */
+export function hiddenInputCellRanges(
+  cells: readonly MarimoNotebookCell[],
+): CellRange[] {
+  return cells
+    .filter((cell) => cell.isCodeHidden)
+    .map((cell) => ({ start: cell.index, end: cell.index + 1 }));
+}
+
+/**
+ * Collapses the code input of `hide_code` cells when a marimo notebook opens.
+ *
+ * VS Code exposes the input-collapsed state as view-only: we can neither read
+ * it nor set it through cell metadata, only fire `notebook.cell.collapseCellInput`.
+ * So this is a one-way sync, not a binding. We collapse a notebook's `hide_code`
+ * cells the first time it becomes active and track it per session, so refocusing
+ * the tab does not re-hide a cell the user expanded to edit. See issue #326.
+ */
+export const HideCodeSyncLive = Layer.scopedDiscard(
+  Effect.gen(function* () {
+    const code = yield* VsCode;
+
+    // Notebooks already collapsed this session; refocusing a tab must not
+    // re-hide a cell the user has since expanded.
+    const collapsed = yield* Ref.make(HashSet.empty<NotebookId>());
+
+    yield* Effect.forkScoped(
+      // Prepend the currently-active editor: onDidChangeActiveNotebookEditor
+      // only emits future changes, so this covers a notebook open at startup.
+      Stream.concat(
+        Stream.fromEffect(code.window.getActiveNotebookEditor()),
+        code.window.activeNotebookEditorChanges(),
+      ).pipe(
+        Stream.runForEach(
+          Effect.fn("HideCodeSync.collapseHiddenCells")(function* (editor) {
+            const notebook = Option.filterMap(editor, (editor) =>
+              MarimoNotebookDocument.tryFrom(editor.notebook),
+            );
+            if (Option.isNone(notebook)) {
+              return;
+            }
+
+            const id = notebook.value.id;
+            if (HashSet.has(yield* Ref.get(collapsed), id)) {
+              return;
+            }
+
+            const ranges = hiddenInputCellRanges(notebook.value.getCells());
+            if (ranges.length === 0) {
+              return;
+            }
+
+            yield* Effect.annotateCurrentSpan({
+              notebook: id,
+              hiddenCells: ranges.length,
+            });
+
+            yield* code.commands
+              .executeCommand("notebook.cell.collapseCellInput", {
+                ranges,
+                document: notebook.value.uri,
+              })
+              .pipe(
+                // Mark only after a successful collapse, so a transient failure
+                // retries on the next activation.
+                Effect.andThen(Ref.update(collapsed, HashSet.add(id))),
+                Effect.catchAll((error) =>
+                  Effect.logWarning("Failed to collapse hidden cells").pipe(
+                    Effect.annotateLogs({ error }),
+                  ),
+                ),
+              );
+          }),
+        ),
+      ),
+    );
+  }),
+);

extension/src/features/Main.ts

@@ -53,6 +53,7 @@ import type { Telemetry } from "../telemetry/Telemetry.ts";
 import { CellMetadataBindingsLive } from "./CellMetadataBindings.ts";
 import { CellStatusBarProviderLive } from "./CellStatusBarProvider.ts";
 import { DebugLayerLive } from "./DebugLayer.ts";
+import { HideCodeSyncLive } from "./HideCodeSync.ts";
 import { MarimoCodeLensProviderLive } from "./MarimoCodeLensProvider.ts";
 import { MarimoFileDetectorLive } from "./MarimoFileDetector.ts";
 import { RegisterCommandsLive } from "./RegisterCommands.ts";
@@ -80,6 +81,7 @@ const MainLive = Layer.empty
     Layer.merge(CellMetadataBindingsLive),
     Layer.merge(ReloadOnConfigChangeLive),
     Layer.merge(ThemeSyncLive),
+    Layer.merge(HideCodeSyncLive),
     Layer.merge(DebugLayerLive),
   )
   .pipe(

extension/src/features/__tests__/HideCodeSync.test.ts

@@ -0,0 +1,132 @@
+import { describe, expect, it } from "@effect/vitest";
+import { Effect, Layer, Option, Ref, TestClock } from "effect";
+
+import {
+  createNotebookCell,
+  createTestNotebookDocument,
+  TestVsCode,
+} from "../../__mocks__/TestVsCode.ts";
+import { MarimoNotebookCell } from "../../schemas/MarimoNotebookDocument.ts";
+import {
+  type CellRange,
+  hiddenInputCellRanges,
+  HideCodeSyncLive,
+} from "../HideCodeSync.ts";
+
+const cell = (index: number, hideCode: boolean) =>
+  MarimoNotebookCell.from(
+    createNotebookCell(
+      createTestNotebookDocument("/test/notebook_mo.py"),
+      {
+        kind: 2,
+        value: "",
+        languageId: "python",
+        metadata: {
+          stableId: `cell-${index}`,
+          options: { hide_code: hideCode },
+        },
+      },
+      index,
+    ),
+  );
+
+describe("hiddenInputCellRanges", () => {
+  it("returns one end-exclusive range per hide_code cell, by index", () => {
+    const cells = [
+      cell(0, false),
+      cell(1, true),
+      cell(2, false),
+      cell(3, true),
+    ];
+    expect(hiddenInputCellRanges(cells)).toEqual([
+      { start: 1, end: 2 },
+      { start: 3, end: 4 },
+    ]);
+  });
+
+  it("returns no ranges when no cell hides its code", () => {
+    expect(hiddenInputCellRanges([cell(0, false)])).toEqual([]);
+  });
+});
+
+const isCellRange = (x: unknown): x is CellRange =>
+  typeof x === "object" &&
+  x !== null &&
+  "start" in x &&
+  typeof x.start === "number" &&
+  "end" in x &&
+  typeof x.end === "number";
+
+const collapseRanges = (arg: unknown): readonly CellRange[] | undefined =>
+  typeof arg === "object" &&
+  arg !== null &&
+  "ranges" in arg &&
+  Array.isArray(arg.ranges) &&
+  arg.ranges.every(isCellRange)
+    ? arg.ranges
+    : undefined;
+
+/** The ranges passed to each `notebook.cell.collapseCellInput` call, in order. */
+const collapses = Effect.fn(function* (vscode: TestVsCode) {
+  const executions = yield* Ref.get(vscode.executions);
+  return executions
+    .filter((e) => e.command === "notebook.cell.collapseCellInput")
+    .map((e) => collapseRanges(e.args[0]));
+});
+
+const withTestCtx = Effect.fn(function* (hideCode: ReadonlyArray<boolean>) {
+  const editor = TestVsCode.makeNotebookEditor("/test/notebook_mo.py", {
+    data: {
+      cells: hideCode.map((hide_code, index) => ({
+        kind: 2,
+        value: "",
+        languageId: "python",
+        metadata: { stableId: `cell-${index}`, options: { hide_code } },
+      })),
+    },
+  });
+  const vscode = yield* TestVsCode.make({
+    initialDocuments: [editor.notebook],
+  });
+  const layer = HideCodeSyncLive.pipe(Layer.provide(vscode.layer));
+  return { vscode, editor, layer };
+});
+
+describe("HideCodeSync", () => {
+  it.scoped(
+    "collapses hide_code cells when a notebook first becomes active",
+    Effect.fn(function* () {
+      const { vscode, editor, layer } = yield* withTestCtx([false, true, true]);
+
+      yield* Effect.gen(function* () {
+        yield* vscode.setActiveNotebookEditor(Option.some(editor));
+        yield* TestClock.adjust("1 millis");
+
+        expect(yield* collapses(vscode)).toEqual([
+          [
+            { start: 1, end: 2 },
+            { start: 2, end: 3 },
+          ],
+        ]);
+      }).pipe(Effect.provide(layer));
+    }),
+  );
+
+  it.scoped(
+    "collapses once and does not re-collapse on tab refocus",
+    Effect.fn(function* () {
+      const { vscode, editor, layer } = yield* withTestCtx([true]);
+
+      yield* Effect.gen(function* () {
+        yield* vscode.setActiveNotebookEditor(Option.some(editor));
+        yield* TestClock.adjust("1 millis");
+        yield* vscode.setActiveNotebookEditor(Option.none());
+        yield* TestClock.adjust("1 millis");
+        yield* vscode.setActiveNotebookEditor(Option.some(editor));
+        yield* TestClock.adjust("1 millis");
+
+        expect(yield* collapses(vscode)).toEqual([[{ start: 0, end: 1 }]]);
+      }).pipe(Effect.provide(layer));
+    }),
+  );
+});

extension/src/schemas/MarimoNotebookDocument.ts

@@ -130,6 +130,22 @@ export class MarimoNotebookCell {
     );
   }
 
+  /**
+   * Whether the cell's code is hidden via `@app.cell(hide_code=True)`.
+   *
+   * marimo stores the decorator's `hide_code` flag in the cell's config; the
+   * LSP deserialize path surfaces it as `metadata.options.hide_code`. VS Code
+   * has no API to read or set the input-collapsed state, so this is the source
+   * of truth the {@link HideCodeSyncLive} feature uses to one-way sync the
+   * editor's collapse state (issue #326).
+   */
+  get isCodeHidden() {
+    return this.metadata.pipe(
+      Option.map((meta) => meta.options?.hide_code === true),
+      Option.getOrElse(() => false),
+    );
+  }
+
   /**
    * The cell's language metadata, if present.
    */