Add missing AI SDK data chunk fields for reconciliation

Author: bendrucker

pydantic_ai_slim/pydantic_ai/ui/vercel_ai/request_types.py

@@ -107,6 +107,7 @@ class DataUIPart(BaseUIPart):
 
     type: Annotated[str, Field(pattern=r'^data-')]
     id: str | None = None
+    """Optional ID for data part reconciliation."""
     data: Any
 
 

pydantic_ai_slim/pydantic_ai/ui/vercel_ai/response_types.py

@@ -16,6 +16,9 @@
 ProviderMetadata = dict[str, dict[str, JSONValue]]
 """Provider metadata."""
 
+FinishReason = Literal['stop', 'length', 'content-filter', 'tool-calls', 'error', 'other', 'unknown'] | None
+"""Reason why the model finished generating."""
+
 
 class BaseChunk(CamelBaseModel, ABC):
     """Abstract base class for response SSE events."""
@@ -145,6 +148,21 @@ class ToolOutputErrorChunk(BaseChunk):
     dynamic: bool | None = None
 
 
+class ToolApprovalRequestChunk(BaseChunk):
+    """Tool approval request chunk for human-in-the-loop approval."""
+
+    type: Literal['tool-approval-request'] = 'tool-approval-request'
+    approval_id: str
+    tool_call_id: str
+
+
+class ToolOutputDeniedChunk(BaseChunk):
+    """Tool output denied chunk when user denies tool execution."""
+
+    type: Literal['tool-output-denied'] = 'tool-output-denied'
+    tool_call_id: str
+
+
 class SourceUrlChunk(BaseChunk):
     """Source URL chunk."""
 
@@ -178,7 +196,11 @@ class DataChunk(BaseChunk):
     """Data chunk with dynamic type."""
 
     type: Annotated[str, Field(pattern=r'^data-')]
+    id: str | None = None
+    """Optional ID for data part reconciliation."""
     data: Any
+    transient: bool | None = None
+    """If true, data is not persisted to the message history."""
 
 
 class StartStepChunk(BaseChunk):
@@ -205,6 +227,8 @@ class FinishChunk(BaseChunk):
     """Finish chunk."""
 
     type: Literal['finish'] = 'finish'
+    finish_reason: FinishReason = None
+    """Reason why the model finished generating."""
     message_metadata: Any | None = None
 
 

tests/test_vercel_ai.py

@@ -1619,6 +1619,38 @@ async def on_complete(run_result: AgentRunResult[Any]) -> AsyncIterator[BaseChun
     )
 
 
+async def test_data_chunk_with_id_and_transient():
+    """Test DataChunk supports optional id and transient fields for AI SDK compatibility."""
+    agent = Agent(model=TestModel())
+
+    request = SubmitMessage(
+        id='foo',
+        messages=[
+            UIMessage(
+                id='bar',
+                role='user',
+                parts=[TextUIPart(text='Hello')],
+            ),
+        ],
+    )
+
+    async def on_complete(run_result: AgentRunResult[Any]) -> AsyncIterator[BaseChunk]:
+        # Yield a data chunk with id for reconciliation
+        yield DataChunk(type='data-task', id='task-123', data={'status': 'complete'})
+        # Yield a transient data chunk (not persisted to history)
+        yield DataChunk(type='data-progress', data={'percent': 100}, transient=True)
+
+    adapter = VercelAIAdapter(agent, request)
+    events = [
+        '[DONE]' if '[DONE]' in event else json.loads(event.removeprefix('data: '))
+        async for event in adapter.encode_stream(adapter.run_stream(on_complete=on_complete))
+    ]
+
+    # Verify the data chunks are present in the events with correct fields
+    assert {'type': 'data-task', 'id': 'task-123', 'data': {'status': 'complete'}} in events
+    assert {'type': 'data-progress', 'data': {'percent': 100}, 'transient': True} in events
+
+
 @pytest.mark.skipif(not starlette_import_successful, reason='Starlette is not installed')
 async def test_adapter_dispatch_request():
     agent = Agent(model=TestModel())