fix schema drift for chat-ui (ai-sdk 6) (marimo-team#8105)

## 📝 Summary

<!--
Provide a concise summary of what this pull request is addressing.

If this PR closes any issues, list them here by number (e.g., Closes
marimo-team#123).
-->
There is a bug with Pydantic-AI where it started returning AI sdk v6
chunks. This has been fixed here
pydantic/pydantic-ai#4166 but we need to do some
patching to ensure we have the fix.

pyproject.toml is only updated for dev & test dependencies.

## 📋 Checklist

- [x] I have read the [contributor
guidelines](https://github.com/marimo-team/marimo/blob/main/CONTRIBUTING.md).
- [ ] For large changes, or changes that affect the public API: this
change was discussed or approved through an issue, on
[Discord](https://marimo.io/discord?ref=pr), or the community
[discussions](https://github.com/marimo-team/marimo/discussions) (Please
provide a link if applicable).
- [x] Tests have been added for the changes made.
- [ ] Documentation has been updated where applicable, including
docstrings for API changes.
- [x] Pull request title is a good summary of the changes - it will be
used in the [release
notes](https://github.com/marimo-team/marimo/releases).

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: Light2Dark

marimo/_ai/llm/_impl.py

@@ -2,12 +2,14 @@
 from __future__ import annotations
 
 import dataclasses
+import json
 import os
 import re
 from typing import TYPE_CHECKING, Any, Callable, Optional, cast
 
 from marimo import _loggers
 from marimo._ai._pydantic_ai_utils import generate_id
+from marimo._plugins.ui._impl.chat.chat import AI_SDK_VERSION, DONE_CHUNK
 from marimo._plugins.utils import remove_none_values
 
 if TYPE_CHECKING:
@@ -789,21 +791,37 @@ def _serialize_vercel_ai_chunk(
     ) -> dict[str, Any] | None:
         """
         Serialize vercel ai chunk to a dictionary. Skip "done" chunks - not part of Vercel AI SDK schema.
-
-        by_alias=True: Use camelCase keys expected by Vercel AI SDK.
-        exclude_none=True: Remove null values which cause validation errors.
+        We use encode as it uses Pydantic-AI's method of serializing dataclasses to JSON.
         """
         try:
-            serialized = chunk.model_dump(
-                mode="json", by_alias=True, exclude_none=True
-            )
+            encoded = chunk.encode(sdk_version=AI_SDK_VERSION)
+            if encoded == DONE_CHUNK:
+                return None
+            result = json.loads(encoded)
+            if not isinstance(result, dict):
+                LOGGER.debug(
+                    "Serialized vercel ai chunk is not a dictionary: %s",
+                    result,
+                )
+            return result  # type: ignore[no-any-return]
+        except TypeError:
+            # Fallback for pydantic-ai < 1.52.0 which doesn't have sdk_version param
+            try:
+                # by_alias=True: Use camelCase keys expected by Vercel AI SDK.
+                # exclude_none=True: Remove null values which cause validation errors.
+                serialized = chunk.model_dump(
+                    mode="json", by_alias=True, exclude_none=True
+                )
+            except Exception as e:
+                LOGGER.error("Error serializing vercel ai chunk: %s", e)
+                return None
+            else:
+                if serialized.get("type") == "done":
+                    return None
+                return serialized
         except Exception as e:
             LOGGER.error("Error serializing vercel ai chunk: %s", e)
             return None
-        else:
-            if serialized.get("type") == "done":
-                return None
-            return serialized
 
     async def _stream_response(
         self, messages: list[ChatMessage], config: ChatModelConfig

marimo/_plugins/ui/_impl/chat/chat.py

@@ -2,9 +2,10 @@
 from __future__ import annotations
 
 import inspect
+import json
 import uuid
 from dataclasses import dataclass
-from typing import Any, Callable, Final, Optional, Union, cast
+from typing import Any, Callable, Final, Literal, Optional, Union, cast
 
 from marimo import _loggers
 from marimo._ai._types import (
@@ -34,6 +35,8 @@
     presence_penalty=0,
 )
 
+# The version of the Vercel AI SDK we use
+AI_SDK_VERSION: Final[Literal[5, 6]] = 5
 DONE_CHUNK: Final[str] = "[DONE]"
 
 
@@ -464,13 +467,16 @@ def handle_chunk(self, chunk: Any) -> None:
             )
 
             if isinstance(chunk, BaseChunk):
-                # by_alias=True: Use camelCase keys expected by Vercel AI SDK.
-                # exclude_none=True: Remove null values which cause validation errors.
-                self.on_send_chunk(
-                    chunk.model_dump(
+                try:
+                    serialized = json.loads(
+                        chunk.encode(sdk_version=AI_SDK_VERSION)
+                    )
+                except TypeError:
+                    # Fallback for pydantic-ai < 1.52.0 which doesn't have sdk_version param
+                    serialized = chunk.model_dump(
                         mode="json", by_alias=True, exclude_none=True
                     )
-                )
+                self.on_send_chunk(serialized)
                 return
 
         # Handle plain text chunks

marimo/_server/ai/providers.py

@@ -25,6 +25,7 @@
     generate_id,
 )
 from marimo._dependencies.dependencies import Dependency, DependencyManager
+from marimo._plugins.ui._impl.chat.chat import AI_SDK_VERSION
 from marimo._server.ai.config import AnyProviderConfig
 from marimo._server.ai.ids import AiModelId
 from marimo._server.ai.tools.tool_manager import get_tool_manager
@@ -166,9 +167,17 @@ async def stream_completion(
         stream_options = stream_options or StreamOptions()
 
         vercel_adapter = self.get_vercel_adapter()
-        adapter = vercel_adapter(
-            agent=agent, run_input=run_input, accept=stream_options.accept
-        )
+        if DependencyManager.pydantic_ai.has_at_version(min_version="1.52.0"):
+            adapter = vercel_adapter(
+                agent=agent,
+                run_input=run_input,
+                accept=stream_options.accept,
+                sdk_version=AI_SDK_VERSION,
+            )
+        else:
+            adapter = vercel_adapter(
+                agent=agent, run_input=run_input, accept=stream_options.accept
+            )
         event_stream = adapter.run_stream()
         return adapter.streaming_response(event_stream)
 

pyproject.toml

@@ -129,7 +129,7 @@ dev = [
     # For linting
     "ruff>=0.13.2",
     # For AI
-    "pydantic-ai-slim[openai]>=1.47.0",
+    "pydantic-ai-slim[openai]>=1.52.0",
 ]
 
 test = [
@@ -186,7 +186,7 @@ test-optional = [
     "anywidget~=0.9.18",
     "ipython~=8.12.3",
     # testing gen ai
-    "pydantic-ai-slim[google,anthropic,bedrock,openai]>=1.47.0",
+    "pydantic-ai-slim[google,anthropic,bedrock,openai]>=1.52.0",
     # - google-auth uses cachetools, and cachetools<5.0.0 uses collections.MutableMapping (removed in Python 3.10)
     "cachetools>=5.0.0",
     "boto3>=1.38.46",
@@ -246,7 +246,7 @@ dependencies = [
     "matplotlib>=3.8.0",
     "sqlglot[rs]>=26.2.0",
     "sqlalchemy>=2.0.40",
-    "pydantic-ai-slim[google,anthropic,bedrock,openai]>=1.47.0",
+    "pydantic-ai-slim[google,anthropic,bedrock,openai]>=1.52.0",
     "loro>=1.5.0",
     "pandas-stubs>=1.5.3.230321",
     "pyiceberg>=0.9.0",

tests/_ai/llm/test_impl.py

@@ -1591,6 +1591,42 @@ def test_pydantic_ai_serialize_vercel_ai_chunk(self) -> None:
             "input": {"query": "test"},
         }
 
+    def test_pydantic_ai_serialize_vercel_ai_chunk_v5(self) -> None:
+        """Test that tool-input-start chunks exclude providerMetadata for SDK v5.
+
+        The Vercel AI SDK v5 schema drifts from v6, so we need to use Pydantic's handling.
+
+        For tool-input-start chunks, providerMetadata must be excluded.
+        See: https://github.com/pydantic/pydantic-ai/pull/4166
+        """
+        from pydantic_ai.ui.vercel_ai.response_types import ToolInputStartChunk
+
+        mock_agent = MagicMock()
+        model = pydantic_ai(mock_agent)
+
+        # Create chunk with providerMetadata (like Google Gemini produces)
+        chunk = ToolInputStartChunk(
+            tool_call_id="tc_1",
+            tool_name="my_tool",
+            provider_metadata={
+                "pydantic_ai": {
+                    "id": "test_id",
+                    "provider_name": "google-gla",
+                    "provider_details": {
+                        "thought_signature": "encrypted_data"
+                    },
+                }
+            },
+        )
+        result = model._serialize_vercel_ai_chunk(chunk)
+
+        # providerMetadata should be excluded for SDK v5 compatibility
+        assert result == {
+            "type": "tool-input-start",
+            "toolCallId": "tc_1",
+            "toolName": "my_tool",
+        }
+
     def test_pydantic_ai_serialize_vercel_ai_chunk_done_type(self) -> None:
         """Test that 'done' type chunks are skipped."""
         from pydantic_ai.ui.vercel_ai.response_types import DoneChunk

tests/_plugins/ui/_impl/chat/test_chat.py

@@ -953,7 +953,6 @@ def on_send_chunk(chunk: dict):
 
     serializer = ChunkSerializer(on_send_chunk=on_send_chunk)
 
-    # Pydantic BaseChunk should be serialized with model_dump
     chunk = TextDeltaChunk(id="text-1", delta="Hello")
     serializer.handle_chunk(chunk)
 
@@ -963,6 +962,51 @@ def on_send_chunk(chunk: dict):
     ]
 
 
+@pytest.mark.skipif(
+    not DependencyManager.pydantic_ai.has(),
+    reason="Pydantic AI is not installed",
+)
+def test_serialize_pydantic_v5():
+    """Test ChunkSerializer excludes providerMetadata from tool-input-start for SDK v5.
+
+    The Vercel AI SDK v5 schema drifts from v6, so we need to use Pydantic's handling.
+
+    Since pydantic-ai uses toolCallId, providerMetadata must be excluded.
+    See: https://github.com/pydantic/pydantic-ai/pull/4166
+    """
+    from pydantic_ai.ui.vercel_ai.response_types import ToolInputStartChunk
+
+    sent_chunks: list[dict] = []
+
+    def on_send_chunk(chunk: dict):
+        sent_chunks.append(chunk)
+
+    serializer = ChunkSerializer(on_send_chunk=on_send_chunk)
+
+    # Create chunk with providerMetadata (like Google Gemini produces)
+    chunk = ToolInputStartChunk(
+        tool_call_id="tc_1",
+        tool_name="my_tool",
+        provider_metadata={
+            "pydantic_ai": {
+                "id": "test_id",
+                "provider_name": "google-gla",
+                "provider_details": {"thought_signature": "encrypted_data"},
+            }
+        },
+    )
+    serializer.handle_chunk(chunk)
+
+    # providerMetadata should be excluded for SDK v5 compatibility
+    assert sent_chunks == [
+        {
+            "type": "tool-input-start",
+            "toolCallId": "tc_1",
+            "toolName": "my_tool",
+        }
+    ]
+
+
 @pytest.mark.skipif(
     not DependencyManager.pydantic_ai.has(),
     reason="Pydantic AI is not installed",