notebook-mcp 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (28) hide show
  1. notebook_mcp-0.1.0/LICENSE +21 -0
  2. notebook_mcp-0.1.0/PKG-INFO +98 -0
  3. notebook_mcp-0.1.0/README.md +75 -0
  4. notebook_mcp-0.1.0/pyproject.toml +52 -0
  5. notebook_mcp-0.1.0/setup.cfg +4 -0
  6. notebook_mcp-0.1.0/src/notebook_mcp/__init__.py +3 -0
  7. notebook_mcp-0.1.0/src/notebook_mcp/analyzer.py +159 -0
  8. notebook_mcp-0.1.0/src/notebook_mcp/ast_analysis.py +74 -0
  9. notebook_mcp-0.1.0/src/notebook_mcp/context_builder.py +17 -0
  10. notebook_mcp-0.1.0/src/notebook_mcp/dependency_graph.py +74 -0
  11. notebook_mcp-0.1.0/src/notebook_mcp/errors.py +18 -0
  12. notebook_mcp-0.1.0/src/notebook_mcp/jupyter_server.py +60 -0
  13. notebook_mcp-0.1.0/src/notebook_mcp/kernel_channels.py +251 -0
  14. notebook_mcp-0.1.0/src/notebook_mcp/models.py +78 -0
  15. notebook_mcp-0.1.0/src/notebook_mcp/notebook_io.py +25 -0
  16. notebook_mcp-0.1.0/src/notebook_mcp/server.py +163 -0
  17. notebook_mcp-0.1.0/src/notebook_mcp/state_engine.py +127 -0
  18. notebook_mcp-0.1.0/src/notebook_mcp/utils.py +11 -0
  19. notebook_mcp-0.1.0/src/notebook_mcp.egg-info/PKG-INFO +98 -0
  20. notebook_mcp-0.1.0/src/notebook_mcp.egg-info/SOURCES.txt +26 -0
  21. notebook_mcp-0.1.0/src/notebook_mcp.egg-info/dependency_links.txt +1 -0
  22. notebook_mcp-0.1.0/src/notebook_mcp.egg-info/entry_points.txt +2 -0
  23. notebook_mcp-0.1.0/src/notebook_mcp.egg-info/requires.txt +12 -0
  24. notebook_mcp-0.1.0/src/notebook_mcp.egg-info/top_level.txt +1 -0
  25. notebook_mcp-0.1.0/tests/test_analyzer_smoke.py +21 -0
  26. notebook_mcp-0.1.0/tests/test_ast_analysis.py +24 -0
  27. notebook_mcp-0.1.0/tests/test_dependency_graph.py +32 -0
  28. notebook_mcp-0.1.0/tests/test_state_engine.py +51 -0
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,98 @@
1
+ Metadata-Version: 2.4
2
+ Name: notebook-mcp
3
+ Version: 0.1.0
4
+ Summary: MCP server that provides notebook-aware context and analysis for .ipynb files
5
+ License-Expression: MIT
6
+ Project-URL: Homepage, https://github.com/your-org/notebook-mcp
7
+ Project-URL: Repository, https://github.com/your-org/notebook-mcp
8
+ Requires-Python: >=3.11
9
+ Description-Content-Type: text/markdown
10
+ License-File: LICENSE
11
+ Requires-Dist: mcp>=1.0.0
12
+ Requires-Dist: nbformat>=5.10.4
13
+ Requires-Dist: pydantic>=2.7.0
14
+ Requires-Dist: httpx>=0.27.0
15
+ Requires-Dist: networkx>=3.3
16
+ Requires-Dist: websockets>=12.0
17
+ Provides-Extra: dev
18
+ Requires-Dist: pytest>=8.2.0; extra == "dev"
19
+ Requires-Dist: ruff>=0.5.0; extra == "dev"
20
+ Requires-Dist: build>=1.2.1; extra == "dev"
21
+ Requires-Dist: twine>=5.1.0; extra == "dev"
22
+ Dynamic: license-file
23
+
24
+ # notebook-mcp
25
+
26
+ MCP server that provides notebook-aware context and analysis for `.ipynb` files.
27
+
28
+ ## Features
29
+
30
+ - Offline notebook analysis (no kernel required)
31
+ - Strip outputs (optional)
32
+ - Stable cell indexing
33
+ - AST-based `defines` / `uses`
34
+ - Best-effort dependency graph
35
+ - Focused context slicing for a target cell
36
+ - Export notebook to a deterministic Python script (best-effort)
37
+
38
+ - Optional Jupyter Server adapter (HTTP)
39
+ - List active sessions
40
+ - Inspect kernel metadata
41
+
42
+ ## Install (dev)
43
+
44
+ ```bash
45
+ python -m pip install -e ".[dev]"
46
+ ```
47
+
48
+ ## Distribution
49
+
50
+ ### Publish Python package (PyPI)
51
+
52
+ 1. Build:
53
+
54
+ ```bash
55
+ python -m pip install -e ".[dev]"
56
+ python -m build
57
+ ```
58
+
59
+ 2. Upload:
60
+
61
+ ```bash
62
+ python -m twine upload dist/*
63
+ ```
64
+
65
+ ### Publish NPM wrapper
66
+
67
+ The Node wrapper lives in `npm/` and publishes as a scoped package.
68
+
69
+ ```bash
70
+ cd npm
71
+ npm publish --access public
72
+ ```
73
+
74
+ ## Run (stdio transport; best for IDE agents)
75
+
76
+ ```bash
77
+ notebook-mcp
78
+ ```
79
+
80
+ ## Run (streamable HTTP transport)
81
+
82
+ ```bash
83
+ set MCP_TRANSPORT=streamable-http
84
+ set MCP_HOST=127.0.0.1
85
+ set MCP_PORT=8000
86
+ notebook-mcp
87
+ ```
88
+
89
+ ## Jupyter Server adapter
90
+
91
+ Set env vars:
92
+
93
+ - `JUPYTER_BASE_URL` e.g. `http://127.0.0.1:8888`
94
+ - `JUPYTER_TOKEN` (if required)
95
+
96
+ ## Notes
97
+
98
+ - Dependency graph and stale detection are best-effort offline. Truthful execution/runtime state requires kernel messaging.
@@ -0,0 +1,75 @@
1
+ # notebook-mcp
2
+
3
+ MCP server that provides notebook-aware context and analysis for `.ipynb` files.
4
+
5
+ ## Features
6
+
7
+ - Offline notebook analysis (no kernel required)
8
+ - Strip outputs (optional)
9
+ - Stable cell indexing
10
+ - AST-based `defines` / `uses`
11
+ - Best-effort dependency graph
12
+ - Focused context slicing for a target cell
13
+ - Export notebook to a deterministic Python script (best-effort)
14
+
15
+ - Optional Jupyter Server adapter (HTTP)
16
+ - List active sessions
17
+ - Inspect kernel metadata
18
+
19
+ ## Install (dev)
20
+
21
+ ```bash
22
+ python -m pip install -e ".[dev]"
23
+ ```
24
+
25
+ ## Distribution
26
+
27
+ ### Publish Python package (PyPI)
28
+
29
+ 1. Build:
30
+
31
+ ```bash
32
+ python -m pip install -e ".[dev]"
33
+ python -m build
34
+ ```
35
+
36
+ 2. Upload:
37
+
38
+ ```bash
39
+ python -m twine upload dist/*
40
+ ```
41
+
42
+ ### Publish NPM wrapper
43
+
44
+ The Node wrapper lives in `npm/` and publishes as a scoped package.
45
+
46
+ ```bash
47
+ cd npm
48
+ npm publish --access public
49
+ ```
50
+
51
+ ## Run (stdio transport; best for IDE agents)
52
+
53
+ ```bash
54
+ notebook-mcp
55
+ ```
56
+
57
+ ## Run (streamable HTTP transport)
58
+
59
+ ```bash
60
+ set MCP_TRANSPORT=streamable-http
61
+ set MCP_HOST=127.0.0.1
62
+ set MCP_PORT=8000
63
+ notebook-mcp
64
+ ```
65
+
66
+ ## Jupyter Server adapter
67
+
68
+ Set env vars:
69
+
70
+ - `JUPYTER_BASE_URL` e.g. `http://127.0.0.1:8888`
71
+ - `JUPYTER_TOKEN` (if required)
72
+
73
+ ## Notes
74
+
75
+ - Dependency graph and stale detection are best-effort offline. Truthful execution/runtime state requires kernel messaging.
@@ -0,0 +1,52 @@
1
+ [project]
2
+ name = "notebook-mcp"
3
+ version = "0.1.0"
4
+ description = "MCP server that provides notebook-aware context and analysis for .ipynb files"
5
+ readme = "README.md"
6
+ requires-python = ">=3.11"
7
+ license = "MIT"
8
+ dependencies = [
9
+ "mcp>=1.0.0",
10
+ "nbformat>=5.10.4",
11
+ "pydantic>=2.7.0",
12
+ "httpx>=0.27.0",
13
+ "networkx>=3.3",
14
+ "websockets>=12.0",
15
+ ]
16
+
17
+ [project.optional-dependencies]
18
+ dev = [
19
+ "pytest>=8.2.0",
20
+ "ruff>=0.5.0",
21
+ "build>=1.2.1",
22
+ "twine>=5.1.0",
23
+ ]
24
+
25
+ [project.urls]
26
+ Homepage = "https://github.com/your-org/notebook-mcp"
27
+ Repository = "https://github.com/your-org/notebook-mcp"
28
+
29
+ [project.scripts]
30
+ notebook-mcp = "notebook_mcp.server:main"
31
+
32
+ [build-system]
33
+ requires = ["setuptools>=69.0.0", "wheel>=0.42.0"]
34
+ build-backend = "setuptools.build_meta"
35
+
36
+ [tool.setuptools]
37
+ package-dir = {"" = "src"}
38
+ include-package-data = true
39
+
40
+ [tool.setuptools.packages.find]
41
+ where = ["src"]
42
+
43
+ [tool.ruff]
44
+ line-length = 100
45
+ target-version = "py311"
46
+
47
+ [tool.ruff.lint]
48
+ select = ["E", "F", "I", "UP", "B"]
49
+
50
+ [tool.pytest.ini_options]
51
+ addopts = "-q"
52
+ pythonpath = ["src"]
@@ -0,0 +1,4 @@
1
+ [egg_info]
2
+ tag_build =
3
+ tag_date = 0
4
+
@@ -0,0 +1,3 @@
1
+ __all__ = ["__version__"]
2
+
3
+ __version__ = "0.1.0"
@@ -0,0 +1,159 @@
1
+ from __future__ import annotations
2
+
3
+ import os
4
+ from functools import lru_cache
5
+ from pathlib import Path
6
+ from typing import Literal
7
+
8
+ from .ast_analysis import summarize_python_source
9
+ from .context_builder import format_cells_as_context
10
+ from .dependency_graph import build_dependency_edges, topo_sort_cells, upstream_slice
11
+ from .errors import CellNotFoundError
12
+ from .models import FocusedContext, NotebookAnalysis, NotebookCell
13
+ from .notebook_io import load_notebook, strip_outputs_inplace
14
+ from .utils import normalize_newlines, sha256_text
15
+
16
+
17
+ def _cell_id(cell: dict, index: int) -> str:
18
+ # Jupyter notebooks may include a stable 'id' field (nbformat 4.5+).
19
+ cid = cell.get("id")
20
+ if isinstance(cid, str) and cid.strip():
21
+ return cid
22
+ return f"cell_{index}" # stable fallback
23
+
24
+
25
+ def analyze_notebook(
26
+ path: str,
27
+ *,
28
+ include_markdown: bool = True,
29
+ strip_outputs: bool = True,
30
+ cell_types: tuple[str, ...] = ("code", "markdown"),
31
+ ) -> NotebookAnalysis:
32
+ mtime = os.path.getmtime(path)
33
+ return _analyze_notebook_cached(
34
+ path,
35
+ mtime,
36
+ include_markdown=include_markdown,
37
+ strip_outputs=strip_outputs,
38
+ cell_types=cell_types,
39
+ )
40
+
41
+
42
+ @lru_cache(maxsize=32)
43
+ def _analyze_notebook_cached(
44
+ path: str,
45
+ mtime: float,
46
+ *,
47
+ include_markdown: bool,
48
+ strip_outputs: bool,
49
+ cell_types: tuple[str, ...],
50
+ ) -> NotebookAnalysis:
51
+ nb = load_notebook(path)
52
+ if strip_outputs:
53
+ strip_outputs_inplace(nb)
54
+
55
+ out_cells: list[NotebookCell] = []
56
+
57
+ for i, cell in enumerate(nb.get("cells", [])):
58
+ ctype = cell.get("cell_type")
59
+ if ctype not in cell_types:
60
+ continue
61
+ if ctype == "markdown" and not include_markdown:
62
+ continue
63
+
64
+ source = normalize_newlines(cell.get("source") or "")
65
+ h = sha256_text(source)
66
+
67
+ defines: list[str] = []
68
+ uses: list[str] = []
69
+ imports: list[str] = []
70
+
71
+ if ctype == "code":
72
+ s = summarize_python_source(source)
73
+ defines = sorted(s.defines)
74
+ uses = sorted(s.uses)
75
+ imports = sorted(s.imports)
76
+
77
+ out_cells.append(
78
+ NotebookCell(
79
+ cell_id=_cell_id(cell, i),
80
+ index=i,
81
+ cell_type=ctype,
82
+ source=source,
83
+ execution_count=cell.get("execution_count"),
84
+ source_hash=h,
85
+ defines=defines,
86
+ uses=uses,
87
+ imports=imports,
88
+ )
89
+ )
90
+
91
+ edges = build_dependency_edges(out_cells)
92
+
93
+ return NotebookAnalysis(
94
+ path=str(Path(path)),
95
+ nbformat=int(nb.get("nbformat", 4)),
96
+ nbformat_minor=int(nb.get("nbformat_minor", 0)),
97
+ cells=out_cells,
98
+ dependency_edges=edges,
99
+ )
100
+
101
+
102
+ def get_focused_context(
103
+ path: str,
104
+ *,
105
+ focus_cell_id: str,
106
+ max_cells: int = 25,
107
+ include_markdown: bool = True,
108
+ ) -> FocusedContext:
109
+ analysis = analyze_notebook(path, include_markdown=include_markdown)
110
+ if focus_cell_id not in {c.cell_id for c in analysis.cells}:
111
+ raise CellNotFoundError(f"Cell not found: {focus_cell_id}")
112
+
113
+ selected_ids = upstream_slice(focus_cell_id, analysis.dependency_edges, max_cells=max_cells)
114
+
115
+ cell_by_id = {c.cell_id: c for c in analysis.cells}
116
+ selected_cells = [cell_by_id[cid] for cid in selected_ids if cid in cell_by_id]
117
+
118
+ # Make output readable: topo order within selected cells when possible.
119
+ selected_edges = [(a, b) for (a, b) in analysis.dependency_edges if a in selected_ids and b in selected_ids]
120
+ selected_cells = topo_sort_cells(selected_cells, selected_edges)
121
+
122
+ context_text = format_cells_as_context(selected_cells)
123
+
124
+ return FocusedContext(
125
+ path=analysis.path,
126
+ focus_cell_id=focus_cell_id,
127
+ selected_cell_ids=[c.cell_id for c in selected_cells],
128
+ context_text=context_text,
129
+ )
130
+
131
+
132
+ def export_notebook_to_script(
133
+ path: str,
134
+ *,
135
+ include_markdown_as_comments: bool = False,
136
+ ) -> str:
137
+ analysis = analyze_notebook(path, include_markdown=True, strip_outputs=True)
138
+
139
+ code_cells: list[NotebookCell] = []
140
+ for c in analysis.cells:
141
+ if c.cell_type == "code":
142
+ code_cells.append(c)
143
+ elif c.cell_type == "markdown" and include_markdown_as_comments:
144
+ # Keep markdown as a docstring-ish block to preserve narrative.
145
+ code_cells.append(
146
+ c.model_copy(update={"cell_type": "code", "source": f"\n\"\"\"\n{c.source}\n\"\"\"\n"})
147
+ )
148
+
149
+ ordered = topo_sort_cells(code_cells, analysis.dependency_edges)
150
+
151
+ parts: list[str] = []
152
+ parts.append(f"# Generated from notebook: {analysis.path}\n")
153
+
154
+ for c in ordered:
155
+ parts.append(f"# --- cell: {c.cell_id} (index={c.index}) ---")
156
+ parts.append(c.source.rstrip())
157
+ parts.append("")
158
+
159
+ return "\n".join(parts).strip() + "\n"
@@ -0,0 +1,74 @@
1
+ from __future__ import annotations
2
+
3
+ import ast
4
+ from dataclasses import dataclass
5
+
6
+
7
+ @dataclass(frozen=True)
8
+ class AstSummary:
9
+ defines: set[str]
10
+ uses: set[str]
11
+ imports: set[str]
12
+
13
+
14
+ class _Analyzer(ast.NodeVisitor):
15
+ def __init__(self) -> None:
16
+ self.defines: set[str] = set()
17
+ self.uses: set[str] = set()
18
+ self.imports: set[str] = set()
19
+
20
+ def visit_Import(self, node: ast.Import) -> None: # noqa: N802
21
+ for alias in node.names:
22
+ name = alias.asname or alias.name.split(".")[0]
23
+ self.defines.add(name)
24
+ self.imports.add(alias.name)
25
+ self.generic_visit(node)
26
+
27
+ def visit_ImportFrom(self, node: ast.ImportFrom) -> None: # noqa: N802
28
+ mod = node.module or ""
29
+ for alias in node.names:
30
+ if alias.name == "*":
31
+ continue
32
+ name = alias.asname or alias.name
33
+ self.defines.add(name)
34
+ self.imports.add(f"{mod}:{alias.name}" if mod else alias.name)
35
+ self.generic_visit(node)
36
+
37
+ def visit_FunctionDef(self, node: ast.FunctionDef) -> None: # noqa: N802
38
+ self.defines.add(node.name)
39
+ self.generic_visit(node)
40
+
41
+ def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> None: # noqa: N802
42
+ self.defines.add(node.name)
43
+ self.generic_visit(node)
44
+
45
+ def visit_ClassDef(self, node: ast.ClassDef) -> None: # noqa: N802
46
+ self.defines.add(node.name)
47
+ self.generic_visit(node)
48
+
49
+ def visit_Name(self, node: ast.Name) -> None: # noqa: N802
50
+ if isinstance(node.ctx, ast.Store):
51
+ self.defines.add(node.id)
52
+ elif isinstance(node.ctx, ast.Load):
53
+ self.uses.add(node.id)
54
+ self.generic_visit(node)
55
+
56
+ def visit_arg(self, node: ast.arg) -> None: # noqa: N802
57
+ # Treat function args as local defines.
58
+ self.defines.add(node.arg)
59
+ self.generic_visit(node)
60
+
61
+
62
+ def summarize_python_source(source: str) -> AstSummary:
63
+ try:
64
+ tree = ast.parse(source)
65
+ except SyntaxError:
66
+ return AstSummary(defines=set(), uses=set(), imports=set())
67
+
68
+ a = _Analyzer()
69
+ a.visit(tree)
70
+
71
+ # If a name is defined, don't treat it as a dependency use.
72
+ uses = a.uses - a.defines
73
+
74
+ return AstSummary(defines=a.defines, uses=uses, imports=a.imports)
@@ -0,0 +1,17 @@
1
+ from __future__ import annotations
2
+
3
+ from .models import NotebookCell
4
+
5
+
6
+ def format_cells_as_context(cells: list[NotebookCell]) -> str:
7
+ parts: list[str] = []
8
+ for c in cells:
9
+ header = f"# --- cell: {c.cell_id} (index={c.index}, type={c.cell_type}, exec={c.execution_count}) ---"
10
+ parts.append(header)
11
+ if c.cell_type == "markdown":
12
+ parts.append(c.source)
13
+ else:
14
+ parts.append(f"```python\n{c.source}\n```")
15
+ parts.append("")
16
+
17
+ return "\n".join(parts).strip() + "\n"
@@ -0,0 +1,74 @@
1
+ from __future__ import annotations
2
+
3
+ from collections import defaultdict
4
+
5
+ import networkx as nx
6
+
7
+ from .models import NotebookCell
8
+
9
+
10
+ def build_dependency_edges(cells: list[NotebookCell]) -> list[tuple[str, str]]:
11
+ last_def: dict[str, str] = {}
12
+ edges: set[tuple[str, str]] = set()
13
+
14
+ for cell in cells:
15
+ for sym in cell.uses:
16
+ src = last_def.get(sym)
17
+ if src and src != cell.cell_id:
18
+ edges.add((src, cell.cell_id))
19
+
20
+ for sym in cell.defines:
21
+ last_def[sym] = cell.cell_id
22
+
23
+ return sorted(edges)
24
+
25
+
26
+ def topo_sort_cells(cells: list[NotebookCell], edges: list[tuple[str, str]]) -> list[NotebookCell]:
27
+ g = nx.DiGraph()
28
+ for c in cells:
29
+ g.add_node(c.cell_id)
30
+ g.add_edges_from(edges)
31
+
32
+ try:
33
+ order = list(nx.topological_sort(g))
34
+ except nx.NetworkXUnfeasible:
35
+ # Cycles are common; fall back to file order.
36
+ return cells
37
+
38
+ index = {c.cell_id: i for i, c in enumerate(cells)}
39
+ order = sorted(order, key=lambda cid: index.get(cid, 10**9))
40
+
41
+ cell_by_id = {c.cell_id: c for c in cells}
42
+ return [cell_by_id[cid] for cid in order if cid in cell_by_id]
43
+
44
+
45
+ def upstream_slice(
46
+ focus_cell_id: str,
47
+ edges: list[tuple[str, str]],
48
+ max_cells: int,
49
+ ) -> list[str]:
50
+ preds: dict[str, set[str]] = defaultdict(set)
51
+ for a, b in edges:
52
+ preds[b].add(a)
53
+
54
+ selected: list[str] = []
55
+ seen: set[str] = set()
56
+ stack: list[str] = [focus_cell_id]
57
+
58
+ while stack and len(selected) < max_cells:
59
+ cid = stack.pop()
60
+ if cid in seen:
61
+ continue
62
+ seen.add(cid)
63
+ selected.append(cid)
64
+
65
+ for p in sorted(preds.get(cid, set())):
66
+ if p not in seen:
67
+ stack.append(p)
68
+
69
+ # Put focus cell last to read naturally.
70
+ if focus_cell_id in selected:
71
+ selected.remove(focus_cell_id)
72
+ selected.append(focus_cell_id)
73
+
74
+ return selected
@@ -0,0 +1,18 @@
1
+ class NotebookMcpError(Exception):
2
+ pass
3
+
4
+
5
+ class NotebookNotFoundError(NotebookMcpError):
6
+ pass
7
+
8
+
9
+ class NotebookParseError(NotebookMcpError):
10
+ pass
11
+
12
+
13
+ class CellNotFoundError(NotebookMcpError):
14
+ pass
15
+
16
+
17
+ class JupyterServerError(NotebookMcpError):
18
+ pass
@@ -0,0 +1,60 @@
1
+ from __future__ import annotations
2
+
3
+ from typing import Any
4
+
5
+ import httpx
6
+
7
+ from .errors import JupyterServerError
8
+ from .models import JupyterKernel, JupyterSession
9
+
10
+
11
+ class JupyterServerClient:
12
+ def __init__(self, base_url: str, token: str | None = None, timeout_s: float = 10.0) -> None:
13
+ self.base_url = base_url.rstrip("/")
14
+ self.token = token
15
+ self._client = httpx.Client(timeout=timeout_s)
16
+
17
+ def _headers(self) -> dict[str, str]:
18
+ if not self.token:
19
+ return {}
20
+ return {"Authorization": f"token {self.token}"}
21
+
22
+ def _get(self, path: str) -> Any:
23
+ url = f"{self.base_url}{path}"
24
+ try:
25
+ r = self._client.get(url, headers=self._headers())
26
+ r.raise_for_status()
27
+ return r.json()
28
+ except Exception as e: # noqa: BLE001
29
+ raise JupyterServerError(f"Jupyter request failed: GET {url}") from e
30
+
31
+ def list_sessions(self) -> list[JupyterSession]:
32
+ data = self._get("/api/sessions")
33
+ out: list[JupyterSession] = []
34
+ for s in data:
35
+ kernel = s.get("kernel") or {}
36
+ out.append(
37
+ JupyterSession(
38
+ id=str(s.get("id") or ""),
39
+ path=s.get("path"),
40
+ name=s.get("name"),
41
+ type=s.get("type"),
42
+ kernel_id=kernel.get("id"),
43
+ raw=s,
44
+ )
45
+ )
46
+ return out
47
+
48
+ def get_kernel(self, kernel_id: str) -> JupyterKernel:
49
+ k = self._get(f"/api/kernels/{kernel_id}")
50
+ return JupyterKernel(
51
+ id=str(k.get("id") or kernel_id),
52
+ name=k.get("name"),
53
+ last_activity=k.get("last_activity"),
54
+ execution_state=k.get("execution_state"),
55
+ connections=k.get("connections"),
56
+ raw=k,
57
+ )
58
+
59
+ def close(self) -> None:
60
+ self._client.close()