├── docs ├── static │ ├── .nojekyll │ ├── img │ │ ├── favicon.ico │ │ ├── block-jewel_black.svg │ │ ├── block-jewel_white.svg │ │ └── logo.svg │ └── demos │ │ └── goose-demo.png ├── src │ ├── pages │ │ ├── markdown-page.md │ │ ├── index.module.css │ │ └── index.tsx │ ├── components │ │ └── HomepageFeatures │ │ │ ├── styles.module.css │ │ │ └── index.tsx │ └── css │ │ └── custom.css ├── tsconfig.json ├── .gitignore ├── docs │ ├── images │ │ ├── block-jewel_black.svg │ │ └── block-jewel_white.svg │ ├── intro.md │ ├── quickstart.md │ ├── usage.md │ ├── architecture.md │ ├── development.md │ └── installation.md ├── sidebars.ts ├── README.md ├── package.json └── docusaurus.config.ts ├── bin ├── pnpm ├── hermit.hcl ├── README.hermit.md ├── activate-hermit ├── activate-hermit.fish └── hermit ├── tests ├── llm_providers │ ├── __init__.py │ ├── config.py │ ├── base.py │ └── claude_code.py ├── __init__.py ├── test_check_server.py ├── test_filtering.py ├── test_llm_tool_calls.py ├── conftest.py ├── test_http_transport.py ├── test_notebook_paths.py └── test_integration.py ├── src └── mcp_jupyter │ ├── __main__.py │ ├── __init__.py │ ├── utils.py │ ├── state.py │ ├── notebook.py │ └── rest_client.py ├── .hooks ├── pre-commit └── pre-commit-py ├── demos └── goose-demo.png ├── Justfile ├── .pre-commit-config.yaml ├── .github └── workflows │ ├── tests.yml │ ├── publish.yml │ └── docs.yml ├── pyproject.toml ├── .gitignore ├── README.md └── LICENSE /docs/static/.nojekyll: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /bin/pnpm: -------------------------------------------------------------------------------- 1 | .pnpm-10.11.0.pkg -------------------------------------------------------------------------------- /tests/llm_providers/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /tests/__init__.py: -------------------------------------------------------------------------------- 1 | # Tests package 2 | -------------------------------------------------------------------------------- /bin/hermit.hcl: -------------------------------------------------------------------------------- 1 | manage-git = false 2 | 3 | github-token-auth { 4 | } 5 | -------------------------------------------------------------------------------- /src/mcp_jupyter/__main__.py: -------------------------------------------------------------------------------- 1 | from mcp_jupyter import main 2 | 3 | main() 4 | -------------------------------------------------------------------------------- /.hooks/pre-commit: -------------------------------------------------------------------------------- 1 | $(git rev-parse --show-toplevel)/.hooks/pre-commit-py 2 | exit $? 3 | -------------------------------------------------------------------------------- /demos/goose-demo.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/block/mcp-jupyter/HEAD/demos/goose-demo.png -------------------------------------------------------------------------------- /docs/static/img/favicon.ico: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/block/mcp-jupyter/HEAD/docs/static/img/favicon.ico -------------------------------------------------------------------------------- /docs/static/demos/goose-demo.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/block/mcp-jupyter/HEAD/docs/static/demos/goose-demo.png -------------------------------------------------------------------------------- /docs/src/pages/markdown-page.md: -------------------------------------------------------------------------------- 1 | --- 2 | title: Markdown page example 3 | --- 4 | 5 | # Markdown page example 6 | 7 | You don't need React to write simple standalone pages. 8 | -------------------------------------------------------------------------------- /docs/src/components/HomepageFeatures/styles.module.css: -------------------------------------------------------------------------------- 1 | .features { 2 | display: flex; 3 | align-items: center; 4 | padding: 2rem 0; 5 | width: 100%; 6 | } 7 | 8 | .featureSvg { 9 | height: 200px; 10 | width: 200px; 11 | } 12 | -------------------------------------------------------------------------------- /docs/tsconfig.json: -------------------------------------------------------------------------------- 1 | { 2 | // This file is not used in compilation. It is here just for a nice editor experience. 3 | "extends": "@docusaurus/tsconfig", 4 | "compilerOptions": { 5 | "baseUrl": "." 6 | }, 7 | "exclude": [".docusaurus", "build"] 8 | } 9 | -------------------------------------------------------------------------------- /bin/README.hermit.md: -------------------------------------------------------------------------------- 1 | # Hermit environment 2 | 3 | This is a [Hermit](https://github.com/cashapp/hermit) bin directory. 4 | 5 | The symlinks in this directory are managed by Hermit and will automatically 6 | download and install Hermit itself as well as packages. These packages are 7 | local to this environment. 8 | -------------------------------------------------------------------------------- /docs/.gitignore: -------------------------------------------------------------------------------- 1 | # Dependencies 2 | /node_modules 3 | 4 | # Production 5 | /build 6 | 7 | # Generated files 8 | .docusaurus 9 | .cache-loader 10 | 11 | # Misc 12 | .DS_Store 13 | .env.local 14 | .env.development.local 15 | .env.test.local 16 | .env.production.local 17 | 18 | npm-debug.log* 19 | yarn-debug.log* 20 | yarn-error.log* 21 | -------------------------------------------------------------------------------- /docs/src/pages/index.module.css: -------------------------------------------------------------------------------- 1 | /** 2 | * CSS files with the .module.css suffix will be treated as CSS modules 3 | * and scoped locally. 4 | */ 5 | 6 | .heroBanner { 7 | padding: 4rem 0; 8 | text-align: center; 9 | position: relative; 10 | overflow: hidden; 11 | } 12 | 13 | @media screen and (max-width: 996px) { 14 | .heroBanner { 15 | padding: 2rem; 16 | } 17 | } 18 | 19 | .buttons { 20 | display: flex; 21 | align-items: center; 22 | justify-content: center; 23 | } 24 | -------------------------------------------------------------------------------- /Justfile: -------------------------------------------------------------------------------- 1 | # Run all CI checks locally (matches GitHub Actions) 2 | ci: test lint format build 3 | @echo "✅ All CI checks passed!" 4 | 5 | # Run tests (matches: uv run pytest tests) 6 | test: 7 | @echo "🧪 Running tests..." 8 | uv run pytest tests 9 | 10 | # Run linting (matches: uvx ruff check) 11 | lint: 12 | @echo "🔍 Running linter..." 13 | uvx ruff check 14 | 15 | # Run format check and fix (matches: uvx ruff format --check but also fixes) 16 | format: 17 | @echo "🎨 Checking and fixing formatting..." 18 | uvx ruff format 19 | 20 | # Build documentation (matches: pnpm build in docs/) 21 | build: 22 | @echo "📚 Building documentation..." 23 | cd docs && pnpm start -------------------------------------------------------------------------------- /bin/activate-hermit: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | # This file must be used with "source bin/activate-hermit" from bash or zsh. 3 | # You cannot run it directly 4 | # 5 | # THIS FILE IS GENERATED; DO NOT MODIFY 6 | 7 | if [ "${BASH_SOURCE-}" = "$0" ]; then 8 | echo "You must source this script: \$ source $0" >&2 9 | exit 33 10 | fi 11 | 12 | BIN_DIR="$(dirname "${BASH_SOURCE[0]:-${(%):-%x}}")" 13 | if "${BIN_DIR}/hermit" noop > /dev/null; then 14 | eval "$("${BIN_DIR}/hermit" activate "${BIN_DIR}/..")" 15 | 16 | if [ -n "${BASH-}" ] || [ -n "${ZSH_VERSION-}" ]; then 17 | hash -r 2>/dev/null 18 | fi 19 | 20 | echo "Hermit environment $("${HERMIT_ENV}"/bin/hermit env HERMIT_ENV) activated" 21 | fi 22 | -------------------------------------------------------------------------------- /.pre-commit-config.yaml: -------------------------------------------------------------------------------- 1 | # See https://pre-commit.com for more information 2 | # See https://pre-commit.com/hooks.html for more hooks 3 | repos: 4 | - repo: https://github.com/pre-commit/pre-commit-hooks 5 | rev: v4.6.0 6 | hooks: 7 | - id: trailing-whitespace 8 | - id: end-of-file-fixer 9 | - id: check-yaml 10 | - id: check-added-large-files 11 | - id: check-json 12 | - id: check-merge-conflict 13 | - id: check-toml 14 | - id: debug-statements 15 | - id: mixed-line-ending 16 | 17 | - repo: https://github.com/astral-sh/ruff-pre-commit 18 | rev: main 19 | hooks: 20 | - id: ruff 21 | args: [--fix] 22 | # Run the formatter 23 | - id: ruff-format 24 | -------------------------------------------------------------------------------- /.github/workflows/tests.yml: -------------------------------------------------------------------------------- 1 | name: Tests 2 | 3 | on: 4 | push: 5 | branches: '*' 6 | pull_request: 7 | branches: '*' 8 | 9 | jobs: 10 | tests: 11 | 12 | runs-on: ubuntu-latest 13 | strategy: 14 | matrix: 15 | python-version: ["3.10", "3.11", "3.12", "3.13"] 16 | 17 | steps: 18 | - uses: actions/checkout@v4 19 | 20 | - name: Install the latest version of uv 21 | uses: astral-sh/setup-uv@v6 22 | with: 23 | version: "latest" 24 | 25 | - name: Ruff 26 | run: | 27 | uvx ruff check 28 | uvx ruff format --check 29 | 30 | - name: Test with Python ${{ matrix.python-version }} 31 | run: | 32 | uv run --python ${{ matrix.python-version }} pytest tests 33 | 34 | - name: Test build 35 | run: uv build 36 | -------------------------------------------------------------------------------- /bin/activate-hermit.fish: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env fish 2 | 3 | # This file must be sourced with "source bin/activate-hermit.fish" from Fish shell. 4 | # You cannot run it directly. 5 | # 6 | # THIS FILE IS GENERATED; DO NOT MODIFY 7 | 8 | if status is-interactive 9 | set BIN_DIR (dirname (status --current-filename)) 10 | 11 | if "$BIN_DIR/hermit" noop > /dev/null 12 | # Source the activation script generated by Hermit 13 | "$BIN_DIR/hermit" activate "$BIN_DIR/.." | source 14 | 15 | # Clear the command cache if applicable 16 | functions -c > /dev/null 2>&1 17 | 18 | # Display activation message 19 | echo "Hermit environment $($HERMIT_ENV/bin/hermit env HERMIT_ENV) activated" 20 | end 21 | else 22 | echo "You must source this script: source $argv[0]" >&2 23 | exit 33 24 | end 25 | -------------------------------------------------------------------------------- /docs/docs/images/block-jewel_black.svg: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /docs/docs/images/block-jewel_white.svg: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /docs/static/img/block-jewel_black.svg: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /docs/static/img/block-jewel_white.svg: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /.hooks/pre-commit-py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # File generated by pre-commit: https://pre-commit.com 3 | # with additional edits to support Square git hooks 4 | import os 5 | import sys 6 | from shutil import which 7 | 8 | INSTALL_PYTHON = "python" 9 | ARGS = [ 10 | "hook-impl", 11 | "--config=.pre-commit-config.yaml", 12 | "--hook-type=pre-commit", 13 | "--skip-on-missing-config", 14 | ] 15 | ARGS.extend(("--hook-dir", os.path.realpath(os.path.dirname(__file__)))) 16 | ARGS.append("--") 17 | ARGS.extend(sys.argv[1:]) 18 | 19 | DNE = "`pre-commit` not found. Make sure to install into your virtualenv before committing" # NOQA 20 | if os.access(INSTALL_PYTHON, os.X_OK): 21 | CMD = [INSTALL_PYTHON, "-mpre_commit"] 22 | elif which("pre-commit"): 23 | CMD = ["pre-commit"] 24 | else: 25 | raise SystemExit(DNE) 26 | 27 | CMD.extend(ARGS) 28 | os.execvp(CMD[0], CMD) 29 | -------------------------------------------------------------------------------- /docs/static/img/logo.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 14 | 15 | 16 | 17 | 18 | 19 | 20 | 21 | 22 | 23 | -------------------------------------------------------------------------------- /docs/sidebars.ts: -------------------------------------------------------------------------------- 1 | import type {SidebarsConfig} from '@docusaurus/plugin-content-docs'; 2 | 3 | // This runs in Node.js - Don't use client-side code here (browser APIs, JSX...) 4 | 5 | /** 6 | * Creating a sidebar enables you to: 7 | - create an ordered group of docs 8 | - render a sidebar for each doc of that group 9 | - provide next/previous navigation 10 | 11 | The sidebars can be generated from the filesystem, or explicitly defined here. 12 | 13 | Create as many sidebars as you want. 14 | */ 15 | const sidebars: SidebarsConfig = { 16 | // By default, Docusaurus generates a sidebar from the docs folder structure 17 | tutorialSidebar: [{type: 'autogenerated', dirName: '.'}], 18 | 19 | // But you can create a sidebar manually 20 | /* 21 | tutorialSidebar: [ 22 | 'intro', 23 | 'hello', 24 | { 25 | type: 'category', 26 | label: 'Tutorial', 27 | items: ['tutorial-basics/create-a-document'], 28 | }, 29 | ], 30 | */ 31 | }; 32 | 33 | export default sidebars; 34 | -------------------------------------------------------------------------------- /docs/README.md: -------------------------------------------------------------------------------- 1 | # MCP Jupyter Documentation 2 | 3 | This is the documentation site for MCP Jupyter Server, built with Docusaurus. 4 | 5 | ## Local Development 6 | 7 | First, activate hermit to get pnpm: 8 | 9 | ```bash 10 | . bin/activate-hermit 11 | ``` 12 | 13 | Then install dependencies and start the development server: 14 | 15 | ```bash 16 | cd docs 17 | pnpm install 18 | pnpm start 19 | ``` 20 | 21 | This command starts a local development server and opens up a browser window. Most changes are reflected live without having to restart the server. 22 | 23 | ## Build 24 | 25 | ```bash 26 | pnpm build 27 | ``` 28 | 29 | This command generates static content into the `build` directory and can be served using any static contents hosting service. 30 | 31 | ## Serve Built Site 32 | 33 | To preview the production build locally: 34 | 35 | ```bash 36 | pnpm serve 37 | ``` 38 | 39 | ## Structure 40 | 41 | - `/docs` - Documentation pages 42 | - `/src/pages` - Custom pages (homepage) 43 | - `/src/components` - React components 44 | - `/static` - Static assets (images, demos) 45 | - `docusaurus.config.ts` - Site configuration -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | [project] 2 | name = "mcp-jupyter" 3 | version = "2.0.2" 4 | description = "MCP Jupyter" 5 | readme = "README.md" 6 | requires-python = ">=3.10" 7 | dependencies = [ 8 | "beautifulsoup4>=4.12.3", 9 | "mcp[cli]>=1.13.1", 10 | "requests>=2.32.3", 11 | "rich", 12 | "jupyter-kernel-client>=0.8.0", 13 | ] 14 | 15 | 16 | [project.scripts] 17 | mcp-jupyter = "mcp_jupyter:main" 18 | 19 | [build-system] 20 | requires = ["hatchling"] 21 | build-backend = "hatchling.build" 22 | 23 | [tool.ruff.format] 24 | docstring-code-format = true 25 | 26 | [tool.ruff.lint] 27 | select = ["I", "D"] 28 | ignore = ["D104", "D100", "D205", "D400"] 29 | 30 | [tool.ruff.lint.isort] 31 | force-sort-within-sections = false 32 | known-first-party = ["block"] 33 | 34 | [tool.ruff.lint.pydocstyle] 35 | convention = "numpy" 36 | 37 | [tool.ruff.lint.per-file-ignores] 38 | "tests/*" = ["D103"] 39 | 40 | [dependency-groups] 41 | dev = [ 42 | "claude-code-sdk>=0.0.21", 43 | "ipykernel>=6.29.5", 44 | "jupyter-collaboration>=3.1.0", 45 | "jupyterlab>=4.3.6", 46 | "pytest>=8.3.5", 47 | "pytest-asyncio>=0.21.0", 48 | "pytest-order>=1.3.0", 49 | "pytest-xdist>=3.6.1", 50 | ] 51 | 52 | [tool.pytest.ini_options] 53 | addopts = "-m 'not llm'" 54 | markers = [ 55 | "llm: marks tests as LLM-generated tool call tests", 56 | ] 57 | filterwarnings = [ 58 | "ignore:.*Jupyter is migrating.*:DeprecationWarning", 59 | "ignore:Exception ignored:pytest.PytestUnraisableExceptionWarning", 60 | ] 61 | -------------------------------------------------------------------------------- /docs/package.json: -------------------------------------------------------------------------------- 1 | { 2 | "name": "docs", 3 | "version": "0.0.0", 4 | "private": true, 5 | "scripts": { 6 | "docusaurus": "docusaurus", 7 | "start": "docusaurus start", 8 | "build": "docusaurus build", 9 | "swizzle": "docusaurus swizzle", 10 | "deploy": "docusaurus deploy", 11 | "clear": "docusaurus clear", 12 | "serve": "docusaurus serve", 13 | "write-translations": "docusaurus write-translations", 14 | "write-heading-ids": "docusaurus write-heading-ids", 15 | "typecheck": "tsc" 16 | }, 17 | "dependencies": { 18 | "@docusaurus/core": "3.7.0", 19 | "@docusaurus/preset-classic": "3.7.0", 20 | "@easyops-cn/docusaurus-search-local": "^0.44.0", 21 | "@mdx-js/react": "^3.0.0", 22 | "clsx": "^2.0.0", 23 | "prism-react-renderer": "^2.3.0", 24 | "react": "^19.0.0", 25 | "react-dom": "^19.0.0" 26 | }, 27 | "devDependencies": { 28 | "@docusaurus/module-type-aliases": "3.7.0", 29 | "@docusaurus/tsconfig": "3.7.0", 30 | "@docusaurus/types": "3.7.0", 31 | "typescript": "~5.6.2" 32 | }, 33 | "browserslist": { 34 | "production": [ 35 | ">0.5%", 36 | "not dead", 37 | "not op_mini all" 38 | ], 39 | "development": [ 40 | "last 3 chrome version", 41 | "last 3 firefox version", 42 | "last 5 safari version" 43 | ] 44 | }, 45 | "engines": { 46 | "node": ">=18.0" 47 | }, 48 | "pnpm": { 49 | "overrides": { 50 | "undici": "5.28.4", 51 | "jiti": "1.20.0" 52 | } 53 | } 54 | } 55 | -------------------------------------------------------------------------------- /bin/hermit: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | # 3 | # THIS FILE IS GENERATED; DO NOT MODIFY 4 | 5 | set -eo pipefail 6 | 7 | export HERMIT_USER_HOME=~ 8 | 9 | if [ -z "${HERMIT_STATE_DIR}" ]; then 10 | case "$(uname -s)" in 11 | Darwin) 12 | export HERMIT_STATE_DIR="${HERMIT_USER_HOME}/Library/Caches/hermit" 13 | ;; 14 | Linux) 15 | export HERMIT_STATE_DIR="${XDG_CACHE_HOME:-${HERMIT_USER_HOME}/.cache}/hermit" 16 | ;; 17 | esac 18 | fi 19 | 20 | export HERMIT_DIST_URL="${HERMIT_DIST_URL:-https://d1abdrezunyhdp.cloudfront.net/square}" 21 | HERMIT_CHANNEL="$(basename "${HERMIT_DIST_URL}")" 22 | export HERMIT_CHANNEL 23 | export HERMIT_EXE=${HERMIT_EXE:-${HERMIT_STATE_DIR}/pkg/hermit@${HERMIT_CHANNEL}/hermit} 24 | 25 | if [ ! -x "${HERMIT_EXE}" ]; then 26 | echo "Bootstrapping ${HERMIT_EXE} from ${HERMIT_DIST_URL}" 1>&2 27 | INSTALL_SCRIPT="$(mktemp)" 28 | # This value must match that of the install script 29 | INSTALL_SCRIPT_SHA256="4b006236f2e5e81939229b377bb355e3608f94d73ff8feccbd5792d1ed5699cd" 30 | if [ "${INSTALL_SCRIPT_SHA256}" = "BYPASS" ]; then 31 | curl -fsSL "${HERMIT_DIST_URL}/install.sh" -o "${INSTALL_SCRIPT}" 32 | else 33 | # Install script is versioned by its sha256sum value 34 | curl -fsSL "${HERMIT_DIST_URL}/install-${INSTALL_SCRIPT_SHA256}.sh" -o "${INSTALL_SCRIPT}" 35 | # Verify install script's sha256sum 36 | openssl dgst -sha256 "${INSTALL_SCRIPT}" | \ 37 | awk -v EXPECTED="$INSTALL_SCRIPT_SHA256" \ 38 | '$2!=EXPECTED {print "Install script sha256 " $2 " does not match " EXPECTED; exit 1}' 39 | fi 40 | /bin/bash "${INSTALL_SCRIPT}" 1>&2 41 | fi 42 | 43 | exec "${HERMIT_EXE}" --level=fatal exec "$0" -- "$@" 44 | -------------------------------------------------------------------------------- /.github/workflows/publish.yml: -------------------------------------------------------------------------------- 1 | name: Publish 2 | 3 | on: 4 | release: 5 | types: [published] 6 | 7 | jobs: 8 | publish: 9 | permissions: 10 | id-token: write 11 | contents: write 12 | 13 | runs-on: ubuntu-latest 14 | environment: 15 | name: pypi 16 | url: https://pypi.org/project/mcp-jupyter/ 17 | steps: 18 | - uses: actions/checkout@v4 19 | 20 | - name: Get current version from pyproject.toml 21 | id: get_version 22 | run: | 23 | echo "VERSION=$(grep -m 1 'version =' "pyproject.toml" | awk -F'"' '{print $2}')" >> $GITHUB_ENV 24 | 25 | - name: Extract tag version 26 | id: extract_tag 27 | run: | 28 | TAG_VERSION=$(echo "${{ github.event.release.tag_name }}" | sed -E 's/v(.*)/\1/') 29 | echo "TAG_VERSION=$TAG_VERSION" >> $GITHUB_ENV 30 | 31 | - name: Check if tag matches version from pyproject.toml 32 | id: check_tag 33 | run: | 34 | if [ "${{ env.TAG_VERSION }}" != "${{ env.VERSION }}" ]; then 35 | echo "::error::Tag version (${{ env.TAG_VERSION }}) does not match version in pyproject.toml (${{ env.VERSION }})." 36 | exit 1 37 | fi 38 | 39 | - name: Install the latest version of uv 40 | uses: astral-sh/setup-uv@v6 41 | with: 42 | version: "latest" 43 | 44 | - name: Build Package 45 | run: uv build 46 | 47 | - name: Upload to GitHub Release 48 | env: 49 | GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} 50 | run: | 51 | gh release upload ${{ github.event.release.tag_name }} dist/* 52 | 53 | - name: Publish package to PyPI 54 | uses: pypa/gh-action-pypi-publish@release/v1 -------------------------------------------------------------------------------- /.github/workflows/docs.yml: -------------------------------------------------------------------------------- 1 | name: Deploy Documentation 2 | 3 | on: 4 | push: 5 | branches: 6 | - main 7 | paths: 8 | - 'docs/**' 9 | pull_request: 10 | branches: 11 | - main 12 | paths: 13 | - 'docs/**' 14 | workflow_dispatch: 15 | 16 | permissions: 17 | contents: read 18 | pages: write 19 | id-token: write 20 | 21 | concurrency: 22 | group: "pages" 23 | cancel-in-progress: false 24 | 25 | jobs: 26 | build: 27 | runs-on: ubuntu-latest 28 | steps: 29 | - name: Checkout 30 | uses: actions/checkout@v4 31 | with: 32 | fetch-depth: 0 33 | 34 | - name: Setup Node.js 35 | uses: actions/setup-node@v4 36 | with: 37 | node-version: '18' 38 | 39 | - name: Setup pnpm 40 | uses: pnpm/action-setup@v4 41 | with: 42 | version: 8 43 | 44 | - name: Install dependencies 45 | run: | 46 | cd docs 47 | pnpm install --no-frozen-lockfile 48 | 49 | - name: Build website 50 | run: | 51 | cd docs 52 | pnpm build 53 | 54 | - name: Setup Pages 55 | uses: actions/configure-pages@v4 56 | 57 | - name: Upload artifact 58 | uses: actions/upload-pages-artifact@v3 59 | with: 60 | path: docs/build 61 | 62 | deploy: 63 | if: github.ref == 'refs/heads/main' && github.event_name != 'pull_request' 64 | environment: 65 | name: github-pages 66 | url: ${{ steps.deployment.outputs.page_url }} 67 | runs-on: ubuntu-latest 68 | needs: build 69 | steps: 70 | - name: Deploy to GitHub Pages 71 | id: deployment 72 | uses: actions/deploy-pages@v4 -------------------------------------------------------------------------------- /src/mcp_jupyter/__init__.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import logging 3 | 4 | from .server import create_server 5 | 6 | logger = logging.getLogger(__name__) 7 | 8 | 9 | def main(): 10 | """MCP Jupyter: Control a Jupyter notebook from MCP.""" 11 | parser = argparse.ArgumentParser( 12 | description="Gives you the ability to control a Jupyter notebook from MCP." 13 | ) 14 | parser.add_argument( 15 | "--transport", 16 | "-t", 17 | choices=["stdio", "http"], 18 | default="stdio", 19 | help="Transport type to use (default: stdio)", 20 | ) 21 | parser.add_argument( 22 | "--port", 23 | "-p", 24 | type=int, 25 | default=8000, 26 | help="Port for HTTP transport (default: 8000)", 27 | ) 28 | parser.add_argument( 29 | "--host", 30 | default="127.0.0.1", 31 | help="Host for HTTP transport (default: 127.0.0.1)", 32 | ) 33 | parser.add_argument( 34 | "--stateless-http", 35 | action="store_true", 36 | help="Enable stateless HTTP mode (no session persistence)", 37 | ) 38 | 39 | args = parser.parse_args() 40 | 41 | # Create the server with appropriate settings 42 | if args.transport == "http": 43 | logger.info( 44 | f"Starting MCP server on {args.host}:{args.port} with HTTP transport" 45 | ) 46 | server = create_server( 47 | host=args.host, port=args.port, stateless_http=args.stateless_http 48 | ) 49 | server.run(transport="streamable-http") 50 | else: 51 | logger.info(f"Starting MCP server with stdio transport") 52 | server = create_server() 53 | server.run(transport="stdio") 54 | 55 | 56 | if __name__ == "__main__": 57 | main() 58 | -------------------------------------------------------------------------------- /docs/src/pages/index.tsx: -------------------------------------------------------------------------------- 1 | import type {ReactNode} from 'react'; 2 | import clsx from 'clsx'; 3 | import Link from '@docusaurus/Link'; 4 | import useDocusaurusContext from '@docusaurus/useDocusaurusContext'; 5 | import Layout from '@theme/Layout'; 6 | import HomepageFeatures from '@site/src/components/HomepageFeatures'; 7 | import Heading from '@theme/Heading'; 8 | 9 | import styles from './index.module.css'; 10 | 11 | function HomepageHeader() { 12 | const {siteConfig} = useDocusaurusContext(); 13 | return ( 14 |
15 |
16 | 17 | {siteConfig.title} 18 | 19 |

