Ensure AG-UI state is isolated between requests. (#2343)

Author: DouweM

docs/ag-ui.md

@@ -82,10 +82,15 @@ The adapter provides full support for
 real-time synchronization between agents and frontend applications.
 
 In the example below we have document state which is shared between the UI and
-server using the [`StateDeps`][pydantic_ai.ag_ui.StateDeps] which implements the
-[`StateHandler`][pydantic_ai.ag_ui.StateHandler] protocol that can be used to automatically
-decode state contained in [`RunAgentInput.state`](https://docs.ag-ui.com/sdk/js/core/types#runagentinput)
-when processing requests.
+server using the [`StateDeps`][pydantic_ai.ag_ui.StateDeps] [dependencies type](./dependencies.md) that can be used to automatically
+validate state contained in [`RunAgentInput.state`](https://docs.ag-ui.com/sdk/js/core/types#runagentinput) using a Pydantic `BaseModel` specified as a generic parameter.
+
+!!! note "Custom dependencies type with AG-UI state"
+    If you want to use your own dependencies type to hold AG-UI state as well as other things, it needs to implements the
+    [`StateHandler`][pydantic_ai.ag_ui.StateHandler] protocol, meaning it needs to be a [dataclass](https://docs.python.org/3/library/dataclasses.html) with a non-optional `state` field. This lets Pydantic AI ensure that state is properly isolated between requests by building a new dependencies object each time.
+
+    If the `state` field's type is a Pydantic `BaseModel` subclass, the raw state dictionary on the request is automatically validated. If not, you can validate the raw value yourself in your dependencies dataclass's `__post_init__` method.
+
 
 ```python {title="ag_ui_state.py" py="3.10"}
 from pydantic import BaseModel

pydantic_ai_slim/pydantic_ai/ag_ui.py

@@ -9,18 +9,25 @@
 import json
 import uuid
 from collections.abc import Iterable, Mapping, Sequence
-from dataclasses import dataclass, field
+from dataclasses import Field, dataclass, field, replace
 from http import HTTPStatus
 from typing import (
+    TYPE_CHECKING,
     Any,
     Callable,
+    ClassVar,
     Final,
     Generic,
     Protocol,
     TypeVar,
     runtime_checkable,
 )
 
+from pydantic_ai.exceptions import UserError
+
+if TYPE_CHECKING:
+    pass
+
 try:
     from ag_ui.core import (
         AssistantMessage,
@@ -288,8 +295,24 @@ async def run(
             if not run_input.messages:
                 raise _NoMessagesError
 
+            raw_state: dict[str, Any] = run_input.state or {}
             if isinstance(deps, StateHandler):
-                deps.state = run_input.state
+                if isinstance(deps.state, BaseModel):
+                    try:
+                        state = type(deps.state).model_validate(raw_state)
+                    except ValidationError as e:  # pragma: no cover
+                        raise _InvalidStateError from e
+                else:
+                    state = raw_state
+
+                deps = replace(deps, state=state)
+            elif raw_state:
+                raise UserError(
+                    f'AG-UI state is provided but `deps` of type `{type(deps).__name__}` does not implement the `StateHandler` protocol: it needs to be a dataclass with a non-optional `state` field.'
+                )
+            else:
+                # `deps` not being a `StateHandler` is OK if there is no state.
+                pass
 
             messages = _messages_from_ag_ui(run_input.messages)
 
@@ -311,7 +334,7 @@ async def run(
             yield encoder.encode(
                 RunErrorEvent(message=e.message, code=e.code),
             )
-        except Exception as e:  # pragma: no cover
+        except Exception as e:
             yield encoder.encode(
                 RunErrorEvent(message=str(e)),
             )
@@ -531,7 +554,11 @@ def _messages_from_ag_ui(messages: list[Message]) -> list[ModelMessage]:
 
 @runtime_checkable
 class StateHandler(Protocol):
-    """Protocol for state handlers in agent runs."""
+    """Protocol for state handlers in agent runs. Requires the class to be a dataclass with a `state` field."""
+
+    # Has to be a dataclass so we can use `replace` to update the state.
+    # From https://github.com/python/typeshed/blob/9ab7fde0a0cd24ed7a72837fcb21093b811b80d8/stdlib/_typeshed/__init__.pyi#L352
+    __dataclass_fields__: ClassVar[dict[str, Field[Any]]]
 
     @property
     def state(self) -> State:
@@ -558,6 +585,7 @@ def state(self, state: State) -> None:
 """Type variable for the state type, which must be a subclass of `BaseModel`."""
 
 
+@dataclass
 class StateDeps(Generic[StateT]):
     """Provides AG-UI state management.
 
@@ -570,42 +598,7 @@ class StateDeps(Generic[StateT]):
     Implements the `StateHandler` protocol.
     """
 
-    def __init__(self, default: StateT) -> None:
-        """Initialize the state with the provided state type."""
-        self._state = default
-
-    @property
-    def state(self) -> StateT:
-        """Get the current state of the agent run.
-
-        Returns:
-            The current run state.
-        """
-        return self._state
-
-    @state.setter
-    def state(self, state: State) -> None:
-        """Set the state of the agent run.
-
-        This method is called to update the state of the agent run with the
-        provided state.
-
-        Implements the `StateHandler` protocol.
-
-        Args:
-            state: The run state, which must be `None` or model validate for the state type.
-
-        Raises:
-            InvalidStateError: If `state` does not validate.
-        """
-        if state is None:
-            # If state is None, we keep the current state, which will be the default state.
-            return
-
-        try:
-            self._state = type(self._state).model_validate(state)
-        except ValidationError as e:  # pragma: no cover
-            raise _InvalidStateError from e
+    state: StateT
 
 
 @dataclass(repr=False)

tests/test_ag_ui.py

@@ -7,6 +7,7 @@
 import json
 import uuid
 from collections.abc import AsyncIterator
+from dataclasses import dataclass
 from http import HTTPStatus
 from typing import Any
 
@@ -17,7 +18,9 @@
 from inline_snapshot import snapshot
 from pydantic import BaseModel
 
+from pydantic_ai._run_context import RunContext
 from pydantic_ai.agent import Agent
+from pydantic_ai.exceptions import UserError
 from pydantic_ai.messages import ModelMessage
 from pydantic_ai.models.function import (
     AgentInfo,
@@ -27,8 +30,9 @@
     DeltaToolCalls,
     FunctionModel,
 )
+from pydantic_ai.models.test import TestModel
 from pydantic_ai.output import OutputDataT
-from pydantic_ai.tools import AgentDepsT
+from pydantic_ai.tools import AgentDepsT, ToolDefinition
 
 from .conftest import IsSameStr
 
@@ -180,7 +184,7 @@ def create_input(
         thread_id=thread_id,
         run_id=uuid_str(),
         messages=list(messages),
-        state=state,
+        state=dict(state) if state else {},
         context=[],
         tools=tools or [],
         forwarded_props=None,
@@ -1050,9 +1054,19 @@ async def stream_function(
 async def test_request_with_state() -> None:
     """Test request with state modification."""
 
+    seen_states: list[int] = []
+
+    async def store_state(
+        ctx: RunContext[StateDeps[StateInt]], tool_defs: list[ToolDefinition]
+    ) -> list[ToolDefinition]:
+        seen_states.append(ctx.deps.state.value)
+        ctx.deps.state.value += 1
+        return tool_defs
+
     agent: Agent[StateDeps[StateInt], str] = Agent(
         model=FunctionModel(stream_function=simple_stream),
         deps_type=StateDeps[StateInt],  # type: ignore[reportUnknownArgumentType]
+        prepare_tools=store_state,
     )
     adapter = _Adapter(agent=agent)
     run_inputs = [
@@ -1074,32 +1088,101 @@ async def test_request_with_state() -> None:
                 id='msg_3',
                 content='Hello, how are you?',
             ),
+        ),
+        create_input(
+            UserMessage(
+                id='msg_4',
+                content='Hello, how are you?',
+            ),
             state=StateInt(value=42),
         ),
     ]
 
-    deps = StateDeps(StateInt())
+    deps = StateDeps(StateInt(value=0))
 
-    last_value = deps.state.value
     for run_input in run_inputs:
         events = list[dict[str, Any]]()
         async for event in adapter.run(run_input, deps=deps):
             events.append(json.loads(event.removeprefix('data: ')))
 
         assert events == simple_result()
-        assert deps.state.value == run_input.state.value if run_input.state is not None else last_value
-        last_value = deps.state.value
+    assert seen_states == snapshot(
+        [
+            41,  # run msg_1, prepare_tools call 1
+            42,  # run msg_1, prepare_tools call 2
+            0,  # run msg_2, prepare_tools call 1
+            1,  # run msg_2, prepare_tools call 2
+            0,  # run msg_3, prepare_tools call 1
+            1,  # run msg_3, prepare_tools call 2
+            42,  # run msg_4, prepare_tools call 1
+            43,  # run msg_4, prepare_tools call 2
+        ]
+    )
+
+
+async def test_request_with_state_without_handler() -> None:
+    agent = Agent(model=FunctionModel(stream_function=simple_stream))
+    adapter = _Adapter(agent=agent)
+    run_input = create_input(
+        UserMessage(
+            id='msg_1',
+            content='Hello, how are you?',
+        ),
+        state=StateInt(value=41),
+    )
+
+    with pytest.raises(
+        UserError,
+        match='AG-UI state is provided but `deps` of type `NoneType` does not implement the `StateHandler` protocol: it needs to be a dataclass with a non-optional `state` field.',
+    ):
+        async for _ in adapter.run(run_input):
+            pass
+
+
+async def test_request_with_state_with_custom_handler() -> None:
+    @dataclass
+    class CustomStateDeps:
+        state: dict[str, Any]
+
+    seen_states: list[dict[str, Any]] = []
+
+    async def store_state(ctx: RunContext[CustomStateDeps], tool_defs: list[ToolDefinition]) -> list[ToolDefinition]:
+        seen_states.append(ctx.deps.state)
+        return tool_defs
+
+    agent: Agent[CustomStateDeps, str] = Agent(
+        model=FunctionModel(stream_function=simple_stream),
+        deps_type=CustomStateDeps,
+        prepare_tools=store_state,
+    )
+    adapter = _Adapter(agent=agent)
+    run_input = create_input(
+        UserMessage(
+            id='msg_1',
+            content='Hello, how are you?',
+        ),
+        state={'value': 42},
+    )
+
+    async for _ in adapter.run(run_input, deps=CustomStateDeps(state={'value': 0})):
+        pass
 
-    assert deps.state.value == 42
+    assert seen_states[-1] == {'value': 42}
 
 
 async def test_concurrent_runs() -> None:
     """Test concurrent execution of multiple runs."""
     import asyncio
 
-    agent = Agent(
-        model=FunctionModel(stream_function=simple_stream),
+    agent: Agent[StateDeps[StateInt], str] = Agent(
+        model=TestModel(),
+        deps_type=StateDeps[StateInt],  # type: ignore[reportUnknownArgumentType]
     )
+
+    @agent.tool
+    async def get_state(ctx: RunContext[StateDeps[StateInt]]) -> int:
+        return ctx.deps.state.value
+
     adapter = _Adapter(agent=agent)
     concurrent_tasks: list[asyncio.Task[list[dict[str, Any]]]] = []
 
@@ -1109,10 +1192,11 @@ async def test_concurrent_runs() -> None:
                 id=f'msg_{i}',
                 content=f'Message {i}',
             ),
+            state=StateInt(value=i),
             thread_id=f'test_thread_{i}',
         )
 
-        task = asyncio.create_task(collect_events_from_adapter(adapter, run_input))
+        task = asyncio.create_task(collect_events_from_adapter(adapter, run_input, deps=StateDeps(StateInt())))
         concurrent_tasks.append(task)
 
     results = await asyncio.gather(*concurrent_tasks)
@@ -1121,9 +1205,23 @@ async def test_concurrent_runs() -> None:
     for i, events in enumerate(results):
         assert events == [
             {'type': 'RUN_STARTED', 'threadId': f'test_thread_{i}', 'runId': (run_id := IsSameStr())},
+            {
+                'type': 'TOOL_CALL_START',
+                'toolCallId': (tool_call_id := IsSameStr()),
+                'toolCallName': 'get_state',
+                'parentMessageId': IsStr(),
+            },
+            {'type': 'TOOL_CALL_END', 'toolCallId': tool_call_id},
+            {
+                'type': 'TOOL_CALL_RESULT',
+                'messageId': IsStr(),
+                'toolCallId': tool_call_id,
+                'content': str(i),
+                'role': 'tool',
+            },
             {'type': 'TEXT_MESSAGE_START', 'messageId': (message_id := IsSameStr()), 'role': 'assistant'},
-            {'type': 'TEXT_MESSAGE_CONTENT', 'messageId': message_id, 'delta': 'success '},
-            {'type': 'TEXT_MESSAGE_CONTENT', 'messageId': message_id, 'delta': '(no tool calls)'},
+            {'type': 'TEXT_MESSAGE_CONTENT', 'messageId': message_id, 'delta': '{"get_s'},
+            {'type': 'TEXT_MESSAGE_CONTENT', 'messageId': message_id, 'delta': 'tate":' + str(i) + '}'},
             {'type': 'TEXT_MESSAGE_END', 'messageId': message_id},
             {'type': 'RUN_FINISHED', 'threadId': f'test_thread_{i}', 'runId': run_id},
         ]