api/controllers/mcp/mcp.py
Rama main · réplica de lectura del repositorio original.
from typing import Any, Literal, get_args
from flask import Response, request
from flask_restx import Resource
from pydantic import BaseModel, Field, RootModel, ValidationError
from sqlalchemy import select
from sqlalchemy.orm import Session
from controllers.common.access_response import mcp_server_not_found_response
from controllers.common.schema import register_response_schema_models, register_schema_model
from controllers.mcp import mcp_ns
from core.db.session_factory import session_factory
from core.mcp import types as mcp_types
from core.mcp.server.streamable_http import handle_mcp_request, negotiate_protocol_version
from fields.base import ResponseModel
from graphon.variables.input_entities import VariableEntity, VariableEntityType
from libs import helper
from models.enums import EndUserType
from models.model import App, AppMCPServer, AppMode, EndUser
from services.app.mcp_server_service import AppMCPServerStatus
_MCP_METHODS = frozenset(
method
for message_type in (mcp_types.ClientRequest, mcp_types.ClientNotification)
for model in get_args(message_type.model_fields["root"].annotation)
for method in get_args(model.model_fields["method"].annotation)
)
class MCPRequestError(Exception):
"""Custom exception for MCP request processing errors"""
def __init__(self, error_code: int, message: str):
self.error_code = error_code
self.message = message
super().__init__(message)
class MCPServerNotFoundError(MCPRequestError):
"""An unavailable server, backing app, or inactive server is one opaque identity failure."""
def __init__(self):
super().__init__(mcp_types.INVALID_REQUEST, "Server Not Found")
class MCPRequestPayload(BaseModel):
jsonrpc: str = Field(description="JSON-RPC version (should be '2.0')")
method: str = Field(description="The method to invoke")
params: dict[str, Any] | None = Field(default=None, description="Parameters for the method")
id: int | str | None = Field(default=None, description="Request ID for tracking responses")
class MCPJSONRPCErrorResponse(ResponseModel):
jsonrpc: Literal["2.0"]
id: int | float | str | None
error: mcp_types.ErrorData
class MCPJSONRPCResponse(RootModel[mcp_types.JSONRPCResponse | MCPJSONRPCErrorResponse]):
pass
register_schema_model(mcp_ns, MCPRequestPayload)
register_response_schema_models(mcp_ns, MCPJSONRPCResponse)
@mcp_ns.route("/server/<string:server_code>/mcp")
class MCPAppApi(Resource):
@mcp_ns.expect(mcp_ns.models[MCPRequestPayload.__name__], validate=False)
@mcp_ns.doc("handle_mcp_request")
@mcp_ns.doc(description="Handle Model Context Protocol (MCP) requests for a specific server")
@mcp_ns.doc(params={"server_code": "Unique identifier for the MCP server"})
@mcp_ns.response(200, "MCP JSON-RPC response", mcp_ns.models[MCPJSONRPCResponse.__name__])
@mcp_ns.response(202, "MCP notification accepted")
@mcp_ns.response(400, "Invalid MCP request or parameters", mcp_ns.models[MCPJSONRPCErrorResponse.__name__])
@mcp_ns.response(404, "Server or app not found", mcp_ns.models[MCPJSONRPCErrorResponse.__name__])
def post(self, server_code: str):
"""Handle MCP requests for a specific server.
Processes JSON-RPC formatted requests according to the Model Context Protocol specification.
Validates the server status and associated app before processing the request.
Args:
server_code: Unique identifier for the MCP server
Returns:
dict: JSON-RPC response from the MCP handler
Raises:
ValidationError: Invalid request format or parameters
"""
# response-contract:ignore MCP route returns Flask Response from JSON-RPC handler
sessions = session_factory.get_session_maker()
try:
# Resolve identity before parsing the body: missing identities have
# one bounded response even for malformed or oversized JSON. Only
# scalar IDs cross this short read session, never its ORM objects.
with sessions.begin() as session:
mcp_server, app = self._get_mcp_server_and_app(server_code, session)
self._validate_server_status(mcp_server)
identity = (mcp_server.id, app.id, app.tenant_id)
except MCPServerNotFoundError:
# Error formatting may read the request body; release DB resources first.
return mcp_server_not_found_response()
# Allowed requests retain the existing full body/protocol parser. The
# 64 KiB cap belongs only to opaque errors, not normal MCP invocations.
args = MCPRequestPayload.model_validate(mcp_ns.payload or {})
request_id: int | str | None = args.id
try:
mcp_request = self._parse_mcp_request(args.model_dump(exclude_none=True))
except MCPRequestError as e:
if "id" not in args.model_fields_set:
return Response("", status=202)
response = helper.compact_generate_response(
{
"jsonrpc": "2.0",
"id": request_id,
"error": {"code": e.error_code, "message": e.message},
}
)
# A legacy HTTP error lets modern clients fall back to initialize.
response.status_code = 400
return response
is_initialize = isinstance(mcp_request.root, mcp_types.InitializeRequest)
header_value = request.headers.get("MCP-Protocol-Version")
protocol_version = negotiate_protocol_version(header_value, is_initialize)
if protocol_version is None:
if isinstance(mcp_request, mcp_types.ClientNotification):
protocol_version = mcp_types.DEFAULT_NEGOTIATED_VERSION
else:
return self._protocol_version_error_response(request_id, header_value)
try:
# Reading/parsing may have taken time. Revalidate in the original
# execution transaction and use fresh, attached ORM records. A code
# reassigned to a different server/App/tenant is not this request's identity.
# A plain session rather than sessions.begin(): app generation may commit this
# session itself (e.g. when it creates a conversation) and keep using it, which a
# begin() context manager rejects as operating on a closed transaction.
with sessions() as session:
mcp_server, app = self._get_mcp_server_and_app(server_code, session)
self._validate_server_status(mcp_server)
if (mcp_server.id, app.id, app.tenant_id) != identity:
raise MCPServerNotFoundError()
user_input_form = self._get_user_input_form(app, session=session)
response = self._process_mcp_message(
mcp_request, request_id, app, mcp_server, user_input_form, session, protocol_version
)
session.commit()
return response
except MCPServerNotFoundError:
return mcp_server_not_found_response()
def _protocol_version_error_response(self, request_id: int | str | None, header_value: str | None) -> Response:
"""Return a JSON-RPC error for an unsupported MCP-Protocol-Version header.
Per JSON-RPC 2.0, an error whose request id is unknown uses a null id, so we echo the
offending request's id directly (None -> null) instead of fabricating a placeholder.
"""
error_data = mcp_types.ErrorData(
code=mcp_types.INVALID_REQUEST,
message=f"Unsupported MCP-Protocol-Version: {header_value}",
)
error_response = {
"jsonrpc": "2.0",
"id": request_id,
"error": error_data.model_dump(by_alias=True, mode="json", exclude_none=True),
}
return helper.compact_generate_response(error_response)
def _get_mcp_server_and_app(self, server_code: str, session: Session) -> tuple[AppMCPServer, App]:
"""Get and validate MCP server and app in one query session"""
mcp_server = session.scalar(select(AppMCPServer).where(AppMCPServer.server_code == server_code).limit(1))
if not mcp_server:
raise MCPServerNotFoundError()
app = session.scalar(
select(App).where(App.id == mcp_server.app_id, App.tenant_id == mcp_server.tenant_id).limit(1)
)
if not app:
raise MCPServerNotFoundError()
return mcp_server, app
def _validate_server_status(self, mcp_server: AppMCPServer):
"""Validate MCP server status"""
if mcp_server.status != AppMCPServerStatus.ACTIVE:
raise MCPServerNotFoundError()
def _process_mcp_message(
self,
mcp_request: mcp_types.ClientRequest | mcp_types.ClientNotification,
request_id: int | str | None,
app: App,
mcp_server: AppMCPServer,
user_input_form: list[VariableEntity],
session: Session,
protocol_version: str,
) -> Response:
"""Process MCP message (notification or request)"""
if isinstance(mcp_request, mcp_types.ClientNotification):
return self._handle_notification(mcp_request)
else:
return self._handle_request(
mcp_request, request_id, app, mcp_server, user_input_form, session, protocol_version
)
def _handle_notification(self, mcp_request: mcp_types.ClientNotification) -> Response:
"""Handle MCP notification"""
# For notifications, only support init notification
if mcp_request.root.method != "notifications/initialized":
raise MCPRequestError(mcp_types.INVALID_REQUEST, "Invalid notification method")
# Return HTTP 202 Accepted for notifications (no response body)
return Response("", status=202, content_type="application/json")
def _handle_request(
self,
mcp_request: mcp_types.ClientRequest,
request_id: int | str | None,
app: App,
mcp_server: AppMCPServer,
user_input_form: list[VariableEntity],
session: Session,
protocol_version: str,
) -> Response:
"""Handle MCP request"""
if request_id is None:
raise MCPRequestError(mcp_types.INVALID_REQUEST, "Request ID is required")
result = self._handle_mcp_request(
app, mcp_server, mcp_request, user_input_form, session, request_id, protocol_version
)
if result is None:
# This shouldn't happen for requests, but handle gracefully
raise MCPRequestError(mcp_types.INTERNAL_ERROR, "No response generated for request")
return helper.compact_generate_response(result.model_dump(by_alias=True, mode="json", exclude_none=True))
def _get_user_input_form(self, app: App, *, session: Session) -> list[VariableEntity]:
"""Get and convert user input form"""
# Get raw user input form based on app mode
if app.mode in {AppMode.ADVANCED_CHAT, AppMode.WORKFLOW}:
workflow = app.workflow_with_session(session=session)
if workflow is None:
raise MCPRequestError(mcp_types.INVALID_REQUEST, "App is unavailable")
raw_user_input_form = workflow.user_input_form(to_old_structure=True)
else:
app_model_config = app.app_model_config_with_session(session=session)
if app_model_config is None:
raise MCPRequestError(mcp_types.INVALID_REQUEST, "App is unavailable")
raw_user_input_form = app_model_config.user_input_form_list
# Convert to VariableEntity objects
try:
return self._convert_user_input_form(raw_user_input_form)
except ValidationError as e:
raise MCPRequestError(mcp_types.INVALID_PARAMS, f"Invalid user_input_form: {str(e)}") from e
def _convert_user_input_form(self, raw_form: list[dict[str, Any]]) -> list[VariableEntity]:
"""Convert raw user input form to VariableEntity objects"""
return [self._create_variable_entity(item) for item in raw_form]
def _create_variable_entity(self, item: dict[str, Any]) -> VariableEntity:
"""Create a single VariableEntity from raw form item"""
variable_type_raw: str = item.get("type", "") or list(item.keys())[0]
try:
variable_type = VariableEntityType(variable_type_raw)
except ValueError as e:
raise MCPRequestError(
mcp_types.INVALID_PARAMS, f"Invalid user_input_form variable type: {variable_type_raw}"
) from e
variable = item[variable_type_raw]
return VariableEntity(
type=variable_type,
variable=variable.get("variable"),
description=variable.get("description") or "",
label=variable.get("label"),
required=variable.get("required", False),
max_length=variable.get("max_length"),
options=variable.get("options") or [],
json_schema=variable.get("json_schema"),
)
def _parse_mcp_request(self, args: dict[str, Any]) -> mcp_types.ClientRequest | mcp_types.ClientNotification:
"""Parse and validate MCP request"""
if args.get("method") not in _MCP_METHODS:
raise MCPRequestError(mcp_types.METHOD_NOT_FOUND, "Method not found")
try:
return mcp_types.ClientRequest.model_validate(args)
except ValidationError:
try:
return mcp_types.ClientNotification.model_validate(args)
except ValidationError as e:
raise MCPRequestError(mcp_types.INVALID_PARAMS, "Invalid MCP request parameters") from e
def _retrieve_end_user(self, tenant_id: str, mcp_server_id: str, session: Session) -> EndUser | None:
"""Get end user in the caller session"""
return session.scalar(
select(EndUser)
.where(EndUser.tenant_id == tenant_id)
.where(EndUser.session_id == mcp_server_id)
.where(EndUser.type == EndUserType.MCP)
.limit(1)
)
def _create_end_user(
self, client_name: str, tenant_id: str, app_id: str, mcp_server_id: str, session: Session
) -> EndUser:
"""Create end user in existing session"""
end_user = EndUser(
tenant_id=tenant_id,
app_id=app_id,
type=EndUserType.MCP,
name=client_name,
session_id=mcp_server_id,
)
session.add(end_user)
session.flush() # Use flush instead of commit to keep transaction open
session.refresh(end_user)
return end_user
def _handle_mcp_request(
self,
app: App,
mcp_server: AppMCPServer,
mcp_request: mcp_types.ClientRequest,
user_input_form: list[VariableEntity],
session: Session,
request_id: int | str,
protocol_version: str,
) -> mcp_types.JSONRPCResponse | mcp_types.JSONRPCError | None:
"""Handle MCP request and return response"""
end_user = self._retrieve_end_user(mcp_server.tenant_id, mcp_server.id, session)
if not end_user and isinstance(mcp_request.root, mcp_types.InitializeRequest):
client_info = mcp_request.root.params.clientInfo
client_name = f"{client_info.name}@{client_info.version}"
end_user = self._create_end_user(client_name, app.tenant_id, app.id, mcp_server.id, session)
return handle_mcp_request(
session, app, mcp_request, user_input_form, mcp_server, end_user, request_id, protocol_version
)