{siteConfig.tagline}

20 |
21 | 24 | Get Started → 25 | 26 | 30 | View on GitHub 31 | 32 |
33 |
34 |
35 | ); 36 | } 37 | 38 | export default function Home(): ReactNode { 39 | const {siteConfig} = useDocusaurusContext(); 40 | return ( 41 | 44 | 45 |
46 | 47 |
48 | 49 | ); 50 | } -------------------------------------------------------------------------------- /tests/llm_providers/config.py: -------------------------------------------------------------------------------- 1 | """Configuration for LLM providers in testing.""" 2 | 3 | import os 4 | from typing import Any, Dict, List 5 | 6 | from .base import LLMProvider 7 | from .claude_code import ClaudeCodeProvider 8 | 9 | 10 | def get_available_providers() -> List[LLMProvider]: 11 | """Get list of available LLM providers based on configuration and environment. 12 | 13 | Returns 14 | ------- 15 | List of available LLM providers 16 | """ 17 | providers = [] 18 | 19 | # Claude Code is always available since it doesn't require API keys 20 | providers.append(ClaudeCodeProvider()) 21 | 22 | # Future providers can be added here with environment checks when implemented 23 | 24 | return providers 25 | 26 | 27 | def get_provider_config() -> Dict[str, Any]: 28 | """Get configuration for LLM providers. 29 | 30 | Returns 31 | ------- 32 | Configuration dictionary 33 | """ 34 | return { 35 | "providers": { 36 | "claude-code": { 37 | "enabled": True, 38 | "requires_api_key": False, 39 | } 40 | }, 41 | } 42 | 43 | 44 | def get_test_providers() -> List[LLMProvider]: 45 | """Get providers that should be used in tests. 46 | 47 | This respects environment variables to enable/disable specific providers. 48 | 49 | Returns 50 | ------- 51 | List of providers to test 52 | """ 53 | config = get_provider_config() 54 | available = get_available_providers() 55 | 56 | # Filter based on configuration 57 | enabled_provider_names = [ 58 | name 59 | for name, provider_config in config["providers"].items() 60 | if provider_config["enabled"] 61 | ] 62 | 63 | return [ 64 | provider for provider in available if provider.name in enabled_provider_names 65 | ] 66 | -------------------------------------------------------------------------------- /tests/llm_providers/base.py: -------------------------------------------------------------------------------- 1 | """Base interface for LLM providers in MCP tool call tests.""" 2 | 3 | from abc import ABC, abstractmethod 4 | from dataclasses import dataclass 5 | from typing import AsyncGenerator, Optional 6 | 7 | 8 | @dataclass 9 | class LLMResponse: 10 | """Response from an LLM provider.""" 11 | 12 | text: str 13 | """The text response from the LLM.""" 14 | 15 | tool_calls_made: int 16 | """Number of tool calls the LLM made.""" 17 | 18 | success: bool 19 | """Whether the LLM successfully completed the task.""" 20 | 21 | error: Optional[str] = None 22 | """Error message if the task failed.""" 23 | 24 | metadata: Optional[dict] = None 25 | """Additional provider-specific metadata.""" 26 | 27 | 28 | class LLMProvider(ABC): 29 | """Base interface for LLM providers that can generate MCP tool calls.""" 30 | 31 | @property 32 | @abstractmethod 33 | def name(self) -> str: 34 | """Return the name of this provider.""" 35 | pass 36 | 37 | @abstractmethod 38 | async def send_task( 39 | self, prompt: str, server_url: str, verbose: bool = False 40 | ) -> AsyncGenerator[str, None]: 41 | """Send a task to the LLM and yield progress messages. 42 | 43 | Args: 44 | prompt: The task prompt to send to the LLM 45 | server_url: The MCP server URL to use 46 | verbose: Whether to show detailed progress 47 | 48 | Yields 49 | ------ 50 | Progress messages as the LLM works 51 | """ 52 | pass 53 | 54 | @abstractmethod 55 | async def get_final_response(self) -> LLMResponse: 56 | """Get the final response after send_task completes. 57 | 58 | Returns 59 | ------- 60 | LLMResponse with the results of the task 61 | """ 62 | pass 63 | 64 | @abstractmethod 65 | async def cleanup(self): 66 | """Clean up any resources used by the provider.""" 67 | pass 68 | -------------------------------------------------------------------------------- /docs/src/components/HomepageFeatures/index.tsx: -------------------------------------------------------------------------------- 1 | import type {ReactNode} from 'react'; 2 | import clsx from 'clsx'; 3 | import Heading from '@theme/Heading'; 4 | import styles from './styles.module.css'; 5 | 6 | type FeatureItem = { 7 | title: string; 8 | Svg: React.ComponentType>; 9 | description: ReactNode; 10 | }; 11 | 12 | const FeatureList: FeatureItem[] = [ 13 | { 14 | title: 'Preserved State', 15 | Svg: () => null, 16 | description: ( 17 | <> 18 | Variables and data remain intact throughout your session. Your AI assistant 19 | can see errors, install packages, and continue where you left off. 20 | 21 | ), 22 | }, 23 | { 24 | title: 'MCP Protocol', 25 | Svg: () => null, 26 | description: ( 27 | <> 28 | Built on the Model Context Protocol, works with any MCP-compatible client 29 | including Goose, Cursor, and more. 30 | 31 | ), 32 | }, 33 | { 34 | title: 'Seamless Integration', 35 | Svg: () => null, 36 | description: ( 37 | <> 38 | Continue your data exploration naturally. Hand off to the AI assistant at 39 | any time to pick up exactly where you left off. 40 | 41 | ), 42 | }, 43 | ]; 44 | 45 | function Feature({title, Svg, description}: FeatureItem) { 46 | return ( 47 |
48 |
49 | 50 |
51 |
52 | {title} 53 |

{description}

