langgraph-api 0.18.0.dev2__py3-none-any.whl → 0.18.0.dev4__py3-none-any.whl

langgraph_api/api/mcp/_handlers.py

@@ -3,29 +3,48 @@ from __future__ import annotations
 import asyncio
 import functools
 import json
-from typing import TYPE_CHECKING, Any, cast
+import re
+from http import HTTPStatus
+from typing import TYPE_CHECKING, Any, assert_never, cast
 
 import structlog
 from langgraph_sdk.client import LangGraphClient, get_client
 from starlette.responses import JSONResponse, Response
 
-from langgraph_api import __version__
-from langgraph_api.api.mcp import _sanitizers
+from langgraph_api.api.mcp import _protocol, _sanitizers
+from langgraph_api.api.mcp._calls import ToolCall
 from langgraph_api.api.mcp._constants import (
     DEFAULT_PAGE_SIZE,
-    ERROR_CODE_INVALID_PARAMS,
+    ERROR_CODE_INVALID_REQUEST,
     ERROR_CODE_METHOD_NOT_FOUND,
-    LATEST_PROTOCOL_VERSION,
+    INVALID_CURSOR_MESSAGE,
+    LATEST_LEGACY_PROTOCOL_VERSION,
+    LEGACY_PROTOCOL_VERSIONS,
     MAX_ASSISTANTS,
-    SUPPORTED_PROTOCOL_VERSIONS,
+    MAX_CURSOR_LENGTH,
+    MAX_PAGE_OFFSET,
+    PAUSED_TOOL_MESSAGE,
 )
+from langgraph_api.api.mcp._outcome import (
+    Completed,
+    Failed,
+    Paused,
+    RunOutcome,
+    parse_outcome,
+)
+from langgraph_api.serde import json_dumpb
 
 if TYPE_CHECKING:
+    from collections.abc import Mapping
+
     from langgraph_api.api.mcp._models import JsonRpcRequest
+    from langgraph_api.cache import JsonValue
     from langgraph_api.route import ApiRequest
 
 logger = structlog.stdlib.get_logger(__name__)
 
+_CURSOR_DIGITS = re.compile(rf"[0-9]{{1,{MAX_CURSOR_LENGTH}}}")
+
 
 @functools.lru_cache(maxsize=1)
 def _client() -> LangGraphClient:
@@ -53,54 +72,43 @@ def handle_get_request() -> Response:
 
 
 async def handle_post_request(request: ApiRequest) -> Response:
-    """Handle HTTP POST requests for JSON-RPC messaging.
-
-    Args:
-        request: The incoming request object
-
-    Returns:
-        Response to the JSON-RPC message
-    """
     body = await request.body()
 
-    # Validate JSON
     try:
         message = json.loads(body)
     except json.JSONDecodeError:
         return create_error_response("Invalid JSON", 400)
 
-    # Validate Accept header
     if not is_valid_accept_header(request):
         return create_error_response(
             "Accept header must include application/json or text/event-stream", 400
         )
 
-    # Validate message format
     if not isinstance(message, dict):
         return create_error_response("Invalid message format.", 400)
 
-    # Determine message type and route to appropriate handler
     id_ = message.get("id")
     method = message.get("method")
+    if (method is not None and not isinstance(method, str)) or (
+        id_ is not None and (isinstance(id_, bool) or not isinstance(id_, (str, int)))
+    ):
+        return _jsonrpc_error(
+            None,
+            _protocol.ProtocolError(
+                "Invalid JSON-RPC method or id", ERROR_CODE_INVALID_REQUEST
+            ),
+        )
 
-    # Check for required jsonrpc field
     if message.get("jsonrpc") != "2.0":
         return create_error_response(
             "Invalid JSON-RPC message. Missing or invalid jsonrpc version.", 400
         )
 
-    # Careful ID checks as the integer 0 is a valid ID
     if id_ is not None and method:
-        # JSON-RPC request
         return await handle_jsonrpc_request(request, cast("JsonRpcRequest", message))
-    elif id_ is not None:
-        # JSON-RPC response
-        return Response(status_code=202)
-    elif method:
-        # JSON-RPC notification
+    elif id_ is not None or method:
         return Response(status_code=202)
     else:
