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:
Hans Aschauer 2026-02-07 08:13:40 +01:00
parent ed58c6ad5a
commit 9cc43e4166
4 changed files with 106 additions and 13 deletions

View file

@ -48,7 +48,14 @@
- `cleanup_idle_kernels()`: Time-based cleanup
- 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
- `ipykernel` removed from server (will be in container image)
- `pyzmq>=27.1.0` for ZMQ support
@ -57,6 +64,7 @@
- [architecture1.md](../docs/architecture1.md) updated with Jupyter Backend section
- [todo.md](../docs/todo.md) Phase 2.2.1 updated with architecture details
- Architecture flow diagram and session mapping explanation
- [JUPYTER_IMPLEMENTATION_STATUS.md](../docs/JUPYTER_IMPLEMENTATION_STATUS.md) tracking document
### 🚧 In Progress / TODO

View file

@ -55,7 +55,9 @@ class JupyterBackend:
memory: Optional[str] = None,
cpu_quota: Optional[int] = 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:
"""
Execute code in stateful session.
@ -71,6 +73,8 @@ class JupyterBackend:
cpu_quota: CPU quota (uses config default if None)
custom_image: Custom image name (uses config default if None)
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:
ExecutionResult with execution output and metadata
@ -106,7 +110,7 @@ class JupyterBackend:
try:
self.session_manager.get_session(session_id)
except SessionError:
# Session doesn't exist, create it
# Session doesn't exist, create it with MCP injection
resource_limits = ResourceLimits(
memory=memory,
cpu_quota=cpu_quota,
@ -117,7 +121,9 @@ class JupyterBackend:
self.session_manager.create_session(
session_id=session_id,
resource_limits=resource_limits,
volumes=volumes
volumes=volumes,
injection_code=injection_code,
bridge_socket_path=bridge_socket_path
)
# Execute in session

View file

@ -99,7 +99,9 @@ class JupyterKernelManager:
def start_kernel(
self,
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:
"""
Start IPython kernel in dedicated container.
@ -108,14 +110,18 @@ class JupyterKernelManager:
1. Generate ZMQ connection info (ports, keys)
2. Create connection file
3. Create container with ipykernel command
4. Start container
5. Wait for kernel to be ready
6. Connect jupyter-client to kernel via ZMQ
7. Verify kernel is responsive
4. Mount bridge socket if provided (for MCP tools)
5. Start container
6. Wait for kernel to be ready
7. Connect jupyter-client to kernel via ZMQ
8. Execute injection code (MCP tools setup) if provided
9. Verify kernel is responsive
Args:
session_id: Session ID this kernel belongs to
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:
kernel_id: Unique identifier for this kernel
@ -132,6 +138,14 @@ class JupyterKernelManager:
connection_file = self._create_connection_file(kernel_id, connection_info)
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
config = ContainerConfig(
image=self.image,
@ -140,7 +154,7 @@ class JupyterKernelManager:
"-f", f"/tmp/kernel-{kernel_id}.json"
],
resource_limits=self.resource_limits,
volumes=volumes or {},
volumes=container_volumes,
# TODO: Port mappings for ZMQ
# TODO: Mount connection file into container
)
@ -164,6 +178,10 @@ class JupyterKernelManager:
if not self._verify_kernel(client):
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
now = datetime.utcnow()
kernel_info = KernelInfo(
@ -478,3 +496,55 @@ _info
return True
except Exception:
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

View file

@ -121,7 +121,9 @@ class SessionManager:
self,
session_id: str,
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:
"""
Create new stateful session.
@ -130,6 +132,8 @@ class SessionManager:
session_id: Unique identifier for session
resource_limits: Resource limits for session
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:
Created Session object
@ -144,8 +148,13 @@ class SessionManager:
# Check max concurrent limit
self._enforce_max_concurrent()
# Start kernel
kernel_id = self.kernel_manager.start_kernel(session_id, volumes=volumes)
# Start kernel with MCP injection if provided
kernel_id = self.kernel_manager.start_kernel(
session_id,
volumes=volumes,
injection_code=injection_code,
bridge_socket_path=bridge_socket_path
)
# Create session
now = datetime.utcnow()