jupytermind 0.3.0

package/src/ai_data_scientist/lifecycle.py

@@ -0,0 +1,221 @@
+"""Cooperative run cancellation and quiescence lifecycle API.
+
+Implements DES-AIDS-039 (REQ-AIDS-051): a process-wide, thread-safe
+registry of run state that lets an external orchestrator request
+cooperative cancellation of a tracked run and wait until all tracked
+activity (cell executions, notebook writes, locks) has settled.
+
+Cancellation is cooperative only: ``request_cancel`` sets a flag that
+``is_cancel_requested`` exposes for callers (such as
+``mcp_gateway.execute_cell``) to check *before* starting new work. It
+cannot interrupt a kernel cell execution already in flight.
+"""
+
+from __future__ import annotations
+
+import threading
+import time
+from dataclasses import dataclass
+from pathlib import Path
+
+_VALID_STATES = frozenset({"running", "cancelling", "cancelled", "completed", "failed"})
+
+
+@dataclass
+class _RunRecord:
+    notebook_path: Path
+    state: str = "running"
+    active_cell_executions: int = 0
+    pending_notebook_writes: int = 0
+    reason: str | None = None
+
+
+@dataclass(frozen=True)
+class RunStatus:
+    """Immutable snapshot of a run's tracked lifecycle state."""
+
+    run_id: str
+    state: str
+    active_cell_executions: int
+    pending_notebook_writes: int
+    locks_held: int
+    notebook_path: str
+    reason: str | None = None
+
+
+class UnknownRunError(KeyError):
+    """Raised when a run_id was never registered via register_run."""
+
+
+_registry: dict[str, _RunRecord] = {}
+_registry_guard = threading.Lock()
+_QUIESCENCE_POLL_INTERVAL_S = 0.01
+
+
+def _get_record(run_id: str) -> _RunRecord:
+    with _registry_guard:
+        record = _registry.get(run_id)
+    if record is None:
+        raise UnknownRunError(f"No run registered with run_id={run_id!r}.")
+    return record
+
+
+def _locks_held_for(notebook_path: Path) -> int:
+    # project_manager keeps one threading.Lock per notebook path; a lock is
+    # "held" iff it currently cannot be acquired without blocking.
+    from ai_data_scientist import project_manager
+
+    lock = project_manager._write_locks.get(notebook_path)
+    if lock is None:
+        return 0
+    acquired = lock.acquire(blocking=False)
+    if acquired:
+        lock.release()
+        return 0
+    return 1
+
+
+# @id CODE-AIDS-060
+# @implements REQ-AIDS-051
+# @design DES-AIDS-039
+def register_run(run_id: str, notebook_path: Path | str) -> None:
+    """Register a new tracked run in state "running"."""
+    with _registry_guard:
+        _registry[run_id] = _RunRecord(notebook_path=Path(notebook_path))
+
+
+# @id CODE-AIDS-061
+# @implements REQ-AIDS-051
+# @design DES-AIDS-039
+def mark_execution_start(run_id: str) -> None:
+    """Record that a cell execution has begun for ``run_id``."""
+    record = _get_record(run_id)
+    with _registry_guard:
+        record.active_cell_executions += 1
+
+
+# @id CODE-AIDS-062
+# @implements REQ-AIDS-051
+# @design DES-AIDS-039
+def mark_execution_end(run_id: str) -> None:
+    """Record that a cell execution has finished for ``run_id``."""
+    record = _get_record(run_id)
+    with _registry_guard:
+        record.active_cell_executions = max(0, record.active_cell_executions - 1)
+
+
+# @id CODE-AIDS-063
+# @implements REQ-AIDS-051
+# @design DES-AIDS-039
+def mark_write_start(run_id: str) -> None:
+    """Record that a notebook write has begun for ``run_id``."""
+    record = _get_record(run_id)
+    with _registry_guard:
+        record.pending_notebook_writes += 1
+
+
+# @id CODE-AIDS-064
+# @implements REQ-AIDS-051
+# @design DES-AIDS-039
+def mark_write_end(run_id: str) -> None:
+    """Record that a notebook write has finished for ``run_id``."""
+    record = _get_record(run_id)
+    with _registry_guard:
+        record.pending_notebook_writes = max(0, record.pending_notebook_writes - 1)
+
+
+# @id CODE-AIDS-065
+# @implements REQ-AIDS-051
+# @design DES-AIDS-039
+def request_cancel(run_id: str, reason: str | None = None) -> None:
+    """Request cooperative cancellation of ``run_id``.
+
+    Idempotent and safe to call on an unknown, already-cancelled, or
+    already-completed/failed run_id: it never raises.
+    """
+    with _registry_guard:
+        record = _registry.get(run_id)
+        if record is None:
+            return
+        if record.state == "running":
+            record.state = "cancelling"
+        record.reason = reason
+
+
+# @id CODE-AIDS-066
+# @implements REQ-AIDS-051
+# @design DES-AIDS-039
+def is_cancel_requested(run_id: str) -> bool:
+    """``True`` iff ``request_cancel`` has been called for ``run_id``."""
+    record = _get_record(run_id)
+    return record.state in ("cancelling", "cancelled")
+
+
+def _finalize(run_id: str, state: str) -> None:
+    record = _get_record(run_id)
+    with _registry_guard:
+        record.state = state
+
+
+# @id CODE-AIDS-067
+# @implements REQ-AIDS-051
+# @design DES-AIDS-039
+def mark_completed(run_id: str) -> None:
+    """Mark ``run_id`` completed, unless it was already cancelling/cancelled."""
+    record = _get_record(run_id)
+    with _registry_guard:
+        if record.state in ("cancelling", "cancelled"):
+            record.state = "cancelled"
+        else:
+            record.state = "completed"
+
+
+# @id CODE-AIDS-068
+# @implements REQ-AIDS-051
+# @design DES-AIDS-039
+def mark_failed(run_id: str) -> None:
+    """Mark ``run_id`` failed."""
+    _finalize(run_id, "failed")
+
+
+# @id CODE-AIDS-069
+# @implements REQ-AIDS-051
+# @design DES-AIDS-039
+def get_run_status(run_id: str) -> RunStatus:
+    """Return a snapshot of ``run_id``'s current tracked lifecycle state."""
+    record = _get_record(run_id)
+    with _registry_guard:
+        return RunStatus(
+            run_id=run_id,
+            state=record.state,
+            active_cell_executions=record.active_cell_executions,
+            pending_notebook_writes=record.pending_notebook_writes,
+            locks_held=_locks_held_for(record.notebook_path),
+            notebook_path=str(record.notebook_path),
+            reason=record.reason,
+        )
+
+
+# @id CODE-AIDS-070
+# @implements REQ-AIDS-051
+# @design DES-AIDS-039
+def wait_for_quiescence(run_id: str, timeout_s: float = 30.0) -> RunStatus:
+    """Poll ``run_id`` until all tracked activity settles or ``timeout_s`` elapses.
+
+    Quiescent means ``active_cell_executions == 0``,
+    ``pending_notebook_writes == 0`` and ``locks_held == 0``. Returns the
+    last observed status either way (callers inspect the returned fields to
+    tell quiescence from a timeout).
+    """
+    deadline = time.monotonic() + timeout_s
+    status = get_run_status(run_id)
+    while (
+        status.active_cell_executions > 0
+        or status.pending_notebook_writes > 0
+        or status.locks_held > 0
+    ):
+        if time.monotonic() >= deadline:
+            return status
+        time.sleep(_QUIESCENCE_POLL_INTERVAL_S)
+        status = get_run_status(run_id)
+    return status

