103 lines
3.7 KiB
Python
103 lines
3.7 KiB
Python
from collections.abc import AsyncIterator
|
|
from datetime import UTC, datetime
|
|
from time import monotonic
|
|
from uuid import uuid4
|
|
|
|
from ..models import AppInfo, ShellInfo
|
|
from ..registry import ShellRegistry
|
|
from ..utils import generate_call_id
|
|
from .messages import ParsedJupyterMessage
|
|
from .shell import JupyterShell
|
|
from .transport import InputRequest
|
|
|
|
|
|
class JupyterApp:
|
|
"""Coordinate named persistent shells backed by Jupyter kernels."""
|
|
|
|
def __init__(self) -> None:
|
|
self.app_id = uuid4().hex
|
|
self.created_at = datetime.now(UTC).isoformat()
|
|
self._started_at = monotonic()
|
|
self._total_calls = 0
|
|
self._shell_registry = ShellRegistry(JupyterShell)
|
|
self._calls: dict[str, JupyterShell] = {}
|
|
self._issued_call_ids: set[str] = set()
|
|
|
|
@property
|
|
def shells(self) -> dict[str, JupyterShell]:
|
|
"""Expose the managed shells for compatibility with the sync app."""
|
|
return self._shell_registry.shells
|
|
|
|
@property
|
|
def last_shell(self) -> str | None:
|
|
"""Return the name of the most recently selected shell."""
|
|
return self._shell_registry.last_shell
|
|
|
|
def _select_shell(self, shell_name: str) -> tuple[str, JupyterShell]:
|
|
return self._shell_registry.select(shell_name)
|
|
|
|
def list_shells(self) -> list[ShellInfo]:
|
|
"""Return metadata for every known Jupyter shell."""
|
|
return self._shell_registry.list_shells()
|
|
|
|
def create_shell(self) -> ShellInfo:
|
|
"""Create a fresh Jupyter shell and return its metadata."""
|
|
return self._shell_registry.create_shell()
|
|
|
|
def get_shell(self, shell_name: str) -> ShellInfo:
|
|
"""Return frontend-safe metadata for one Jupyter shell."""
|
|
return self._shell_registry.get_shell(shell_name)
|
|
|
|
def get_last_shell(self) -> ShellInfo:
|
|
"""Return metadata for the most recently selected Jupyter shell."""
|
|
return self._shell_registry.get_last_shell()
|
|
|
|
def info(self) -> AppInfo:
|
|
"""Return runtime metadata without starting any kernels."""
|
|
return AppInfo(
|
|
app_id=self.app_id,
|
|
created_at=self.created_at,
|
|
uptime_seconds=monotonic() - self._started_at,
|
|
shell_count=len(self.shells),
|
|
last_shell=self.last_shell,
|
|
total_calls=self._total_calls,
|
|
)
|
|
|
|
def _new_call_id(self) -> str:
|
|
"""Generate a readable call ID that is unique for this app."""
|
|
call_id = generate_call_id()
|
|
while call_id in self._issued_call_ids:
|
|
call_id = generate_call_id()
|
|
self._issued_call_ids.add(call_id)
|
|
return call_id
|
|
|
|
async def run_code_stream(
|
|
self,
|
|
code: str,
|
|
shell_name: str = "last",
|
|
) -> AsyncIterator[ParsedJupyterMessage | InputRequest]:
|
|
"""Execute code and yield its subprocess messages."""
|
|
_, shell = self._select_shell(shell_name)
|
|
call_id = self._new_call_id()
|
|
self._total_calls += 1
|
|
self._calls[call_id] = shell
|
|
try:
|
|
async for message in shell.run_cell_stream(code, call_id=call_id):
|
|
yield message
|
|
finally:
|
|
self._calls.pop(call_id, None)
|
|
|
|
async def reply_to_input(self, call_id: str, value: str) -> None:
|
|
"""Reply to an input request emitted by a running call."""
|
|
try:
|
|
shell = self._calls[call_id]
|
|
except KeyError as error:
|
|
raise KeyError(f"Call {call_id} does not exist") from error
|
|
await shell.reply_to_input(call_id, value)
|
|
|
|
async def shutdown(self) -> None:
|
|
"""Stop all kernel processes owned by this app."""
|
|
awaitables = [shell.shutdown() for shell in self.shells.values()]
|
|
for awaitable in awaitables:
|
|
await awaitable
|