langgraph-api 0.15.0.dev1__py3-none-any.whl → 0.15.0.dev2__py3-none-any.whl

langgraph_api/__init__.py

@@ -1 +1 @@
-__version__ = "0.15.0.dev1"
+__version__ = "0.15.0.dev2"

langgraph_api/api/a2a.py

@@ -16,6 +16,7 @@ import binascii
 import functools
 import os
 import uuid
+from dataclasses import dataclass, replace
 from datetime import UTC, datetime
 from typing import Any, Literal, NotRequired, cast
 
@@ -27,6 +28,15 @@ from starlette.responses import JSONResponse, Response
 from typing_extensions import TypedDict
 
 from langgraph_api import __version__
+from langgraph_api.api.a2ui import (
+    A2UI_EXTENSION_URI,
+    A2UI_MIME_TYPE,
+    A2UIValidationError,
+    is_a2ui_mime_type,
+    normalize_metadata,
+    validate_client_metadata,
+    validate_payload,
+)
 from langgraph_api.config import A2A_ALLOWED_TOOL_CALL_RESULTS
 from langgraph_api.metadata import USER_API_URL
 from langgraph_api.route import ApiRequest, ApiRoute
@@ -103,6 +113,73 @@ class InvalidAgentResponseError(ValueError):
     """Raised when public agent output cannot be represented as an A2A part."""
 
 
+@dataclass(frozen=True)
+class A2UIContext:
+    enabled: bool
+    active: bool
+    metadata: dict[str, Any]
+    extensions: list[str]
+    drop_invalid: bool = False
+
+
+_A2UI_DISABLED_CONTEXT = A2UIContext(False, False, {}, [])
+
+
+def _a2ui_enabled(assistant: dict[str, Any]) -> bool:
+    metadata = assistant.get("metadata")
+    a2a_metadata = metadata.get("a2a") if isinstance(metadata, dict) else None
+    return isinstance(a2a_metadata, dict) and a2a_metadata.get("a2ui") is True
+
+
+def _a2ui_context_for_run(
+    run: dict[str, Any], *, enabled: bool | None = None
+) -> A2UIContext:
+    metadata = run.get("metadata")
+    a2a_metadata = metadata.get("a2a") if isinstance(metadata, dict) else None
+    if not isinstance(a2a_metadata, dict):
+        a2a_metadata = {}
+    return A2UIContext(
+        enabled=a2a_metadata.get("a2ui") is True if enabled is None else enabled,
+        active=a2a_metadata.get("a2ui_active") is True,
+        metadata={},
+        extensions=[],
+    )
+
+
+def _a2ui_run_metadata(
+    assistant: dict[str, Any], context: A2UIContext
+) -> dict[str, Any] | None:
+    if not context.enabled:
+        return None
+    metadata = assistant.get("metadata")
+    a2a_metadata = metadata.get("a2a") if isinstance(metadata, dict) else None
+    return {
+        "a2a": {
+            **(a2a_metadata if isinstance(a2a_metadata, dict) else {}),
+            "a2ui_active": context.active,
+        }
+    }
+
+
+def _a2ui_message_context(
+    message: dict[str, Any], assistant: dict[str, Any]
+) -> A2UIContext:
+    metadata = message.get("metadata") or {}
+    extensions = message.get("extensions") or []
+    if (
+        not isinstance(metadata, dict)
+        or not isinstance(extensions, list)
+        or not all(isinstance(extension, str) for extension in extensions)
+    ):
+        raise A2UIValidationError("Invalid A2A message metadata or extensions")
+    return A2UIContext(
+        enabled=_a2ui_enabled(assistant),
+        active=A2UI_EXTENSION_URI in extensions,
+        metadata=metadata,
+        extensions=extensions,
+    )
+
+
 # ============================================================================
 # Legacy (v0.x) format helpers
 # ============================================================================