package/src/ai_data_scientist/mcp_gateway.py

@@ -0,0 +1,113 @@
+"""Jupyter MCP execution gateway.
+
+Implements DES-AIDS-004: routes all analysis code execution exclusively
+through the configured Jupyter MCP client, enforces execution timeouts,
+and classifies unavailability/timeout failures so no partial cell is ever
+persisted to the notebook.
+"""
+
+from __future__ import annotations
+
+import concurrent.futures
+from typing import Protocol
+
+import nbformat
+
+from ai_data_scientist.project_manager import ProjectHandle, enqueue_write, next_execution_count
+
+DEFAULT_TIMEOUT_MS = 30000
+
+
+class MCPUnavailableError(RuntimeError):
+    """The configured Jupyter MCP server or kernel could not be reached."""
+
+
+class MCPExecutionTimeoutError(RuntimeError):
+    """Execution did not complete within the configured timeout."""
+
+
+class MCPClient(Protocol):
+    """Minimal contract any Jupyter MCP transport must satisfy."""
+
+    def execute(self, code: str) -> dict: ...
+
+
+# @id CODE-AIDS-003
+# @implements REQ-AIDS-003
+# @design DES-AIDS-004
+# @id CODE-AIDS-031
+# @implements REQ-AIDS-031
+# @design DES-AIDS-004
+# @id CODE-AIDS-048
+# @implements REQ-AIDS-042
+# @design DES-AIDS-030
+def execute_cell(client: MCPClient, code: str, timeout_ms: int = DEFAULT_TIMEOUT_MS) -> dict:
+    """Execute ``code`` exclusively through ``client`` with a hard timeout.
+
+    Raises ``MCPExecutionTimeoutError`` if execution exceeds ``timeout_ms``
+    and ``MCPUnavailableError`` if the client cannot reach the MCP server
+    or kernel.
+
+    The timeout is non-blocking for the caller: detection uses
+    ``concurrent.futures.wait`` (not ``future.result(timeout=...)`` inside a
+    ``with ThreadPoolExecutor`` block, whose ``__exit__`` would otherwise
+    perform an implicit ``shutdown(wait=True)`` and block the caller until
+    the still-running worker finishes). The executor is shut down with
+    ``wait=False`` so a worker that keeps running past the deadline cannot
+    delay the caller; its eventual result is discarded and never written to
+    the notebook, since only the value returned here is ever persisted.
+    """
+    timeout_s = timeout_ms / 1000
+    executor = concurrent.futures.ThreadPoolExecutor(max_workers=1)
+    try:
+        future = executor.submit(client.execute, code)
+        try:
+            done, _pending = concurrent.futures.wait({future}, timeout=timeout_s)
+            if future not in done:
+                raise MCPExecutionTimeoutError(
+                    f"Execution exceeded the configured {timeout_ms}ms timeout "
+                    f"(実行が設定タイムアウト{timeout_ms}msを超えました)."
+                )
+            return future.result()
+        except ConnectionError as exc:
+            raise MCPUnavailableError(
+                "The configured Jupyter MCP server or kernel is unreachable "
+                "(設定されたJupyter MCPサーバー/カーネルに接続できません)."
+            ) from exc
+    finally:
+        executor.shutdown(wait=False)
+
+
+# @id CODE-AIDS-030
+# @implements REQ-AIDS-030
+# @design DES-AIDS-004
+def run_and_record(
+    client: MCPClient,
+    handle: ProjectHandle,
+    code: str,
+    timeout_ms: int = DEFAULT_TIMEOUT_MS,
+) -> dict:
+    """Execute ``code`` and, only on success, append a cell to the notebook.
+
+    On failure the classified error propagates and no cell is written,
+    satisfying REQ-AIDS-030 (unavailability) and REQ-AIDS-031 (timeout).
+    """
+    result = execute_cell(client, code, timeout_ms=timeout_ms)
+    stamped_count: dict[str, int] = {}
+
+    def add_cell(nb):
+        execution_count = next_execution_count(nb)
+        stamped_count["value"] = execution_count
+        cell = nbformat.v4.new_code_cell(code)
+        cell["execution_count"] = execution_count
+        cell["outputs"] = [
+            nbformat.v4.new_output(
+                "execute_result",
+                data={"text/plain": str(result.get("output", ""))},
+                execution_count=execution_count,
+            )
+        ]
+        nb.cells.append(cell)
+
+    enqueue_write(handle, add_cell)
+    return {**result, "execution_count": stamped_count["value"]}