-        # Invalid message format
         return create_error_response(
             "Invalid message format. A message is to be either a JSON-RPC "
             "request, response, or notification."
@@ -146,114 +154,85 @@ async def handle_jsonrpc_request(
     request: ApiRequest,
     message: JsonRpcRequest,
 ) -> Response:
-    """Handle JSON-RPC requests (messages with both id and method).
-
-    Args:
-        request: The incoming request object
-        message: The parsed JSON-RPC message
-
-    Returns:
-        Response to the request
-    """
-    method = message["method"]
-    params = message.get("params", {})
-
-    if method == "initialize":
-        result_or_error = handle_initialize_request(message)
-    elif method == "ping":
-        result_or_error = {"result": {}}
-    elif method == "tools/list":
-        result_or_error = await handle_tools_list(request, params)
-    elif method == "tools/call":
-        result_or_error = await handle_tools_call(request, params)
-    else:
-        result_or_error = {
-            "error": {
-                "code": ERROR_CODE_METHOD_NOT_FOUND,
-                "message": f"Method not found: {method}",
-            }
-        }
-
-    # Process the result or error output
-    exists = {"error", "result"} - set(result_or_error.keys())
-    if len(exists) != 1:
-        raise AssertionError(
-            "Internal server error. Invalid response in MCP protocol implementation."
+    try:
+        call = _protocol.parse_call(request, message)
+    except _protocol.ProtocolError as error:
+        return _jsonrpc_error(message["id"], error)
+    try:
+        result = await _dispatch(request, call)
+    except _protocol.ProtocolError as error:
+        return _jsonrpc_error(
+            call.id, error, status=call.era.http_status(error), headers=call.era.headers
         )
+    return _result_response(call, result)
+
+
+async def _dispatch(request: ApiRequest, call: _protocol.McpCall) -> dict[str, Any]:
+    match call.era, call.method:
+        case _protocol.Modern(), "server/discover":
+            return _protocol.discover()
+        case _protocol.Legacy(), "initialize":
+            return handle_initialize_request(call.message)
+        case _protocol.Legacy(), "ping":
+            return {}
+        case _, "tools/list":
+            return await handle_tools_list(request, call.params)
+        case _, "tools/call":
+            return await handle_tools_call(request, call)
+        case _:
+            raise _protocol.ProtocolError(
+                f"Method not found: {call.method}", ERROR_CODE_METHOD_NOT_FOUND
+            )
 
+
+def _result_response(call: _protocol.McpCall, result: dict[str, Any]) -> Response:
     return JSONResponse(
         {
             "jsonrpc": "2.0",
-            "id": message["id"],
-            **result_or_error,
-        }
+            "id": call.id,
+            "result": call.era.result(call.method, result),
+        },
+        headers=call.era.headers,
     )
 
 
-def _negotiate_protocol_version(requested: str) -> str:
-    """Negotiate MCP protocol version with the client.
+def _jsonrpc_error(
+    id_: str | int | None,
+    error: _protocol.ProtocolError,
+    *,
+    status: HTTPStatus = HTTPStatus.BAD_REQUEST,
+    headers: Mapping[str, str] | None = None,
+) -> Response:
+    return JSONResponse(
+        {"jsonrpc": "2.0", "id": id_, "error": error.error},
+        status_code=status,
+        headers=headers,
+    )
 
-    Returns the requested version if supported, otherwise the latest.
-    """
-    if requested in SUPPORTED_PROTOCOL_VERSIONS:
+
+def _negotiate_protocol_version(requested: str) -> str:
+    if requested in LEGACY_PROTOCOL_VERSIONS:
         return requested
-    return LATEST_PROTOCOL_VERSION
+    return LATEST_LEGACY_PROTOCOL_VERSION
 
 
 def handle_initialize_request(message: JsonRpcRequest) -> dict[str, Any]:
-    """Handle initialize requests to establish protocol version.
-
-    Negotiates the protocol version with the client. If the client requests
-    a supported version, the server echoes it back. Otherwise the server
-    responds with the latest version it supports and the client may disconnect
-    if it cannot work with that version.
-
-    Args:
-        message: The JSON-RPC request message
-
-    Returns:
-        Response with negotiated protocol details
-    """
-    params = message.get("params", {})
-    requested = params.get("protocolVersion", LATEST_PROTOCOL_VERSION)
-    negotiated = _negotiate_protocol_version(requested)
-
+    params = _protocol.params_of(message)
+    requested = params.get("protocolVersion", LATEST_LEGACY_PROTOCOL_VERSION)
     return {
-        "result": {
-            "protocolVersion": negotiated,
-            "capabilities": {
-                "tools": {
-                    "listChanged": False,
-                }
-            },
-            "serverInfo": {"name": "LangGraph", "version": __version__},
-        }
+        "protocolVersion": _negotiate_protocol_version(requested),
+        "capabilities": _protocol.SERVER_CAPABILITIES,
+        "serverInfo": _protocol.SERVER_INFO,
     }
 
 
 async def handle_tools_list(
     request: ApiRequest, params: dict[str, Any]
 ) -> dict[str, Any]:
-    """Handle tools/list request to get available assistants as tools.
-
-    Args:
-        request: The incoming request object. Used for propagating any headers
-                 for authentication purposes.
-        params: The parameters for the tools/list request
-
-    Returns:
-        Dictionary containing list of available tools
-    """
     client = _client()
 
-    try:
-        cursor = params.get("cursor", 0)
-        cursor = int(cursor)
-    except ValueError:
-        cursor = 0
+    cursor = _page_offset(params.get("cursor", 0))
 
-    # Get assistants from the API
-    # For now set a large limit to get all assistants
     assistants = await client.assistants.search(
         offset=cursor, limit=DEFAULT_PAGE_SIZE, headers=request.headers
     )
@@ -263,9 +242,8 @@ async def handle_tools_list(
     else:
         next_cursor = None
 
-    # Deduplicate by normalized name, preserving original for the title field
     seen_names: set[str] = set()
-    unique: list[tuple[Any, str]] = []  # (assistant, normalized_name)
+    unique: list[tuple[Any, str]] = []
     for assistant in assistants:
         normalized = _sanitizers.normalize_name(assistant["name"])
         if normalized in seen_names:
@@ -279,10 +257,6 @@ async def handle_tools_list(
         unique.append((assistant, normalized))
 
     async def _get_tool(assistant: Any, normalized_name: str) -> dict[str, Any] | None:
-        """Get tool definition for an assistant.
-
-        Returns None if schema generation fails (e.g., non-serializable state).
-        """
         try:
             schemas = await client.assistants.get_schemas(
                 assistant["assistant_id"], headers=request.headers
@@ -292,7 +266,6 @@ async def handle_tools_list(
                 raise ValueError("Input schema generation failed")
             if input_schema:
                 _sanitizers.simplify_mcp_schema_inplace(input_schema)
-            # MCP spec requires inputSchema to be a JSON Schema with type "object"
             input_schema.setdefault("type", "object")
             return {
                 "name": normalized_name,
@@ -317,99 +290,76 @@ async def handle_tools_list(
             return None
 
     tools_or_none = await asyncio.gather(*(_get_tool(a, n) for a, n in unique))
-    # Filter out assistants that failed schema generation
     tools = [t for t in tools_or_none if t is not None]
 
     result = {"tools": tools}
-
     if next_cursor is not None:
-        result["nextCursor"] = next_cursor
+        result["nextCursor"] = str(next_cursor)
+    return result
 
-    return {
-        "result": result,
-    }
+
+def _page_offset(cursor: object) -> int:
+    match cursor:
+        case int() if not isinstance(cursor, bool):
+            offset = cursor
+        case str() if _CURSOR_DIGITS.fullmatch(cursor):
+            offset = int(cursor)
+        case _:
+            raise _protocol.ProtocolError(INVALID_CURSOR_MESSAGE)
+    if not 0 <= offset <= MAX_PAGE_OFFSET:
+        raise _protocol.ProtocolError(INVALID_CURSOR_MESSAGE)
+    return offset
 
 
 async def handle_tools_call(
-    request: ApiRequest, params: dict[str, Any]
+    request: ApiRequest, call: _protocol.McpCall
 ) -> dict[str, Any]:
-    """Handle tools/call request to execute an assistant.
-
-    Args:
-        request: The incoming request
-        params: The parameters for the tool call
-
-    Returns:
-        The result of the tool execution
-    """
-    client = _client()
-
-    tool_name = params.get("name")
-
-    if not tool_name:
-        return {
-            "error": {
-                "code": ERROR_CODE_INVALID_PARAMS,
-                "message": f"Unknown tool: {tool_name}",
-            },
-        }
-
-    arguments = params.get("arguments", {})
-    context = params.get("context")
-    assistants = await client.assistants.search(
-        limit=MAX_ASSISTANTS, headers=request.headers
-    )
-    matching_assistant = []
-    for assistant in assistants:
-        if (
-            normalized_name := _sanitizers.normalize_name(assistant["name"])
-        ) and normalized_name == tool_name:
-            matching_assistant.append(assistant)
-
-    num_assistants = len(matching_assistant)
-
-    if num_assistants == 0:
-        return {
-            "error": {
-                "code": ERROR_CODE_INVALID_PARAMS,
-                "message": f"Unknown tool: {tool_name}",
-            },
-        }
-    elif num_assistants > 1:
-        return {
-            "error": {
-                "code": ERROR_CODE_INVALID_PARAMS,
-                "message": "Multiple tools found with the same name.",
-            },
-        }
-    else:
-        tool_name = matching_assistant[0]["assistant_id"]
-
-    value = await client.runs.wait(
+    tool_call = ToolCall.of(call.params)
+    assistant_id = await _assistant_serving(request, tool_call.name)
+    value = await _client().runs.wait(
         thread_id=None,
-        assistant_id=tool_name,
-        input=arguments,
-        context=context,
+        assistant_id=assistant_id,
+        input=tool_call.arguments,
+        context=tool_call.context,
         headers=request.headers,
         raise_error=False,
     )
+    return _tool_result(parse_outcome(value))
+
 
-    if isinstance(value, dict) and "__error__" in value:
-        # This is a run-time error in the tool.
-        return {
-            "result": {
+async def _assistant_serving(request: ApiRequest, tool_name: str) -> str:
+    assistants = await _client().assistants.search(
+        limit=MAX_ASSISTANTS, headers=request.headers
+    )
+    matching = [
+        assistant
+        for assistant in assistants
+        if _sanitizers.normalize_name(assistant["name"]) == tool_name
+    ]
+    if not matching:
+        raise _protocol.ProtocolError(f"Unknown tool: {tool_name}")
+    if len(matching) > 1:
+        raise _protocol.ProtocolError("Multiple tools found with the same name.")
+    return matching[0]["assistant_id"]
+
+
+def _tool_result(outcome: RunOutcome) -> dict[str, JsonValue]:
+    match outcome:
+        case Completed(output=output):
+            return {"content": [_text_content(json_dumpb(output).decode())]}
+        case Failed(message=message):
+            return {"isError": True, "content": [_text_content(message)]}
+        case Paused(interrupts=interrupts):
+            return {
                 "isError": True,
                 "content": [
-                    {"type": "text", "text": value["__error__"]["error"]},
+                    _text_content(PAUSED_TOOL_MESSAGE),
+                    _text_content(json_dumpb(interrupts).decode()),
                 ],
             }
-        }
+        case _:
+            assert_never(outcome)
 
-    # All good, return the result
-    return {
-        "result": {
-            "content": [
-                {"type": "text", "text": repr(value)},
-            ]
-        }
-    }
+
+def _text_content(text: str) -> dict[str, str]:
+    return {"type": "text", "text": text}

langgraph_api/api/mcp/_outcome.py

@@ -0,0 +1,43 @@
+from dataclasses import dataclass
+
+from langgraph_api.api.mcp._constants import RUN_FAILED_MESSAGE
+from langgraph_api.cache import JsonValue
+from langgraph_api.schema import Interrupt
+
+
+@dataclass(frozen=True, slots=True)
+class Completed:
+    output: JsonValue
+
+
+@dataclass(frozen=True, slots=True)
+class Failed:
+    message: str
+
+
+@dataclass(frozen=True, slots=True)
+class Paused:
+    interrupts: tuple[Interrupt, ...]
+
+
+RunOutcome = Completed | Failed | Paused
+
+
+def parse_outcome(value: JsonValue) -> RunOutcome:
+    match value:
+        case {"__error__": error, **rest} if not rest:
+            return Failed(_error_message(error))
+        case {"__interrupt__": [_, *_] as interrupts}:
+            return Paused(tuple(interrupts))
+        case _:
+            return Completed(value)
+
+
+def _error_message(error: JsonValue) -> str:
+    match error:
+        case {"error": str() as name, "message": str() as message} if message:
+            return f"{name}: {message}"
+        case {"error": str() as name}:
+            return name
+        case _:
+            return RUN_FAILED_MESSAGE