@@ -284,7 +361,15 @@ async def _validate_supports_messages(
 
     # Validate messages field only if there are text or file parts
     has_message_parts = any(
-        ("text" in part) or ("file" in part) for part in parts if isinstance(part, dict)
+        ("text" in part)
+        or ("file" in part)
+        or (
+            "data" in part
+            and isinstance(part.get("metadata"), dict)
+            and is_a2ui_mime_type(part["metadata"].get("mimeType"))
+        )
+        for part in parts
+        if isinstance(part, dict)
     )
     if has_message_parts:
         input_schema = schemas.get("input_schema") or schemas.get("state_schema")
@@ -506,12 +591,14 @@ def _process_a2a_message_parts(
     parts: list[dict[str, Any]],
     message_role: str,
     message_id: str,
+    *,
+    a2ui: A2UIContext,
 ) -> dict[str, Any]:
     """Convert A2A message parts to LangChain messages format.
 
-    Text-only messages keep today's per-part string content. When any file part
-    is present, emit one consolidated message whose ``content`` is an ordered
-    list of LangChain content blocks (text + file/image/audio/video).
+    Text-only messages keep today's per-part string content. When a file or
+    A2UI part is present, emit one consolidated message whose ``content`` is an
+    ordered list of LangChain content blocks.
 
     Args:
         parts: List of A2A message parts
@@ -526,7 +613,19 @@ def _process_a2a_message_parts(
     messages = []
     additional_data = {}
     content_blocks: list[dict[str, Any]] = []
-    has_file_parts = any(isinstance(part, dict) and "file" in part for part in parts)
+    validate_client_metadata(a2ui.metadata)
+    has_content_blocks = any(
+        isinstance(part, dict)
+        and (
+            "file" in part
+            or (
+                "data" in part
+                and isinstance(part.get("metadata"), dict)
+                and is_a2ui_mime_type(part["metadata"].get("mimeType"))
+            )
+        )
+        for part in parts
+    )
     langgraph_role = "human" if message_role == "ROLE_USER" else "assistant"
 
     for part in parts:
@@ -544,7 +643,7 @@ def _process_a2a_message_parts(
             )
 
         if "text" in part:
-            if has_file_parts:
+            if has_content_blocks:
                 content_blocks.append({"type": "text", "text": part["text"]})
             else:
                 messages.append(
@@ -559,8 +658,26 @@ def _process_a2a_message_parts(
             content_blocks.append(_a2a_file_part_to_content_block(part))
 
         elif "data" in part:
-            # Data parts become structured input parameters
             part_data = part.get("data", {})
+            part_metadata = (
+                part.get("metadata") if isinstance(part.get("metadata"), dict) else {}
+            )
+            if is_a2ui_mime_type(part_metadata.get("mimeType")):
+                if not a2ui.enabled:
+                    raise A2UIValidationError("A2UI is not enabled for this agent")
+                if not a2ui.active:
+                    raise A2UIValidationError("The A2UI extension must be activated")
+                validate_payload(part_data, direction="client")
+                content_blocks.append(
+                    {
+                        "type": "data",
+                        "data": part_data,
+                        "metadata": normalize_metadata(part_metadata),
+                    }
+                )
+                continue
+
+            # Non-A2UI data parts remain structured top-level input parameters.
             if not isinstance(part_data, dict):
                 raise ValueError(
                     "DataPart must contain a JSON object in the 'data' field"
@@ -573,7 +690,7 @@ def _process_a2a_message_parts(
                 "A2A agents support 'text', 'data', and 'file' parts only."
             )
 
-    if has_file_parts:
+    if has_content_blocks:
         messages.append(
             {
                 "role": langgraph_role,
@@ -582,6 +699,14 @@ def _process_a2a_message_parts(
             }
         )
 
+    if a2ui.metadata or a2ui.extensions:
+        a2a_data = {
+            **({"metadata": a2ui.metadata} if a2ui.metadata else {}),
+            **({"extensions": a2ui.extensions} if a2ui.extensions else {}),
+        }
+        for graph_message in messages:
+            graph_message["additional_kwargs"] = {"a2a": a2a_data}
+
     if not messages and not additional_data:
         raise ValueError(
             "Message must contain at least one valid text, data, or file part"
@@ -707,7 +832,12 @@ def _content_block_to_a2a_file_part(block: dict[str, Any]) -> dict[str, Any]:
     return {"kind": "file", "file": file_obj}
 
 
-def _content_to_a2a_parts(content: Any) -> list[dict[str, Any]]:
+def _content_to_a2a_parts(
+    content: Any,
+    *,
+    a2ui: A2UIContext,
+    a2ui_direction: Literal["client", "server"] = "server",
+) -> list[dict[str, Any]]:
     """Convert public LangChain message content to ordered A2A parts."""
 
     if isinstance(content, str):
@@ -732,6 +862,39 @@ def _content_to_a2a_parts(content: Any) -> list[dict[str, Any]]:
             text = block.get("text")
             if isinstance(text, str):
                 text_parts.append(text)
+        elif isinstance(block, dict) and block.get("type") == "data":
+            metadata = (
+                block.get("metadata") if isinstance(block.get("metadata"), dict) else {}
+            )
+            if not is_a2ui_mime_type(metadata.get("mimeType")):
+                continue
+            if not a2ui.enabled:
+                if a2ui.drop_invalid:
+                    continue
+                raise A2UIValidationError(
+                    "Agent emitted A2UI while the extension was disabled"
+                )
+            if not a2ui.active:
+                continue
+            data = block.get("data")
+            try:
+                validate_payload(data, direction=a2ui_direction)
+            except A2UIValidationError:
+                if not a2ui.drop_invalid:
+                    raise
+                logger.warning(
+                    "Dropping invalid A2UI block from task history",
+                    direction=a2ui_direction,
+                )
+                continue
+            flush_text()
+            parts.append(
+                {
+                    "kind": "data",
+                    "data": data,
+                    "metadata": normalize_metadata(metadata),
+                }
+            )
         elif (
             isinstance(block, dict) and block.get("type") in _MEDIA_CONTENT_BLOCK_TYPES
         ):
@@ -746,7 +909,7 @@ def _a2a_parts_have_content(parts: list[dict[str, Any]]) -> bool:
     """Return whether converted parts contain public text or file content."""
 
     return any(
-        part.get("kind") == "file"
+        part.get("kind") in {"data", "file"}
         or (part.get("kind") == "text" and bool(part.get("text")))
         for part in parts
     )
@@ -772,7 +935,9 @@ def _content_to_text(content: Any) -> str:
     return str(content)
 
 
-def _extract_a2a_response(result: dict[str, Any]) -> list[dict[str, Any]]:
+def _extract_a2a_response(
+    result: dict[str, Any], *, a2ui: A2UIContext
+) -> list[dict[str, Any]]:
     """Extract public A2A parts from the last assistant response.
 
     Args:
@@ -802,7 +967,10 @@ def _extract_a2a_response(result: dict[str, Any]) -> list[dict[str, Any]]:
         if not isinstance(message, dict):
             continue
         if _a2a_role_for_message(message) == "ROLE_AGENT" and "content" in message:
-            parts = _content_to_a2a_parts(message["content"])
+            parts = _content_to_a2a_parts(
+                message["content"],
+                a2ui=a2ui,
+            )
             if _a2a_parts_have_content(parts):
                 return parts
 
@@ -1135,6 +1303,7 @@ def _convert_messages_to_a2a_format(
     context_id: str,
     *,
     history_length: int | None = None,
+    a2ui: A2UIContext,
 ) -> list[dict[str, Any]]:
     """Convert LangChain messages to A2A message format.
 
@@ -1154,6 +1323,8 @@ def _convert_messages_to_a2a_format(
         List of A2A messages
     """
 
+    history_a2ui = replace(a2ui, drop_invalid=True)
+
     # Convert each LangChain message to A2A format
     a2a_messages = []
     for msg in messages:
@@ -1169,7 +1340,11 @@ def _convert_messages_to_a2a_format(
                 a2a_role = "ROLE_AGENT"
                 parts = [_tool_results_data_part([tool_result])]
             else:
-                parts = _content_to_a2a_parts(msg.get("content", ""))
+                parts = _content_to_a2a_parts(
+                    msg.get("content", ""),
+                    a2ui=history_a2ui,
+                    a2ui_direction=("client" if a2a_role == "ROLE_USER" else "server"),
+                )
                 if a2a_role == "ROLE_AGENT" and not _a2a_parts_have_content(parts):
                     continue
                 if not parts:
@@ -1197,6 +1372,7 @@ async def _create_task_response(
     context_id: str,
     result: dict[str, Any],
     assistant_id: str,
+    a2ui: A2UIContext,
     history_length: int | None = None,
 ) -> dict[str, Any]:
     """Create A2A Task response structure for both success and failure cases.
@@ -1218,6 +1394,7 @@ async def _create_task_response(
         task_id,
         context_id,
         history_length=history_length,
+        a2ui=a2ui,
     )
 
     base_task: dict[str, Any] = {
@@ -1257,7 +1434,13 @@ async def _create_task_response(
             "timestamp": datetime.now(UTC).isoformat(),
         }
         base_task["artifacts"] = [
-            _create_response_artifact(_extract_a2a_response(result), assistant_id)
+            _create_response_artifact(
+                _extract_a2a_response(
+                    result,
+                    a2ui=a2ui,
+                ),
+                assistant_id,
+            )
         ]
 
     return {"result": {"task": base_task}}
@@ -1510,7 +1693,7 @@ async def handle_jsonrpc_request(
     elif method == "SendMessage":
         result_or_error = await handle_message_send(request, params, assistant_id)
     elif method == "GetTask":
-        result_or_error = await handle_tasks_get(request, params)
+        result_or_error = await handle_tasks_get(request, params, assistant_id)
     elif method == "CancelTask":
         result_or_error = await handle_tasks_cancel(request, params)
     elif method == "ListTasks":
@@ -1656,8 +1839,21 @@ async def handle_message_send(
 
         # Process A2A message parts into LangChain messages format
         try:
+            a2ui = _a2ui_message_context(message, assistant)
             message_role = _normalize_input_role(message.get("role", "ROLE_USER"))
-            input_content = _process_a2a_message_parts(parts, message_role, message_id)
+            input_content = _process_a2a_message_parts(
+                parts,
+                message_role,
+                message_id,
+                a2ui=a2ui,
+            )
+        except A2UIValidationError as e:
+            return {
+                "error": {
+                    "code": ERROR_CODE_INVALID_PARAMS,
+                    "message": str(e),
+                }
+            }
         except ValueError as e:
             return {
                 "error": {
@@ -1714,6 +1910,7 @@ async def handle_message_send(
                 input=input_content,
                 command=command,
                 context=run_context,
+                metadata=_a2ui_run_metadata(assistant, a2ui),
                 if_not_exists="create",
                 headers=request.headers,
             )
@@ -1740,6 +1937,7 @@ async def handle_message_send(
             result=result,
             assistant_id=assistant_id,
             history_length=history_length,
+            a2ui=a2ui,
         )
 
     except InvalidAgentResponseError as e:
@@ -1784,7 +1982,7 @@ async def _get_historical_messages_for_task(
 
 
 async def handle_tasks_get(
-    request: ApiRequest, params: dict[str, Any]
+    request: ApiRequest, params: dict[str, Any], assistant_id: str
 ) -> dict[str, Any]:
     """Handle tasks/get requests to retrieve task status.
 
@@ -1843,7 +2041,8 @@ async def handle_tasks_get(
 
         try:
             # TODO: fix the N+1 query issue
-            run_info, thread_info = await asyncio.gather(
+            assistant, run_info, thread_info = await asyncio.gather(
+                _get_assistant(assistant_id, request.headers),
                 client.runs.get(
                     thread_id=context_id,
                     run_id=run_id,
@@ -1896,7 +2095,12 @@ async def handle_tasks_get(
                     task_id,
                     context_id,
                     history_length=history_length,
+                    a2ui=_a2ui_context_for_run(
+                        run_info, enabled=_a2ui_enabled(assistant)
+                    ),
                 )
+            except A2UIValidationError:
+                raise
             except InvalidAgentResponseError as e:
                 return {
                     "error": {
@@ -2225,6 +2429,7 @@ async def handle_list_tasks(
 
         # Collect all runs from matching threads
         all_tasks: list[dict[str, Any]] = []
+        a2ui_by_task_id: dict[str, A2UIContext] = {}
         for tid in thread_ids:
             try:
                 runs = await client.runs.list(
@@ -2267,6 +2472,7 @@ async def handle_list_tasks(
                     task["artifacts"] = []
 
                 all_tasks.append(task)
+                a2ui_by_task_id[task_id] = _a2ui_context_for_run(run)
 
         # Sort by timestamp descending (newest first)
         all_tasks.sort(
@@ -2290,6 +2496,7 @@ async def handle_list_tasks(
                     task["id"],
                     task["contextId"],
                     history_length=history_length,
+                    a2ui=a2ui_by_task_id[task["id"]],
                 )
 
             await asyncio.gather(*(populate_history(task) for task in page_tasks))
@@ -2445,6 +2652,25 @@ async def generate_agent_card(request: ApiRequest, assistant_id: str) -> dict[st
         assistant.get("description") or f"{assistant_name} assistant"
     )
     input_modes, output_modes = _resolve_agent_card_modes(assistant)
+    extensions: list[dict[str, Any]] = []
+    if _a2ui_enabled(assistant):
+        if A2UI_MIME_TYPE not in input_modes:
+            input_modes.append(A2UI_MIME_TYPE)
+        if A2UI_MIME_TYPE not in output_modes:
+            output_modes.append(A2UI_MIME_TYPE)
+        extensions.append(
+            {
+                "uri": A2UI_EXTENSION_URI,
+                "description": "Ability to render A2UI v0.9",
+                "required": False,
+                "params": {
+                    "v0.9": {
+                        "supportedCatalogIds": [],
+                        "acceptsInlineCatalogs": False,
+                    }
+                },
+            }
+        )
 
     # For now, each assistant has one main skill - itself
     skills = [
@@ -2488,6 +2714,14 @@ async def generate_agent_card(request: ApiRequest, assistant_id: str) -> dict[st
 
     agent_url = f"{base_url}{agent_path}"
 
+    capabilities: dict[str, Any] = {
+        "streaming": True,
+        "pushNotifications": False,  # Not implemented yet
+        "stateTransitionHistory": False,
+    }
+    if extensions:
+        capabilities["extensions"] = extensions
+
     return {
         "name": assistant_name,
         "description": assistant_description,
@@ -2495,15 +2729,11 @@ async def generate_agent_card(request: ApiRequest, assistant_id: str) -> dict[st
         "supportedInterfaces": [
             {
                 "url": agent_url,
-                "protocolBinding": "jsonrpc",
+                "protocolBinding": "JSONRPC",
                 "protocolVersion": "1.0",
             },
         ],
-        "capabilities": {
-            "streaming": True,
-            "pushNotifications": False,  # Not implemented yet
-            "stateTransitionHistory": False,
-        },
+        "capabilities": capabilities,
         "defaultInputModes": input_modes,
         "defaultOutputModes": output_modes,
         "skills": skills,
@@ -2680,10 +2910,27 @@ async def handle_message_stream(
 
             # Process A2A message parts into LangChain messages format
             try:
+                a2ui = _a2ui_message_context(message, assistant)
                 message_role = _normalize_input_role(message.get("role", "ROLE_USER"))
                 input_content = _process_a2a_message_parts(
-                    parts, message_role, message_id
+                    parts,
+                    message_role,
+                    message_id,
+                    a2ui=a2ui,
+                )
+            except A2UIValidationError as e:
+                yield (
+                    b"message",
+                    {
+                        "jsonrpc": "2.0",
+                        "id": rpc_id,
+                        "error": {
+                            "code": ERROR_CODE_INVALID_PARAMS,
+                            "message": str(e),
+                        },
+                    },
                 )
+                return
             except ValueError as e:
                 yield (
                     b"message",
@@ -2775,6 +3022,7 @@ async def handle_message_stream(
                 input=input_content,
                 command=command,
                 context=run_context,
+                metadata=_a2ui_run_metadata(assistant, a2ui),
                 headers=request.headers,
             )
 
@@ -2825,7 +3073,26 @@ async def handle_message_stream(
                             final_parts: list[dict[str, Any]] | None = None
                             if isinstance(result, dict):
                                 try:
-                                    final_parts = _extract_a2a_response(result)
+                                    final_parts = _extract_a2a_response(
+                                        result,
+                                        a2ui=a2ui,
+                                    )
+                                except A2UIValidationError:
+                                    await logger.aexception(
+                                        "Agent emitted an invalid A2UI payload"
+                                    )
+                                    yield (
+                                        b"message",
+                                        {
+                                            "jsonrpc": "2.0",
+                                            "id": rpc_id,
+                                            "error": {
+                                                "code": ERROR_CODE_INTERNAL_ERROR,
+                                                "message": "Internal server error",
+                                            },
+                                        },
+                                    )
+                                    return
                                 except InvalidAgentResponseError as e:
                                     yield (
                                         b"message",
@@ -3031,7 +3298,24 @@ async def handle_message_stream(
             # its plain terminal status-update.
             if fallback_state == "TASK_STATE_COMPLETED" and isinstance(result, dict):
                 try:
-                    fallback_parts = _extract_a2a_response(result)
+                    fallback_parts = _extract_a2a_response(
+                        result,
+                        a2ui=a2ui,
+                    )
+                except A2UIValidationError:
+                    await logger.aexception("Agent emitted an invalid A2UI payload")
+                    yield (
+                        b"message",
+                        {
+                            "jsonrpc": "2.0",
+                            "id": rpc_id,
+                            "error": {
+                                "code": ERROR_CODE_INTERNAL_ERROR,
+                                "message": "Internal server error",
+                            },
+                        },
+                    )
+                    return
                 except InvalidAgentResponseError as e:
                     yield (
                         b"message",
@@ -3206,6 +3490,12 @@ a2a_routes = [
         handle_a2a_assistant_endpoint,
         methods=["GET", "POST", "DELETE"],
     ),
+    ApiRoute(
+        "/a2a/{assistant_id}/",
+        handle_a2a_assistant_endpoint,
+        methods=["GET", "POST", "DELETE"],
+        include_in_schema=False,
+    ),
     # Per-assistant agent card (multi-tenant pattern)
     ApiRoute(
         "/a2a/{assistant_id}/.well-known/agent-card.json",