package/src/ai_data_scientist/mcp_runtime.py

@@ -0,0 +1,194 @@
+"""Jupyter MCP runtime lifecycle manager.
+
+Implements DES-AIDS-025: on the first execution request with no healthy
+runtime registered, installs (if needed) and starts JupyterLab and the
+jupyter-mcp-server inside the project's managed Python environment, bound to
+127.0.0.1 with an automatically selected free port and a randomly generated
+token, persists that runtime's identity to a machine-local state file so a
+separate later invocation reuses it instead of starting a duplicate one, and
+exposes explicit status/stop operations.
+"""
+
+from __future__ import annotations
+
+import dataclasses
+import json
+import time
+from pathlib import Path
+from typing import Protocol
+
+from ai_data_scientist.mcp_gateway import MCPUnavailableError
+
+DEFAULT_STARTUP_TIMEOUT_MS = 30000
+_HEALTH_POLL_INTERVAL_S = 0.2
+DEFAULT_STATE_PATH = Path.home() / ".cache" / "ai-data-scientist" / "mcp_runtime.json"
+
+
+@dataclasses.dataclass(frozen=True)
+class RuntimeInfo:
+    """Identity of a running Jupyter MCP runtime (REQ-AIDS-034).
+
+    Two separate (port, token) pairs are tracked because JupyterLab and the
+    jupyter-mcp-server (streamable-http transport) are independent processes
+    with independent listeners and independent authentication: JUPYTER_TOKEN
+    guards the Jupyter server itself, while the MCP token guards the MCP
+    server's own HTTP endpoint that the concrete client (DES-AIDS-026) calls.
+    """
+
+    jupyter_pid: int
+    jupyter_port: int
+    jupyter_token: str
+    mcp_server_pid: int
+    mcp_port: int
+    mcp_token: str
+
+
+# @id CODE-AIDS-050
+# @implements REQ-AIDS-041
+# @design DES-AIDS-029
+@dataclasses.dataclass(frozen=True)
+class StopResult:
+    """Outcome of a ``stop(..., wait=True)`` call (REQ-AIDS-041).
+
+    ``still_running`` lists any recorded PIDs that were still alive when the
+    bounded poll timed out; empty when every recorded process exited before
+    the timeout.
+    """
+
+    still_running: tuple[int, ...]
+
+
+class RuntimeLauncher(Protocol):
+    """Process-management contract any Jupyter MCP runtime backend must satisfy.
+
+    Kept as an injectable Protocol (mirroring mcp_gateway.MCPClient) so the
+    lifecycle logic below is fully unit-testable without actually installing
+    or starting JupyterLab/jupyter-mcp-server.
+    """
+
+    def start_jupyter(self) -> tuple[int, int, str]: ...
+    def start_mcp_server(self, jupyter_port: int, jupyter_token: str) -> tuple[int, int, str]: ...
+    def is_healthy(self, info: RuntimeInfo) -> bool: ...
+    def terminate(self, pid: int) -> None: ...
+    def is_process_alive(self, pid: int) -> bool: ...
+
+
+def _read_state(state_path: Path) -> RuntimeInfo | None:
+    if not state_path.exists():
+        return None
+    try:
+        data = json.loads(state_path.read_text(encoding="utf-8"))
+        return RuntimeInfo(**data)
+    except (json.JSONDecodeError, TypeError, KeyError):
+        return None
+
+
+def _write_state(state_path: Path, info: RuntimeInfo) -> None:
+    state_path.parent.mkdir(parents=True, exist_ok=True)
+    state_path.write_text(json.dumps(dataclasses.asdict(info)), encoding="utf-8")
+
+
+def _clear_state(state_path: Path) -> None:
+    state_path.unlink(missing_ok=True)
+
+
+# @id CODE-AIDS-039
+# @implements REQ-AIDS-034, REQ-AIDS-035, REQ-AIDS-036
+# @design DES-AIDS-025
+def ensure_runtime(
+    launcher: RuntimeLauncher,
+    state_path: Path = DEFAULT_STATE_PATH,
+    timeout_ms: int = DEFAULT_STARTUP_TIMEOUT_MS,
+) -> RuntimeInfo:
+    """Return a healthy Jupyter MCP runtime, starting one only if needed.
+
+    Reuses a still-running runtime recorded in ``state_path`` (REQ-AIDS-035).
+    Otherwise starts a new localhost-only runtime with an auto-selected port
+    and a random token (REQ-AIDS-034). Raises ``MCPUnavailableError`` without
+    registering a partially started runtime if it does not become healthy
+    within ``timeout_ms`` (REQ-AIDS-036).
+    """
+    existing = _read_state(state_path)
+    if existing is not None and launcher.is_healthy(existing):
+        return existing
+    if existing is not None:
+        _clear_state(state_path)
+
+    jupyter_pid, jupyter_port, jupyter_token = launcher.start_jupyter()
+    mcp_server_pid, mcp_port, mcp_token = launcher.start_mcp_server(jupyter_port, jupyter_token)
+    candidate = RuntimeInfo(
+        jupyter_pid=jupyter_pid,
+        jupyter_port=jupyter_port,
+        jupyter_token=jupyter_token,
+        mcp_server_pid=mcp_server_pid,
+        mcp_port=mcp_port,
+        mcp_token=mcp_token,
+    )
+
+    deadline = time.monotonic() + timeout_ms / 1000
+    while time.monotonic() < deadline:
+        if launcher.is_healthy(candidate):
+            _write_state(state_path, candidate)
+            return candidate
+        time.sleep(_HEALTH_POLL_INTERVAL_S)
+
+    launcher.terminate(jupyter_pid)
+    launcher.terminate(mcp_server_pid)
+    raise MCPUnavailableError(
+        "Jupyter MCP runtime did not become healthy within "
+        f"{timeout_ms}ms (Jupyter MCPランタイムが{timeout_ms}ms以内に起動しませんでした)."
+    )
+
+
+# @id CODE-AIDS-040
+# @implements REQ-AIDS-037
+# @design DES-AIDS-025
+def status(launcher: RuntimeLauncher, state_path: Path = DEFAULT_STATE_PATH) -> RuntimeInfo | None:
+    """Return the currently running, healthy runtime, or None if none is running."""
+    existing = _read_state(state_path)
+    if existing is None:
+        return None
+    if launcher.is_healthy(existing):
+        return existing
+    _clear_state(state_path)
+    return None
+
+
+# @id CODE-AIDS-041
+# @implements REQ-AIDS-037
+# @design DES-AIDS-025
+def stop(
+    launcher: RuntimeLauncher,
+    state_path: Path = DEFAULT_STATE_PATH,
+    wait: bool = False,
+    timeout_s: float = 5.0,
+    poll_interval_s: float = 0.2,
+) -> StopResult | None:
+    """Terminate the recorded runtime's processes and clear the state file.
+
+    With the default ``wait=False`` this is fire-and-forget (unchanged
+    behavior): it sends SIGTERM and returns ``None`` immediately without
+    confirming the processes actually exited. With ``wait=True``
+    (REQ-AIDS-041), it polls ``launcher.is_process_alive`` for both recorded
+    PIDs every ``poll_interval_s`` seconds until neither is alive or
+    ``timeout_s`` elapses, then returns a ``StopResult`` recording any PIDs
+    still alive at that point (empty when both exited cleanly).
+    """
+    existing = _read_state(state_path)
+    if existing is None:
+        return StopResult(still_running=()) if wait else None
+    launcher.terminate(existing.jupyter_pid)
+    launcher.terminate(existing.mcp_server_pid)
+    _clear_state(state_path)
+
+    if not wait:
+        return None
+
+    pids = [existing.jupyter_pid, existing.mcp_server_pid]
+    deadline = time.monotonic() + timeout_s
+    while time.monotonic() < deadline:
+        pids = [pid for pid in pids if launcher.is_process_alive(pid)]
+        if not pids:
+            break
+        time.sleep(poll_interval_s)
+    return StopResult(still_running=tuple(pids))