54 |
55 |
56 | ); 57 | } 58 | 59 | export default function HomepageFeatures(): ReactNode { 60 | return ( 61 |
62 |
63 |
64 | {FeatureList.map((props, idx) => ( 65 | 66 | ))} 67 |
68 |
69 |
70 | ); 71 | } -------------------------------------------------------------------------------- /docs/docs/intro.md: -------------------------------------------------------------------------------- 1 | --- 2 | sidebar_position: 1 3 | slug: / 4 | --- 5 | 6 | # Introduction 7 | 8 | :::warning API Compatibility Notice 9 | This project is currently focused on MCP (Model Context Protocol) usage. There are **no API compatibility guarantees** between versions as the interface is actively evolving. Breaking changes may occur in any release. 10 | ::: 11 | 12 | MCP Jupyter Server allows you to use AI assistants like [Goose](https://block.github.io/goose/) or Cursor to pair with you in JupyterLab notebooks where the state of your variables is preserved by the JupyterLab Kernel. 13 | 14 | ## Why MCP Jupyter? 15 | 16 | The key advantage of MCP Jupyter is **state preservation**. This allows you to: 17 | 18 | - Work with your AI assistant in a notebook where variables and data remain intact 19 | - Let the AI see errors and install missing packages automatically 20 | - Do data exploration yourself, then hand off to the agent to continue 21 | - Maintain context throughout your entire session 22 | 23 | ## How It Works 24 | 25 | MCP Jupyter acts as a bridge between MCP-compatible AI clients and your JupyterLab server: 26 | 27 | ``` 28 | AI Client (Goose/Cursor) <--> MCP Jupyter Server <--> JupyterLab Kernel 29 | | 30 | Preserved State 31 | (Variables/Data) 32 | ``` 33 | 34 | This architecture ensures that your notebook state is maintained throughout the session, enabling seamless collaboration between you and your AI assistant. 35 | 36 | ## Key Features 37 | 38 | - **State Preservation**: Variables and data persist across interactions 39 | - **Error Handling**: AI can see and respond to errors in real-time 40 | - **Package Management**: Automatic package installation when needed 41 | - **Seamless Handoff**: Switch between manual and AI-assisted work anytime 42 | - **MCP Protocol**: Works with any MCP-compatible client 43 | - **Optimized Architecture**: 4 consolidated tools following MCP best practices 44 | - **Workflow-oriented Design**: Tools match AI collaboration patterns vs API endpoints 45 | 46 | ## Next Steps 47 | 48 | - [Quickstart Guide →](/docs/quickstart) 49 | - [Installation →](/docs/installation) 50 | - [Architecture →](/docs/architecture) 51 | - [Usage Examples →](/docs/usage) -------------------------------------------------------------------------------- /docs/src/css/custom.css: -------------------------------------------------------------------------------- 1 | /** 2 | * Any CSS included here will be global. The classic template 3 | * bundles Infima by default. Infima is a CSS framework designed to 4 | * work well for content-centric websites. 5 | */ 6 | 7 | /* You can override the default Infima variables here. */ 8 | :root { 9 | --ifm-color-primary: #0066ff; 10 | --ifm-color-primary-dark: #0052cc; 11 | --ifm-color-primary-darker: #004bb5; 12 | --ifm-color-primary-darkest: #003d99; 13 | --ifm-color-primary-light: #1a75ff; 14 | --ifm-color-primary-lighter: #3385ff; 15 | --ifm-color-primary-lightest: #66a3ff; 16 | --ifm-code-font-size: 95%; 17 | --docusaurus-highlighted-code-line-bg: rgba(0, 0, 0, 0.1); 18 | --ifm-hero-background-color: #f5f7fa; 19 | } 20 | 21 | /* For readability concerns, you should choose a lighter palette in dark mode. */ 22 | [data-theme='dark'] { 23 | --ifm-color-primary: #4d94ff; 24 | --ifm-color-primary-dark: #2e7eff; 25 | --ifm-color-primary-darker: #1f73ff; 26 | --ifm-color-primary-darkest: #005ce6; 27 | --ifm-color-primary-light: #6ca6ff; 28 | --ifm-color-primary-lighter: #7ab0ff; 29 | --ifm-color-primary-lightest: #99c2ff; 30 | --docusaurus-highlighted-code-line-bg: rgba(0, 0, 0, 0.3); 31 | --ifm-hero-background-color: #1a1a1a; 32 | } 33 | 34 | /* Hero customization */ 35 | .hero--primary { 36 | background-color: var(--ifm-hero-background-color); 37 | color: var(--ifm-font-color-base); 38 | } 39 | 40 | .hero__title { 41 | font-size: 3.5rem; 42 | font-weight: 700; 43 | margin-bottom: 1rem; 44 | } 45 | 46 | .hero__subtitle { 47 | font-size: 1.5rem; 48 | font-weight: 400; 49 | opacity: 0.8; 50 | max-width: 600px; 51 | margin: 0 auto 2rem; 52 | } 53 | 54 | /* Button styling similar to Goose */ 55 | .button { 56 | border-radius: 8px; 57 | font-weight: 600; 58 | padding: 12px 24px; 59 | transition: all 0.2s ease; 60 | } 61 | 62 | .button:hover { 63 | transform: translateY(-2px); 64 | box-shadow: 0 4px 12px rgba(0, 0, 0, 0.15); 65 | } 66 | 67 | /* Feature section styling */ 68 | .features { 69 | padding: 4rem 0; 70 | background-color: var(--ifm-background-surface-color); 71 | } 72 | 73 | .features h3 { 74 | font-size: 1.5rem; 75 | margin-bottom: 1rem; 76 | } 77 | 78 | /* Navbar styling */ 79 | .navbar__title { 80 | font-weight: 700; 81 | font-size: 1.25rem; 82 | } 83 | 84 | /* Code blocks */ 85 | pre { 86 | border-radius: 8px; 87 | } 88 | 89 | /* Responsive improvements */ 90 | @media (max-width: 768px) { 91 | .hero__title { 92 | font-size: 2.5rem; 93 | } 94 | 95 | .hero__subtitle { 96 | font-size: 1.25rem; 97 | } 98 | } -------------------------------------------------------------------------------- /tests/test_check_server.py: -------------------------------------------------------------------------------- 1 | """Test that check_server query type doesn't try to access notebooks.""" 2 | 3 | from unittest.mock import Mock, patch 4 | 5 | import pytest 6 | 7 | from mcp_jupyter.rest_client import NotebookClient 8 | from mcp_jupyter.server import query_notebook 9 | 10 | 11 | def test_check_server_without_notebook(jupyter_server): 12 | """Test that check_server works without a notebook path.""" 13 | # check_server should work even with a non-existent notebook path 14 | # because it doesn't actually access the notebook 15 | result = query_notebook( 16 | "non_existent_notebook", # This notebook doesn't exist 17 | "check_server", 18 | server_url=jupyter_server, 19 | ) 20 | assert result == "Jupyter server is running" 21 | 22 | 23 | def test_list_sessions_without_notebook(jupyter_server): 24 | """Test that list_sessions works without accessing a specific notebook.""" 25 | # list_sessions should work with any notebook path 26 | # because it doesn't actually access the notebook 27 | result = query_notebook( 28 | "any_notebook_name", # This notebook doesn't need to exist 29 | "list_sessions", 30 | server_url=jupyter_server, 31 | ) 32 | assert isinstance(result, list) 33 | 34 | 35 | def test_get_default_kernel_info(jupyter_server): 36 | """Test that kernel info can be retrieved from the server.""" 37 | # Create a client to test kernel info retrieval 38 | client = NotebookClient( 39 | server_url=jupyter_server, 40 | notebook_path="test_kernel.ipynb", # Doesn't need to exist for this test 41 | token="BLOCK", 42 | ) 43 | 44 | # Test getting kernel info 45 | kernelspec, language_info = client._get_default_kernel_info() 46 | 47 | # Verify kernelspec structure 48 | assert isinstance(kernelspec, dict) 49 | assert "display_name" in kernelspec 50 | assert "language" in kernelspec 51 | assert "name" in kernelspec 52 | 53 | # Verify language_info structure 54 | assert isinstance(language_info, dict) 55 | assert "name" in language_info 56 | 57 | # Verify content makes sense for a typical Jupyter setup 58 | assert kernelspec["language"] == "python" 59 | assert language_info["name"] == "python" 60 | assert "python" in kernelspec["display_name"].lower() 61 | 62 | 63 | def test_reconnect_interval_configured(): 64 | """Test that KernelClient is initialized with reconnect_interval=5.""" 65 | with patch("mcp_jupyter.server.KernelClient") as mock_kernel_client: 66 | with patch("mcp_jupyter.server.get_kernel_id", return_value="mock-kernel-id"): 67 | # Mock the kernel instance 68 | mock_instance = Mock() 69 | mock_instance.kernel_id = "mock-kernel-id" 70 | mock_kernel_client.return_value = mock_instance 71 | 72 | # Import and call get_kernel 73 | from mcp_jupyter.server import get_kernel 74 | 75 | get_kernel("test_notebook.ipynb", server_url="http://localhost:8888") 76 | 77 | # Verify KernelClient was called with reconnect_interval in client_kwargs 78 | mock_kernel_client.assert_called_once() 79 | call_kwargs = mock_kernel_client.call_args.kwargs 80 | assert "client_kwargs" in call_kwargs 81 | assert call_kwargs["client_kwargs"] == {"reconnect_interval": 5} 82 | -------------------------------------------------------------------------------- /docs/docs/quickstart.md: -------------------------------------------------------------------------------- 1 | --- 2 | sidebar_position: 2 3 | --- 4 | 5 | # Quickstart 6 | 7 | Get up and running with MCP Jupyter in minutes. 8 | 9 | ## Requirements 10 | 11 | - [UV](https://docs.astral.sh/uv/) - Required for installation 12 | - JupyterLab with `ipykernel` (required) 13 | - `jupyter-collaboration` extension (recommended for automatic sync) 14 | - An MCP-compatible client (e.g., [Goose](https://block.github.io/goose/), Cursor) 15 | 16 | ## Installation 17 | 18 | MCP Jupyter Server uses stdio and can be added to any MCP client with: 19 | 20 | ```bash 21 | uvx mcp-jupyter 22 | ``` 23 | 24 | ## Quick Setup 25 | 26 | ### 1. Start Jupyter Server 27 | 28 | First, set up and start your Jupyter server: 29 | 30 | import Tabs from '@theme/Tabs'; 31 | import TabItem from '@theme/TabItem'; 32 | 33 | 34 | 35 | 36 | ```bash 37 | # Using uv venv 38 | uv venv 39 | source .venv/bin/activate # On Windows: .venv\Scripts\activate 40 | uv pip install jupyterlab jupyter-collaboration ipykernel 41 | jupyter lab --port 8888 --IdentityProvider.token BLOCK --ip 0.0.0.0 42 | ``` 43 | 44 | ```bash 45 | # Using uv project 46 | uv init jupyter-workspace 47 | cd jupyter-workspace 48 | uv add jupyterlab jupyter-collaboration ipykernel 49 | uv run jupyter lab --port 8888 --IdentityProvider.token BLOCK --ip 0.0.0.0 50 | ``` 51 | 52 | 53 | 54 | 55 | ```bash 56 | # Using uv venv 57 | uv venv 58 | source .venv/bin/activate # On Windows: .venv\Scripts\activate 59 | uv pip install jupyterlab ipykernel 60 | jupyter lab --port 8888 --IdentityProvider.token BLOCK --ip 0.0.0.0 61 | ``` 62 | 63 | ```bash 64 | # Using uv project 65 | uv init jupyter-workspace 66 | cd jupyter-workspace 67 | uv add jupyterlab ipykernel 68 | uv run jupyter lab --port 8888 --IdentityProvider.token BLOCK --ip 0.0.0.0 69 | ``` 70 | 71 | :::caution Manual Sync Required 72 | Without `jupyter-collaboration`, you must manually sync changes: 73 | - **Agent makes changes** → Reload Notebook from Disk (File menu) 74 | - **You make edits** → Save (Ctrl+S / Cmd+S) 75 | ::: 76 | 77 | 78 | 79 | 80 | :::tip 81 | The server expects a token for authentication. If the `TOKEN` environment variable is not set, it defaults to "BLOCK". 82 | ::: 83 | 84 | ### 2. Configure Your MCP Client 85 | 86 | #### For Goose 87 | 88 | Add MCP Jupyter to your Goose configuration: 89 | 90 | ```bash 91 | goose session --with-extension "uvx mcp-jupyter" 92 | ``` 93 | 94 | #### For Other Clients 95 | 96 | Add the following to your MCP client configuration: 97 | 98 | ```json 99 | { 100 | "mcpServers": { 101 | "jupyter": { 102 | "command": "uvx", 103 | "args": ["mcp-jupyter"] 104 | } 105 | } 106 | } 107 | ``` 108 | 109 | ### 3. Start Using 110 | 111 | Once configured, you can: 112 | 113 | 1. Create or open a notebook through your AI assistant 114 | 2. Execute code cells with preserved state 115 | 3. Let the AI handle errors and install packages 116 | 4. Switch between manual and AI-assisted work seamlessly 117 | 118 | ## Example Session 119 | 120 | ```python 121 | # Your AI assistant can help you: 122 | # 1. Load and explore data 123 | # 2. Visualize results 124 | # 3. Debug errors 125 | # 4. Install missing packages 126 | # All while preserving your notebook state! 127 | ``` 128 | 129 | ## Next Steps 130 | 131 | - [Detailed Installation Guide →](/docs/installation) 132 | - [Architecture →](/docs/architecture) 133 | - [Usage Examples →](/docs/usage) 134 | - [Development Setup →](/docs/development) 135 | -------------------------------------------------------------------------------- /tests/llm_providers/claude_code.py: -------------------------------------------------------------------------------- 1 | """Claude Code provider for MCP tool call testing.""" 2 | 3 | from typing import AsyncGenerator 4 | 5 | from claude_code_sdk import AssistantMessage, TextBlock, query 6 | 7 | from .base import LLMProvider, LLMResponse 8 | 9 | 10 | class ClaudeCodeProvider(LLMProvider): 11 | """Provider that uses Claude Code SDK to generate MCP tool calls.""" 12 | 13 | def __init__(self): 14 | self._responses = [] 15 | self._message_count = 0 16 | self._verbose = False 17 | 18 | @property 19 | def name(self) -> str: 20 | """Return the provider name.""" 21 | return "claude-code" 22 | 23 | async def send_task( 24 | self, prompt: str, server_url: str, verbose: bool = False 25 | ) -> AsyncGenerator[str, None]: 26 | """Send task to Claude Code and yield progress messages.""" 27 | self._verbose = verbose 28 | self._responses = [] 29 | self._message_count = 0 30 | 31 | # Create the full task prompt 32 | full_prompt = f"""Using the MCP Jupyter server at {server_url}, please: 33 | 34 | {prompt} 35 | 36 | Make sure to execute the cells so we can see the output.""" 37 | 38 | if verbose: 39 | yield f"Sending query to Claude Code: {full_prompt}" 40 | yield "=" * 80 41 | yield "Claude Code is working..." 42 | yield "=" * 80 43 | 44 | # Send query to Claude Code and process responses 45 | async for message in query(prompt=full_prompt): 46 | self._message_count += 1 47 | 48 | # Filter out internal metadata messages 49 | message_type = type(message).__name__ 50 | if message_type == "ResultMessage": 51 | continue 52 | 53 | if verbose: 54 | yield f"\n--- Message {self._message_count} ---" 55 | yield f"Message type: {message_type}" 56 | 57 | if isinstance(message, AssistantMessage): 58 | for i, block in enumerate(message.content): 59 | if verbose: 60 | yield f"Block {i + 1} type: {type(block).__name__}" 61 | 62 | if isinstance(block, TextBlock): 63 | self._responses.append(block.text) 64 | if verbose: 65 | yield f"Claude: {block.text}" 66 | else: 67 | if verbose: 68 | yield f"Other content: {block}" 69 | else: 70 | if verbose: 71 | # Only show meaningful system messages, not all internal ones 72 | if message_type not in ["SystemMessage"]: 73 | yield f"Other message: {message}" 74 | 75 | if verbose: 76 | yield "=" * 80 77 | yield f"Claude Code finished. Total messages: {self._message_count}" 78 | yield "=" * 80 79 | 80 | async def get_final_response(self) -> LLMResponse: 81 | """Get the final response from Claude Code.""" 82 | success = len(self._responses) > 0 83 | full_text = "\n".join(self._responses) 84 | 85 | return LLMResponse( 86 | text=full_text, 87 | tool_calls_made=self._message_count, 88 | success=success, 89 | error=None if success else "No responses received from Claude Code", 90 | metadata={ 91 | "provider": "claude-code", 92 | "message_count": self._message_count, 93 | "response_count": len(self._responses), 94 | }, 95 | ) 96 | 97 | async def cleanup(self): 98 | """Clean up Claude Code provider resources.""" 99 | self._responses = [] 100 | self._message_count = 0 101 | -------------------------------------------------------------------------------- /tests/test_filtering.py: -------------------------------------------------------------------------------- 1 | """Tests for filtering functions that remove verbose output data.""" 2 | 3 | import pytest 4 | 5 | from mcp_jupyter.utils import filter_image_outputs 6 | 7 | 8 | class TestFilterImageOutputs: 9 | """Test filter_image_outputs function from utils.py.""" 10 | 11 | def test_filter_png_image(self): 12 | """Test filtering of PNG image data.""" 13 | outputs = [ 14 | { 15 | "output_type": "display_data", 16 | "data": { 17 | "text/plain": ["
"], 18 | "image/png": "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg==", 19 | }, 20 | } 21 | ] 22 | 23 | filtered = filter_image_outputs(outputs) 24 | 25 | assert len(filtered) == 1 26 | assert "image/png" not in filtered[0]["data"] 27 | assert "text/plain" in filtered[0]["data"] 28 | text_plain = filtered[0]["data"]["text/plain"] 29 | assert isinstance(text_plain, list) 30 | assert "Image generated (PNG format)" in "".join(text_plain) 31 | 32 | def test_filter_multiple_image_formats(self): 33 | """Test filtering of multiple image formats.""" 34 | outputs = [ 35 | { 36 | "output_type": "execute_result", 37 | "data": { 38 | "text/plain": [""], 39 | "image/png": "base64_png_data_here", 40 | "image/jpeg": "base64_jpeg_data_here", 41 | "image/svg+xml": "...", 42 | }, 43 | } 44 | ] 45 | 46 | filtered = filter_image_outputs(outputs) 47 | 48 | assert len(filtered) == 1 49 | data = filtered[0]["data"] 50 | assert "image/png" not in data 51 | assert "image/jpeg" not in data 52 | assert "image/svg+xml" not in data 53 | assert "text/plain" in data 54 | text_plain = data["text/plain"] 55 | assert isinstance(text_plain, list) 56 | assert "Image generated (PNG, JPEG, SVG+XML format)" in "".join(text_plain) 57 | 58 | def test_preserve_non_image_data(self): 59 | """Test that non-image data is preserved.""" 60 | outputs = [ 61 | {"output_type": "stream", "name": "stdout", "text": ["Hello World\n"]}, 62 | { 63 | "output_type": "execute_result", 64 | "data": {"text/plain": ["42"], "text/html": ["42"]}, 65 | }, 66 | ] 67 | 68 | filtered = filter_image_outputs(outputs) 69 | 70 | assert len(filtered) == 2 71 | # Stream output should be unchanged 72 | assert filtered[0] == outputs[0] 73 | # Execute result without images should be unchanged 74 | assert filtered[1] == outputs[1] 75 | 76 | def test_no_data_field(self): 77 | """Test outputs without data field.""" 78 | outputs = [ 79 | { 80 | "output_type": "display_data" 81 | # No data field 82 | } 83 | ] 84 | 85 | filtered = filter_image_outputs(outputs) 86 | 87 | assert len(filtered) == 1 88 | assert filtered[0] == outputs[0] 89 | 90 | def test_create_text_plain_when_missing(self): 91 | """Test creating text/plain field when it doesn't exist.""" 92 | outputs = [ 93 | {"output_type": "display_data", "data": {"image/png": "base64_data_here"}} 94 | ] 95 | 96 | filtered = filter_image_outputs(outputs) 97 | 98 | assert len(filtered) == 1 99 | assert "image/png" not in filtered[0]["data"] 100 | assert filtered[0]["data"]["text/plain"] == "Image generated (PNG format)" 101 | -------------------------------------------------------------------------------- /tests/test_llm_tool_calls.py: -------------------------------------------------------------------------------- 1 | """Test that an LLM can generate correct MCP tool calls for notebook operations.""" 2 | 3 | import pytest 4 | 5 | from mcp_jupyter.server import query_notebook 6 | 7 | from .llm_providers.config import get_test_providers 8 | 9 | LLM_PROVIDERS = get_test_providers() 10 | 11 | 12 | def get_test_task() -> str: 13 | return """ 14 | 1. Check if the Jupyter server is running 15 | 2. Create a new notebook called 'math_functions' 16 | 3. Add a function that calculates the area of a circle 17 | 4. Add another function that calculates the factorial of a number 18 | 5. Edit the first cell to add a proper docstring and test the function 19 | """ 20 | 21 | 22 | @pytest.mark.llm 23 | @pytest.mark.asyncio 24 | @pytest.mark.parametrize("provider", LLM_PROVIDERS, ids=lambda p: p.name) 25 | async def test_llm_generates_correct_tool_calls(llm_jupyter_server, provider): 26 | """Test that an LLM provider can generate correct MCP tool calls.""" 27 | print(f"\n{'=' * 80}") 28 | print(f"Testing {provider.name.upper()} provider") 29 | print(f"{'=' * 80}") 30 | 31 | try: 32 | task_prompt = get_test_task() 33 | 34 | async for progress_message in provider.send_task( 35 | prompt=task_prompt, server_url=llm_jupyter_server, verbose=True 36 | ): 37 | print(progress_message) 38 | 39 | response = await provider.get_final_response() 40 | 41 | print(f"\n{'=' * 80}") 42 | print(f"{provider.name.upper()} RESPONSE SUMMARY") 43 | print(f"{'=' * 80}") 44 | print(f"Success: {response.success}") 45 | print(f"Tool calls made: {response.tool_calls_made}") 46 | print(f"Error: {response.error or 'None'}") 47 | print(f"Response text length: {len(response.text)} characters") 48 | if response.metadata: 49 | print(f"Metadata: {response.metadata}") 50 | 51 | assert response.success, f"{provider.name} should have completed successfully" 52 | assert response.tool_calls_made > 0, ( 53 | f"{provider.name} should have made tool calls" 54 | ) 55 | 56 | all_cells = query_notebook( 57 | "math_functions", "view_source", server_url=llm_jupyter_server 58 | ) 59 | assert len(all_cells) >= 2, f"Expected at least 2 cells, got {len(all_cells)}" 60 | 61 | function_cells = [ 62 | cell 63 | for cell in all_cells 64 | if cell.get("cell_type") == "code" and "def " in cell.get("source", "") 65 | ] 66 | assert len(function_cells) >= 2, ( 67 | f"Expected at least 2 function definitions, got {len(function_cells)}" 68 | ) 69 | 70 | source_text = " ".join([cell.get("source", "") for cell in all_cells]) 71 | assert "area" in source_text.lower() or "circle" in source_text.lower(), ( 72 | "Expected circle area function" 73 | ) 74 | assert "factorial" in source_text.lower(), "Expected factorial function" 75 | 76 | executed_cells = [ 77 | cell for cell in all_cells if cell.get("execution_count") is not None 78 | ] 79 | assert len(executed_cells) > 0, "Expected at least one cell to be executed" 80 | 81 | first_code_cell = next( 82 | (cell for cell in all_cells if cell.get("cell_type") == "code"), None 83 | ) 84 | if first_code_cell: 85 | source = first_code_cell.get("source", "") 86 | has_docstring = '"""' in source or "'''" in source or 'r"""' in source 87 | if has_docstring: 88 | print("✓ Found docstring in first cell as requested") 89 | 90 | print(f"\n{'=' * 80}") 91 | print(f"{provider.name.upper()} SUCCESS SUMMARY") 92 | print(f"{'=' * 80}") 93 | print(f"✓ Successfully created notebook with {len(all_cells)} cells") 94 | print(f"✓ Found {len(function_cells)} function definitions") 95 | print(f"✓ Found {len(executed_cells)} executed cells") 96 | print(f"✓ {provider.name} successfully generated MCP tool calls!") 97 | 98 | finally: 99 | await provider.cleanup() 100 | -------------------------------------------------------------------------------- /docs/docusaurus.config.ts: -------------------------------------------------------------------------------- 1 | import {themes as prismThemes} from 'prism-react-renderer'; 2 | import type {Config} from '@docusaurus/types'; 3 | import type * as Preset from '@docusaurus/preset-classic'; 4 | 5 | // This runs in Node.js - Don't use client-side code here (browser APIs, JSX...) 6 | 7 | const config: Config = { 8 | title: 'MCP Jupyter', 9 | tagline: 'Your AI assistant in JupyterLab, preserving notebook state seamlessly', 10 | favicon: 'img/block-jewel_black.svg', 11 | 12 | // Set the production url of your site here 13 | url: 'https://block.github.io', 14 | // Set the // pathname under which your site is served 15 | // For GitHub pages deployment, it is often '//' 16 | baseUrl: '/mcp-jupyter/', 17 | 18 | // GitHub pages deployment config. 19 | // If you aren't using GitHub pages, you don't need these. 20 | organizationName: 'block', // Usually your GitHub org/user name. 21 | projectName: 'mcp-jupyter', // Usually your repo name. 22 | 23 | onBrokenLinks: 'throw', 24 | onBrokenMarkdownLinks: 'warn', 25 | 26 | // Even if you don't use internationalization, you can use this field to set 27 | // useful metadata like html lang. For example, if your site is Chinese, you 28 | // may want to replace "en" with "zh-Hans". 29 | i18n: { 30 | defaultLocale: 'en', 31 | locales: ['en'], 32 | }, 33 | 34 | presets: [ 35 | [ 36 | 'classic', 37 | { 38 | docs: { 39 | sidebarPath: './sidebars.ts', 40 | // Please change this to your repo. 41 | // Remove this to remove the "edit this page" links. 42 | editUrl: 43 | 'https://github.com/block/mcp-jupyter/tree/main/docs/', 44 | }, 45 | blog: false, 46 | theme: { 47 | customCss: './src/css/custom.css', 48 | }, 49 | } satisfies Preset.Options, 50 | ], 51 | ], 52 | 53 | themes: [ 54 | // Local search theme 55 | [ 56 | require.resolve("@easyops-cn/docusaurus-search-local"), 57 | { 58 | hashed: true, 59 | language: ["en"], 60 | highlightSearchTermsOnTargetPage: true, 61 | explicitSearchResultPath: true, 62 | }, 63 | ], 64 | ], 65 | 66 | themeConfig: { 67 | // Replace with your project's social card 68 | image: 'img/docusaurus-social-card.jpg', 69 | navbar: { 70 | title: 'MCP Jupyter', 71 | logo: { 72 | alt: 'MCP Jupyter Logo', 73 | src: 'img/block-jewel_black.svg', 74 | srcDark: 'img/block-jewel_white.svg', 75 | }, 76 | items: [ 77 | { 78 | type: 'docSidebar', 79 | sidebarId: 'tutorialSidebar', 80 | position: 'left', 81 | label: 'Docs', 82 | }, 83 | { 84 | to: '/docs/quickstart', 85 | label: 'Quickstart', 86 | position: 'left', 87 | }, 88 | { 89 | type: 'search', 90 | position: 'right', 91 | }, 92 | { 93 | href: 'https://github.com/block/mcp-jupyter', 94 | label: 'GitHub', 95 | position: 'right', 96 | }, 97 | ], 98 | }, 99 | footer: { 100 | style: 'dark', 101 | links: [ 102 | { 103 | title: 'Docs', 104 | items: [ 105 | { 106 | label: 'Quickstart', 107 | to: '/docs/quickstart', 108 | }, 109 | { 110 | label: 'Installation', 111 | to: '/docs/installation', 112 | }, 113 | ], 114 | }, 115 | { 116 | title: 'Community', 117 | items: [ 118 | { 119 | label: 'GitHub', 120 | href: 'https://github.com/block/mcp-jupyter', 121 | }, 122 | { 123 | label: 'Goose', 124 | href: 'https://block.github.io/goose/', 125 | }, 126 | ], 127 | }, 128 | ], 129 | copyright: `Copyright © ${new Date().getFullYear()} Block, Inc.`, 130 | }, 131 | prism: { 132 | theme: prismThemes.github, 133 | darkTheme: prismThemes.dracula, 134 | }, 135 | } satisfies Preset.ThemeConfig, 136 | }; 137 | 138 | export default config; 139 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Distribution / packaging 10 | .Python 11 | build/ 12 | develop-eggs/ 13 | dist/ 14 | downloads/ 15 | eggs/ 16 | .eggs/ 17 | .hermit/ 18 | lib/ 19 | lib64/ 20 | parts/ 21 | sdist/ 22 | var/ 23 | wheels/ 24 | share/python-wheels/ 25 | *.egg-info/ 26 | .installed.cfg 27 | *.egg 28 | MANIFEST 29 | 30 | # PyInstaller 31 | # Usually these files are written by a python script from a template 32 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 33 | *.manifest 34 | *.spec 35 | 36 | # Installer logs 37 | pip-log.txt 38 | pip-delete-this-directory.txt 39 | 40 | # Unit test / coverage reports 41 | htmlcov/ 42 | .tox/ 43 | .nox/ 44 | .coverage 45 | .coverage.* 46 | .cache 47 | nosetests.xml 48 | coverage.xml 49 | *.cover 50 | *.py,cover 51 | .hypothesis/ 52 | .pytest_cache/ 53 | cover/ 54 | 55 | # Translations 56 | *.mo 57 | *.pot 58 | 59 | # Django stuff: 60 | *.log 61 | local_settings.py 62 | db.sqlite3 63 | db.sqlite3-journal 64 | 65 | # Flask stuff: 66 | instance/ 67 | .webassets-cache 68 | 69 | # Scrapy stuff: 70 | .scrapy 71 | 72 | # Sphinx documentation 73 | docs/_build/ 74 | 75 | # PyBuilder 76 | .pybuilder/ 77 | target/ 78 | 79 | # Jupyter Notebook 80 | .ipynb_checkpoints 81 | 82 | # IPython 83 | profile_default/ 84 | ipython_config.py 85 | 86 | # pyenv 87 | # For a library or package, you might want to ignore these files since the code is 88 | # intended to run in multiple environments; otherwise, check them in: 89 | # .python-version 90 | 91 | # pipenv 92 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 93 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 94 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 95 | # install all needed dependencies. 96 | #Pipfile.lock 97 | 98 | # UV 99 | # Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control. 100 | # This is especially recommended for binary packages to ensure reproducibility, and is more 101 | # commonly ignored for libraries. 102 | uv.lock 103 | 104 | # poetry 105 | # Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. 106 | # This is especially recommended for binary packages to ensure reproducibility, and is more 107 | # commonly ignored for libraries. 108 | # https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control 109 | #poetry.lock 110 | 111 | # pdm 112 | # Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control. 113 | #pdm.lock 114 | # pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it 115 | # in version control. 116 | # https://pdm.fming.dev/latest/usage/project/#working-with-version-control 117 | .pdm.toml 118 | .pdm-python 119 | .pdm-build/ 120 | 121 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm 122 | __pypackages__/ 123 | 124 | # Celery stuff 125 | celerybeat-schedule 126 | celerybeat.pid 127 | 128 | # SageMath parsed files 129 | *.sage.py 130 | 131 | # Environments 132 | .env 133 | .venv 134 | env/ 135 | venv/ 136 | ENV/ 137 | env.bak/ 138 | venv.bak/ 139 | 140 | # Spyder project settings 141 | .spyderproject 142 | .spyproject 143 | 144 | # Rope project settings 145 | .ropeproject 146 | 147 | # mkdocs documentation 148 | /site 149 | 150 | # mypy 151 | .mypy_cache/ 152 | .dmypy.json 153 | dmypy.json 154 | 155 | # Pyre type checker 156 | .pyre/ 157 | 158 | # pytype static type analyzer 159 | .pytype/ 160 | 161 | # Cython debug symbols 162 | cython_debug/ 163 | 164 | # PyCharm 165 | # JetBrains specific template is maintained in a separate JetBrains.gitignore that can 166 | # be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore 167 | # and can be added to the global gitignore or merged into this file. For a more nuclear 168 | # option (not recommended) you can uncomment the following to ignore the entire idea folder. 169 | #.idea/ 170 | 171 | # Ruff stuff: 172 | .ruff_cache/ 173 | 174 | # PyPI configuration file 175 | .pypirc 176 | -------------------------------------------------------------------------------- /docs/docs/usage.md: -------------------------------------------------------------------------------- 1 | --- 2 | sidebar_position: 4 3 | --- 4 | 5 | # Usage 6 | 7 | Learn how to use MCP Jupyter effectively with your AI assistant. 8 | 9 | ## Basic Usage 10 | 11 | ### Creating a New Notebook 12 | 13 | Ask your AI assistant to create a notebook: 14 | 15 | ``` 16 | "Create a new notebook called data_analysis.ipynb" 17 | ``` 18 | 19 | The AI will: 20 | 1. Create the notebook file 21 | 2. Start a kernel 22 | 3. Be ready for your commands 23 | 24 | ### Working with Existing Notebooks 25 | 26 | ``` 27 | "Open the notebook experiments/model_training.ipynb" 28 | ``` 29 | 30 | Your AI assistant will connect to the existing notebook and preserve all current state. 31 | 32 | ## Key Features 33 | 34 | ### State Preservation 35 | All variables, data, and models remain available throughout your session. Work with large datasets without reloading, and hand off complex objects between you and the AI. 36 | 37 | ### Automatic Error Recovery 38 | The AI sees execution errors in real-time and can automatically install missing packages, fix syntax issues, or suggest corrections. 39 | 40 | ### Seamless Collaboration 41 | Switch between manual exploration and AI assistance at any point. The AI builds on your work, and you can take over whenever needed. 42 | 43 | ### Smart Package Management 44 | Missing dependencies are automatically detected and installed, so your workflow isn't interrupted by import errors. 45 | 46 | ## Common Use Cases 47 | 48 | MCP Jupyter excels at collaborative data work. Here are popular use cases: 49 | 50 | ### Data Analysis & Exploration 51 | - **Data cleaning & profiling**: "Handle missing values, outliers, and analyze data quality" 52 | - **Exploratory analysis**: "Show me key patterns, distributions, and statistical summaries" 53 | - **Trend analysis**: "Plot time series trends with seasonality and correlations" 54 | 55 | ### Machine Learning & Modeling 56 | - **End-to-end ML pipeline**: "Prepare data, engineer features, and compare multiple algorithms" 57 | - **Model optimization**: "Tune hyperparameters and evaluate performance comprehensively" 58 | - **Experiment analysis**: "Analyze A/B tests and statistical significance" 59 | 60 | ### Data Visualization & Reporting 61 | - **Automated visualization**: "Create appropriate charts and statistical plots for this data" 62 | - **Custom dashboards**: "Build interactive visualizations and reports" 63 | - **Anomaly detection**: "Identify and visualize unusual patterns" 64 | 65 | ### Research & Advanced Analysis 66 | - **Hypothesis testing**: "Test statistical differences and relationships between variables" 67 | - **Cohort & behavioral analysis**: "Track user patterns and segment analysis over time" 68 | - **Concept exploration**: "Demonstrate and compare different analytical methods" 69 | 70 | ### Workflow Automation 71 | - **Data pipelines**: "Create repeatable ETL processes and data validation workflows" 72 | - **Report automation**: "Generate recurring analysis reports with charts and summaries" 73 | - **Code assistance**: "Debug analysis code and explain complex statistical concepts" 74 | 75 | ## Best Practices 76 | 77 | ### 1. Clear Instructions 78 | 79 | Be specific about what you want: 80 | - ❌ "Analyze the data" 81 | - ✅ "Perform exploratory data analysis focusing on customer segments and seasonal patterns" 82 | 83 | ### 2. Specify Cell Types Clearly 84 | 85 | Help the AI choose the right cell type and operation: 86 | - **For code**: "Add a code cell that loads the data" 87 | - **For markdown**: "Create a markdown cell with the project title and description" 88 | - **For mixed content**: "Add a markdown cell explaining the analysis, then add code to implement it" 89 | 90 | ### 3. Handle Operation Errors 91 | 92 | Common AI mistakes and corrections: 93 | - ❌ AI says "edit_markdown" → ✅ Should be `operation="add_markdown"` or `operation="edit_markdown"` 94 | - ❌ Putting ASCII art in code cells → ✅ "Put that ASCII art in a markdown cell instead" 95 | - ❌ IndentationError on non-code content → ✅ "That content belongs in markdown, not code" 96 | 97 | ### 4. Iterative Refinement 98 | 99 | Work iteratively with the AI: 100 | ``` 101 | 1. "Load and preview the customer data" 102 | 2. Review the output 103 | 3. "Focus on customers from the last quarter" 104 | 4. "Now segment them by purchase frequency" 105 | ``` 106 | 107 | ### 5. State Management 108 | 109 | - Keep important variables in the global namespace 110 | - Use descriptive variable names 111 | - Periodically check available variables with `dir()` or `locals()` 112 | 113 | ### 6. Error Recovery 114 | 115 | When errors occur: 116 | - Let the AI see and handle the error 117 | - Clarify cell type if there's confusion: "That should be markdown, not code" 118 | - Provide context if needed 119 | - The AI will install packages or fix issues automatically 120 | 121 | ## Demo Example 122 | 123 | ![MCP Jupyter Demo](/demos/goose-demo.png) 124 | 125 | [View the generated notebook →](https://github.com/block/mcp-jupyter/blob/main/demos/demo.ipynb) 126 | 127 | ## Tips and Tricks 128 | 129 | 1. **Use Markdown cells**: Ask the AI to document its analysis 130 | 2. **Save checkpoints**: Periodically save important state 131 | 3. **Combine approaches**: Use AI for boilerplate, manually tune details 132 | 4. **Leverage errors**: Let errors guide package installation 133 | 5. **Incremental development**: Build complex analyses step by step 134 | 135 | ## Next Steps 136 | 137 | - [Development Guide →](/docs/development) 138 | -------------------------------------------------------------------------------- /src/mcp_jupyter/utils.py: -------------------------------------------------------------------------------- 1 | import hashlib 2 | import json 3 | import logging 4 | import os 5 | import re 6 | import signal 7 | import subprocess 8 | import time 9 | from collections import defaultdict 10 | from contextlib import contextmanager 11 | from functools import wraps 12 | from typing import List, Optional, Union 13 | 14 | import requests 15 | from jupyter_kernel_client import KernelClient 16 | from mcp.server.fastmcp import FastMCP 17 | from mcp.shared.exceptions import McpError 18 | from mcp.types import INTERNAL_ERROR, INVALID_PARAMS, ErrorData 19 | from rich.console import Console 20 | from rich.logging import RichHandler 21 | 22 | TOKEN = os.getenv("TOKEN", "BLOCK") 23 | 24 | 25 | def _ensure_ipynb_extension(notebook_path: str) -> str: 26 | """Ensure the notebook path has the .ipynb extension. 27 | 28 | Args: 29 | notebook_path: Path to a notebook file 30 | 31 | Returns 32 | ------- 33 | str: The notebook path with .ipynb extension 34 | """ 35 | if not notebook_path.endswith(".ipynb"): 36 | return f"{notebook_path}.ipynb" 37 | return notebook_path 38 | 39 | 40 | def filter_image_outputs(outputs: List[dict]) -> List[dict]: 41 | """Filter out base64 images and replace with text indicators. 42 | 43 | Args: 44 | outputs: List of output dictionaries from cell execution 45 | 46 | Returns 47 | ------- 48 | List[dict]: Filtered outputs with images replaced by text indicators 49 | """ 50 | filtered_outputs = [] 51 | 52 | for output in outputs: 53 | # Create a copy of the output to avoid modifying the original 54 | filtered_output = output.copy() 55 | 56 | # Check for image data in display_data or execute_result outputs 57 | if output.get("output_type") in ["display_data", "execute_result"]: 58 | data = output.get("data", {}) 59 | if data: 60 | # Create a copy of data to avoid modifying the original 61 | filtered_data = data.copy() 62 | 63 | # Check for various image formats and replace with text indicators 64 | image_types = ["image/png", "image/jpeg", "image/svg+xml", "image/gif"] 65 | images_found = [] 66 | 67 | for img_type in image_types: 68 | if img_type in filtered_data: 69 | # Remove the base64 image data 70 | del filtered_data[img_type] 71 | images_found.append(img_type.split("/")[1].upper()) 72 | 73 | if images_found: 74 | # Add a text indicator for the removed images 75 | image_indicator = ( 76 | f"Image generated ({', '.join(images_found)} format)" 77 | ) 78 | if "text/plain" in filtered_data: 79 | # If there's already text/plain content (like "
"), 80 | # append the indicator to show image was filtered 81 | existing_text = filtered_data["text/plain"] 82 | 83 | # Handle both string and list formats 84 | if isinstance(existing_text, list): 85 | existing_text_str = "".join(existing_text) 86 | else: 87 | existing_text_str = existing_text 88 | 89 | if "Figure" in existing_text_str or "Axes" in existing_text_str: 90 | # Keep the existing figure description and add our indicator 91 | if isinstance(existing_text, list): 92 | filtered_data["text/plain"] = existing_text + [ 93 | f" - {image_indicator}" 94 | ] 95 | else: 96 | filtered_data["text/plain"] = ( 97 | existing_text + f" - {image_indicator}" 98 | ) 99 | else: 100 | # For other text, append on new line 101 | if isinstance(existing_text, list): 102 | filtered_data["text/plain"] = existing_text + [ 103 | f"\n{image_indicator}" 104 | ] 105 | else: 106 | filtered_data["text/plain"] = ( 107 | existing_text + f"\n{image_indicator}" 108 | ) 109 | else: 110 | # Create new text/plain output 111 | filtered_data["text/plain"] = image_indicator 112 | 113 | filtered_output["data"] = filtered_data 114 | 115 | filtered_outputs.append(filtered_output) 116 | 117 | return filtered_outputs 118 | 119 | 120 | def extract_output(output: dict) -> str: 121 | """Extract output from a Jupyter notebook cell. 122 | 123 | Args: 124 | output: Output dictionary from cell execution 125 | 126 | Returns 127 | ------- 128 | str: The extracted output text. For different output types: 129 | - display_data: returns data["text/plain"] 130 | - execute_result: returns data["text/plain"] 131 | - stream: returns text 132 | - error: returns traceback 133 | - other: returns empty string 134 | 135 | Raises 136 | ------ 137 | KeyError: If required keys are missing from the output dictionary 138 | """ 139 | if output["output_type"] == "display_data": 140 | return output["data"]["text/plain"] 141 | elif output["output_type"] == "execute_result": 142 | return output["data"]["text/plain"] 143 | elif output["output_type"] == "stream": 144 | return output["text"] 145 | elif output["output_type"] == "error": 146 | return output["traceback"] 147 | else: 148 | return "" 149 | -------------------------------------------------------------------------------- /docs/docs/architecture.md: -------------------------------------------------------------------------------- 1 | --- 2 | sidebar_position: 5 3 | --- 4 | 5 | # Architecture 6 | 7 | Understanding the design and structure of MCP Jupyter. 8 | 9 | ## Synchronization Architecture 10 | 11 | MCP Jupyter uses a **hybrid approach** combining REST API operations with RTC (Real-Time Collaboration) infrastructure: 12 | 13 | ### Agent Operations (REST-based) 14 | ``` 15 | AI Agent → MCP Server → REST Client → Jupyter REST API → Notebook File 16 | ``` 17 | 18 | - **MCP server operations** use Jupyter's REST API for reliability 19 | - Changes are **explicitly saved** to the notebook file on disk 20 | - Provides **consistent, stateful** modifications 21 | - Handles **error recovery** and **validation** cleanly 22 | 23 | ### User Collaboration (RTC-enabled) 24 | ``` 25 | User Browser ← RTC WebSocket ← Jupyter Server ← Notebook File → REST API → MCP Agent 26 | ``` 27 | 28 | - **User's JupyterLab interface** uses RTC (Real-Time Collaboration) via WebSockets 29 | - **Automatic synchronization** keeps user's view up-to-date with agent changes 30 | - **Real-time updates** appear in user's browser without refresh 31 | - **Collaborative editing** maintains consistency across multiple clients 32 | 33 | ### Benefits of This Architecture 34 | 35 | 1. **Reliability**: REST operations provide atomic, validated changes 36 | 2. **Real-time sync**: Users see agent changes immediately via RTC 37 | 3. **State preservation**: Kernel state and variables persist across operations 38 | 4. **Error handling**: REST provides clear error responses for debugging 39 | 5. **Compatibility**: Works with existing Jupyter infrastructure 40 | 41 | This design allows the **agent to make reliable changes** while ensuring **users stay synchronized** through Jupyter's built-in collaboration features. 42 | 43 | ## Tool Design 44 | 45 | MCP Jupyter uses **4 consolidated tools** (reduced from 11): 46 | 47 | ### 1. `query_notebook` - Read Operations 48 | All read-only operations for querying notebook information: 49 | - `view_source` - View cell source code (single cell or all) 50 | - `check_server` - Check if Jupyter server is accessible 51 | - `list_sessions` - List all notebook sessions 52 | - `get_position_index` - Get cell index by execution count or cell ID 53 | 54 | Uses `query_type` parameter to specify which operation to perform. 55 | 56 | ### 2. `modify_notebook_cells` - Cell Operations 57 | All cell modification operations: 58 | - `add_code` - Add and optionally execute code cells 59 | - `edit_code` - Edit existing code cells 60 | - `add_markdown` - Add markdown cells 61 | - `edit_markdown` - Edit existing markdown cells 62 | - `delete` - Delete cells 63 | 64 | Uses `operation` parameter to specify which action to perform. 65 | 66 | ### 3. `execute_notebook_code` - Execution Operations 67 | All code execution operations: 68 | - `execute_cell` - Execute existing code cells 69 | - `install_packages` - Install packages using uv pip 70 | 71 | Uses `execution_type` parameter to specify the type of execution. 72 | 73 | ### 4. `setup_notebook` - Initialization 74 | Notebook setup and kernel connection: 75 | - Creates new notebooks if needed 76 | - Connects to existing kernels 77 | - Manages notebook sessions 78 | 79 | ## Key Components 80 | 81 | ### 1. MCP Server (`server.py`) 82 | - Handles MCP protocol with consolidated tools 83 | - Manages parameter routing and validation 84 | - Provides float-to-int conversion for position_index 85 | - Routes requests to appropriate internal functions 86 | 87 | ### 2. REST Notebook Client (`rest_client.py`) 88 | - **Primary interface** for MCP server operations 89 | - Uses Jupyter's REST API for reliable cell modifications 90 | - Handles notebook CRUD operations (create, read, update, delete) 91 | - Provides compatibility layer with existing RTC-based code 92 | - Ensures consistent state management through explicit save operations 93 | 94 | ### 3. Notebook Manager (`notebook.py`) 95 | - Creates and manages notebooks on Jupyter server 96 | - Handles kernel lifecycle management 97 | - Manages notebook sessions 98 | - Handles directory creation for nested paths 99 | 100 | ### 4. State Tracker (`state.py`) 101 | - Tracks notebook state changes 102 | - Manages state consistency between operations 103 | - Provides decorators: 104 | - `@state_dependent` - Validates state before operations 105 | - `@refreshes_state` - Updates state after operations 106 | 107 | ### 5. Utilities (`utils.py`) 108 | - Helper functions for path handling 109 | - Parameter extraction and validation 110 | - File extension management (.ipynb) 111 | - Image filtering for output optimization 112 | 113 | ## Data Flow 114 | 115 | ``` 116 | AI Client → MCP Server → Tool Router → Internal Function → REST Client → Jupyter Server 117 | ↓ ↓ 118 | Parameter Validation Explicit Save 119 | ↓ ↓ 120 | State Management RTC Broadcast 121 | ↓ ↓ 122 | Response Formatting User Browser Update 123 | ``` 124 | 125 | ## State Management 126 | 127 | MCP Jupyter maintains notebook state consistency through: 128 | 129 | 1. **State Hashing** - Tracks changes to notebook content 130 | 2. **Dependency Decorators** - Ensures operations use current state 131 | 3. **Server URL Tracking** - Maps notebooks to their Jupyter servers 132 | 4. **Kernel Management** - Maintains kernel connections 133 | 134 | ## Adding New Operations 135 | 136 | To extend functionality, add new operations to existing tools: 137 | 138 | ```python 139 | # In query_notebook 140 | if query_type == "my_new_query": 141 | return _query_my_new_operation(notebook_path, **parameters) 142 | 143 | # In modify_notebook_cells 144 | if operation == "my_new_operation": 145 | return _modify_my_new_operation(notebook_path, **parameters) 146 | ``` 147 | 148 | ## Error Handling 149 | 150 | - **Parameter Validation** - Type checking and conversion 151 | - **Connection Errors** - Jupyter server connectivity issues 152 | - **State Mismatches** - Notebook state inconsistencies 153 | - **Execution Failures** - Kernel execution problems 154 | 155 | All errors provide actionable messages to help users recover. 156 | -------------------------------------------------------------------------------- /docs/docs/development.md: -------------------------------------------------------------------------------- 1 | --- 2 | sidebar_position: 6 3 | --- 4 | 5 | # Development 6 | 7 | Set up MCP Jupyter for development and contribution. 8 | 9 | ## Development Setup 10 | 11 | ### 1. Clone the Repository 12 | 13 | ```bash 14 | mkdir ~/Development 15 | cd ~/Development 16 | git clone https://github.com/block/mcp-jupyter.git 17 | cd mcp-jupyter 18 | ``` 19 | 20 | ### 2. Create Development Environment 21 | 22 | ```bash 23 | # Sync all dependencies including dev tools 24 | uv sync 25 | ``` 26 | 27 | ### 3. Run Tests 28 | 29 | ```bash 30 | # Run all tests 31 | uv run pytest tests/ 32 | 33 | # Run with coverage 34 | uv run pytest --cov=mcp_jupyter tests/ 35 | 36 | # Run specific test file 37 | uv run pytest tests/test_integration.py 38 | 39 | # Run LLM tool call generation tests 40 | uv run pytest -m llm -v 41 | 42 | # Run all tests except LLM tests (default behavior) 43 | uv run pytest -v 44 | ``` 45 | 46 | ## Using Development Version 47 | 48 | ### With Goose 49 | 50 | For development, use the local installation: 51 | 52 | ```bash 53 | goose session --with-extension "uv run --directory $(pwd) mcp-jupyter" 54 | ``` 55 | 56 | This allows you to make changes and test them immediately by restarting Goose. 57 | 58 | ### With Other Clients 59 | 60 | Update your MCP configuration to point to your local installation: 61 | 62 | ```json 63 | { 64 | "mcpServers": { 65 | "jupyter": { 66 | "command": "uv", 67 | "args": ["run", "--directory", "/path/to/mcp-jupyter", "mcp-jupyter"], 68 | "env": { 69 | "TOKEN": "your-token-here" 70 | } 71 | } 72 | } 73 | } 74 | ``` 75 | 76 | ## Project Structure 77 | 78 | ``` 79 | mcp-jupyter/ 80 | ├── src/ 81 | │ └── mcp_jupyter/ 82 | │ ├── __init__.py 83 | │ ├── __main__.py # Entry point 84 | │ ├── server.py # MCP server implementation 85 | │ ├── notebook.py # Notebook operations 86 | │ ├── jupyter.py # Jupyter integration 87 | │ ├── state.py # State management 88 | │ └── utils.py # Utilities 89 | ├── tests/ 90 | │ ├── test_integration.py # Integration tests with real Jupyter server 91 | │ ├── test_notebook_paths.py # Unit tests for notebook path handling 92 | │ ├── test_llm_tool_calls.py # LLM tool call generation tests 93 | │ └── llm_providers/ # LLM provider architecture 94 | │ ├── base.py # Base provider interface 95 | │ ├── claude_code.py # Claude Code provider 96 | │ └── config.py # Provider configuration 97 | ├── demos/ 98 | │ ├── demo.ipynb 99 | │ └── goose-demo.png 100 | ├── docs/ # Documentation site 101 | ├── pyproject.toml 102 | └── README.md 103 | ``` 104 | 105 | ## Making Changes 106 | 107 | ### Code Style 108 | 109 | We use `ruff` for linting and formatting: 110 | 111 | ```bash 112 | # Format code 113 | uv run ruff format . 114 | 115 | # Check linting 116 | uv run ruff check . 117 | 118 | # Fix linting issues 119 | uv run ruff check --fix . 120 | ``` 121 | 122 | ### Testing Changes 123 | 124 | 1. **Unit Tests**: Test individual functions 125 | 2. **Integration Tests**: Test with real Jupyter server 126 | 3. **LLM Tests**: Test how well LLMs generate MCP tool calls 127 | 4. **Manual Testing**: Test with your MCP client 128 | 129 | Example test: 130 | 131 | ```python 132 | def test_notebook_creation(): 133 | """Test creating a new notebook.""" 134 | notebook_path = "test_notebook.ipynb" 135 | cells = ["import pandas as pd", "print('Hello, World!')"] 136 | 137 | create_new_notebook(notebook_path, cells, server_url, token) 138 | 139 | assert check_notebook_exists(notebook_path, server_url, token) 140 | ``` 141 | 142 | ### LLM Evaluation 143 | 144 | The project includes comprehensive testing for how well different LLMs can generate MCP tool calls from natural language prompts. 145 | 146 | #### Test Architecture 147 | 148 | - **Pluggable providers**: Easy to add new LLMs (Claude Code, Gemini, OpenAI, etc.) 149 | - **Standardized interface**: All providers implement the same `LLMProvider` interface 150 | - **Parameterized tests**: Same test validates all providers consistently 151 | - **Real-time monitoring**: Watch LLMs generate tool calls with verbose output 152 | 153 | #### Running LLM Tests 154 | 155 | ```bash 156 | # Run LLM tool call generation tests 157 | uv run pytest -m llm -v 158 | 159 | # See LLM working in real-time (shows detailed progress) 160 | uv run pytest -m llm -v -s 161 | 162 | # Test specific provider 163 | uv run pytest -k "claude-code" -m llm -v 164 | ``` 165 | 166 | #### What Gets Tested 167 | 168 | Each LLM provider is validated on: 169 | 170 | 1. **Understanding natural language prompts** about Jupyter tasks 171 | 2. **Generating correct MCP tool calls** (`query_notebook`, `setup_notebook`, etc.) 172 | 3. **Successfully executing the calls** to create notebooks with expected content 173 | 4. **Error handling** when operations fail 174 | 175 | #### Adding New LLM Providers 176 | 177 | 1. **Create provider class** in `tests/llm_providers/`: 178 | ```python 179 | class MyLLMProvider(LLMProvider): 180 | @property 181 | def name(self) -> str: 182 | return "my-llm" 183 | 184 | async def send_task(self, prompt: str, server_url: str, verbose: bool = False): 185 | # Implement LLM interaction 186 | pass 187 | 188 | async def get_final_response(self) -> LLMResponse: 189 | # Return results with success metrics 190 | pass 191 | ``` 192 | 193 | 2. **Update configuration** in `tests/llm_providers/config.py`: 194 | ```python 195 | if os.getenv("MY_LLM_API_KEY"): 196 | from .my_llm import MyLLMProvider 197 | providers.append(MyLLMProvider()) 198 | ``` 199 | 200 | 3. **Test automatically**: Provider included when environment variables are set 201 | 202 | This makes it easy to validate and compare how different LLMs perform at MCP tool call generation. 203 | 204 | ## Debugging 205 | 206 | ### Using VS Code 207 | 208 | 1. Create `.vscode/launch.json`: 209 | 210 | ```json 211 | { 212 | "version": "0.2.0", 213 | "configurations": [ 214 | { 215 | "name": "Debug MCP Jupyter", 216 | "type": "python", 217 | "request": "launch", 218 | "module": "mcp_jupyter", 219 | "justMyCode": true, 220 | "env": { 221 | "PYTHONPATH": "${workspaceFolder}", 222 | "TOKEN": "BLOCK" 223 | } 224 | } 225 | ] 226 | } 227 | ``` 228 | 229 | 2. Set breakpoints in the code 230 | 3. Run with F5 231 | 232 | 233 | ## Contributing 234 | 235 | ### 1. Fork and Branch 236 | 237 | ```bash 238 | git checkout -b feature/your-feature-name 239 | ``` 240 | 241 | ### 2. Make Changes 242 | 243 | - Follow the code style 244 | - Add tests for new features 245 | - Update documentation 246 | 247 | ### 3. Test Thoroughly 248 | 249 | ```bash 250 | # Run tests 251 | uv run pytest tests/ 252 | 253 | # Check formatting 254 | uv run ruff format --check . 255 | 256 | # Check types 257 | uv run mypy src/mcp_jupyter 258 | ``` 259 | 260 | ### 4. Submit PR 261 | 262 | 1. Push to your fork 263 | 2. Create pull request 264 | 3. Describe changes clearly 265 | 4. Link any related issues 266 | 267 | ## Next Steps 268 | 269 | - [Architecture →](/docs/architecture) 270 | - [Usage Guide →](/docs/usage) 271 | - [Installation →](/docs/installation) 272 | -------------------------------------------------------------------------------- /docs/docs/installation.md: -------------------------------------------------------------------------------- 1 | --- 2 | sidebar_position: 3 3 | --- 4 | 5 | # Installation 6 | 7 | Detailed installation instructions for MCP Jupyter Server. 8 | 9 | ## Prerequisites 10 | 11 | ### 1. Install UV 12 | 13 | UV is required for running MCP Jupyter. Install it using one of these methods: 14 | 15 | ```bash 16 | # macOS/Linux 17 | curl -LsSf https://astral.sh/uv/install.sh | sh 18 | 19 | # Windows 20 | powershell -c "irm https://astral.sh/uv/install.ps1 | iex" 21 | 22 | # Or via pip 23 | pip install uv 24 | ``` 25 | 26 | ### 2. JupyterLab Setup 27 | 28 | MCP Jupyter requires a running JupyterLab server with `ipykernel`. We recommend also installing `jupyter-collaboration` for real-time synchronization. 29 | 30 | import Tabs from '@theme/Tabs'; 31 | import TabItem from '@theme/TabItem'; 32 | 33 | 34 | 35 | 36 | Real-time collaboration enables automatic synchronization between the AI agent and your notebook interface. Changes sync instantly without manual reloading. 37 | 38 | ```bash 39 | # Using uv venv 40 | uv venv 41 | source .venv/bin/activate # On Windows: .venv\Scripts\activate 42 | uv pip install jupyterlab jupyter-collaboration ipykernel 43 | 44 | # Optional: Install additional packages 45 | uv pip install numpy pandas matplotlib 46 | ``` 47 | 48 | ```bash 49 | # Using uv project 50 | uv init jupyter-workspace 51 | cd jupyter-workspace 52 | uv add jupyterlab jupyter-collaboration ipykernel 53 | 54 | # Optional: Add additional packages 55 | uv add numpy pandas matplotlib 56 | ``` 57 | 58 | 59 | 60 | 61 | You can use MCP Jupyter without `jupyter-collaboration`, but you'll need to manually sync changes between the agent and JupyterLab. 62 | 63 | ```bash 64 | # Using uv venv 65 | uv venv 66 | source .venv/bin/activate # On Windows: .venv\Scripts\activate 67 | uv pip install jupyterlab ipykernel 68 | 69 | # Optional: Install additional packages 70 | uv pip install numpy pandas matplotlib 71 | ``` 72 | 73 | ```bash 74 | # Using uv project 75 | uv init jupyter-workspace 76 | cd jupyter-workspace 77 | uv add jupyterlab ipykernel 78 | 79 | # Optional: Add additional packages 80 | uv add numpy pandas matplotlib 81 | ``` 82 | 83 | :::warning Important Workflow Differences 84 | When running **without** `jupyter-collaboration`: 85 | - **After the agent makes changes**: Use JupyterLab's "Reload Notebook from Disk" command (File → Reload Notebook from Disk) to see the updates 86 | - **After you edit in the notebook**: Save your changes (Ctrl+S / Cmd+S) so the agent can see them 87 | 88 | With RTC enabled, these manual steps are **not needed** - changes sync automatically. 89 | ::: 90 | 91 | 92 | 93 | 94 | ## Server Installation 95 | 96 | ### Using UV (Recommended) 97 | 98 | The simplest way to use MCP Jupyter is via uvx: 99 | 100 | ```bash 101 | uvx mcp-jupyter 102 | ``` 103 | 104 | This command will automatically download and run the latest version. 105 | 106 | ### Transport Modes 107 | 108 | MCP Jupyter supports two transport protocols: 109 | 110 | - **stdio** (default) - Standard input/output communication, ideal for IDE integration 111 | - **http** - Streamable HTTP transport with session management, ideal for web clients and remote access 112 | 113 | #### Use Cases for HTTP Transport 114 | 115 | - **Serverless deployments**: Host the MCP server in cloud environments (AWS Lambda, Google Cloud Functions, etc.) 116 | - **Remote access**: Connect to the server from different machines or networks 117 | - **Web integrations**: Build web-based AI assistants that connect to the MCP server 118 | - **Multi-user environments**: Deploy a central MCP server that multiple clients can connect to 119 | - **Stateless operations**: Use `--stateless-http` for environments where session persistence isn't needed or desired 120 | 121 | #### Using HTTP Transport 122 | 123 | Start the server with HTTP transport: 124 | 125 | ```bash 126 | # HTTP transport on default port 8000 127 | uvx mcp-jupyter --transport http 128 | 129 | # HTTP transport on custom port 130 | uvx mcp-jupyter --transport http --port 8090 131 | 132 | # HTTP transport in stateless mode (no session persistence) 133 | uvx mcp-jupyter --transport http --port 8090 --stateless-http 134 | ``` 135 | 136 | ### From Source 137 | 138 | For development or customization: 139 | 140 | ```bash 141 | # Clone the repository 142 | git clone https://github.com/block/mcp-jupyter.git 143 | cd mcp-jupyter 144 | 145 | # Sync all dependencies including dev tools 146 | uv sync 147 | ``` 148 | 149 | ## Starting the Jupyter Server 150 | 151 | ### Basic Setup 152 | 153 | ```bash 154 | jupyter lab --port 8888 --IdentityProvider.token BLOCK --ip 0.0.0.0 155 | ``` 156 | 157 | ### Custom Configuration 158 | 159 | For production use or specific setups: 160 | 161 | ```bash 162 | # Set custom token 163 | export TOKEN=your-secure-token 164 | jupyter lab --port 8888 --IdentityProvider.token $TOKEN 165 | 166 | # Use config file 167 | jupyter lab --config=/path/to/jupyter_config.py 168 | ``` 169 | 170 | ## Client Configuration 171 | 172 | ### Goose 173 | 174 | Add to your Goose session: 175 | 176 | ```bash 177 | goose session --with-extension "uvx mcp-jupyter" 178 | ``` 179 | 180 | Or for development: 181 | 182 | ```bash 183 | goose session --with-extension "uv run /path/to/mcp-jupyter/.venv/bin/mcp-jupyter" 184 | ``` 185 | 186 | ### Cursor 187 | 188 | #### Option 1: stdio Transport (Recommended for IDE) 189 | 190 | Add to your `.cursor/mcp.json`: 191 | 192 | ```json 193 | { 194 | "mcpServers": { 195 | "jupyter": { 196 | "command": "uvx", 197 | "args": ["mcp-jupyter"], 198 | "env": { 199 | "TOKEN": "your-token-here" 200 | } 201 | } 202 | } 203 | } 204 | ``` 205 | 206 | #### Option 2: HTTP Transport 207 | 208 | For HTTP transport, first start the server separately: 209 | 210 | ```bash 211 | uvx mcp-jupyter --transport http --port 8090 212 | ``` 213 | 214 | Then configure Cursor's `.cursor/mcp.json`: 215 | 216 | ```json 217 | { 218 | "mcpServers": { 219 | "jupyter-http": { 220 | "url": "http://localhost:8090/mcp/" // ⚠️ Trailing slash is REQUIRED 221 | } 222 | } 223 | } 224 | ``` 225 | 226 | :::warning Important 227 | The trailing slash (`/mcp/`) is **required** for Cursor to connect properly to the HTTP endpoint. 228 | ::: 229 | 230 | ### Other MCP Clients 231 | 232 | The general pattern for any MCP client: 233 | 234 | ```json 235 | { 236 | "command": "uvx", 237 | "args": ["mcp-jupyter"], 238 | "stdio": true 239 | } 240 | ``` 241 | 242 | ## Troubleshooting 243 | 244 | ### Common Issues 245 | 246 | 1. **"Jupyter server is not accessible"** 247 | - Ensure Jupyter is running on the expected port 248 | - Check firewall settings 249 | - Verify the token matches 250 | 251 | 2. **"No kernel found"** 252 | - Make sure you have opened a notebook in Jupyter 253 | - Check that ipykernel is installed 254 | - Verify the notebook path is correct 255 | 256 | 3. **Package installation fails** 257 | - Ensure your virtual environment has pip 258 | - Check write permissions 259 | - Verify internet connectivity 260 | 261 | ### Debug Mode 262 | 263 | Enable debug logging: 264 | 265 | ```bash 266 | export MCP_JUPYTER_DEBUG=1 267 | uvx mcp-jupyter 268 | ``` 269 | 270 | ## Next Steps 271 | 272 | - [Usage Guide →](/docs/usage) 273 | - [Development Setup →](/docs/development) 274 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Jupyter MCP Server 2 | 3 | > **⚠️ API Compatibility Notice**: This project is currently focused on MCP (Model Context Protocol) usage. There are **no API compatibility guarantees** between versions as the interface is actively evolving. Breaking changes may occur in any release. 4 | 5 | Jupyter MCP Server allows you to use tools like [Goose](https://block.github.io/goose/) or Cursor to pair with you in a JupyterLab notebook where the state of your variables is preserved by the JupyterLab Kernel. This enables seamless collaboration where agents can install packages, fix errors, and hand off to you for data exploration at any time. 6 | 7 | **Architecture**: Uses Jupyter's REST API for reliable agent operations while maintaining real-time user synchronization through RTC. See [Architecture Documentation](docs/docs/architecture.md) for detailed technical information. 8 | 9 | ## Key Features 10 | 11 | - **4 Consolidated MCP Tools** (reduced from 11): 12 | - `query_notebook` - All read-only operations (view source, check server, etc.) 13 | - `modify_notebook_cells` - All cell modifications (add, edit, delete cells) 14 | - `execute_notebook_code` - All execution operations (run cells, install packages) 15 | - `setup_notebook` - Notebook initialization and kernel connection 16 | - **Workflow-oriented design** optimized for AI agent collaboration 17 | - **State preservation** across notebook sessions 18 | - **Automatic parameter validation** with float-to-int conversion 19 | 20 | This works with any client that supports MCP but will focus on using Goose for the examples. 21 | 22 | ## Requirements 23 | You will need [UV](https://docs.astral.sh/uv/) is required to be installed. 24 | 25 | ## Installation 26 | This MCP server supports multiple transport modes and can be added to client with the command `uvx mcp-jupyter`. 27 | 28 | ### Transport Modes 29 | 30 | The server supports two transport protocols: 31 | - **stdio** (default) - Standard input/output communication, ideal for local IDE integrations 32 | - **http** - Streamable HTTP transport with session management, enabling serverless deployments and remote access 33 | 34 | #### Use Cases for HTTP Transport 35 | - **Serverless deployments**: Host the MCP server in cloud environments (AWS Lambda, Google Cloud Functions, etc.) 36 | - **Remote access**: Connect to the server from different machines or networks 37 | - **Web integrations**: Build web-based AI assistants that connect to the MCP server 38 | - **Stateless operations**: Use `--stateless-http` for environments where session persistence isn't needed 39 | 40 | To use a specific transport: 41 | ```bash 42 | # Default stdio transport 43 | uvx mcp-jupyter 44 | 45 | # HTTP transport on custom port (stateful - maintains session) 46 | uvx mcp-jupyter --transport http --port 8080 47 | 48 | # HTTP transport in stateless mode (no session persistence) 49 | uvx mcp-jupyter --transport http --port 8080 --stateless-http 50 | ``` 51 | 52 | ### Using HTTP Transport with Cursor 53 | 54 | To connect Cursor to an HTTP MCP server: 55 | 56 | 1. Start the server separately: 57 | ```bash 58 | uvx mcp-jupyter --transport http --port 8090 59 | ``` 60 | 61 | 2. Configure Cursor's `.cursor/mcp.json`: 62 | ```json 63 | { 64 | "mcpServers": { 65 | "notebook-http": { 66 | "url": "http://localhost:8090/mcp/" // ⚠️ Trailing slash is REQUIRED 67 | } 68 | } 69 | } 70 | ``` 71 | 72 | **Important:** The trailing slash (`/mcp/`) is required for Cursor to connect properly to the HTTP endpoint. 73 | 74 | ## Usage 75 | 76 | ### Start Jupyter 77 | The server expects that a server is already running on a port that is available to the client. If the environmental variable TOKEN is not set, it will default to "BLOCK". 78 | 79 | **Option 1: With Real-Time Collaboration (Recommended)** 80 | 81 | Real-time collaboration (`jupyter-collaboration`) enables automatic synchronization between the AI agent and your notebook interface. Changes made by the agent appear instantly in your browser. 82 | 83 | ```bash 84 | # Using uv venv 85 | uv venv 86 | source .venv/bin/activate # On Windows: .venv\Scripts\activate 87 | uv pip install jupyterlab jupyter-collaboration ipykernel 88 | jupyter lab --port 8888 --IdentityProvider.token BLOCK --ip 0.0.0.0 89 | 90 | # OR using uv project 91 | uv init jupyter-workspace && cd jupyter-workspace 92 | uv add jupyterlab jupyter-collaboration ipykernel 93 | uv run jupyter lab --port 8888 --IdentityProvider.token BLOCK --ip 0.0.0.0 94 | ``` 95 | 96 | **Option 2: Without Real-Time Collaboration** 97 | 98 | You can use MCP Jupyter without `jupyter-collaboration`, but you'll need to manually sync changes: 99 | 100 | ```bash 101 | # Using uv venv 102 | uv venv 103 | source .venv/bin/activate # On Windows: .venv\Scripts\activate 104 | uv pip install jupyterlab ipykernel 105 | jupyter lab --port 8888 --IdentityProvider.token BLOCK --ip 0.0.0.0 106 | 107 | # OR using uv project 108 | uv init jupyter-workspace && cd jupyter-workspace 109 | uv add jupyterlab ipykernel 110 | uv run jupyter lab --port 8888 --IdentityProvider.token BLOCK --ip 0.0.0.0 111 | ``` 112 | 113 | > **Important Workflow Differences Without RTC:** 114 | > - **After the agent makes changes**: Use the "Reload Notebook from Disk" command in JupyterLab to see the updates 115 | > - **After you edit in the notebook**: Click Save (Ctrl+S / Cmd+S) so the agent can see your changes 116 | > 117 | > With RTC enabled, these manual steps are not needed - changes sync automatically. 118 | 119 | ### Goose Usage 120 | 121 | Here's a demonstration of the tool in action: 122 | 123 | ![MCP Jupyter Demo](demos/goose-demo.png) 124 | 125 | You can view the Generated notebook here: [View Demo Notebook](demos/demo.ipynb) 126 | 127 | ## Development 128 | Steps remain similar except you will need to clone this mcp-jupyter repository and use that for the server instead of the precompiled version. 129 | 130 | ### MCP Server 131 | 132 | 1. Clone and setup the repository: 133 | ```bash 134 | mkdir ~/Development 135 | cd ~/Development 136 | git clone https://github.com/block/mcp-jupyter.git 137 | cd mcp-jupyter 138 | 139 | # Sync all dependencies 140 | uv sync 141 | ``` 142 | 143 | Using editable mode allows you to make changes to the server and only have you need to restart Goose, etc. 144 | `goose session --with-extension "uv run --directory $(pwd) mcp-jupyter"` 145 | 146 | ## LLM Evaluation 147 | 148 | This project includes a comprehensive testing infrastructure for validating how well different LLMs can generate MCP tool calls from natural language prompts. 149 | 150 | ### Test Architecture 151 | 152 | The LLM testing system uses a pluggable provider architecture: 153 | 154 | - **`LLMProvider`**: Abstract base class that all providers implement 155 | - **`LLMResponse`**: Standardized response format with success metrics and metadata 156 | - **Parameterized tests**: Same test runs against all available providers 157 | 158 | ### Current Providers 159 | 160 | - **ClaudeCodeProvider**: Uses the Claude Code SDK (no API key required) 161 | 162 | ### Running LLM Tests 163 | 164 | ```bash 165 | # Run LLM tool call generation tests 166 | uv run pytest -m llm -v 167 | 168 | # See LLM working in real-time (shows detailed progress) 169 | uv run pytest -m llm -v -s 170 | 171 | # Run all tests except LLM tests (default behavior) 172 | uv run pytest -v 173 | ``` 174 | 175 | ### What the Tests Validate 176 | 177 | Each LLM provider is tested on its ability to: 178 | 179 | 1. **Understand natural language prompts** about Jupyter notebook tasks 180 | 2. **Generate correct MCP tool calls** (`query_notebook`, `setup_notebook`, `modify_notebook_cells`) 181 | 3. **Successfully execute the calls** to create notebooks with expected content 182 | 4. **Handle errors gracefully** when operations fail 183 | 184 | ### Adding New Providers 185 | 186 | To add a new LLM provider: 187 | 188 | 1. **Implement the interface**: 189 | ```python 190 | # tests/llm_providers/my_llm.py 191 | from .base import LLMProvider, LLMResponse 192 | 193 | class MyLLMProvider(LLMProvider): 194 | @property 195 | def name(self) -> str: 196 | return "my-llm" 197 | 198 | async def send_task(self, prompt: str, server_url: str, verbose: bool = False): 199 | # Implement LLM interaction 200 | pass 201 | 202 | async def get_final_response(self) -> LLMResponse: 203 | # Return standardized response 204 | pass 205 | 206 | async def cleanup(self): 207 | # Clean up resources 208 | pass 209 | ``` 210 | 211 | 2. **Update configuration**: 212 | ```python 213 | # tests/llm_providers/config.py - add to get_available_providers() 214 | if os.getenv("MY_LLM_API_KEY"): 215 | from .my_llm import MyLLMProvider 216 | providers.append(MyLLMProvider()) 217 | ``` 218 | 219 | 3. **Test automatically**: Your provider will be included in parameterized tests when its environment variables are set. 220 | 221 | This infrastructure makes it easy to validate and compare how different LLMs perform at generating MCP tool calls for Jupyter notebook automation. 222 | -------------------------------------------------------------------------------- /tests/conftest.py: -------------------------------------------------------------------------------- 1 | """Shared test configuration with single warm server for session reuse.""" 2 | 3 | import asyncio 4 | import os 5 | import shutil 6 | import signal 7 | import subprocess 8 | import time 9 | from pathlib import Path 10 | 11 | import pytest 12 | 13 | import mcp_jupyter.server 14 | 15 | # Fixtures for MCP Jupyter integration tests 16 | 17 | # Constants 18 | SERVER_PORT = 9999 19 | TOKEN = "BLOCK" 20 | 21 | # LLM test constants 22 | LLM_SERVER_PORT = 10000 23 | 24 | 25 | async def _check_server_health(server_url: str, token: str) -> bool: 26 | """Check if Jupyter server is healthy using asyncio.""" 27 | try: 28 | # Use asyncio to make HTTP request without external dependencies 29 | reader, writer = await asyncio.wait_for( 30 | asyncio.open_connection("localhost", int(server_url.split(":")[-1])), 31 | timeout=1.0, 32 | ) 33 | 34 | # Send HTTP request 35 | request = f"GET /api/sessions HTTP/1.1\r\nHost: localhost\r\nAuthorization: token {token}\r\nConnection: close\r\n\r\n" 36 | writer.write(request.encode()) 37 | await writer.drain() 38 | 39 | # Read response 40 | response = await asyncio.wait_for(reader.read(1024), timeout=1.0) 41 | writer.close() 42 | await writer.wait_closed() 43 | 44 | # Check if we got a 200 OK response 45 | return b"200 OK" in response 46 | 47 | except (asyncio.TimeoutError, ConnectionRefusedError, OSError): 48 | return False 49 | 50 | 51 | async def _delete_notebook(server_url: str, notebook_name: str, token: str) -> bool: 52 | """Delete a notebook using asyncio.""" 53 | try: 54 | # Use asyncio to make HTTP DELETE request 55 | reader, writer = await asyncio.wait_for( 56 | asyncio.open_connection("localhost", int(server_url.split(":")[-1])), 57 | timeout=2.0, 58 | ) 59 | 60 | # Send HTTP DELETE request 61 | request = f"DELETE /api/contents/{notebook_name} HTTP/1.1\r\nHost: localhost\r\nAuthorization: token {token}\r\nConnection: close\r\n\r\n" 62 | writer.write(request.encode()) 63 | await writer.drain() 64 | 65 | # Read response 66 | response = await asyncio.wait_for(reader.read(1024), timeout=2.0) 67 | writer.close() 68 | await writer.wait_closed() 69 | 70 | # Check if we got a successful response (2xx) 71 | return b"200 " in response or b"204 " in response 72 | 73 | except (asyncio.TimeoutError, ConnectionRefusedError, OSError): 74 | return False 75 | 76 | 77 | def _start_jupyter_server(port: int, test_dir_name: str, server_type: str = ""): 78 | """Start a Jupyter server with given configuration. 79 | 80 | Args: 81 | port: Port number for the server 82 | test_dir_name: Name of the test directory 83 | server_type: Optional prefix for log messages (e.g., "LLM ") 84 | 85 | Returns 86 | ------- 87 | Server URL string 88 | """ 89 | test_notebooks_dir = Path(test_dir_name) 90 | server_url = f"http://localhost:{port}" 91 | 92 | # Clean up potential leftovers from previous failed runs 93 | if test_notebooks_dir.exists(): 94 | shutil.rmtree(test_notebooks_dir) 95 | 96 | # Create a directory for test notebooks 97 | test_notebooks_dir.mkdir(exist_ok=True) 98 | 99 | # Start the Jupyter server process using uv run 100 | jupyter_cmd = [ 101 | "uv", 102 | "run", 103 | "jupyter", 104 | "lab", 105 | f"--port={port}", 106 | f"--IdentityProvider.token={TOKEN}", 107 | "--ip=0.0.0.0", 108 | "--no-browser", 109 | "--ServerApp.disable_check_xsrf=True", # Skip XSRF checks for faster startup 110 | "--ServerApp.allow_origin='*'", # Allow all origins 111 | "--LabServerApp.open_browser=False", # Ensure no browser attempts 112 | f"--ServerApp.root_dir={test_notebooks_dir.absolute()}", # Set root directory 113 | ] 114 | 115 | # Add --allow-root flag if running as root (handle systems without geteuid) 116 | try: 117 | if hasattr(os, "geteuid") and os.geteuid() == 0: # Check if running as root 118 | jupyter_cmd.append("--allow-root") 119 | except (AttributeError, OSError): 120 | # On systems without geteuid (Windows) or other issues, add --allow-root anyway 121 | jupyter_cmd.append("--allow-root") 122 | 123 | # Start the Jupyter server process 124 | print(f"Starting {server_type}Jupyter server on port {port}...") 125 | server_process = subprocess.Popen( 126 | jupyter_cmd, 127 | stdout=subprocess.PIPE, 128 | stderr=subprocess.PIPE, 129 | text=True, 130 | preexec_fn=os.setsid, 131 | ) 132 | 133 | # Wait for the server to start with optimized polling 134 | max_retries = 30 # More retries for reliability 135 | retry_interval = 0.25 # Check every 250ms for faster detection 136 | initial_wait = 0.5 # Brief initial delay 137 | 138 | time.sleep(initial_wait) 139 | 140 | for attempt in range(max_retries): 141 | try: 142 | # Check server health using asyncio 143 | is_healthy = asyncio.run(_check_server_health(server_url, TOKEN)) 144 | if is_healthy: 145 | print( 146 | f"{server_type}Jupyter server started successfully (attempt {attempt + 1})" 147 | ) 148 | break 149 | except Exception: 150 | pass 151 | time.sleep(retry_interval) 152 | if attempt % 8 == 0: # Print every 2 seconds 153 | print( 154 | f"Waiting for {server_type.lower()}server to start... (attempt {attempt + 1}/{max_retries})" 155 | ) 156 | else: 157 | # Server didn't start in time, kill the process and raise an exception 158 | try: 159 | os.killpg(os.getpgid(server_process.pid), signal.SIGTERM) 160 | except ProcessLookupError: 161 | pass # Process already terminated 162 | stdout, stderr = server_process.communicate() 163 | print(f"{server_type}Jupyter server stdout: {stdout}") 164 | print(f"{server_type}Jupyter server stderr: {stderr}") 165 | pytest.fail(f"{server_type}Jupyter server failed to start in time") 166 | 167 | # Reset notebook state hash at session start 168 | try: 169 | from mcp_jupyter.server import NotebookState 170 | 171 | NotebookState.contents_hash = "" 172 | NotebookState.notebook_server_urls = {} 173 | except ImportError: 174 | print("Warning: Could not import NotebookState, state management disabled") 175 | 176 | return server_url, server_process, test_notebooks_dir 177 | 178 | 179 | def _cleanup_jupyter_server( 180 | server_process, test_notebooks_dir: Path, server_type: str = "" 181 | ): 182 | """Clean up a Jupyter server and its test directory. 183 | 184 | Args: 185 | server_process: The subprocess.Popen server process 186 | test_notebooks_dir: Path to the test directory to remove 187 | server_type: Optional prefix for log messages (e.g., "LLM ") 188 | """ 189 | # Cleanup: kill the Jupyter server process and all its children 190 | print(f"Shutting down {server_type}Jupyter server") 191 | try: 192 | os.killpg(os.getpgid(server_process.pid), signal.SIGTERM) 193 | server_process.wait(timeout=5) 194 | except ProcessLookupError: 195 | print(f"{server_type}Server process already terminated.") 196 | except subprocess.TimeoutExpired: 197 | print(f"{server_type}Server process did not terminate gracefully, killing.") 198 | os.killpg(os.getpgid(server_process.pid), signal.SIGKILL) 199 | server_process.wait() 200 | 201 | # Remove the entire test directory and its contents 202 | print(f"Removing {server_type.lower()}test directory: {test_notebooks_dir}") 203 | if test_notebooks_dir.exists(): 204 | shutil.rmtree(test_notebooks_dir) 205 | 206 | 207 | @pytest.fixture(scope="session") 208 | def jupyter_server(): 209 | """Session-scoped Jupyter server that stays warm throughout all tests.""" 210 | server_url, server_process, test_notebooks_dir = _start_jupyter_server( 211 | SERVER_PORT, "test_notebooks_session", "session " 212 | ) 213 | 214 | yield server_url 215 | 216 | # Clean up any lingering kernel connections 217 | try: 218 | if mcp_jupyter.server.kernel is not None: 219 | print("Cleaning up kernel connection...") 220 | mcp_jupyter.server.kernel.stop() 221 | except Exception as e: 222 | print(f"Error cleaning up kernel: {e}") 223 | 224 | _cleanup_jupyter_server(server_process, test_notebooks_dir, "session ") 225 | 226 | 227 | @pytest.fixture 228 | def test_notebook(jupyter_server): 229 | """Create a test notebook with some initial cells for testing.""" 230 | try: 231 | from mcp_jupyter.server import setup_notebook 232 | except ImportError: 233 | pytest.skip("mcp_jupyter package not available") 234 | 235 | notebook_name = "test_tools_notebook" 236 | 237 | # Create an empty notebook 238 | result = setup_notebook(notebook_name, server_url=jupyter_server) 239 | 240 | # Add initial cells using modify_notebook_cells 241 | from mcp_jupyter.server import modify_notebook_cells 242 | 243 | modify_notebook_cells( 244 | notebook_name, "add_code", "# Initial cell\nprint('Hello from initial cell')" 245 | ) 246 | 247 | modify_notebook_cells( 248 | notebook_name, 249 | "add_code", 250 | "def add(a, b):\n return a + b\n\nprint(add(2, 3))", 251 | ) 252 | 253 | # Small delay to ensure notebook is fully saved and available 254 | time.sleep(0.2) 255 | 256 | yield f"{notebook_name}.ipynb" 257 | 258 | # Cleanup: delete the test notebook after test 259 | try: 260 | asyncio.run(_delete_notebook(jupyter_server, f"{notebook_name}.ipynb", TOKEN)) 261 | # Reset notebook state after deletion 262 | try: 263 | from mcp_jupyter.server import NotebookState 264 | 265 | NotebookState.contents_hash = "" 266 | except ImportError: 267 | pass # State management not available 268 | except Exception: 269 | pass # Ignore cleanup errors 270 | 271 | 272 | @pytest.fixture(scope="session") 273 | def llm_jupyter_server(): 274 | """Session-scoped Jupyter server for LLM tests that stays warm throughout all tests.""" 275 | server_url, server_process, test_notebooks_dir = _start_jupyter_server( 276 | LLM_SERVER_PORT, "test_notebooks_llm", "LLM " 277 | ) 278 | 279 | yield server_url 280 | 281 | _cleanup_jupyter_server(server_process, test_notebooks_dir, "LLM ") 282 | -------------------------------------------------------------------------------- /src/mcp_jupyter/state.py: -------------------------------------------------------------------------------- 1 | import hashlib 2 | import json 3 | import logging 4 | import time 5 | from functools import wraps 6 | from typing import Optional 7 | 8 | import requests 9 | from mcp.shared.exceptions import McpError 10 | from mcp.types import INTERNAL_ERROR, ErrorData 11 | 12 | from .utils import TOKEN, _ensure_ipynb_extension 13 | 14 | # Setup logger 15 | logger = logging.getLogger(__name__) 16 | 17 | 18 | class NotebookState: 19 | """Tracks the state of the notebook (or more precisely, Goose's knowledge of the 20 | notebook state). 21 | 22 | This class provides class methods for: 23 | - Tracking notebook content hash based on a user-provided Jupyter server. 24 | - Checking for notebook changes against the user-provided server. 25 | - Decorators for state-dependent and state-refreshing operations. 26 | - Managing server URLs associated with notebook paths (relative to the server root). 27 | 28 | Notebook paths are assumed to be relative to the root directory of the 29 | user-provided Jupyter server. 30 | 31 | Raises 32 | ------ 33 | McpError: When state-dependent operations are attempted with outdated state 34 | """ 35 | 36 | contents_hash: str = "" 37 | # Add a dictionary to store server URLs for each notebook path 38 | notebook_server_urls: dict = {} 39 | 40 | @classmethod 41 | def get_server_url(cls, notebook_path: str) -> str: 42 | """Get the server URL associated with a notebook path (relative to server root). 43 | 44 | Args: 45 | notebook_path: Path to the notebook file, relative to the Jupyter server root. 46 | 47 | Returns 48 | ------- 49 | str: The server URL associated with this notebook, or the default URL 50 | """ 51 | notebook_path = _ensure_ipynb_extension(notebook_path) 52 | return cls.notebook_server_urls.get(notebook_path, "http://localhost:8888") 53 | 54 | @classmethod 55 | def set_server_url(cls, notebook_path: str, server_url: str): 56 | """Associate a server URL with a notebook path (relative to server root). 57 | 58 | Args: 59 | notebook_path: Path to the notebook file, relative to the Jupyter server root. 60 | server_url: Server URL to associate with this notebook 61 | """ 62 | notebook_path = _ensure_ipynb_extension(notebook_path) 63 | cls.notebook_server_urls[notebook_path] = server_url 64 | logger.info(f"Associated notebook {notebook_path} with server URL {server_url}") 65 | 66 | @classmethod 67 | def _get_new_hash(cls, notebook_path: str, server_url: str = None) -> str: 68 | """Hash the notebook contents from the user-provided Jupyter server. 69 | 70 | Args: 71 | notebook_path: Path to the notebook file (.ipynb extension will be added if missing), 72 | relative to the Jupyter server root. 73 | server_url: The server URL to use. Defaults to None. 74 | 75 | Returns 76 | ------- 77 | str: SHA-256 hash of the notebook contents 78 | 79 | Raises 80 | ------ 81 | RequestException: If unable to fetch notebook contents 82 | IOError: If unable to read notebook file 83 | """ 84 | # Ensure the notebook path has the .ipynb extension 85 | notebook_path = _ensure_ipynb_extension(notebook_path) 86 | 87 | # Use the stored server_url if none is provided 88 | if server_url is None: 89 | server_url = cls.get_server_url(notebook_path) 90 | 91 | # Empirically this seems to be long enough that autosaving kicks in after a change. 92 | # Also tried fetching source in other ways and triggering saves to avoid sleeping, but 93 | # haven't found a good way to get this working reliably. 94 | time.sleep(1.5) 95 | 96 | response = requests.get( 97 | f"{server_url}/api/contents/{notebook_path}", 98 | headers={"Authorization": f"token {TOKEN}"}, 99 | ) 100 | 101 | response.raise_for_status() 102 | notebook_content = response.json()["content"] 103 | 104 | return hashlib.sha256(json.dumps(notebook_content).encode()).hexdigest() 105 | 106 | @classmethod 107 | def update_hash( 108 | cls, notebook_path: str, server_url: str = None, caller: Optional[str] = None 109 | ): 110 | """Update the stored hash of notebook contents from the user-provided server. 111 | 112 | Args: 113 | notebook_path: Path to the notebook file (.ipynb extension will be added if missing), 114 | relative to the Jupyter server root. 115 | server_url: The server URL to use. Defaults to None. 116 | caller: Optional name of calling function for logging 117 | 118 | Returns 119 | ------- 120 | None 121 | 122 | Raises 123 | ------ 124 | RequestException: If unable to fetch notebook contents 125 | IOError: If unable to read notebook file 126 | """ 127 | # Ensure the notebook path has the .ipynb extension 128 | notebook_path = _ensure_ipynb_extension(notebook_path) 129 | 130 | old_hash = cls.contents_hash 131 | cls.contents_hash = cls._get_new_hash(notebook_path, server_url) 132 | prefix = f"[{caller}] " if caller else "" 133 | if old_hash != cls.contents_hash: 134 | logger.info( 135 | f"{prefix}Updated notebook hash from {old_hash!r} to " 136 | f"{cls.contents_hash!r}." 137 | ) 138 | else: 139 | logger.info( 140 | f"{prefix}No change in notebook hash (still {cls.contents_hash!r})." 141 | ) 142 | 143 | @classmethod 144 | def check_for_changes(cls, notebook_path: str, server_url: str = None) -> dict: 145 | """Check if the notebook has changed on the user-provided server. 146 | 147 | Args: 148 | notebook_path: Path to the notebook file (.ipynb extension will be added if missing), 149 | relative to the Jupyter server root. 150 | server_url: The server URL to use. Defaults to None. 151 | 152 | Returns 153 | ------- 154 | dict: Contains: 155 | - has_changed: bool indicating if notebook changed 156 | - new_hash: str new content hash 157 | - old_hash: str previous content hash 158 | 159 | Raises 160 | ------ 161 | RequestException: If unable to fetch notebook contents 162 | IOError: If unable to read notebook file 163 | """ 164 | # Ensure the notebook path has the .ipynb extension 165 | notebook_path = _ensure_ipynb_extension(notebook_path) 166 | 167 | hashed = cls._get_new_hash(notebook_path, server_url) 168 | has_changed = cls.contents_hash and hashed != cls.contents_hash 169 | return { 170 | "has_changed": has_changed, 171 | "new_hash": hashed, 172 | "old_hash": cls.contents_hash, 173 | } 174 | 175 | @classmethod 176 | def state_dependent(cls, func): 177 | """Decorate functions that goose should only use if it knows the current 178 | state of the notebook. 179 | State_dependent functions will raise an error if the notebook has changed since the last 180 | the contents_hash attribute was updated in order to encourage goose to view the notebook 181 | source. After the wrapped function executes, the contents_hash will be updated to 182 | reflect the new state. 183 | 184 | Args: 185 | func: The function to decorate. 186 | 187 | Returns 188 | ------- 189 | The decorated function. 190 | 191 | Raises 192 | ------ 193 | OutdatedStateError: If the notebook has changed since the last time the contents_hash 194 | attribute was updated. 195 | """ 196 | 197 | @wraps(func) 198 | def wrapper(*args, **kwargs): 199 | # Extract notebook_path from args/kwargs 200 | if args and len(args) > 0: 201 | notebook_path = args[0] 202 | elif "notebook_path" in kwargs: 203 | notebook_path = kwargs["notebook_path"] 204 | else: 205 | raise ValueError( 206 | "notebook_path (relative to server root) must be provided as first argument or keyword argument" 207 | ) 208 | 209 | # Get server_url from kwargs or use default 210 | server_url = kwargs.get("server_url", None) 211 | 212 | # Ensure the notebook path has the .ipynb extension and update the args/kwargs 213 | notebook_path = _ensure_ipynb_extension(notebook_path) 214 | if args and len(args) > 0: 215 | args = list(args) 216 | args[0] = notebook_path 217 | args = tuple(args) 218 | else: 219 | kwargs["notebook_path"] = notebook_path 220 | 221 | changes = cls.check_for_changes(notebook_path, server_url) 222 | if changes["has_changed"]: 223 | raise McpError( 224 | ErrorData( 225 | code=INTERNAL_ERROR, 226 | message=f"Notebook has changed since you last saw it ({changes['old_hash']}" 227 | f" --> {changes['new_hash']}). Use your view_source tool to " 228 | "update your knowledge of the notebook contents. If you are trying to " 229 | "edit/delete a specific cell, you can likely start by viewing just that " 230 | "cell rather than the whole notebook.", 231 | ) 232 | ) 233 | 234 | result = func(*args, **kwargs) 235 | cls.update_hash(notebook_path, server_url, caller=func.__name__) 236 | return result 237 | 238 | return wrapper 239 | 240 | @classmethod 241 | def refreshes_state(cls, func): 242 | """Decorate functions that update goose's knowledge of the notebook but 243 | do not require checking state before execution. I.e. use `state_dependent` for functions 244 | that goose should not execute unless it knows the notebook state; use `refreshes_state` if 245 | you merely want execution of the function to refresh goose's knowledge of current state. 246 | 247 | Args: 248 | func: The function to decorate. 249 | 250 | Returns 251 | ------- 252 | The decorated function. 253 | """ 254 | 255 | @wraps(func) 256 | def wrapper(*args, **kwargs): 257 | # Extract notebook_path from args/kwargs 258 | if args and len(args) > 0: 259 | notebook_path = args[0] 260 | elif "notebook_path" in kwargs: 261 | notebook_path = kwargs["notebook_path"] 262 | else: 263 | raise ValueError( 264 | "notebook_path (relative to server root) must be provided as first argument or keyword argument" 265 | ) 266 | 267 | # Get server_url from kwargs or use default 268 | server_url = kwargs.get("server_url", None) 269 | 270 | # Ensure the notebook path has the .ipynb extension and update the args/kwargs 271 | notebook_path = _ensure_ipynb_extension(notebook_path) 272 | if args and len(args) > 0: 273 | args = list(args) 274 | args[0] = notebook_path 275 | args = tuple(args) 276 | else: 277 | kwargs["notebook_path"] = notebook_path 278 | 279 | result = func(*args, **kwargs) 280 | cls.update_hash(notebook_path, server_url, caller=func.__name__) 281 | return result 282 | 283 | return wrapper 284 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Apache License 2 | Version 2.0, January 2004 3 | http://www.apache.org/licenses/ 4 | 5 | TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 6 | 7 | 1. Definitions. 8 | 9 | "License" shall mean the terms and conditions for use, reproduction, 10 | and distribution as defined by Sections 1 through 9 of this document. 11 | 12 | "Licensor" shall mean the copyright owner or entity authorized by 13 | the copyright owner that is granting the License. 14 | 15 | "Legal Entity" shall mean the union of the acting entity and all 16 | other entities that control, are controlled by, or are under common 17 | control with that entity. For the purposes of this definition, 18 | "control" means (i) the power, direct or indirect, to cause the 19 | direction or management of such entity, whether by contract or 20 | otherwise, or (ii) ownership of fifty percent (50%) or more of the 21 | outstanding shares, or (iii) beneficial ownership of such entity. 22 | 23 | "You" (or "Your") shall mean an individual or Legal Entity 24 | exercising permissions granted by this License. 25 | 26 | "Source" form shall mean the preferred form for making modifications, 27 | including but not limited to software source code, documentation 28 | source, and configuration files. 29 | 30 | "Object" form shall mean any form resulting from mechanical 31 | transformation or translation of a Source form, including but 32 | not limited to compiled object code, generated documentation, 33 | and conversions to other media types. 34 | 35 | "Work" shall mean the work of authorship, whether in Source or 36 | Object form, made available under the License, as indicated by a 37 | copyright notice that is included in or attached to the work 38 | (an example is provided in the Appendix below). 39 | 40 | "Derivative Works" shall mean any work, whether in Source or Object 41 | form, that is based on (or derived from) the Work and for which the 42 | editorial revisions, annotations, elaborations, or other modifications 43 | represent, as a whole, an original work of authorship. For the purposes 44 | of this License, Derivative Works shall not include works that remain 45 | separable from, or merely link (or bind by name) to the interfaces of, 46 | the Work and Derivative Works thereof. 47 | 48 | "Contribution" shall mean any work of authorship, including 49 | the original version of the Work and any modifications or additions 50 | to that Work or Derivative Works thereof, that is intentionally 51 | submitted to Licensor for inclusion in the Work by the copyright owner 52 | or by an individual or Legal Entity authorized to submit on behalf of 53 | the copyright owner. For the purposes of this definition, "submitted" 54 | means any form of electronic, verbal, or written communication sent 55 | to the Licensor or its representatives, including but not limited to 56 | communication on electronic mailing lists, source code control systems, 57 | and issue tracking systems that are managed by, or on behalf of, the 58 | Licensor for the purpose of discussing and improving the Work, but 59 | excluding communication that is conspicuously marked or otherwise 60 | designated in writing by the copyright owner as "Not a Contribution." 61 | 62 | "Contributor" shall mean Licensor and any individual or Legal Entity 63 | on behalf of whom a Contribution has been received by Licensor and 64 | subsequently incorporated within the Work. 65 | 66 | 2. Grant of Copyright License. Subject to the terms and conditions of 67 | this License, each Contributor hereby grants to You a perpetual, 68 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 69 | copyright license to reproduce, prepare Derivative Works of, 70 | publicly display, publicly perform, sublicense, and distribute the 71 | Work and such Derivative Works in Source or Object form. 72 | 73 | 3. Grant of Patent License. Subject to the terms and conditions of 74 | this License, each Contributor hereby grants to You a perpetual, 75 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 76 | (except as stated in this section) patent license to make, have made, 77 | use, offer to sell, sell, import, and otherwise transfer the Work, 78 | where such license applies only to those patent claims licensable 79 | by such Contributor that are necessarily infringed by their 80 | Contribution(s) alone or by combination of their Contribution(s) 81 | with the Work to which such Contribution(s) was submitted. If You 82 | institute patent litigation against any entity (including a 83 | cross-claim or counterclaim in a lawsuit) alleging that the Work 84 | or a Contribution incorporated within the Work constitutes direct 85 | or contributory patent infringement, then any patent licenses 86 | granted to You under this License for that Work shall terminate 87 | as of the date such litigation is filed. 88 | 89 | 4. Redistribution. You may reproduce and distribute copies of the 90 | Work or Derivative Works thereof in any medium, with or without 91 | modifications, and in Source or Object form, provided that You 92 | meet the following conditions: 93 | 94 | (a) You must give any other recipients of the Work or 95 | Derivative Works a copy of this License; and 96 | 97 | (b) You must cause any modified files to carry prominent notices 98 | stating that You changed the files; and 99 | 100 | (c) You must retain, in the Source form of any Derivative Works 101 | that You distribute, all copyright, patent, trademark, and 102 | attribution notices from the Source form of the Work, 103 | excluding those notices that do not pertain to any part of 104 | the Derivative Works; and 105 | 106 | (d) If the Work includes a "NOTICE" text file as part of its 107 | distribution, then any Derivative Works that You distribute must 108 | include a readable copy of the attribution notices contained 109 | within such NOTICE file, excluding those notices that do not 110 | pertain to any part of the Derivative Works, in at least one 111 | of the following places: within a NOTICE text file distributed 112 | as part of the Derivative Works; within the Source form or 113 | documentation, if provided along with the Derivative Works; or, 114 | within a display generated by the Derivative Works, if and 115 | wherever such third-party notices normally appear. The contents 116 | of the NOTICE file are for informational purposes only and 117 | do not modify the License. You may add Your own attribution 118 | notices within Derivative Works that You distribute, alongside 119 | or as an addendum to the NOTICE text from the Work, provided 120 | that such additional attribution notices cannot be construed 121 | as modifying the License. 122 | 123 | You may add Your own copyright statement to Your modifications and 124 | may provide additional or different license terms and conditions 125 | for use, reproduction, or distribution of Your modifications, or 126 | for any such Derivative Works as a whole, provided Your use, 127 | reproduction, and distribution of the Work otherwise complies with 128 | the conditions stated in this License. 129 | 130 | 5. Submission of Contributions. Unless You explicitly state otherwise, 131 | any Contribution intentionally submitted for inclusion in the Work 132 | by You to the Licensor shall be under the terms and conditions of 133 | this License, without any additional terms or conditions. 134 | Notwithstanding the above, nothing herein shall supersede or modify 135 | the terms of any separate license agreement you may have executed 136 | with Licensor regarding such Contributions. 137 | 138 | 6. Trademarks. This License does not grant permission to use the trade 139 | names, trademarks, service marks, or product names of the Licensor, 140 | except as required for reasonable and customary use in describing the 141 | origin of the Work and reproducing the content of the NOTICE file. 142 | 143 | 7. Disclaimer of Warranty. Unless required by applicable law or 144 | agreed to in writing, Licensor provides the Work (and each 145 | Contributor provides its Contributions) on an "AS IS" BASIS, 146 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or 147 | implied, including, without limitation, any warranties or conditions 148 | of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A 149 | PARTICULAR PURPOSE. You are solely responsible for determining the 150 | appropriateness of using or redistributing the Work and assume any 151 | risks associated with Your exercise of permissions under this License. 152 | 153 | 8. Limitation of Liability. In no event and under no legal theory, 154 | whether in tort (including negligence), contract, or otherwise, 155 | unless required by applicable law (such as deliberate and grossly 156 | negligent acts) or agreed to in writing, shall any Contributor be 157 | liable to You for damages, including any direct, indirect, special, 158 | incidental, or consequential damages of any character arising as a 159 | result of this License or out of the use or inability to use the 160 | Work (including but not limited to damages for loss of goodwill, 161 | work stoppage, computer failure or malfunction, or any and all 162 | other commercial damages or losses), even if such Contributor 163 | has been advised of the possibility of such damages. 164 | 165 | 9. Accepting Warranty or Additional Liability. While redistributing 166 | the Work or Derivative Works thereof, You may choose to offer, 167 | and charge a fee for, acceptance of support, warranty, indemnity, 168 | or other liability obligations and/or rights consistent with this 169 | License. However, in accepting such obligations, You may act only 170 | on Your own behalf and on Your sole responsibility, not on behalf 171 | of any other Contributor, and only if You agree to indemnify, 172 | defend, and hold each Contributor harmless for any liability 173 | incurred by, or claims asserted against, such Contributor by reason 174 | of your accepting any such warranty or additional liability. 175 | 176 | END OF TERMS AND CONDITIONS 177 | 178 | APPENDIX: How to apply the Apache License to your work. 179 | 180 | To apply the Apache License to your work, attach the following 181 | boilerplate notice, with the fields enclosed by brackets "[]" 182 | replaced with your own identifying information. (Don't include 183 | the brackets!) The text should be enclosed in the appropriate 184 | comment syntax for the file format. We also recommend that a 185 | file or class name and description of purpose be included on the 186 | same "printed page" as the copyright notice for easier 187 | identification within third-party archives. 188 | 189 | Licensed under the Apache License, Version 2.0 (the "License"); 190 | you may not use this file except in compliance with the License. 191 | You may obtain a copy of the License at 192 | 193 | http://www.apache.org/licenses/LICENSE-2.0 194 | 195 | Copyright 2025 Block, Inc. 196 | 197 | Unless required by applicable law or agreed to in writing, software 198 | distributed under the License is distributed on an "AS IS" BASIS, 199 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 200 | See the License for the specific language governing permissions and 201 | limitations under the License. -------------------------------------------------------------------------------- /tests/test_http_transport.py: -------------------------------------------------------------------------------- 1 | """Tests for HTTP transport functionality.""" 2 | 3 | import threading 4 | import time 5 | from unittest.mock import MagicMock, patch 6 | 7 | import pytest 8 | import requests 9 | 10 | from mcp_jupyter.server import create_server 11 | 12 | 13 | def wait_for_server(url, timeout=10, poll_interval=0.5): 14 | """Wait for server to be ready by polling an endpoint. 15 | 16 | Args: 17 | url: Server URL to check 18 | timeout: Maximum time to wait in seconds 19 | poll_interval: Time between checks in seconds 20 | 21 | Returns 22 | ------- 23 | True if server is ready, False if timeout reached 24 | """ 25 | start_time = time.time() 26 | while time.time() - start_time < timeout: 27 | try: 28 | response = requests.get(url, timeout=1) 29 | if response.status_code in [ 30 | 200, 31 | 404, 32 | 405, 33 | ]: # Any response means server is up 34 | return True 35 | except requests.exceptions.RequestException: 36 | pass 37 | time.sleep(poll_interval) 38 | return False 39 | 40 | 41 | class TestHTTPTransport: 42 | """Test HTTP transport functionality.""" 43 | 44 | def test_create_server_with_default_settings(self): 45 | """Test server creation with default settings.""" 46 | server = create_server() 47 | assert server is not None 48 | assert server.name == "notebook" 49 | 50 | def test_create_server_with_custom_host_port(self): 51 | """Test server creation with custom host and port.""" 52 | server = create_server(host="0.0.0.0", port=9090) 53 | assert server is not None 54 | assert server.name == "notebook" 55 | 56 | def test_create_server_with_stateless_mode(self): 57 | """Test server creation with stateless HTTP mode.""" 58 | server = create_server(stateless_http=True) 59 | assert server is not None 60 | assert server.name == "notebook" 61 | 62 | def test_http_server_startup(self): 63 | """Test that HTTP server can start and respond to requests.""" 64 | server = create_server(port=8081) 65 | 66 | # Start server in a separate thread 67 | server_thread = threading.Thread( 68 | target=lambda: server.run(transport="streamable-http"), daemon=True 69 | ) 70 | server_thread.start() 71 | 72 | # Wait for server to be ready 73 | if not wait_for_server("http://127.0.0.1:8081/", timeout=10): 74 | pytest.skip("HTTP server did not start in time") 75 | 76 | try: 77 | # Test that server responds to HTTP requests 78 | response = requests.post( 79 | "http://127.0.0.1:8081/mcp", 80 | headers={ 81 | "Content-Type": "application/json", 82 | "Accept": "application/json, text/event-stream", 83 | }, 84 | json={ 85 | "jsonrpc": "2.0", 86 | "method": "initialize", 87 | "params": { 88 | "protocolVersion": "2025-06-18", 89 | "capabilities": {"tools": {}}, 90 | "clientInfo": {"name": "test-client", "version": "1.0.0"}, 91 | }, 92 | "id": 1, 93 | }, 94 | timeout=5, 95 | ) 96 | 97 | # Check that we get a response (even if it's an error due to missing session) 98 | assert response.status_code in [200, 202, 400, 406] 99 | 100 | except requests.exceptions.RequestException: 101 | pytest.skip("HTTP server did not start in time") 102 | 103 | def test_tools_registered_correctly(self): 104 | """Test that all tools are registered when server is created.""" 105 | # The server is already created at module import time with tools registered 106 | # Just verify that the tools are accessible 107 | server = create_server() 108 | assert server is not None 109 | assert server.name == "notebook" 110 | 111 | # Since FastMCP is initialized at module level with decorators, 112 | # we can't easily mock it. The test for server creation is sufficient. 113 | 114 | def test_server_singleton_behavior(self): 115 | """Test that create_server returns the same instance when called multiple times.""" 116 | # Reset the global mcp to None 117 | import mcp_jupyter.server 118 | 119 | original_mcp = mcp_jupyter.server.mcp 120 | mcp_jupyter.server.mcp = None 121 | 122 | try: 123 | server1 = create_server() 124 | server2 = create_server() 125 | 126 | # Should return the same instance 127 | assert server1 is server2 128 | finally: 129 | # Restore original state 130 | mcp_jupyter.server.mcp = original_mcp 131 | 132 | 133 | class TestCLIArguments: 134 | """Test CLI argument parsing.""" 135 | 136 | @patch("argparse.ArgumentParser.parse_args") 137 | @patch("mcp_jupyter.create_server") 138 | @patch("sys.argv", ["mcp-jupyter"]) 139 | def test_cli_default_transport(self, mock_create, mock_parse): 140 | """Test default transport is stdio.""" 141 | from mcp_jupyter import main 142 | 143 | mock_server = MagicMock() 144 | mock_create.return_value = mock_server 145 | 146 | mock_args = MagicMock() 147 | mock_args.transport = "stdio" 148 | mock_args.port = 8000 149 | mock_args.host = "127.0.0.1" 150 | mock_args.stateless_http = False 151 | mock_parse.return_value = mock_args 152 | 153 | # This would normally run the server, but we're mocking it 154 | try: 155 | main() 156 | except SystemExit: 157 | pass 158 | 159 | mock_server.run.assert_called_once_with(transport="stdio") 160 | 161 | @patch("argparse.ArgumentParser.parse_args") 162 | @patch("mcp_jupyter.create_server") 163 | @patch("sys.argv", ["mcp-jupyter", "--transport", "http", "--port", "8080"]) 164 | def test_cli_http_transport(self, mock_create, mock_parse): 165 | """Test HTTP transport argument.""" 166 | from mcp_jupyter import main 167 | 168 | mock_server = MagicMock() 169 | mock_create.return_value = mock_server 170 | 171 | mock_args = MagicMock() 172 | mock_args.transport = "http" 173 | mock_args.port = 8080 174 | mock_args.host = "127.0.0.1" 175 | mock_args.stateless_http = False 176 | mock_parse.return_value = mock_args 177 | 178 | try: 179 | main() 180 | except SystemExit: 181 | pass 182 | 183 | # Check that streamable-http transport is used 184 | mock_server.run.assert_called_once_with(transport="streamable-http") 185 | mock_create.assert_called_once_with( 186 | host="127.0.0.1", port=8080, stateless_http=False 187 | ) 188 | 189 | @patch("argparse.ArgumentParser.parse_args") 190 | @patch("mcp_jupyter.create_server") 191 | @patch("sys.argv", ["mcp-jupyter", "--transport", "http", "--stateless-http"]) 192 | def test_cli_stateless_mode(self, mock_create, mock_parse): 193 | """Test stateless HTTP mode argument.""" 194 | from mcp_jupyter import main 195 | 196 | mock_server = MagicMock() 197 | mock_create.return_value = mock_server 198 | 199 | mock_args = MagicMock() 200 | mock_args.transport = "http" 201 | mock_args.port = 8000 202 | mock_args.host = "127.0.0.1" 203 | mock_args.stateless_http = True 204 | mock_parse.return_value = mock_args 205 | 206 | try: 207 | main() 208 | except SystemExit: 209 | pass 210 | 211 | mock_create.assert_called_once_with( 212 | host="127.0.0.1", port=8000, stateless_http=True 213 | ) 214 | 215 | 216 | class TestHTTPEndpoints: 217 | """Test HTTP endpoint behavior.""" 218 | 219 | def test_http_initialize_endpoint(self): 220 | """Test the initialize endpoint with proper headers.""" 221 | server = create_server(port=8082) 222 | 223 | # Start server in a separate thread 224 | server_thread = threading.Thread( 225 | target=lambda: server.run(transport="streamable-http"), daemon=True 226 | ) 227 | server_thread.start() 228 | 229 | # Wait for server to be ready 230 | if not wait_for_server("http://127.0.0.1:8082/", timeout=10): 231 | pytest.skip("HTTP server did not start in time") 232 | 233 | try: 234 | # Test initialize endpoint 235 | response = requests.post( 236 | "http://127.0.0.1:8082/mcp", 237 | headers={ 238 | "Content-Type": "application/json", 239 | "Accept": "application/json, text/event-stream", 240 | }, 241 | json={ 242 | "jsonrpc": "2.0", 243 | "method": "initialize", 244 | "params": { 245 | "protocolVersion": "2025-06-18", 246 | "capabilities": {"tools": {}}, 247 | "clientInfo": {"name": "test-client", "version": "1.0.0"}, 248 | }, 249 | "id": 1, 250 | }, 251 | timeout=5, 252 | stream=True, 253 | ) 254 | 255 | assert response.status_code == 200 256 | 257 | # Parse SSE response 258 | content = response.text 259 | assert "event: message" in content or "jsonrpc" in content 260 | 261 | # Check for session ID header 262 | assert ( 263 | "mcp-session-id" in response.headers 264 | or "Mcp-Session-Id" in response.headers 265 | ) 266 | 267 | except requests.exceptions.RequestException as e: 268 | pytest.skip(f"HTTP server did not respond properly: {e}") 269 | 270 | def test_http_missing_accept_headers(self): 271 | """Test that server rejects requests without proper Accept headers.""" 272 | server = create_server(port=8083) 273 | 274 | # Start server in a separate thread 275 | server_thread = threading.Thread( 276 | target=lambda: server.run(transport="streamable-http"), daemon=True 277 | ) 278 | server_thread.start() 279 | 280 | # Wait for server to be ready 281 | if not wait_for_server("http://127.0.0.1:8083/", timeout=10): 282 | pytest.skip("HTTP server did not start in time") 283 | 284 | try: 285 | # Test with missing Accept header 286 | response = requests.post( 287 | "http://127.0.0.1:8083/mcp", 288 | headers={ 289 | "Content-Type": "application/json", 290 | # Missing Accept header 291 | }, 292 | json={ 293 | "jsonrpc": "2.0", 294 | "method": "initialize", 295 | "params": { 296 | "protocolVersion": "2025-06-18", 297 | "capabilities": {"tools": {}}, 298 | "clientInfo": {"name": "test-client", "version": "1.0.0"}, 299 | }, 300 | "id": 1, 301 | }, 302 | timeout=5, 303 | ) 304 | 305 | # Should reject with 406 Not Acceptable 306 | assert response.status_code == 406 307 | 308 | # Check error message 309 | error_data = response.json() 310 | assert "error" in error_data 311 | assert "Not Acceptable" in error_data["error"]["message"] 312 | 313 | except requests.exceptions.RequestException as e: 314 | pytest.skip(f"HTTP server did not respond properly: {e}") 315 | -------------------------------------------------------------------------------- /src/mcp_jupyter/notebook.py: -------------------------------------------------------------------------------- 1 | """ 2 | Jupyter notebook management module for MCP. 3 | 4 | This module provides functions to create, open, and interact with Jupyter notebooks. 5 | 6 | Important note about paths: 7 | --------------------------- 8 | All notebook_path parameters throughout this module are assumed to be relative to the 9 | Jupyter server root directory, not absolute paths on the local filesystem. 10 | 11 | For example, if your Jupyter server is running with root directory at '/home/user/jupyter', 12 | and you want to access a notebook at '/home/user/jupyter/examples/demo.ipynb', you would 13 | use notebook_path='examples/demo.ipynb'. 14 | 15 | This is particularly important when using port forwarding to access a Jupyter server running 16 | on a different machine. The paths are always evaluated relative to the server's root, not 17 | the client's filesystem. 18 | """ 19 | 20 | import json 21 | import logging 22 | import os 23 | import time 24 | from typing import Any, Dict, List, Optional 25 | 26 | import requests 27 | from mcp.shared.exceptions import McpError 28 | from mcp.types import INTERNAL_ERROR, ErrorData 29 | 30 | from .utils import _ensure_ipynb_extension 31 | 32 | # Setup notebook logger 33 | logger = logging.getLogger(__name__) 34 | 35 | 36 | def check_notebook_exists(notebook_path: str, server_url: str, token: str) -> bool: 37 | """Check if a notebook exists at the given path. 38 | 39 | Args: 40 | notebook_path: Path to the notebook, relative to Jupyter server root 41 | server_url: Jupyter server URL 42 | token: Authentication token 43 | 44 | Returns 45 | ------- 46 | bool: True if notebook exists, False otherwise 47 | """ 48 | notebook_path = _ensure_ipynb_extension(notebook_path) 49 | 50 | try: 51 | response = requests.get( 52 | f"{server_url}/api/contents/{notebook_path}", 53 | headers={"Authorization": f"token {token}"}, 54 | ) 55 | return response.status_code == 200 56 | except requests.RequestException: 57 | return False 58 | 59 | 60 | def create_new_notebook(notebook_path: str, server_url: str, token: str) -> None: 61 | """Create a new empty Jupyter notebook. 62 | 63 | Args: 64 | notebook_path: Path where to create the notebook, relative to Jupyter server root. 65 | Intermediate directories will be created if they don't exist. 66 | server_url: Jupyter server URL 67 | token: Authentication token 68 | 69 | Raises 70 | ------ 71 | McpError: If notebook creation fails 72 | """ 73 | notebook_path = _ensure_ipynb_extension(notebook_path) 74 | directory_path = os.path.dirname(notebook_path) 75 | 76 | # 1. Ensure the target directory exists on the server 77 | if directory_path: # Only proceed if there's a directory component 78 | try: 79 | # Check if directory exists 80 | dir_check_response = requests.get( 81 | f"{server_url}/api/contents/{directory_path}", 82 | headers={"Authorization": f"token {token}"}, 83 | ) 84 | # If directory doesn't exist (404), create it 85 | if dir_check_response.status_code == 404: 86 | logger.info(f"Directory {directory_path} not found, creating it.") 87 | dir_create_response = requests.put( 88 | f"{server_url}/api/contents/{directory_path}", 89 | headers={ 90 | "Authorization": f"token {token}", 91 | "Content-Type": "application/json", 92 | }, 93 | json={"type": "directory"}, 94 | ) 95 | dir_create_response.raise_for_status() # Raise exception if directory creation fails 96 | logger.info(f"Successfully created directory {directory_path}") 97 | # Raise exception for other non-successful status codes 98 | elif dir_check_response.status_code >= 400: 99 | dir_check_response.raise_for_status() 100 | 101 | except requests.RequestException as e: 102 | logger.error(f"Error checking or creating directory {directory_path}: {e}") 103 | if hasattr(e, "response") and e.response is not None: 104 | error_message = ( 105 | f"Server error: {e.response.status_code} - {e.response.text}" 106 | ) 107 | else: 108 | error_message = str(e) 109 | raise McpError( 110 | ErrorData( 111 | code=INTERNAL_ERROR, 112 | message=f"Could not ensure directory exists: {error_message}", 113 | ) 114 | ) 115 | 116 | # 2. Create the notebook file 117 | # Create a new notebook with default empty structure 118 | notebook_content = { 119 | "type": "notebook", 120 | "content": { 121 | "metadata": { 122 | "kernelspec": { 123 | "name": "python3", 124 | "display_name": "Python 3", 125 | "language": "python", 126 | }, 127 | "language_info": {"name": "python", "version": "3.8"}, 128 | }, 129 | "nbformat": 4, 130 | "nbformat_minor": 5, 131 | "cells": [], 132 | }, 133 | } 134 | 135 | # Notebook starts empty - use modify_notebook_cells to add content 136 | 137 | # Make the API request to create the notebook 138 | try: 139 | create_response = requests.put( 140 | f"{server_url}/api/contents/{notebook_path}", 141 | headers={ 142 | "Authorization": f"token {token}", 143 | "Content-Type": "application/json", 144 | }, 145 | json=notebook_content, 146 | ) 147 | 148 | create_response.raise_for_status() 149 | logger.info(f"Created new notebook at {notebook_path}") 150 | except requests.RequestException as e: 151 | logger.error(f"Error creating notebook: {e}") 152 | if hasattr(e, "response") and e.response is not None: 153 | error_message = ( 154 | f"Server error: {e.response.status_code} - {e.response.text}" 155 | ) 156 | else: 157 | error_message = str(e) 158 | 159 | raise McpError( 160 | ErrorData( 161 | code=INTERNAL_ERROR, 162 | message=f"Could not create notebook: {error_message}", 163 | ) 164 | ) 165 | 166 | 167 | def list_notebook_sessions(server_url: str, token: str) -> List[Dict[str, Any]]: 168 | """List all notebook sessions. 169 | 170 | Args: 171 | server_url: Jupyter server URL 172 | token: Authentication token 173 | 174 | Returns 175 | ------- 176 | List[Dict[str, Any]]: List of session data 177 | 178 | Raises 179 | ------ 180 | McpError: If unable to list sessions 181 | """ 182 | try: 183 | response = requests.get( 184 | f"{server_url}/api/sessions", headers={"Authorization": f"token {token}"} 185 | ) 186 | response.raise_for_status() 187 | return response.json() 188 | except requests.RequestException as e: 189 | logger.error(f"Error listing notebook sessions: {e}") 190 | raise McpError( 191 | ErrorData( 192 | code=INTERNAL_ERROR, 193 | message=f"Could not list notebook sessions: {str(e)}", 194 | ) 195 | ) 196 | 197 | 198 | def start_notebook_kernel( 199 | notebook_path: str, server_url: str, token: str 200 | ) -> Dict[str, Any]: 201 | """Start a kernel for a notebook. 202 | 203 | Args: 204 | notebook_path: Path to the notebook, relative to Jupyter server root 205 | server_url: Jupyter server URL 206 | token: Authentication token 207 | 208 | Returns 209 | ------- 210 | Dict[str, Any]: Kernel session information 211 | 212 | Raises 213 | ------ 214 | McpError: If kernel creation fails 215 | """ 216 | notebook_path = _ensure_ipynb_extension(notebook_path) 217 | 218 | # Check if a kernel is already running for this notebook 219 | sessions = list_notebook_sessions(server_url, token) 220 | session_exists = any(session["path"] == notebook_path for session in sessions) 221 | 222 | if session_exists: 223 | logger.info(f"Kernel already running for notebook {notebook_path}") 224 | return next(session for session in sessions if session["path"] == notebook_path) 225 | 226 | # Start a new kernel session for this notebook 227 | session_data = { 228 | "kernel": {"name": "python3"}, 229 | "name": os.path.basename(notebook_path), 230 | "path": notebook_path, 231 | "type": "notebook", 232 | } 233 | 234 | try: 235 | kernel_response = requests.post( 236 | f"{server_url}/api/sessions", 237 | headers={ 238 | "Authorization": f"token {token}", 239 | "Content-Type": "application/json", 240 | }, 241 | json=session_data, 242 | ) 243 | 244 | kernel_response.raise_for_status() 245 | logger.info(f"Started kernel for notebook {notebook_path}") 246 | return kernel_response.json() 247 | except requests.RequestException as e: 248 | logger.error(f"Error starting kernel: {e}") 249 | raise McpError( 250 | ErrorData( 251 | code=INTERNAL_ERROR, 252 | message=f"Could not start kernel for notebook: {str(e)}", 253 | ) 254 | ) 255 | 256 | 257 | def get_notebook_info( 258 | notebook_path: str, server_url: str, token: str 259 | ) -> Dict[str, Any]: 260 | """Get information about a notebook. 261 | 262 | Args: 263 | notebook_path: Path to the notebook, relative to Jupyter server root 264 | server_url: Jupyter server URL 265 | token: Authentication token 266 | 267 | Returns 268 | ------- 269 | Dict[str, Any]: Notebook information 270 | 271 | Raises 272 | ------ 273 | McpError: If unable to get notebook info 274 | """ 275 | notebook_path = _ensure_ipynb_extension(notebook_path) 276 | 277 | try: 278 | response = requests.get( 279 | f"{server_url}/api/contents/{notebook_path}", 280 | headers={"Authorization": f"token {token}"}, 281 | ) 282 | response.raise_for_status() 283 | return response.json() 284 | except requests.RequestException as e: 285 | logger.error(f"Error getting notebook info: {e}") 286 | raise McpError( 287 | ErrorData( 288 | code=INTERNAL_ERROR, 289 | message=f"Could not get notebook info: {str(e)}", 290 | ) 291 | ) 292 | 293 | 294 | def prepare_notebook( 295 | notebook_path: str, 296 | server_url: str = "http://localhost:8888", 297 | token: str = None, 298 | ) -> Dict[str, Any]: 299 | """Prepare notebook for use and start kernel. 300 | 301 | Creates an empty notebook if it doesn't exist. To add content, use the 302 | modify_notebook_cells operations after creation. 303 | 304 | Args: 305 | notebook_path: Path to the notebook, relative to Jupyter server root 306 | server_url: Jupyter server URL 307 | token: Authentication token (defaults to env var) 308 | 309 | Returns 310 | ------- 311 | Dict[str, Any]: Notebook information with status message 312 | 313 | Raises 314 | ------ 315 | McpError: If notebook preparation fails 316 | """ 317 | if token is None: 318 | token = os.getenv("TOKEN", "BLOCK") 319 | 320 | notebook_path = _ensure_ipynb_extension(notebook_path) 321 | 322 | try: 323 | # Check if notebook exists 324 | notebook_exists = check_notebook_exists(notebook_path, server_url, token) 325 | 326 | # Create notebook if it doesn't exist 327 | if not notebook_exists: 328 | create_new_notebook(notebook_path, server_url, token) 329 | notebook_created = True 330 | else: 331 | notebook_created = False 332 | logger.info(f"Notebook {notebook_path} already exists") 333 | 334 | # Start kernel 335 | start_notebook_kernel(notebook_path, server_url, token) 336 | 337 | # Get notebook info 338 | info = get_notebook_info(notebook_path, server_url, token) 339 | 340 | # Add message about notebook creation status 341 | if notebook_exists: 342 | info["message"] = f"Notebook {notebook_path} already exists" 343 | else: 344 | info["message"] = f"Notebook {notebook_path} created" 345 | 346 | return info 347 | 348 | except Exception as e: 349 | if not isinstance(e, McpError): 350 | logger.error(f"Error preparing notebook: {e}") 351 | raise McpError( 352 | ErrorData( 353 | code=INTERNAL_ERROR, message=f"Could not prepare notebook: {e}" 354 | ) 355 | ) 356 | raise e 357 | -------------------------------------------------------------------------------- /tests/test_notebook_paths.py: -------------------------------------------------------------------------------- 1 | """Tests for Jupyter notebook path handling, especially with remote Jupyter servers.""" 2 | 3 | import os 4 | import unittest 5 | from unittest.mock import MagicMock, call, patch 6 | 7 | import requests 8 | from mcp.shared.exceptions import McpError 9 | 10 | from mcp_jupyter.notebook import ( 11 | check_notebook_exists, 12 | create_new_notebook, 13 | get_notebook_info, 14 | prepare_notebook, 15 | start_notebook_kernel, 16 | ) 17 | from mcp_jupyter.utils import _ensure_ipynb_extension 18 | 19 | # Mock responses 20 | mock_success_response = MagicMock() 21 | mock_success_response.status_code = 200 22 | mock_success_response.json.return_value = {"name": "test.ipynb", "path": "test.ipynb"} 23 | 24 | mock_404_response = MagicMock() 25 | mock_404_response.status_code = 404 26 | 27 | mock_create_response = MagicMock() 28 | mock_create_response.status_code = 201 # Typically 201 for created 29 | 30 | 31 | @patch("requests.put", return_value=mock_create_response) 32 | @patch("requests.get", return_value=mock_success_response) 33 | class TestNotebookPaths(unittest.TestCase): 34 | """Test notebook path handling functionality.""" 35 | 36 | def setUp(self): 37 | """Set up test fixtures, if any.""" 38 | # Access mocks via self if needed, but they are reset automatically 39 | # by the decorator context usually. Explicit reset for clarity. 40 | # We access them via the arguments injected into the test methods. 41 | pass # No assignments needed here, mocks are handled per test method 42 | 43 | def test_check_notebook_exists_correct_path(self, mock_get, mock_put): 44 | """Test that check_notebook_exists uses the relative path.""" 45 | # Setup specific mock behavior for this test 46 | mock_get.return_value = mock_success_response 47 | mock_get.reset_mock() # Reset before use in test 48 | mock_put.reset_mock() 49 | 50 | # Test 51 | notebook_path = "subfolder/my_notebook.ipynb" 52 | server_url = "http://remote.server:8888" 53 | token = "test-token" 54 | result = check_notebook_exists(notebook_path, server_url, token) 55 | 56 | # Assert 57 | mock_get.assert_called_once_with( 58 | f"{server_url}/api/contents/{notebook_path}", 59 | headers={"Authorization": f"token {token}"}, 60 | ) 61 | self.assertTrue(result) 62 | 63 | def test_create_notebook_correct_path(self, mock_get, mock_put): 64 | """Test that create_new_notebook uses the relative path.""" 65 | # Setup 66 | # Simulate directory already exists for this test case 67 | mock_get.return_value = mock_success_response 68 | mock_put.return_value = mock_create_response 69 | mock_get.reset_mock() 70 | mock_put.reset_mock() 71 | 72 | # Test with a path that includes subdirectories 73 | notebook_path = "subfolder/my_notebook" # Function adds .ipynb 74 | expected_notebook_path_ext = "subfolder/my_notebook.ipynb" 75 | directory_path = "subfolder" 76 | cells = ["print('Hello')"] 77 | server_url = "http://remote.server:8888" 78 | token = "test-token" 79 | 80 | # Execute 81 | create_new_notebook(notebook_path, server_url, token) 82 | 83 | # Assert 84 | # Check that GET was called to check the directory 85 | mock_get.assert_called_once_with( 86 | f"{server_url}/api/contents/{directory_path}", 87 | headers={"Authorization": f"token {token}"}, 88 | ) 89 | # Check PUT was called once (only for the notebook file) 90 | mock_put.assert_called_once() 91 | put_call_args = mock_put.call_args 92 | self.assertEqual( 93 | put_call_args[0][0], 94 | f"{server_url}/api/contents/{expected_notebook_path_ext}", 95 | ) 96 | self.assertEqual(put_call_args[1]["headers"]["Authorization"], f"token {token}") 97 | 98 | def test_create_notebook_creates_directory(self, mock_get, mock_put): 99 | """Test that create_new_notebook creates the directory if it doesn't exist.""" 100 | # Setup 101 | notebook_path = "new_folder/my_test_notebook" # Function adds .ipynb 102 | expected_notebook_path_ext = "new_folder/my_test_notebook.ipynb" 103 | directory_path = "new_folder" 104 | server_url = "http://remote.server:8888" 105 | token = "test-token" 106 | cells = ["print('hello dir')"] 107 | 108 | # Configure mocks for this test: 109 | mock_get.return_value = mock_404_response 110 | mock_put.return_value = mock_create_response 111 | # Ensure raise_for_status on the PUT response mock doesn't raise 112 | mock_put.return_value.raise_for_status = MagicMock() 113 | mock_get.reset_mock() 114 | mock_put.reset_mock() 115 | 116 | # Execute 117 | create_new_notebook(notebook_path, server_url, token) 118 | 119 | # Assert 120 | # Check that GET was called to check the directory 121 | mock_get.assert_called_once_with( 122 | f"{server_url}/api/contents/{directory_path}", 123 | headers={"Authorization": f"token {token}"}, 124 | ) 125 | 126 | # Check that PUT was called twice 127 | self.assertEqual(mock_put.call_count, 2) 128 | 129 | # Get the actual calls made to the mock 130 | put_calls = mock_put.call_args_list 131 | 132 | # Assert the directory creation call details 133 | dir_call_expected = call( 134 | f"{server_url}/api/contents/{directory_path}", 135 | headers={ 136 | "Authorization": f"token {token}", 137 | "Content-Type": "application/json", 138 | }, 139 | json={"type": "directory"}, 140 | ) 141 | self.assertEqual(put_calls[0], dir_call_expected) 142 | 143 | # Assert the notebook creation call details (check URL, token, type) 144 | notebook_call_actual = put_calls[1] 145 | self.assertEqual( 146 | notebook_call_actual[0][0], 147 | f"{server_url}/api/contents/{expected_notebook_path_ext}", 148 | ) # Check URL arg 149 | self.assertEqual( 150 | notebook_call_actual[1]["headers"]["Authorization"], f"token {token}" 151 | ) # Check token in kwargs['headers'] 152 | self.assertEqual( 153 | notebook_call_actual[1]["json"]["type"], "notebook" 154 | ) # Check type in kwargs['json'] 155 | 156 | # Ensure raise_for_status was called on the mock responses 157 | self.assertEqual(mock_create_response.raise_for_status.call_count, 2) 158 | 159 | def test_get_notebook_info_correct_path(self, mock_get, mock_put): 160 | """Test that get_notebook_info uses the relative path.""" 161 | # Setup specific mock behavior 162 | mock_get.return_value = mock_success_response 163 | mock_get.reset_mock() 164 | mock_put.reset_mock() 165 | 166 | # Test 167 | notebook_path = "another/folder/info_test.ipynb" 168 | server_url = "http://host.com:9999" 169 | token = "info-token" 170 | get_notebook_info(notebook_path, server_url, token) 171 | 172 | # Assert 173 | mock_get.assert_called_once_with( 174 | f"{server_url}/api/contents/{notebook_path}", 175 | headers={"Authorization": f"token {token}"}, 176 | ) 177 | 178 | def test_prepare_notebook_end_to_end(self, mock_get, mock_put): 179 | """Test the prepare_notebook function creates and starts kernel.""" 180 | # Setup Mocks for prepare_notebook sequence: 181 | # 1. check_notebook_exists (GET /api/contents/...) -> 404 182 | # 2. create_new_notebook -> directory check (GET /api/contents/...) -> 404 183 | # 3. create_new_notebook -> directory create (PUT /api/contents/...) -> 201 184 | # 4. create_new_notebook -> notebook create (PUT /api/contents/...) -> 201 185 | # 5. list_notebook_sessions (GET /api/sessions) -> [] (no existing session) 186 | # 6. start_notebook_kernel -> session create (POST /api/sessions) -> 200 (or 201) 187 | # 7. get_notebook_info (GET /api/contents/...) -> 200 188 | 189 | mock_post_response = MagicMock() # For POST /api/sessions 190 | mock_post_response.status_code = 201 191 | mock_post_response.json.return_value = { 192 | "id": "kernel-123", 193 | "path": "prep/end_to_end.ipynb", 194 | "kernel": {"id": "k1"}, 195 | } 196 | mock_post_response.raise_for_status = MagicMock() 197 | 198 | # Use side_effect for multiple calls with different responses 199 | mock_get.side_effect = [ 200 | mock_404_response, # 1. check_notebook_exists 201 | mock_404_response, # 2. create_new_notebook (dir check) 202 | MagicMock(status_code=200, json=lambda: []), # 5. list_notebook_sessions 203 | mock_success_response, # 7. get_notebook_info 204 | ] 205 | # PUT is called twice (dir, then notebook) 206 | mock_put.return_value = mock_create_response 207 | mock_put.return_value.raise_for_status = MagicMock() 208 | 209 | mock_get.reset_mock() 210 | mock_put.reset_mock() 211 | 212 | # Patch requests.post specifically for this test 213 | with patch("requests.post", return_value=mock_post_response) as mock_post: 214 | # Test 215 | notebook_path = "prep/end_to_end" 216 | expected_notebook_path_ext = "prep/end_to_end.ipynb" 217 | server_url = "http://prepare.it:8888" 218 | token = "prep-token" 219 | cells = ["import os"] 220 | result = prepare_notebook(notebook_path, server_url, token) 221 | 222 | # Assertions 223 | # Check GET calls 224 | get_calls = mock_get.call_args_list 225 | self.assertEqual(len(get_calls), 4) 226 | self.assertEqual( 227 | get_calls[0], 228 | call( 229 | f"{server_url}/api/contents/{expected_notebook_path_ext}", 230 | headers={"Authorization": f"token {token}"}, 231 | ), 232 | ) # Check exists 233 | self.assertEqual( 234 | get_calls[1], 235 | call( 236 | f"{server_url}/api/contents/prep", 237 | headers={"Authorization": f"token {token}"}, 238 | ), 239 | ) # Create dir check 240 | self.assertEqual( 241 | get_calls[2], 242 | call( 243 | f"{server_url}/api/sessions", 244 | headers={"Authorization": f"token {token}"}, 245 | ), 246 | ) # List sessions 247 | self.assertEqual( 248 | get_calls[3], 249 | call( 250 | f"{server_url}/api/contents/{expected_notebook_path_ext}", 251 | headers={"Authorization": f"token {token}"}, 252 | ), 253 | ) # Get info 254 | 255 | # Check PUT calls (dir + notebook) 256 | self.assertEqual(mock_put.call_count, 2) 257 | put_calls = mock_put.call_args_list 258 | self.assertEqual( 259 | put_calls[0], 260 | call( 261 | f"{server_url}/api/contents/prep", 262 | headers={ 263 | "Authorization": f"token {token}", 264 | "Content-Type": "application/json", 265 | }, 266 | json={"type": "directory"}, 267 | ), 268 | ) 269 | self.assertEqual( 270 | put_calls[1][0][0], 271 | f"{server_url}/api/contents/{expected_notebook_path_ext}", 272 | ) # Check URL of notebook PUT 273 | self.assertEqual(put_calls[1][1]["headers"]["Authorization"], f"token {token}") 274 | 275 | # Check POST call (start kernel) 276 | mock_post.assert_called_once() 277 | post_call_args = mock_post.call_args 278 | self.assertEqual(post_call_args[0][0], f"{server_url}/api/sessions") 279 | self.assertEqual(post_call_args[1]["json"]["path"], expected_notebook_path_ext) 280 | self.assertEqual( 281 | post_call_args[1]["headers"]["Authorization"], f"token {token}" 282 | ) 283 | 284 | # Check result content 285 | self.assertIn("message", result) 286 | self.assertTrue(result["message"].endswith("created")) 287 | self.assertEqual( 288 | result["path"], "test.ipynb" 289 | ) # From mock_success_response json 290 | 291 | def test_start_kernel_correct_path(self, mock_get, mock_put): 292 | """Test that start_notebook_kernel uses the relative path and checks sessions.""" 293 | # Setup Mocks 294 | # 1. list_notebook_sessions (GET /api/sessions) -> [] (no existing session) 295 | # 2. start_notebook_kernel -> session create (POST /api/sessions) -> 200 (or 201) 296 | 297 | mock_post_response = MagicMock() 298 | mock_post_response.status_code = 201 299 | mock_post_response.json.return_value = { 300 | "id": "kernel-xyz", 301 | "path": "start/me/up.ipynb", 302 | "kernel": {"id": "k2"}, 303 | } 304 | mock_post_response.raise_for_status = MagicMock() 305 | 306 | # Mock GET for list_notebook_sessions to return empty list 307 | mock_get.return_value = MagicMock(status_code=200, json=lambda: []) 308 | mock_get.return_value.raise_for_status = MagicMock() # Mock raise_for_status 309 | 310 | mock_get.reset_mock() 311 | mock_put.reset_mock() 312 | 313 | # Patch requests.post specifically for this test 314 | with patch("requests.post", return_value=mock_post_response) as mock_post: 315 | # Test 316 | notebook_path = "start/me/up" 317 | expected_notebook_path_ext = "start/me/up.ipynb" 318 | server_url = "http://start.it:8888" 319 | token = "start-token" 320 | start_notebook_kernel(notebook_path, server_url, token) 321 | 322 | # Assertions 323 | # Check GET call (list sessions) 324 | mock_get.assert_called_once_with( 325 | f"{server_url}/api/sessions", headers={"Authorization": f"token {token}"} 326 | ) 327 | 328 | # Check POST call (create session) 329 | mock_post.assert_called_once() 330 | post_call_args = mock_post.call_args 331 | self.assertEqual(post_call_args[0][0], f"{server_url}/api/sessions") # URL 332 | self.assertEqual( 333 | post_call_args[1]["headers"]["Authorization"], f"token {token}" 334 | ) # Header 335 | self.assertEqual( 336 | post_call_args[1]["json"]["path"], expected_notebook_path_ext 337 | ) # Body path 338 | self.assertEqual(post_call_args[1]["json"]["type"], "notebook") # Body type 339 | 340 | 341 | if __name__ == "__main__": 342 | unittest.main() 343 | -------------------------------------------------------------------------------- /src/mcp_jupyter/rest_client.py: -------------------------------------------------------------------------------- 1 | """Simple REST-based notebook client for Jupyter. 2 | 3 | This module provides a clean, reliable interface to Jupyter notebooks using 4 | only REST API calls. It avoids the complexity of RTC/WebSocket connections 5 | while still benefiting from server-side collaborative features. 6 | 7 | TODO: Add batch operations context manager for better performance with large notebooks. 8 | Could implement `with client.batch() as batch:` pattern to group operations. 9 | """ 10 | 11 | import json 12 | import logging 13 | from typing import Any, Dict, List, Optional 14 | 15 | import requests 16 | 17 | # Setup logger 18 | logger = logging.getLogger(__name__) 19 | 20 | 21 | class NotebookClient: 22 | """Simple REST-based Jupyter notebook client. 23 | 24 | This client uses only REST API calls to interact with Jupyter notebooks, 25 | avoiding the complexity and reliability issues of WebSocket/RTC connections. 26 | The server-side RTC infrastructure still provides collaboration benefits 27 | for other clients. It maintains the same interface as the 28 | jupyter_nbmodel_client.NbModelClient RTC client. 29 | """ 30 | 31 | def __init__( 32 | self, server_url: str, notebook_path: str, token: Optional[str] = None 33 | ): 34 | """Initialize the notebook client. 35 | 36 | Args: 37 | server_url: Base Jupyter server URL (e.g., http://localhost:8888) 38 | notebook_path: Path to notebook relative to server root 39 | token: Authentication token for the server 40 | """ 41 | self._server_url = server_url.rstrip("/") 42 | self._notebook_path = notebook_path 43 | self._token = token 44 | 45 | @property 46 | def connected(self) -> bool: 47 | """Check if client is connected.""" 48 | try: 49 | response = requests.get( 50 | f"{self._server_url}/api/contents", 51 | headers=self._make_request_headers(), 52 | timeout=10, 53 | ) 54 | return response.status_code == 200 55 | except: 56 | return False 57 | 58 | @property 59 | def cells(self) -> List[Dict[str, Any]]: 60 | """Get all cells as a list of dictionaries.""" 61 | return self._get_notebook_content().get("cells", []) 62 | 63 | @property 64 | def _doc(self): 65 | """Compatibility property to access cells like the old RTC client. 66 | 67 | This provides compatibility with code that accesses notebook._doc.ycells.to_py(). 68 | """ 69 | 70 | class CellsWrapper: 71 | def __init__(self, cells): 72 | self._cells = cells 73 | 74 | def to_py(self): 75 | return self._cells 76 | 77 | def __iter__(self): 78 | return iter(self._cells) 79 | 80 | def __len__(self): 81 | return len(self._cells) 82 | 83 | def __getitem__(self, index): 84 | return self._cells[index] 85 | 86 | class DocWrapper: 87 | def __init__(self, cells): 88 | self.ycells = CellsWrapper(cells) 89 | 90 | return DocWrapper(self.cells) 91 | 92 | def _get_notebook_content(self) -> Dict[str, Any]: 93 | """Get current notebook content from server.""" 94 | try: 95 | response = requests.get( 96 | f"{self._server_url}/api/contents/{self._notebook_path}", 97 | headers=self._make_request_headers(), 98 | timeout=10, 99 | ) 100 | 101 | if response.status_code == 404: 102 | self._create_empty_notebook() 103 | return self._get_notebook_content() 104 | else: 105 | response.raise_for_status() 106 | data = response.json() 107 | return data.get("content", {}) 108 | 109 | except Exception as e: 110 | logger.error(f"Failed to get notebook content: {e}") 111 | raise 112 | 113 | def _make_request_headers(self) -> Dict[str, str]: 114 | """Create headers for REST requests.""" 115 | headers = {"Content-Type": "application/json"} 116 | if self._token: 117 | headers["Authorization"] = f"token {self._token}" 118 | return headers 119 | 120 | def _get_default_kernel_info(self) -> tuple[Dict[str, Any], Dict[str, Any]]: 121 | """Get default kernel specification from the Jupyter server. 122 | 123 | Returns 124 | ------- 125 | tuple: (kernelspec, language_info) dictionaries 126 | 127 | Raises 128 | ------ 129 | requests.RequestException: If unable to get kernel specs from server 130 | """ 131 | # Get available kernel specs 132 | response = requests.get( 133 | f"{self._server_url}/api/kernelspecs", 134 | headers=self._make_request_headers(), 135 | timeout=10, 136 | ) 137 | response.raise_for_status() 138 | data = response.json() 139 | 140 | # Get default kernel name (usually 'python3' or similar) 141 | default_kernel_name = data.get("default", "python3") 142 | kernelspecs = data.get("kernelspecs", {}) 143 | 144 | if default_kernel_name not in kernelspecs: 145 | raise RuntimeError( 146 | f"Default kernel '{default_kernel_name}' not found in available kernelspecs" 147 | ) 148 | 149 | spec = kernelspecs[default_kernel_name]["spec"] 150 | kernelspec = { 151 | "display_name": spec.get("display_name", "Python 3"), 152 | "language": spec.get("language", "python"), 153 | "name": default_kernel_name, 154 | } 155 | 156 | # Extract language info 157 | language_info = { 158 | "name": spec.get("language", "python"), 159 | } 160 | 161 | # Add version if available in metadata 162 | metadata = spec.get("metadata", {}) 163 | if "interpreter" in metadata: 164 | interpreter = metadata["interpreter"] 165 | if "version" in interpreter: 166 | language_info["version"] = interpreter["version"] 167 | 168 | return kernelspec, language_info 169 | 170 | def connect(self) -> None: 171 | """Connect to the Jupyter server.""" 172 | try: 173 | response = requests.get( 174 | f"{self._server_url}/api/contents", 175 | headers=self._make_request_headers(), 176 | timeout=10, 177 | ) 178 | response.raise_for_status() 179 | logger.info(f"✅ Connected to Jupyter server at {self._server_url}") 180 | 181 | except Exception as e: 182 | logger.error(f"❌ Failed to connect to Jupyter server: {e}") 183 | raise 184 | 185 | def disconnect(self) -> None: 186 | """Disconnect from the server.""" 187 | logger.info("Disconnected from Jupyter server") 188 | 189 | def _create_empty_notebook(self) -> None: 190 | """Create an empty notebook on the server.""" 191 | kernelspec, language_info = self._get_default_kernel_info() 192 | 193 | empty_notebook = { 194 | "cells": [], 195 | "metadata": { 196 | "kernelspec": kernelspec, 197 | "language_info": language_info, 198 | }, 199 | "nbformat": 4, 200 | "nbformat_minor": 5, 201 | } 202 | 203 | notebook_data = {"type": "notebook", "content": empty_notebook} 204 | 205 | response = requests.put( 206 | f"{self._server_url}/api/contents/{self._notebook_path}", 207 | headers=self._make_request_headers(), 208 | data=json.dumps(notebook_data), 209 | timeout=10, 210 | ) 211 | response.raise_for_status() 212 | logger.info(f"Created empty notebook: {self._notebook_path}") 213 | 214 | def _save_notebook(self, notebook_content: Dict[str, Any]) -> None: 215 | """Save notebook content to the server.""" 216 | notebook_data = {"type": "notebook", "content": notebook_content} 217 | 218 | response = requests.put( 219 | f"{self._server_url}/api/contents/{self._notebook_path}", 220 | headers=self._make_request_headers(), 221 | data=json.dumps(notebook_data), 222 | timeout=10, 223 | ) 224 | response.raise_for_status() 225 | logger.debug("Notebook saved successfully") 226 | 227 | def add_code_cell(self, content: str) -> int: 228 | """Add a code cell at the end of the notebook. 229 | 230 | Args: 231 | content: Source code for the cell 232 | 233 | Returns 234 | ------- 235 | Position index where the cell was inserted 236 | """ 237 | notebook_content = self._get_notebook_content() 238 | 239 | cell = { 240 | "cell_type": "code", 241 | "source": content, 242 | "metadata": {}, 243 | "outputs": [], 244 | "execution_count": None, 245 | } 246 | 247 | notebook_content["cells"].append(cell) 248 | self._save_notebook(notebook_content) 249 | return len(notebook_content["cells"]) - 1 250 | 251 | def insert_code_cell(self, position: int, content: str) -> None: 252 | """Insert a code cell at a specific position. 253 | 254 | Args: 255 | position: Position to insert at (0-indexed) 256 | content: Source code for the cell 257 | """ 258 | notebook_content = self._get_notebook_content() 259 | 260 | cell = { 261 | "cell_type": "code", 262 | "source": content, 263 | "metadata": {}, 264 | "outputs": [], 265 | "execution_count": None, 266 | } 267 | 268 | # Ensure position is within bounds 269 | max_position = len(notebook_content["cells"]) 270 | if position > max_position: 271 | position = max_position 272 | 273 | notebook_content["cells"].insert(position, cell) 274 | self._save_notebook(notebook_content) 275 | 276 | def add_markdown_cell(self, content: str) -> int: 277 | """Add a markdown cell at the end of the notebook. 278 | 279 | Args: 280 | content: Markdown content for the cell 281 | 282 | Returns 283 | ------- 284 | Position index where the cell was inserted 285 | """ 286 | notebook_content = self._get_notebook_content() 287 | 288 | cell = {"cell_type": "markdown", "source": content, "metadata": {}} 289 | 290 | notebook_content["cells"].append(cell) 291 | self._save_notebook(notebook_content) 292 | return len(notebook_content["cells"]) - 1 293 | 294 | def insert_markdown_cell(self, position: int, content: str) -> None: 295 | """Insert a markdown cell at a specific position. 296 | 297 | Args: 298 | position: Position to insert at (0-indexed) 299 | content: Markdown content for the cell 300 | """ 301 | notebook_content = self._get_notebook_content() 302 | 303 | cell = {"cell_type": "markdown", "source": content, "metadata": {}} 304 | 305 | # Ensure position is within bounds 306 | max_position = len(notebook_content["cells"]) 307 | if position > max_position: 308 | position = max_position 309 | 310 | notebook_content["cells"].insert(position, cell) 311 | self._save_notebook(notebook_content) 312 | 313 | def edit_cell(self, position: int, new_content: str) -> None: 314 | """Edit the content of a cell at the specified position. 315 | 316 | Args: 317 | position: Position of the cell to edit (0-indexed) 318 | new_content: New content for the cell 319 | """ 320 | notebook_content = self._get_notebook_content() 321 | cells = notebook_content["cells"] 322 | if position >= len(cells): 323 | raise IndexError(f"Cell index {position} out of range") 324 | 325 | cells[position]["source"] = new_content 326 | self._save_notebook(notebook_content) 327 | 328 | def delete_cell(self, position: int) -> None: 329 | """Delete a cell at the specified position. 330 | 331 | Args: 332 | position: Position of the cell to delete (0-indexed) 333 | """ 334 | notebook_content = self._get_notebook_content() 335 | cells = notebook_content["cells"] 336 | if position >= len(cells): 337 | raise IndexError(f"Cell index {position} out of range") 338 | 339 | cells.pop(position) 340 | self._save_notebook(notebook_content) 341 | 342 | def execute_cell(self, position: int, kernel_client) -> Dict[str, Any]: 343 | """Execute a code cell using the provided kernel client. 344 | 345 | Args: 346 | position: Position of the cell to execute (0-indexed) 347 | kernel_client: The kernel client to use for execution 348 | 349 | Returns 350 | ------- 351 | Execution results containing outputs and execution count 352 | """ 353 | notebook_content = self._get_notebook_content() 354 | cells = notebook_content["cells"] 355 | if position >= len(cells): 356 | raise IndexError(f"Cell index {position} out of range") 357 | 358 | cell = cells[position] 359 | if cell["cell_type"] != "code": 360 | raise ValueError(f"Cell at position {position} is not a code cell") 361 | 362 | # Execute using the existing kernel client 363 | result = kernel_client.execute(cell["source"]) 364 | 365 | # Update cell with results 366 | cell["outputs"] = result.get("outputs", []) 367 | cell["execution_count"] = result.get("execution_count") 368 | self._save_notebook(notebook_content) 369 | 370 | return result 371 | 372 | def get_cell(self, position: int) -> Dict[str, Any]: 373 | """Get a cell by position. 374 | 375 | Args: 376 | position: Position of the cell (0-indexed) 377 | 378 | Returns 379 | ------- 380 | Cell data as dictionary 381 | """ 382 | notebook_content = self._get_notebook_content() 383 | cells = notebook_content["cells"] 384 | if position >= len(cells): 385 | raise IndexError(f"Cell index {position} out of range") 386 | 387 | return cells[position] 388 | 389 | def refresh(self) -> None: 390 | """Refresh notebook content from server to detect external changes.""" 391 | pass 392 | 393 | def __getitem__(self, position: int) -> Dict[str, Any]: 394 | """Get a cell by position using bracket notation. 395 | 396 | Args: 397 | position: Position of the cell (0-indexed) 398 | 399 | Returns 400 | ------- 401 | Cell data as dictionary 402 | """ 403 | return self.get_cell(position) 404 | 405 | def __setitem__(self, position: int, cell_data: Dict[str, Any]) -> None: 406 | """Set cell content using bracket notation. 407 | 408 | Args: 409 | position: Position of the cell (0-indexed) 410 | cell_data: New cell data dictionary 411 | """ 412 | notebook_content = self._get_notebook_content() 413 | cells = notebook_content["cells"] 414 | if position >= len(cells): 415 | raise IndexError(f"Cell index {position} out of range") 416 | 417 | # Update the cell in place 418 | cells[position] = cell_data 419 | self._save_notebook(notebook_content) 420 | -------------------------------------------------------------------------------- /tests/test_integration.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | 3 | from mcp_jupyter.server import ( 4 | execute_notebook_code, 5 | modify_notebook_cells, 6 | query_notebook, 7 | setup_notebook, 8 | ) 9 | 10 | # TOKEN constant 11 | TOKEN = "BLOCK" 12 | 13 | 14 | # Server fixtures are now provided by conftest.py 15 | 16 | 17 | # Test notebook fixture is now provided by conftest.py 18 | 19 | 20 | def test_notebook_creation(jupyter_server): 21 | """Test notebook creation functionality.""" 22 | notebook_name = "test_creation" 23 | 24 | # Create a new notebook - specify server_url on creation 25 | result = setup_notebook(notebook_name, server_url=jupyter_server) 26 | assert result is not None 27 | assert "message" in result 28 | assert result["message"] == f"Notebook {notebook_name}.ipynb created" 29 | 30 | # Try creating the same notebook again - no need to specify server_url 31 | result = setup_notebook(notebook_name) 32 | assert result["message"] == f"Notebook {notebook_name}.ipynb already exists" 33 | 34 | 35 | def test_add_code_cell(jupyter_server, test_notebook): 36 | """Test adding a code cell to a notebook.""" 37 | # Add a simple code cell - test_notebook path is already relative to root 38 | result = modify_notebook_cells( 39 | test_notebook, "add_code", "x = 10\nprint(f'x = {x}')" 40 | ) 41 | 42 | # Verify execution results 43 | assert "execution_count" in result 44 | assert "outputs" in result 45 | assert len(result["outputs"]) > 0 46 | assert result["status"] == "ok" 47 | 48 | # Check output content 49 | output_text = "" 50 | for output in result["outputs"]: 51 | if output["output_type"] == "stream": 52 | output_text += output["text"] 53 | 54 | assert "x = 10" in output_text 55 | 56 | # Test adding cell at specific position (not at the bottom) 57 | # First get current cell count 58 | all_cells_before = query_notebook(test_notebook, "view_source") 59 | initial_count = len(all_cells_before) 60 | 61 | # Add a cell at position 1 (second position) 62 | positioned_result = modify_notebook_cells( 63 | test_notebook, 64 | "add_code", 65 | "y = 20\nprint(f'y = {y}')", 66 | position_index=1, 67 | execute=False, 68 | ) 69 | 70 | # Verify the cell was added (should be empty result since execute=False) 71 | assert positioned_result == {} 72 | 73 | # Check that the cell was inserted at the correct position 74 | all_cells_after = query_notebook(test_notebook, "view_source") 75 | assert len(all_cells_after) == initial_count + 1 76 | 77 | # Verify the cell at position 1 contains our new content 78 | inserted_cell = all_cells_after[1] 79 | assert inserted_cell["cell_type"] == "code" 80 | assert "y = 20" in inserted_cell["source"] 81 | assert inserted_cell["execution_count"] is None # Not executed 82 | 83 | 84 | def test_add_markdown_cell(jupyter_server, test_notebook): 85 | """Test adding a markdown cell to a notebook.""" 86 | # Add a markdown cell - test_notebook path is already relative to root 87 | result = modify_notebook_cells( 88 | test_notebook, 89 | "add_markdown", 90 | "# Test Markdown\nThis is a *markdown* cell with **formatting**.", 91 | ) 92 | 93 | # Verify result 94 | assert result["message"] == "Markdown cell added" 95 | assert not result["error"] 96 | 97 | 98 | def test_edit_markdown_cell(jupyter_server, test_notebook): 99 | """Test editing a markdown cell.""" 100 | # First add a markdown cell to edit 101 | add_result = modify_notebook_cells( 102 | test_notebook, 103 | "add_markdown", 104 | cell_content="# Original Title\n\nOriginal content.", 105 | ) 106 | assert add_result["message"] == "Markdown cell added" 107 | 108 | # Get all cells to find the markdown cell we just added 109 | all_cells = query_notebook(test_notebook, "view_source") 110 | 111 | # Find the markdown cell by content 112 | markdown_position = None 113 | for i, cell in enumerate(all_cells): 114 | if cell.get("cell_type") == "markdown" and "Original Title" in cell.get( 115 | "source", "" 116 | ): 117 | markdown_position = i 118 | break 119 | 120 | assert markdown_position is not None, ( 121 | "Could not find the markdown cell we just added" 122 | ) 123 | 124 | # Edit the markdown cell 125 | edit_result = modify_notebook_cells( 126 | test_notebook, 127 | "edit_markdown", 128 | cell_content="# Updated Title\n\nThis content has been updated!", 129 | position_index=markdown_position, 130 | ) 131 | 132 | # Verify edit result 133 | assert edit_result["message"] == "Markdown cell edited" 134 | assert not edit_result["error"] 135 | 136 | # Verify the cell was actually changed 137 | updated_cell = query_notebook( 138 | test_notebook, "view_source", position_index=markdown_position 139 | ) 140 | assert "Updated Title" in updated_cell["source"] 141 | assert "This content has been updated!" in updated_cell["source"] 142 | assert updated_cell["cell_type"] == "markdown" 143 | 144 | 145 | def test_view_source(jupyter_server, test_notebook): 146 | """Test viewing notebook source.""" 147 | # View all cells - test_notebook path is already relative to root 148 | all_cells = query_notebook(test_notebook, "view_source") 149 | 150 | # Verify we got a list of cells 151 | assert isinstance(all_cells, list) 152 | assert len(all_cells) >= 2 # Should have at least our 2 initial cells 153 | 154 | # Find the cell with the add function by content, not by execution count 155 | cell_with_add_function = None 156 | for cell in all_cells: 157 | if cell.get("source") and "def add(a, b):" in cell.get("source"): 158 | cell_with_add_function = cell 159 | break 160 | 161 | assert cell_with_add_function is not None 162 | execution_count = cell_with_add_function.get("execution_count") 163 | 164 | # Now view just that specific cell by execution count (if it has one) 165 | if execution_count is not None: 166 | specific_cell = query_notebook( 167 | test_notebook, "view_source", execution_count=int(execution_count) 168 | ) 169 | assert isinstance(specific_cell, dict) 170 | assert "def add(a, b):" in specific_cell["source"] 171 | else: 172 | # If no execution count (cell might not have been executed yet), 173 | # find the cell by position instead 174 | position = None 175 | for i, cell in enumerate(all_cells): 176 | if cell.get("source") and "def add(a, b):" in cell.get("source"): 177 | position = i 178 | break 179 | 180 | if position is not None: 181 | specific_cell = query_notebook( 182 | test_notebook, "view_source", position_index=position 183 | ) 184 | assert isinstance(specific_cell, dict) 185 | assert "def add(a, b):" in specific_cell["source"] 186 | 187 | 188 | def test_get_position_index(jupyter_server, test_notebook): 189 | """Test getting the position index of a cell.""" 190 | # First, explicitly execute a cell to ensure we have at least one with an execution count 191 | # Add a cell we can easily identify 192 | modify_notebook_cells( 193 | test_notebook, 194 | "add_code", 195 | "# Test cell for get_position_index\nprint('Hello from test cell')", 196 | ) 197 | 198 | # Now get all cells 199 | all_cells = query_notebook(test_notebook, "view_source") 200 | 201 | # Find our cell either by content or by execution count 202 | position_to_find = None 203 | cell_id_to_find = None 204 | 205 | for i, cell in enumerate(all_cells): 206 | if cell.get("source") and "Test cell for get_position_index" in cell.get( 207 | "source" 208 | ): 209 | position_to_find = i 210 | cell_id_to_find = cell.get("id") 211 | execution_count = cell.get("execution_count") 212 | break 213 | 214 | assert position_to_find is not None, "Could not find our test cell" 215 | 216 | # Try to get position by content (using cell_id) 217 | if cell_id_to_find: 218 | position_by_id = query_notebook( 219 | test_notebook, "get_position_index", cell_id=cell_id_to_find 220 | ) 221 | assert position_by_id == position_to_find 222 | 223 | # If we have an execution count, test that path too 224 | if execution_count is not None: 225 | position_by_exec = query_notebook( 226 | test_notebook, "get_position_index", execution_count=int(execution_count) 227 | ) 228 | assert position_by_exec == position_to_find 229 | 230 | # If we don't have an execution count, just log a message 231 | else: 232 | print("Cell has no execution_count, skipping that part of the test") 233 | 234 | 235 | def test_edit_code_cell(jupyter_server, test_notebook): 236 | """Test editing a code cell.""" 237 | # First, view all cells to find the one we want to edit 238 | all_cells = query_notebook(test_notebook, "view_source") 239 | 240 | # Find the cell with the add function by content 241 | position_index = None 242 | for i, cell in enumerate(all_cells): 243 | if cell.get("source") and "def add(a, b):" in cell.get("source"): 244 | position_index = i 245 | break 246 | 247 | # If we didn't find the add function cell, use the first code cell 248 | if position_index is None: 249 | for i, cell in enumerate(all_cells): 250 | if cell.get("cell_type") == "code": 251 | position_index = i 252 | break 253 | 254 | assert position_index is not None, "Could not find a code cell to edit" 255 | 256 | # Edit the cell 257 | modified_code = "def multiply(a, b):\n return a * b\n\nprint(multiply(3, 4))" 258 | result = modify_notebook_cells( 259 | test_notebook, "edit_code", modified_code, position_index 260 | ) 261 | 262 | # Verify execution results 263 | assert "execution_count" in result 264 | assert "outputs" in result 265 | assert result["status"] == "ok" 266 | 267 | # Check output content 268 | output_text = "" 269 | for output in result["outputs"]: 270 | if output["output_type"] == "stream": 271 | output_text += output["text"] 272 | 273 | assert "12" in output_text # 3 * 4 = 12 274 | 275 | # Verify the cell was actually changed 276 | updated_cell = query_notebook( 277 | test_notebook, "view_source", position_index=position_index 278 | ) 279 | assert "def multiply(a, b):" in updated_cell["source"] 280 | 281 | # Test editing by execution_count (two-step workflow) 282 | # Get the execution_count of the cell we just edited 283 | execution_count = updated_cell.get("execution_count") 284 | assert execution_count is not None, ( 285 | "Cell should have execution_count after being executed" 286 | ) 287 | 288 | # Step 1: Find position by execution_count 289 | found_position = query_notebook( 290 | test_notebook, "get_position_index", execution_count=execution_count 291 | ) 292 | assert found_position == position_index, ( 293 | "execution_count lookup should return same position" 294 | ) 295 | 296 | # Step 2: Edit the cell using the found position 297 | modified_code2 = "def divide(a, b):\n return a / b\n\nprint(divide(12, 3))" 298 | result2 = modify_notebook_cells( 299 | test_notebook, "edit_code", modified_code2, found_position 300 | ) 301 | 302 | # Verify the second edit worked 303 | assert "execution_count" in result2 304 | assert result2["status"] == "ok" 305 | 306 | # Check output content 307 | output_text2 = "" 308 | for output in result2["outputs"]: 309 | if output["output_type"] == "stream": 310 | output_text2 += output["text"] 311 | assert "4.0" in output_text2 # 12 / 3 = 4.0 312 | 313 | # Verify the cell content changed again 314 | final_cell = query_notebook( 315 | test_notebook, "view_source", position_index=position_index 316 | ) 317 | assert "def divide(a, b):" in final_cell["source"] 318 | 319 | 320 | def test_execute_cell(jupyter_server, test_notebook): 321 | """Test executing a cell.""" 322 | # First add a cell without executing it - no need to specify server_url 323 | result = modify_notebook_cells( 324 | test_notebook, 325 | "add_code", 326 | "result = 5 ** 2\nprint(f'5 squared is {result}')", 327 | execute=False, 328 | ) 329 | 330 | # When execute=False, we get position_index back, not a result dict 331 | # Get all cells to find the last one (which should be the one we just added) 332 | all_cells = query_notebook(test_notebook, "view_source") 333 | position_index = len(all_cells) - 1 334 | 335 | # Now execute it - no need to specify server_url 336 | result = execute_notebook_code(test_notebook, "execute_cell", position_index) 337 | 338 | # Verify execution results 339 | assert "execution_count" in result 340 | assert "outputs" in result 341 | assert result["status"] == "ok" 342 | 343 | # Check output content 344 | output_text = "" 345 | for output in result["outputs"]: 346 | if output["output_type"] == "stream": 347 | output_text += output["text"] 348 | 349 | assert "5 squared is 25" in output_text 350 | 351 | 352 | def test_delete_cell(jupyter_server, test_notebook): 353 | """Test deleting a cell.""" 354 | # Add a cell that we'll delete - no need to specify server_url 355 | modify_notebook_cells(test_notebook, "add_code", "# This cell will be deleted") 356 | 357 | # Get all cells to find the last one (which should be the one we just added) 358 | all_cells = query_notebook(test_notebook, "view_source") 359 | position_index = len(all_cells) - 1 360 | 361 | # Delete the cell - no need to specify server_url 362 | result = modify_notebook_cells( 363 | test_notebook, "delete", position_index=position_index 364 | ) 365 | 366 | # Verify result 367 | assert result["message"] == "Cell deleted" 368 | assert not result["error"] 369 | 370 | # Verify the cell was actually deleted 371 | updated_cells = query_notebook(test_notebook, "view_source") 372 | assert len(updated_cells) == len(all_cells) - 1 373 | 374 | 375 | def test_install_packages(jupyter_server, test_notebook): 376 | """Test installing packages.""" 377 | # Install a small, common package - no need to specify server_url 378 | result = execute_notebook_code( 379 | test_notebook, "install_packages", package_names="pyyaml" 380 | ) 381 | 382 | # Just verify we got a string response 383 | assert isinstance(result, str) 384 | assert "pyyaml" in result 385 | 386 | # Verify we can import the package - no need to specify server_url 387 | import_result = modify_notebook_cells( 388 | test_notebook, "add_code", "import yaml\nprint('PyYAML successfully imported')" 389 | ) 390 | 391 | # Check output content 392 | output_text = "" 393 | for output in import_result["outputs"]: 394 | if output["output_type"] == "stream": 395 | output_text += output["text"] 396 | 397 | assert "successfully imported" in output_text 398 | 399 | 400 | def test_check_jupyter_server(jupyter_server): 401 | """Test that check_jupyter_server correctly verifies server connectivity.""" 402 | # We still need to specify server_url here since this function doesn't use notebook_path 403 | result = query_notebook("", "check_server", server_url=jupyter_server) 404 | assert result == "Jupyter server is running" 405 | 406 | 407 | def test_complex_code_execution(jupyter_server, test_notebook): 408 | """Test executing more complex code with multiple outputs.""" 409 | # Add a cell with multiple print statements and a calculation - no need to specify server_url 410 | code = """ 411 | import math 412 | 413 | def calculate_circle_properties(radius): 414 | area = math.pi * radius ** 2 415 | circumference = 2 * math.pi * radius 416 | return area, circumference 417 | 418 | radius = 5 419 | area, circumference = calculate_circle_properties(radius) 420 | 421 | print(f"Radius: {radius}") 422 | print(f"Area: {area:.2f}") 423 | print(f"Circumference: {circumference:.2f}") 424 | """ 425 | 426 | result = modify_notebook_cells(test_notebook, "add_code", code) 427 | 428 | # Verify execution results 429 | assert result["status"] == "ok" 430 | 431 | # Check output content 432 | output_text = "" 433 | for output in result["outputs"]: 434 | if output["output_type"] == "stream": 435 | output_text += output["text"] 436 | 437 | assert "Radius: 5" in output_text 438 | assert "Area: 78.54" in output_text 439 | assert "Circumference: 31.42" in output_text 440 | 441 | 442 | def test_notebook_creation_with_new_directory(jupyter_server): 443 | """Test that creating a notebook in a non-existent directory works.""" 444 | import requests 445 | 446 | dir_name = "new_dir_integration" 447 | notebook_base_name = "my_subdir_notebook" 448 | # Path relative to the server root (where jupyter lab was started) 449 | relative_notebook_path = f"{dir_name}/{notebook_base_name}" 450 | 451 | # 1. Attempt to create the notebook (this should also create the directory) 452 | creation_result = setup_notebook(relative_notebook_path, server_url=jupyter_server) 453 | assert "message" in creation_result 454 | assert "created" in creation_result["message"] # Check it was created 455 | 456 | # 2. Verify the directory exists via API 457 | try: 458 | dir_response = requests.get( 459 | f"{jupyter_server}/api/contents/{dir_name}", 460 | headers={"Authorization": f"token {TOKEN}"}, 461 | ) 462 | dir_response.raise_for_status() # Should be 200 OK 463 | dir_data = dir_response.json() 464 | assert dir_data["type"] == "directory" 465 | assert dir_data["name"] == dir_name 466 | except requests.RequestException as e: 467 | pytest.fail(f"Failed to verify directory existence via API: {e}") 468 | 469 | # 3. Verify the notebook file exists via API 470 | try: 471 | nb_response = requests.get( 472 | f"{jupyter_server}/api/contents/{relative_notebook_path}.ipynb", 473 | headers={"Authorization": f"token {TOKEN}"}, 474 | ) 475 | nb_response.raise_for_status() # Should be 200 OK 476 | nb_data = nb_response.json() 477 | assert nb_data["type"] == "notebook" 478 | assert nb_data["name"] == f"{notebook_base_name}.ipynb" 479 | except requests.RequestException as e: 480 | pytest.fail(f"Failed to verify notebook existence via API: {e}") 481 | --------------------------------------------------------------------------------