ci(run-notebook): replace whole-cell bash partition with nbclient

.github/actions/run-notebook/action.yml

@@ -1,5 +1,5 @@
 name: "Run Notebook"
-description: "Run a notebook"
+description: "Run a notebook end-to-end against a real Jupyter kernel via nbclient"
 
 inputs:
   notebook:
@@ -26,34 +26,18 @@ runs:
       with:
         python-version: '3.11'
 
-    - name: Install dependencies
+    - name: Install runner dependencies
       shell: bash
       run: |
         pip install --upgrade pip
-        pip install nbformat
+        pip install nbformat nbclient ipykernel
 
-    - id: convert
+    - name: Run the notebook
       shell: bash
-      name: Convert notebook into tmpdir script
       run: |
-        python .github/actions/run-notebook/convert-notebook.py ${{ inputs.notebook }}
-
-    - name: View the run script
-      shell: bash
-      run: |
-        cat ${{ steps.convert.outputs.script_path }}
-
-    - name: View converted notebook content
-      shell: bash
-      run: |
-        cat ${{ steps.convert.outputs.notebook_path }}
-
-    - name: Run the converted notebook
-      shell: bash
-      run: |
-        bash ${{ steps.convert.outputs.script_path }}
+        python .github/actions/run-notebook/run-notebook.py "${{ inputs.notebook }}"
       env:
         PINECONE_API_KEY: ${{ inputs.PINECONE_API_KEY }}
         OPENAI_API_KEY: ${{ inputs.OPENAI_API_KEY }}
         HF_TOKEN: ${{ inputs.HF_TOKEN }}
-        ANTHROPIC_API_KEY: ${{ inputs.ANTHROPIC_API_KEY }}
+        ANTHROPIC_API_KEY: ${{ inputs.ANTHROPIC_API_KEY }}

.github/actions/run-notebook/run-notebook.py

@@ -0,0 +1,72 @@
+#!/usr/bin/env python
+"""Execute a notebook end-to-end against a real Jupyter kernel.
+
+Replaces the cell-partitioning approach in convert-notebook.py with a
+straight nbclient drive. This means:
+
+  - !pip install and other shell magics run via the kernel's normal magic
+    handling (the kernel spawns a subshell), exactly as in Colab/Jupyter Lab.
+  - A code cell may freely mix `!pip install` and Python imports.
+  - Errors surface with cell and traceback context, not as
+    `line N: import: command not found`.
+
+Notebook deps are installed into the same Python environment as the
+runner (the kernel and runner share an interpreter), so `!pip install foo`
+in cell 1 means cell 2 can `import foo`.
+
+Usage:
+  run-notebook.py <notebook-path>
+
+Exits non-zero on any cell failure.
+"""
+
+import os
+import sys
+
+import nbformat
+from nbclient import NotebookClient
+from nbclient.exceptions import CellExecutionError, CellTimeoutError, DeadKernelError
+
+CELL_TIMEOUT = int(os.environ.get("NOTEBOOK_CELL_TIMEOUT", "600"))
+
+
+def main() -> int:
+    if len(sys.argv) != 2:
+        print("usage: run-notebook.py <notebook-path>", file=sys.stderr)
+        return 2
+
+    notebook_path = sys.argv[1]
+    print(f"Executing {notebook_path}")
+
+    nb = nbformat.read(notebook_path, as_version=4)
+    client = NotebookClient(
+        nb,
+        timeout=CELL_TIMEOUT,
+        kernel_name="python3",
+        resources={"metadata": {"path": os.path.dirname(notebook_path) or "."}},
+    )
+
+    try:
+        client.execute()
+    except CellTimeoutError as exc:
+        print(
+            f"\nCell exceeded timeout of {CELL_TIMEOUT}s "
+            f"(override via NOTEBOOK_CELL_TIMEOUT):\n{exc}",
+            file=sys.stderr,
+        )
+        return 1
+    except DeadKernelError as exc:
+        print(
+            f"\nKernel died during execution (OOM/segfault?):\n{exc}", file=sys.stderr
+        )
+        return 1
+    except CellExecutionError as exc:
+        print(f"\nCell execution failed:\n{exc}", file=sys.stderr)
+        return 1
+
+    print(f"PASS — {notebook_path}")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())