package/src/ai_data_scientist/mcp_transport.py

@@ -0,0 +1,53 @@
+"""Real streamable-http transport for the Jupyter MCP server.
+
+Implements the ``_Transport`` callable jupyter_mcp_client.JupyterMCPClient
+expects, using the real ``mcp`` Python SDK. Verified interactively against a
+live jupyter-mcp-server 2.2.3: the installed SDK exposes
+``mcp.client.streamable_http.streamable_http_client`` (not
+``streamablehttp_client``), yields a 2-tuple ``(read, write)``, accepts auth
+headers only via an injected ``httpx.AsyncClient``, and each discovered
+``Tool`` exposes its JSON schema as ``input_schema`` (snake_case).
+"""
+
+from __future__ import annotations
+
+import asyncio
+
+import httpx
+from mcp import ClientSession
+from mcp.client.streamable_http import streamable_http_client
+
+_TOOL_NAME = "execute_code"
+
+
+async def _execute_code_async(port: int, token: str, code: str) -> dict:
+    url = f"http://127.0.0.1:{port}/mcp"
+    headers = {"Authorization": "Bearer " + token}
+    async with (
+        httpx.AsyncClient(headers=headers) as http_client,
+        streamable_http_client(url, http_client=http_client) as (read, write),
+        ClientSession(read, write) as session,
+    ):
+        await session.initialize()
+        result = await session.call_tool(_TOOL_NAME, {"code": code})
+        if result.is_error:
+            text = "".join(block.text for block in result.content if hasattr(block, "text"))
+            return {"status": "error", "output": text}
+        outputs = (result.structured_content or {}).get("outputs")
+        if outputs is not None:
+            output = "\n".join(str(item) for item in outputs)
+        else:
+            output = "".join(block.text for block in result.content if hasattr(block, "text"))
+        return {"status": "ok", "output": output}
+
+
+# @id CODE-AIDS-045
+# @implements REQ-AIDS-038
+# @design DES-AIDS-026
+def execute_code(port: int, token: str, code: str) -> dict:
+    """Run ``code`` on a live jupyter-mcp-server and return its result.
+
+    Matches the ``_Transport`` Protocol expected by JupyterMCPClient:
+    ``(port, token, code) -> {"status": "ok"|"error", "output": str}``.
+    """
+    return asyncio.run(_execute_code_async(port, token, code))