diff --git a/docs/JUPYTER_IMPLEMENTATION_STATUS.md b/docs/JUPYTER_IMPLEMENTATION_STATUS.md index e2ec736..b57dc0d 100644 --- a/docs/JUPYTER_IMPLEMENTATION_STATUS.md +++ b/docs/JUPYTER_IMPLEMENTATION_STATUS.md @@ -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 diff --git a/src/mcp_forge/execution/jupyter/backend.py b/src/mcp_forge/execution/jupyter/backend.py index 64d38a5..d51b3ec 100644 --- a/src/mcp_forge/execution/jupyter/backend.py +++ b/src/mcp_forge/execution/jupyter/backend.py @@ -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 diff --git a/src/mcp_forge/execution/jupyter/kernel.py b/src/mcp_forge/execution/jupyter/kernel.py index eb85845..9aa913a 100644 --- a/src/mcp_forge/execution/jupyter/kernel.py +++ b/src/mcp_forge/execution/jupyter/kernel.py @@ -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 diff --git a/src/mcp_forge/execution/jupyter/sessions.py b/src/mcp_forge/execution/jupyter/sessions.py index d04d229..5b0a494 100644 --- a/src/mcp_forge/execution/jupyter/sessions.py +++ b/src/mcp_forge/execution/jupyter/sessions.py @@ -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()