Implement support for input_required status

Author: LucaButBoring

src/mcp/server/lowlevel/async_operations.py

@@ -172,6 +172,32 @@ def cancel_session_operations(self, session_id: str) -> int:
 
         return canceled_count
 
+    def mark_input_required(self, token: str) -> bool:
+        """Mark operation as requiring input from client."""
+        operation = self._operations.get(token)
+        if not operation:
+            return False
+
+        # Can only move to input_required from submitted or working states
+        if operation.status not in ("submitted", "working"):
+            return False
+
+        operation.status = "input_required"
+        return True
+
+    def mark_input_completed(self, token: str) -> bool:
+        """Mark operation as no longer requiring input, return to working state."""
+        operation = self._operations.get(token)
+        if not operation:
+            return False
+
+        # Can only move from input_required back to working
+        if operation.status != "input_required":
+            return False
+
+        operation.status = "working"
+        return True
+
     async def start_cleanup_task(self) -> None:
         """Start the background cleanup task."""
         if self._cleanup_task is not None:

src/mcp/server/lowlevel/server.py

@@ -641,6 +641,35 @@ async def _handle_cancelled_notification(self, notification: types.CancelledNoti
         logger.debug(f"Received cancellation notification for request {request_id}")
         self.handle_cancelled_notification(request_id)
 
+    def send_request_for_operation(self, token: str, request: types.ServerRequest) -> None:
+        """Send a request associated with an async operation."""
+        # Mark operation as requiring input
+        if self.async_operations.mark_input_required(token):
+            # Add operation token to request
+            if hasattr(request.root, "params") and request.root.params is not None:
+                if not hasattr(request.root.params, "operation") or request.root.params.operation is None:
+                    # Create operation field if it doesn't exist
+                    operation_data = types.RequestParams.Operation(token=token)
+                    request.root.params.operation = operation_data
+            logger.debug(f"Marked operation {token} as input_required and added to request")
+
+    def send_notification_for_operation(self, token: str, notification: types.ServerNotification) -> None:
+        """Send a notification associated with an async operation."""
+        # Mark operation as requiring input
+        if self.async_operations.mark_input_required(token):
+            # Add operation token to notification
+            if hasattr(notification.root, "params") and notification.root.params is not None:
+                if not hasattr(notification.root.params, "operation") or notification.root.params.operation is None:
+                    # Create operation field if it doesn't exist
+                    operation_data = types.NotificationParams.Operation(token=token)
+                    notification.root.params.operation = operation_data
+            logger.debug(f"Marked operation {token} as input_required and added to notification")
+
+    def complete_request_for_operation(self, token: str) -> None:
+        """Mark that a request for an operation has been completed."""
+        if self.async_operations.mark_input_completed(token):
+            logger.debug(f"Marked operation {token} as no longer requiring input")
+
     async def run(
         self,
         read_stream: MemoryObjectReceiveStream[SessionMessage | Exception],
@@ -741,6 +770,16 @@ async def _handle_request(
                 )
                 response = await handler(req)
 
+                # Handle operation token in response (for input_required operations)
+                if (
+                    hasattr(req, "params")
+                    and req.params is not None
+                    and hasattr(req.params, "operation")
+                    and req.params.operation is not None
+                ):
+                    operation_token = req.params.operation.token
+                    self.complete_request_for_operation(operation_token)
+
                 # Track async operations for cancellation
                 if isinstance(req, types.CallToolRequest):
                     result = response.root

tests/server/test_cancellation_logic.py

@@ -239,3 +239,221 @@ async def run_handler():
         assert isinstance(response, types.ServerResult)
         payload_result = cast(types.GetOperationPayloadResult, response.root)
         assert payload_result.result == result
+
+
+class TestCancellationLogic:
+    """Test cancellation logic for async operations."""
+
+    def test_handle_cancelled_notification(self):
+        """Test handling of cancelled notifications."""
+        manager = AsyncOperationManager()
+        server = Server("Test", async_operations=manager)
+
+        # Create an operation
+        operation = manager.create_operation("test_tool", {"arg": "value"}, "session1")
+
+        # Track the operation with a request ID
+        request_id = "req_123"
+        server._request_to_operation[request_id] = operation.token
+
+        # Handle cancellation
+        server.handle_cancelled_notification(request_id)
+
+        # Verify operation was cancelled
+        cancelled_op = manager.get_operation(operation.token)
+        assert cancelled_op is not None
+        assert cancelled_op.status == "canceled"
+
+        # Verify mapping was cleaned up
+        assert request_id not in server._request_to_operation
+
+    def test_cancelled_notification_handler(self):
+        """Test the async cancelled notification handler."""
+        manager = AsyncOperationManager()
+        server = Server("Test", async_operations=manager)
+
+        # Create an operation
+        operation = manager.create_operation("test_tool", {"arg": "value"}, "session1")
+
+        # Track the operation with a request ID
+        request_id = "req_456"
+        server._request_to_operation[request_id] = operation.token
+
+        # Create cancelled notification
+        notification = types.CancelledNotification(params=types.CancelledNotificationParams(requestId=request_id))
+
+        # Handle the notification
+        import asyncio
+
+        asyncio.run(server._handle_cancelled_notification(notification))
+
+        # Verify operation was cancelled
+        cancelled_op = manager.get_operation(operation.token)
+        assert cancelled_op is not None
+        assert cancelled_op.status == "canceled"
+
+    def test_validate_operation_token_cancelled(self):
+        """Test that cancelled operations are rejected."""
+        manager = AsyncOperationManager()
+        server = Server("Test", async_operations=manager)
+
+        # Create and cancel an operation
+        operation = manager.create_operation("test_tool", {"arg": "value"}, "session1")
+        manager.cancel_operation(operation.token)
+
+        # Verify that accessing cancelled operation raises error
+        with pytest.raises(McpError) as exc_info:
+            server._validate_operation_token(operation.token)
+
+        assert exc_info.value.error.code == -32602
+        assert "cancelled" in exc_info.value.error.message.lower()
+
+    def test_nonexistent_request_id_cancellation(self):
+        """Test cancellation of non-existent request ID."""
+        server = Server("Test")
+
+        # Should not raise error for non-existent request ID
+        server.handle_cancelled_notification("nonexistent_request")
+
+        # Verify no operations were affected
+        assert len(server._request_to_operation) == 0
+
+
+class TestInputRequiredBehavior:
+    """Test input_required status handling for async operations."""
+
+    def test_mark_input_required(self):
+        """Test marking operation as requiring input."""
+        manager = AsyncOperationManager()
+
+        # Create operation in submitted state
+        operation = manager.create_operation("test_tool", {"arg": "value"}, "session1")
+        assert operation.status == "submitted"
+
+        # Mark as input required
+        result = manager.mark_input_required(operation.token)
+        assert result is True
+
+        # Verify status changed
+        updated_op = manager.get_operation(operation.token)
+        assert updated_op is not None
+        assert updated_op.status == "input_required"
+
+    def test_mark_input_required_from_working(self):
+        """Test marking working operation as requiring input."""
+        manager = AsyncOperationManager()
+
+        # Create and mark as working
+        operation = manager.create_operation("test_tool", {"arg": "value"}, "session1")
+        manager.mark_working(operation.token)
+        assert operation.status == "working"
+
+        # Mark as input required
+        result = manager.mark_input_required(operation.token)
+        assert result is True
+        assert operation.status == "input_required"
+
+    def test_mark_input_required_invalid_states(self):
+        """Test that input_required can only be set from valid states."""
+        manager = AsyncOperationManager()
+
+        # Test from completed state
+        operation = manager.create_operation("test_tool", {"arg": "value"}, "session1")
+        manager.complete_operation(operation.token, types.CallToolResult(content=[]))
+
+        result = manager.mark_input_required(operation.token)
+        assert result is False
+        assert operation.status == "completed"
+
+    def test_mark_input_completed(self):
+        """Test marking input as completed."""
+        manager = AsyncOperationManager()
+
+        # Create operation and mark as input required
+        operation = manager.create_operation("test_tool", {"arg": "value"}, "session1")
+        manager.mark_input_required(operation.token)
+        assert operation.status == "input_required"
+
+        # Mark input as completed
+        result = manager.mark_input_completed(operation.token)
+        assert result is True
+        assert operation.status == "working"
+
+    def test_mark_input_completed_invalid_state(self):
+        """Test that input can only be completed from input_required state."""
+        manager = AsyncOperationManager()
+
+        # Create operation in submitted state
+        operation = manager.create_operation("test_tool", {"arg": "value"}, "session1")
+        assert operation.status == "submitted"
+
+        # Try to mark input completed from wrong state
+        result = manager.mark_input_completed(operation.token)
+        assert result is False
+        assert operation.status == "submitted"
+
+    def test_nonexistent_token_operations(self):
+        """Test input_required operations on nonexistent tokens."""
+        manager = AsyncOperationManager()
+
+        # Test with fake token
+        assert manager.mark_input_required("fake_token") is False
+        assert manager.mark_input_completed("fake_token") is False
+
+    def test_server_send_request_for_operation(self):
+        """Test server method for sending requests with operation tokens."""
+        manager = AsyncOperationManager()
+        server = Server("Test", async_operations=manager)
+
+        # Create operation
+        operation = manager.create_operation("test_tool", {"arg": "value"}, "session1")
+        manager.mark_working(operation.token)
+
+        # Create a mock request
+        request = types.ServerRequest(
+            types.CreateMessageRequest(
+                params=types.CreateMessageRequestParams(
+                    messages=[types.SamplingMessage(role="user", content=types.TextContent(type="text", text="test"))],
+                    maxTokens=100,
+                )
+            )
+        )
+
+        # Send request for operation
+        server.send_request_for_operation(operation.token, request)
+
+        # Verify operation status changed
+        updated_op = manager.get_operation(operation.token)
+        assert updated_op is not None
+        assert updated_op.status == "input_required"
+
+    def test_server_complete_request_for_operation(self):
+        """Test server method for completing requests."""
+        manager = AsyncOperationManager()
+        server = Server("Test", async_operations=manager)
+
+        # Create operation and mark as input required
+        operation = manager.create_operation("test_tool", {"arg": "value"}, "session1")
+        manager.mark_input_required(operation.token)
+
+        # Complete request for operation
+        server.complete_request_for_operation(operation.token)
+
+        # Verify operation status changed back to working
+        updated_op = manager.get_operation(operation.token)
+        assert updated_op is not None
+        assert updated_op.status == "working"
+
+    def test_input_required_is_terminal_check(self):
+        """Test that input_required is not considered a terminal state."""
+        manager = AsyncOperationManager()
+
+        # Create operation and mark as input required
+        operation = manager.create_operation("test_tool", {"arg": "value"}, "session1")
+        manager.mark_input_required(operation.token)
+
+        # Verify it's not terminal
+        assert not operation.is_terminal
+
+        # Verify it doesn't expire while in input_required state
+        assert not operation.is_expired