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
|
||||
- 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
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue