fix(vercel): Align dump_messages output with Vercel spec (#4196)

Co-authored-by: Claude Opus 4.5 <noreply@anthropic.com>

Author: dsfaccini

pydantic_ai_slim/pydantic_ai/ui/vercel_ai/_adapter.py

@@ -4,7 +4,7 @@
 
 import json
 import uuid
-from collections.abc import Sequence
+from collections.abc import Callable, Sequence
 from dataclasses import KW_ONLY, dataclass
 from functools import cached_property
 from typing import TYPE_CHECKING, Any, Literal, cast
@@ -38,12 +38,15 @@
 from ...tools import AgentDepsT, DeferredToolResults, ToolDenied
 from .. import MessagesBuilder, UIAdapter
 from ._event_stream import VercelAIEventStream
-from ._utils import dump_provider_metadata, iter_metadata_chunks, iter_tool_approval_responses, load_provider_metadata
+from ._utils import (
+    dump_provider_metadata,
+    iter_metadata_chunks,
+    iter_tool_approval_responses,
+    load_provider_metadata,
+    tool_return_output,
+)
 from .request_types import (
     DataUIPart,
-    DynamicToolInputAvailablePart,
-    DynamicToolOutputAvailablePart,
-    DynamicToolOutputErrorPart,
     DynamicToolUIPart,
     FileUIPart,
     ProviderMetadata,
@@ -84,6 +87,36 @@
 request_data_ta: TypeAdapter[RequestData] = TypeAdapter(RequestData)
 
 
+def _generate_message_id(
+    msg: ModelRequest | ModelResponse, role: Literal['system', 'user', 'assistant'], message_index: int
+) -> str:
+    """Generate a deterministic message ID based on message content and position.
+
+    Priority order:
+    1. For `ModelResponse` with `provider_response_id` set, use '{provider_response_id}-{message_index}'.
+    2. For any message with run_id set, use '{run_id}-{message_index}'.
+    3. Fallback: UUID5 from 'timestamp-kind-role-message_index'.
+    """
+    if isinstance(msg, ModelResponse) and msg.provider_response_id:
+        return f'{msg.provider_response_id}-{message_index}'
+    if msg.run_id:
+        return f'{msg.run_id}-{message_index}'
+    ts_str = msg.timestamp.isoformat() if msg.timestamp else ''
+    return str(uuid.uuid5(uuid.NAMESPACE_OID, f'{ts_str}-{msg.kind}-{role}-{message_index}'))
+
+
+def _safe_args_as_dict(part: ToolCallPart | BuiltinToolCallPart) -> dict[str, Any] | str:
+    """Safely convert tool call args to dict, falling back to JSON string on parse failure.
+
+    In practice, incomplete tool calls don't reach dump_messages(), but this provides
+    defensive handling for edge cases like interrupted streaming or invalid JSON.
+    """
+    try:
+        return part.args_as_dict()
+    except (ValueError, AssertionError):
+        return part.args_as_json_str()
+
+
 @dataclass
 class VercelAIAdapter(UIAdapter[RequestData, UIMessage, BaseChunk, AgentDepsT, OutputDataT]):
     """UI adapter for the Vercel AI protocol."""
@@ -485,21 +518,20 @@ def _dump_response_message(
                             ToolOutputErrorPart(
                                 type=tool_name,
                                 tool_call_id=part.tool_call_id,
-                                input=part.args_as_json_str(),
+                                input=_safe_args_as_dict(part),
                                 error_text=error_text,
                                 state='output-error',
                                 provider_executed=True,
                                 call_provider_metadata=combined_provider_meta,
                             )
                         )
                     else:
-                        content = builtin_return.model_response_str()
                         ui_parts.append(
                             ToolOutputAvailablePart(
                                 type=tool_name,
                                 tool_call_id=part.tool_call_id,
-                                input=part.args_as_json_str(),
-                                output=content,
+                                input=_safe_args_as_dict(part),
+                                output=tool_return_output(builtin_return),
                                 state='output-available',
                                 provider_executed=True,
                                 call_provider_metadata=combined_provider_meta,
@@ -513,7 +545,7 @@ def _dump_response_message(
                         ToolInputAvailablePart(
                             type=tool_name,
                             tool_call_id=part.tool_call_id,
-                            input=part.args_as_json_str(),
+                            input=_safe_args_as_dict(part),
                             state='input-available',
                             provider_executed=True,
                             call_provider_metadata=call_provider_metadata,
@@ -524,16 +556,17 @@ def _dump_response_message(
                 call_provider_metadata = dump_provider_metadata(
                     id=part.id, provider_name=part.provider_name, provider_details=part.provider_details
                 )
+                tool_type = f'tool-{part.tool_name}'
 
                 if isinstance(tool_result, ToolReturnPart):
-                    content = tool_result.model_response_str()
                     ui_parts.append(
-                        DynamicToolOutputAvailablePart(
-                            tool_name=part.tool_name,
+                        ToolOutputAvailablePart(
+                            type=tool_type,
                             tool_call_id=part.tool_call_id,
-                            input=part.args_as_json_str(),
-                            output=content,
+                            input=_safe_args_as_dict(part),
+                            output=tool_return_output(tool_result),
                             state='output-available',
+                            provider_executed=False,
                             call_provider_metadata=call_provider_metadata,
                         )
                     )
@@ -542,22 +575,24 @@ def _dump_response_message(
                 elif isinstance(tool_result, RetryPromptPart):
                     error_text = tool_result.model_response()
                     ui_parts.append(
-                        DynamicToolOutputErrorPart(
-                            tool_name=part.tool_name,
+                        ToolOutputErrorPart(
+                            type=tool_type,
                             tool_call_id=part.tool_call_id,
-                            input=part.args_as_json_str(),
+                            input=_safe_args_as_dict(part),
                             error_text=error_text,
                             state='output-error',
+                            provider_executed=False,
                             call_provider_metadata=call_provider_metadata,
                         )
                     )
                 else:
                     ui_parts.append(
-                        DynamicToolInputAvailablePart(
-                            tool_name=part.tool_name,
+                        ToolInputAvailablePart(
+                            type=tool_type,
                             tool_call_id=part.tool_call_id,
-                            input=part.args_as_json_str(),
+                            input=_safe_args_as_dict(part),
                             state='input-available',
+                            provider_executed=False,
                             call_provider_metadata=call_provider_metadata,
                         )
                     )
@@ -570,11 +605,19 @@ def _dump_response_message(
     def dump_messages(
         cls,
         messages: Sequence[ModelMessage],
+        *,
+        generate_message_id: Callable[[ModelRequest | ModelResponse, Literal['system', 'user', 'assistant'], int], str]
+        | None = None,
     ) -> list[UIMessage]:
         """Transform Pydantic AI messages into Vercel AI messages.
 
         Args:
             messages: A sequence of ModelMessage objects to convert
+            generate_message_id: Optional custom function to generate message IDs. If provided,
+                it receives the message, the role ('system', 'user', or 'assistant'), and the
+                message index (incremented per UIMessage appended), and should return a unique
+                string ID. If not provided, uses `provider_response_id` for responses,
+                run_id-based IDs for messages with run_id, or a deterministic UUID5 fallback.
 
         Returns:
             A list of UIMessage objects in Vercel AI format
@@ -589,23 +632,34 @@ def dump_messages(
                     elif isinstance(part, RetryPromptPart) and part.tool_name:
                         tool_results[part.tool_call_id] = part
 
+        id_generator = generate_message_id or _generate_message_id
         result: list[UIMessage] = []
+        message_index = 0
 
         for msg in messages:
             if isinstance(msg, ModelRequest):
                 system_ui_parts, user_ui_parts = cls._dump_request_message(msg)
                 if system_ui_parts:
-                    result.append(UIMessage(id=str(uuid.uuid4()), role='system', parts=system_ui_parts))
+                    result.append(
+                        UIMessage(id=id_generator(msg, 'system', message_index), role='system', parts=system_ui_parts)
+                    )
+                    message_index += 1
 
                 if user_ui_parts:
-                    result.append(UIMessage(id=str(uuid.uuid4()), role='user', parts=user_ui_parts))
+                    result.append(
+                        UIMessage(id=id_generator(msg, 'user', message_index), role='user', parts=user_ui_parts)
+                    )
+                    message_index += 1
 
             elif isinstance(  # pragma: no branch
                 msg, ModelResponse
             ):
                 ui_parts: list[UIMessagePart] = cls._dump_response_message(msg, tool_results)
                 if ui_parts:  # pragma: no branch
-                    result.append(UIMessage(id=str(uuid.uuid4()), role='assistant', parts=ui_parts))
+                    result.append(
+                        UIMessage(id=id_generator(msg, 'assistant', message_index), role='assistant', parts=ui_parts)
+                    )
+                    message_index += 1
             else:
                 assert_never(msg)
 

pydantic_ai_slim/pydantic_ai/ui/vercel_ai/_event_stream.py

@@ -11,7 +11,6 @@
 from pydantic_core import to_json
 
 from ...messages import (
-    BaseToolReturnPart,
     BuiltinToolCallPart,
     BuiltinToolReturnPart,
     FilePart,
@@ -30,7 +29,7 @@
 from ...run import AgentRunResultEvent
 from ...tools import AgentDepsT, DeferredToolRequests
 from .. import UIEventStream
-from ._utils import dump_provider_metadata, iter_metadata_chunks, iter_tool_approval_responses
+from ._utils import dump_provider_metadata, iter_metadata_chunks, iter_tool_approval_responses, tool_return_output
 from .request_types import RequestData
 from .response_types import (
     BaseChunk,
@@ -260,7 +259,7 @@ async def handle_builtin_tool_call_end(self, part: BuiltinToolCallPart) -> Async
     async def handle_builtin_tool_return(self, part: BuiltinToolReturnPart) -> AsyncIterator[BaseChunk]:
         yield ToolOutputAvailableChunk(
             tool_call_id=part.tool_call_id,
-            output=self._tool_return_output(part),
+            output=tool_return_output(part),
             provider_executed=True,
         )
 
@@ -278,7 +277,7 @@ async def handle_function_tool_result(self, event: FunctionToolResultEvent) -> A
         elif isinstance(part, RetryPromptPart):
             yield ToolOutputErrorChunk(tool_call_id=tool_call_id, error_text=part.model_response())
         else:
-            yield ToolOutputAvailableChunk(tool_call_id=tool_call_id, output=self._tool_return_output(part))
+            yield ToolOutputAvailableChunk(tool_call_id=tool_call_id, output=tool_return_output(part))
 
         # ToolOutputAvailableChunk/ToolOutputErrorChunk.output may hold user parts
         # (e.g. text, images) that Vercel AI does not currently have chunk types for.
@@ -289,8 +288,3 @@ async def handle_function_tool_result(self, event: FunctionToolResultEvent) -> A
         if isinstance(part, ToolReturnPart):
             for chunk in iter_metadata_chunks(part):
                 yield chunk
-
-    def _tool_return_output(self, part: BaseToolReturnPart) -> Any:
-        output = part.model_response_object()
-        # Unwrap the return value from the output dictionary if it exists
-        return output.get('return_value', output)

pydantic_ai_slim/pydantic_ai/ui/vercel_ai/_utils.py

@@ -3,7 +3,7 @@
 from collections.abc import Iterable, Iterator
 from typing import Any
 
-from pydantic_ai.messages import ProviderDetailsDelta, ToolReturnPart
+from pydantic_ai.messages import BaseToolReturnPart, ProviderDetailsDelta, ToolReturnPart
 from pydantic_ai.ui.vercel_ai.request_types import (
     DynamicToolInputAvailablePart,
     DynamicToolInputStreamingPart,
@@ -29,6 +29,16 @@
 PROVIDER_METADATA_KEY = 'pydantic_ai'
 
 
+def tool_return_output(part: BaseToolReturnPart) -> Any:
+    """Extract the return value from a tool return part.
+
+    If the model response object contains a 'return_value' key, return its value,
+    otherwise return the entire output dict. This matches the streaming output format.
+    """
+    output = part.model_response_object()
+    return output.get('return_value', output)
+
+
 def load_provider_metadata(provider_metadata: ProviderMetadata | None) -> dict[str, Any]:
     """Load the Pydantic AI metadata from the provider metadata."""
     return provider_metadata.get(PROVIDER_METADATA_KEY, {}) if provider_metadata else {}