Add MCP injection support to Jupyter backend
- start_kernel() now accepts injection_code and bridge_socket_path - Bridge Unix domain socket mounted as volume in container - Injection code executed once on kernel startup (silent, no history) - New _execute_injection_code() helper method with error handling - Parameters passed through SessionManager.create_session() - Parameters passed through JupyterBackend.execute() - Updated JUPYTER_IMPLEMENTATION_STATUS.md This mirrors the MCP injection pattern from the simple executor, allowing MCP tools to be available in Jupyter sessions.
This commit is contained in:
parent
ed58c6ad5a
commit
9cc43e4166
4 changed files with 106 additions and 13 deletions
|
|
@ -48,7 +48,14 @@
|
||||||
- `cleanup_idle_kernels()`: Time-based cleanup
|
- `cleanup_idle_kernels()`: Time-based cleanup
|
||||||
- Activity timestamp updates
|
- Activity timestamp updates
|
||||||
|
|
||||||
6. **Dependencies**
|
6. **MCP Integration**
|
||||||
|
- `start_kernel()` accepts `injection_code` and `bridge_socket_path`
|
||||||
|
- Bridge socket mounted as volume in container (Unix domain socket)
|
||||||
|
- Injection code executed once on kernel startup (silent, no history)
|
||||||
|
- Error handling for injection failures
|
||||||
|
- Passed through SessionManager and JupyterBackend
|
||||||
|
|
||||||
|
7. **Dependencies**
|
||||||
- `jupyter-client>=8.8.0` added to server dependencies
|
- `jupyter-client>=8.8.0` added to server dependencies
|
||||||
- `ipykernel` removed from server (will be in container image)
|
- `ipykernel` removed from server (will be in container image)
|
||||||
- `pyzmq>=27.1.0` for ZMQ support
|
- `pyzmq>=27.1.0` for ZMQ support
|
||||||
|
|
@ -57,6 +64,7 @@
|
||||||
- [architecture1.md](../docs/architecture1.md) updated with Jupyter Backend section
|
- [architecture1.md](../docs/architecture1.md) updated with Jupyter Backend section
|
||||||
- [todo.md](../docs/todo.md) Phase 2.2.1 updated with architecture details
|
- [todo.md](../docs/todo.md) Phase 2.2.1 updated with architecture details
|
||||||
- Architecture flow diagram and session mapping explanation
|
- Architecture flow diagram and session mapping explanation
|
||||||
|
- [JUPYTER_IMPLEMENTATION_STATUS.md](../docs/JUPYTER_IMPLEMENTATION_STATUS.md) tracking document
|
||||||
|
|
||||||
### 🚧 In Progress / TODO
|
### 🚧 In Progress / TODO
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -55,7 +55,9 @@ class JupyterBackend:
|
||||||
memory: Optional[str] = None,
|
memory: Optional[str] = None,
|
||||||
cpu_quota: Optional[int] = None,
|
cpu_quota: Optional[int] = None,
|
||||||
custom_image: Optional[str] = None,
|
custom_image: Optional[str] = None,
|
||||||
volumes: Optional[Dict[str, dict]] = None
|
volumes: Optional[Dict[str, dict]] = None,
|
||||||
|
injection_code: Optional[str] = None,
|
||||||
|
bridge_socket_path: Optional[str] = None
|
||||||
) -> ExecutionResult:
|
) -> ExecutionResult:
|
||||||
"""
|
"""
|
||||||
Execute code in stateful session.
|
Execute code in stateful session.
|
||||||
|
|
@ -71,6 +73,8 @@ class JupyterBackend:
|
||||||
cpu_quota: CPU quota (uses config default if None)
|
cpu_quota: CPU quota (uses config default if None)
|
||||||
custom_image: Custom image name (uses config default if None)
|
custom_image: Custom image name (uses config default if None)
|
||||||
volumes: Volume mounts dict
|
volumes: Volume mounts dict
|
||||||
|
injection_code: Optional MCP tool injection code (executed once at session start)
|
||||||
|
bridge_socket_path: Optional path to MCP bridge socket for mounting
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
ExecutionResult with execution output and metadata
|
ExecutionResult with execution output and metadata
|
||||||
|
|
@ -106,7 +110,7 @@ class JupyterBackend:
|
||||||
try:
|
try:
|
||||||
self.session_manager.get_session(session_id)
|
self.session_manager.get_session(session_id)
|
||||||
except SessionError:
|
except SessionError:
|
||||||
# Session doesn't exist, create it
|
# Session doesn't exist, create it with MCP injection
|
||||||
resource_limits = ResourceLimits(
|
resource_limits = ResourceLimits(
|
||||||
memory=memory,
|
memory=memory,
|
||||||
cpu_quota=cpu_quota,
|
cpu_quota=cpu_quota,
|
||||||
|
|
@ -117,7 +121,9 @@ class JupyterBackend:
|
||||||
self.session_manager.create_session(
|
self.session_manager.create_session(
|
||||||
session_id=session_id,
|
session_id=session_id,
|
||||||
resource_limits=resource_limits,
|
resource_limits=resource_limits,
|
||||||
volumes=volumes
|
volumes=volumes,
|
||||||
|
injection_code=injection_code,
|
||||||
|
bridge_socket_path=bridge_socket_path
|
||||||
)
|
)
|
||||||
|
|
||||||
# Execute in session
|
# Execute in session
|
||||||
|
|
|
||||||
|
|
@ -99,7 +99,9 @@ class JupyterKernelManager:
|
||||||
def start_kernel(
|
def start_kernel(
|
||||||
self,
|
self,
|
||||||
session_id: str,
|
session_id: str,
|
||||||
volumes: Optional[Dict[str, dict]] = None
|
volumes: Optional[Dict[str, dict]] = None,
|
||||||
|
injection_code: Optional[str] = None,
|
||||||
|
bridge_socket_path: Optional[str] = None
|
||||||
) -> str:
|
) -> str:
|
||||||
"""
|
"""
|
||||||
Start IPython kernel in dedicated container.
|
Start IPython kernel in dedicated container.
|
||||||
|
|
@ -108,14 +110,18 @@ class JupyterKernelManager:
|
||||||
1. Generate ZMQ connection info (ports, keys)
|
1. Generate ZMQ connection info (ports, keys)
|
||||||
2. Create connection file
|
2. Create connection file
|
||||||
3. Create container with ipykernel command
|
3. Create container with ipykernel command
|
||||||
4. Start container
|
4. Mount bridge socket if provided (for MCP tools)
|
||||||
5. Wait for kernel to be ready
|
5. Start container
|
||||||
6. Connect jupyter-client to kernel via ZMQ
|
6. Wait for kernel to be ready
|
||||||
7. Verify kernel is responsive
|
7. Connect jupyter-client to kernel via ZMQ
|
||||||
|
8. Execute injection code (MCP tools setup) if provided
|
||||||
|
9. Verify kernel is responsive
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
session_id: Session ID this kernel belongs to
|
session_id: Session ID this kernel belongs to
|
||||||
volumes: Optional volume mounts
|
volumes: Optional volume mounts
|
||||||
|
injection_code: Optional MCP tool injection code to execute at startup
|
||||||
|
bridge_socket_path: Optional path to MCP bridge socket for mounting
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
kernel_id: Unique identifier for this kernel
|
kernel_id: Unique identifier for this kernel
|
||||||
|
|
@ -132,6 +138,14 @@ class JupyterKernelManager:
|
||||||
connection_file = self._create_connection_file(kernel_id, connection_info)
|
connection_file = self._create_connection_file(kernel_id, connection_info)
|
||||||
|
|
||||||
try:
|
try:
|
||||||
|
# Set up volumes (user volumes + bridge socket if provided)
|
||||||
|
container_volumes = volumes.copy() if volumes else {}
|
||||||
|
if bridge_socket_path:
|
||||||
|
container_volumes[bridge_socket_path] = {
|
||||||
|
"bind": bridge_socket_path,
|
||||||
|
"mode": "rw"
|
||||||
|
}
|
||||||
|
|
||||||
# Create container with ipykernel
|
# Create container with ipykernel
|
||||||
config = ContainerConfig(
|
config = ContainerConfig(
|
||||||
image=self.image,
|
image=self.image,
|
||||||
|
|
@ -140,7 +154,7 @@ class JupyterKernelManager:
|
||||||
"-f", f"/tmp/kernel-{kernel_id}.json"
|
"-f", f"/tmp/kernel-{kernel_id}.json"
|
||||||
],
|
],
|
||||||
resource_limits=self.resource_limits,
|
resource_limits=self.resource_limits,
|
||||||
volumes=volumes or {},
|
volumes=container_volumes,
|
||||||
# TODO: Port mappings for ZMQ
|
# TODO: Port mappings for ZMQ
|
||||||
# TODO: Mount connection file into container
|
# TODO: Mount connection file into container
|
||||||
)
|
)
|
||||||
|
|
@ -164,6 +178,10 @@ class JupyterKernelManager:
|
||||||
if not self._verify_kernel(client):
|
if not self._verify_kernel(client):
|
||||||
raise KernelError(f"Kernel {kernel_id} not responsive")
|
raise KernelError(f"Kernel {kernel_id} not responsive")
|
||||||
|
|
||||||
|
# Execute injection code if provided (MCP tools setup)
|
||||||
|
if injection_code:
|
||||||
|
self._execute_injection_code(client, injection_code, kernel_id)
|
||||||
|
|
||||||
# Register kernel
|
# Register kernel
|
||||||
now = datetime.utcnow()
|
now = datetime.utcnow()
|
||||||
kernel_info = KernelInfo(
|
kernel_info = KernelInfo(
|
||||||
|
|
@ -478,3 +496,55 @@ _info
|
||||||
return True
|
return True
|
||||||
except Exception:
|
except Exception:
|
||||||
return False
|
return False
|
||||||
|
|
||||||
|
def _execute_injection_code(
|
||||||
|
self,
|
||||||
|
client: BlockingKernelClient,
|
||||||
|
injection_code: str,
|
||||||
|
kernel_id: str
|
||||||
|
) -> None:
|
||||||
|
"""
|
||||||
|
Execute MCP tool injection code on kernel startup.
|
||||||
|
|
||||||
|
This runs once when the kernel starts to set up MCP tools.
|
||||||
|
Unlike regular code execution, we don't capture output.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
client: Connected kernel client
|
||||||
|
injection_code: Python code to inject (MCP tools setup)
|
||||||
|
kernel_id: Kernel ID for error messages
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
KernelError: If injection code fails to execute
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
# Execute injection code silently
|
||||||
|
_msg_id = client.execute(injection_code, silent=True, store_history=False)
|
||||||
|
|
||||||
|
# Wait for execution to complete
|
||||||
|
timeout = 10 # Injection should be fast
|
||||||
|
while True:
|
||||||
|
try:
|
||||||
|
msg = client.get_iopub_msg(timeout=timeout)
|
||||||
|
msg_type = msg['header']['msg_type']
|
||||||
|
|
||||||
|
if msg_type == 'error':
|
||||||
|
content = msg['content']
|
||||||
|
error_msg = '\n'.join(content.get('traceback', [str(content)]))
|
||||||
|
raise KernelError(
|
||||||
|
f"MCP injection failed in kernel {kernel_id}: {error_msg}"
|
||||||
|
)
|
||||||
|
|
||||||
|
elif msg_type == 'status':
|
||||||
|
if msg['content']['execution_state'] == 'idle':
|
||||||
|
break # Injection complete
|
||||||
|
|
||||||
|
except zmq.error.Again:
|
||||||
|
break # Timeout, assume success
|
||||||
|
|
||||||
|
except KernelError:
|
||||||
|
raise
|
||||||
|
except Exception as e:
|
||||||
|
raise KernelError(
|
||||||
|
f"Failed to execute MCP injection code in kernel {kernel_id}: {e}"
|
||||||
|
) from e
|
||||||
|
|
|
||||||
|
|
@ -121,7 +121,9 @@ class SessionManager:
|
||||||
self,
|
self,
|
||||||
session_id: str,
|
session_id: str,
|
||||||
resource_limits: ResourceLimits,
|
resource_limits: ResourceLimits,
|
||||||
volumes: Optional[Dict[str, dict]] = None
|
volumes: Optional[Dict[str, dict]] = None,
|
||||||
|
injection_code: Optional[str] = None,
|
||||||
|
bridge_socket_path: Optional[str] = None
|
||||||
) -> Session:
|
) -> Session:
|
||||||
"""
|
"""
|
||||||
Create new stateful session.
|
Create new stateful session.
|
||||||
|
|
@ -130,6 +132,8 @@ class SessionManager:
|
||||||
session_id: Unique identifier for session
|
session_id: Unique identifier for session
|
||||||
resource_limits: Resource limits for session
|
resource_limits: Resource limits for session
|
||||||
volumes: Optional volume mounts
|
volumes: Optional volume mounts
|
||||||
|
injection_code: Optional MCP tool injection code to execute at startup
|
||||||
|
bridge_socket_path: Optional path to MCP bridge socket for mounting
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
Created Session object
|
Created Session object
|
||||||
|
|
@ -144,8 +148,13 @@ class SessionManager:
|
||||||
# Check max concurrent limit
|
# Check max concurrent limit
|
||||||
self._enforce_max_concurrent()
|
self._enforce_max_concurrent()
|
||||||
|
|
||||||
# Start kernel
|
# Start kernel with MCP injection if provided
|
||||||
kernel_id = self.kernel_manager.start_kernel(session_id, volumes=volumes)
|
kernel_id = self.kernel_manager.start_kernel(
|
||||||
|
session_id,
|
||||||
|
volumes=volumes,
|
||||||
|
injection_code=injection_code,
|
||||||
|
bridge_socket_path=bridge_socket_path
|
||||||
|
)
|
||||||
|
|
||||||
# Create session
|
# Create session
|
||||||
now = datetime.utcnow()
|
now = datetime.utcnow()
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue