chore: import upstream snapshot with attribution

python/assistant-stream-hello-world/.gitignore

@@ -0,0 +1 @@
+.vercel

python/assistant-stream-hello-world/api/chat/completions/index.py

@@ -0,0 +1,20 @@
+from assistant_stream import create_run, RunController
+from assistant_stream.serialization import DataStreamResponse
+
+from fastapi import FastAPI
+from fastapi.middleware.cors import CORSMiddleware
+
+import asyncio
+
+app = FastAPI()
+app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"])
+
+
+@app.post("/api/chat/completions")
+async def chat_completions():
+    async def run(controller: RunController):
+        controller.append_text("Hello ")
+        await asyncio.sleep(1)
+        controller.append_text("world.")
+
+    return DataStreamResponse(create_run(run))

python/assistant-stream-hello-world/requirements.txt

@@ -0,0 +1,2 @@
+assistant-stream==0.0.31
+fastapi==0.128.0

python/assistant-stream/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Simon Farshid
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

python/assistant-stream/README.md

@@ -0,0 +1 @@
+## assistant-stream

python/assistant-stream/README_LANGGRAPH.md

@@ -0,0 +1,122 @@
+# LangGraph Integration for assistant-stream
+
+This document describes the LangGraph integration for the assistant-stream package.
+
+## Installation
+
+To use the LangGraph integration, install the assistant-stream package with the langgraph extra:
+
+```bash
+pip install assistant-stream[langgraph]
+```
+
+This will install the required dependencies including `langchain-core`.
+
+## Usage
+
+The integration exposes `append_langgraph_event`, which folds the events produced by LangGraph's native streaming into the state managed by a `RunController`. It is meant to consume `graph.astream(..., stream_mode=["messages", "updates"])` output directly: the event type is the LangGraph stream mode name, and the payload is the raw chunk LangGraph yields for that mode.
+
+### Function Signature
+
+```python
+def append_langgraph_event(
+    state: Any,
+    _namespace: Any,
+    type: str,
+    payload: Any,
+) -> None
+```
+
+### Parameters
+
+- **state**: The state to mutate, normally `controller.state`. It is read and written with dictionary-style access.
+- **_namespace**: The LangGraph namespace for the event. It is accepted for forward compatibility and is currently unused.
+- **type**: The LangGraph stream mode that produced the event, either `"messages"` or `"updates"`.
+- **payload**: The raw chunk LangGraph yields for that stream mode, described below.
+
+### Event Types
+
+#### Message events (`type="messages"`)
+
+The payload is the `(message, metadata)` tuple that LangGraph yields in `"messages"` mode, where `message` is a single `BaseMessage` (often an `AIMessageChunk`) and `metadata` is currently unused.
+
+The function will:
+
+- Create a `messages` list in the state if it does not exist.
+- Convert the message to a plain dict with `model_dump()`.
+- Merge into an existing message when the `id` (or `tool_call_id`) matches: for an `AIMessageChunk` the message is merged with `add_ai_message_chunks`, then patched into state with granular `set` / `append-text` operations where possible. This lets streaming text and tool-call argument chunks update only the field that changed instead of sending the whole message again. If the shape cannot be represented safely as object-stream operations, the helper falls back to replacing the message.
+- Append the message when no existing id matches.
+
+```python
+from langchain_core.messages import AIMessageChunk
+
+# one chunk yielded by stream_mode="messages"
+chunk = (AIMessageChunk(content="Hello", id="msg1"), {})
+append_langgraph_event(controller.state, namespace, "messages", chunk)
+```
+
+#### Updates events (`type="updates"`)
+
+The payload is the `{node_name: {channel: value}}` dict that LangGraph yields in `"updates"` mode.
+
+The function will:
+
+- Write each channel value directly onto the state (`state[channel] = value`), so the node name is not retained.
+- Skip the `messages` channel, since messages are handled by message events.
+- Skip a node whose value is not a dict.
+
+```python
+updates = {"weather_agent": {"status": "completed", "temperature": 72}}
+append_langgraph_event(controller.state, namespace, "updates", updates)
+# state now contains {"status": "completed", "temperature": 72}
+```
+
+### Notes
+
+- The state holds plain JSON values (lists, dicts, str, int, bool, None); LangChain messages are converted with `model_dump()` before they are stored.
+- Event types other than `"messages"` and `"updates"` are ignored.
+
+## Example Integration
+
+The events come straight from `graph.astream`, so a run callback forwards each one to `append_langgraph_event`:
+
+```python
+from assistant_stream import RunController, create_run
+from assistant_stream.modules.langgraph import append_langgraph_event
+from assistant_stream.serialization import DataStreamResponse
+
+
+async def run_callback(controller: RunController):
+    async for namespace, event_type, chunk in graph.astream(
+        {"messages": input_messages},
+        stream_mode=["messages", "updates"],
+        subgraphs=True,
+    ):
+        append_langgraph_event(controller.state, namespace, event_type, chunk)
+
+
+stream = create_run(run_callback, state={})
+response = DataStreamResponse(stream)
+```
+
+As the assistant message streams in, its chunks merge by id, so `controller.state["messages"]` ends with a single assistant message:
+
+```python
+# controller.state
+# {
+#     "messages": [
+#         {"type": "human", "content": "What is the weather?", "id": "user1"},
+#         {"type": "ai", "content": "The weather is sunny today.", "id": "ai1"},
+#     ],
+# }
+```
+
+See `python/assistant-transport-backend-langgraph` for a complete server built on this pattern, including subgraph state via `get_tool_call_subgraph_state`.
+
+## Testing
+
+The integration is covered by unit tests that exercise `append_langgraph_event` against a real state proxy:
+
+```bash
+uv run pytest tests/test_langgraph.py
+```

python/assistant-stream/pyproject.toml

@@ -0,0 +1,30 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "assistant-stream"
+version = "0.0.34"
+description = ""
+authors = [
+  { name="Simon Farshid", email="simon@assistant-ui.com" },
+]
+readme = "README.md"
+requires-python = ">=3.10,<4.0"
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "License :: OSI Approved :: MIT License",
+    "Operating System :: OS Independent",
+]
+dependencies = [
+    "starlette>=0.37.2",
+]
+
+[project.optional-dependencies]
+langgraph = ["langchain-core>=0.3.0"]
+redis = ["redis>=5"]
+dev = ["pytest<8", "redis>=5"]
+
+[project.urls]
+Homepage = "https://github.com/assistant-ui/assistant-ui"
+Issues = "https://github.com/assistant-ui/assistant-ui/issues"

python/assistant-stream/src/assistant_stream/__init__.py

@@ -0,0 +1,20 @@
+from assistant_stream.serialization.assistant_stream_response import (
+    AssistantStreamResponse,
+)
+from assistant_stream.create_run import (
+    create_run,
+    RunController,
+)
+
+try:
+    from assistant_stream.modules.langgraph import append_langgraph_event, get_tool_call_subgraph_state
+
+    __all__ = [
+        "AssistantStreamResponse",
+        "create_run",
+        "RunController",
+        "append_langgraph_event",
+        "get_tool_call_subgraph_state",
+    ]
+except ImportError:
+    __all__ = ["AssistantStreamResponse", "create_run", "RunController"]

python/assistant-stream/src/assistant_stream/assistant_stream_chunk.py

@@ -0,0 +1,99 @@
+from dataclasses import dataclass
+from typing import Any, Dict, List, Literal, Optional, TypedDict, Union
+
+
+# Define the data classes for different chunk types
+@dataclass
+class TextDeltaChunk:
+    text_delta: str
+    type: str = "text-delta"
+    parent_id: Optional[str] = None
+
+
+@dataclass
+class ReasoningDeltaChunk:
+    reasoning_delta: str
+    type: str = "reasoning-delta"
+    parent_id: Optional[str] = None
+
+
+@dataclass
+class ToolCallBeginChunk:
+    tool_call_id: str
+    tool_name: str
+    type: str = "tool-call-begin"
+    parent_id: Optional[str] = None
+
+
+@dataclass
+class ToolCallDeltaChunk:
+    tool_call_id: str
+    args_text_delta: str
+    type: str = "tool-call-delta"
+
+
+@dataclass
+class ToolResultChunk:
+    tool_call_id: str
+    result: Any
+    artifact: Any | None = None
+    is_error: bool = False
+    type: str = "tool-result"
+
+
+@dataclass
+class DataChunk:
+    data: Any
+    type: str = "data"
+
+
+@dataclass
+class ErrorChunk:
+    error: str
+    type: str = "error"
+
+
+# Define ObjectStream operation types as TypedDict
+class ObjectStreamSetOperation(TypedDict):
+    path: List[str]
+    value: Any
+    type: Literal["set"]
+
+
+class ObjectStreamAppendTextOperation(TypedDict):
+    path: List[str]
+    value: str
+    type: Literal["append-text"]
+
+
+ObjectStreamOperation = Union[ObjectStreamSetOperation, ObjectStreamAppendTextOperation]
+
+
+@dataclass
+class UpdateStateChunk:
+    operations: List[ObjectStreamOperation]
+    type: str = "update-state"
+
+
+@dataclass
+class SourceChunk:
+    id: str
+    url: str
+    source_type: str = "url"
+    title: Optional[str] = None
+    type: str = "source"
+    parent_id: Optional[str] = None
+
+
+# Define the union type for AssistantStreamChunk
+AssistantStreamChunk = Union[
+    TextDeltaChunk,
+    ReasoningDeltaChunk,
+    ToolCallBeginChunk,
+    ToolCallDeltaChunk,
+    ToolResultChunk,
+    DataChunk,
+    ErrorChunk,
+    UpdateStateChunk,
+    SourceChunk,
+]

python/assistant-stream/src/assistant_stream/create_run.py

@@ -0,0 +1,275 @@
+import asyncio
+import logging
+from typing import Any, AsyncGenerator, Callable, Coroutine, List, Optional, Sequence, Union
+from assistant_stream.assistant_stream_chunk import (
+    AssistantStreamChunk,
+    TextDeltaChunk,
+    ReasoningDeltaChunk,
+    ToolResultChunk,
+    DataChunk,
+    ErrorChunk,
+    SourceChunk,
+    ToolCallBeginChunk,
+)
+from assistant_stream.modules.tool_call import (
+    create_tool_call,
+    ToolCallController,
+    generate_openai_style_tool_call_id,
+)
+from assistant_stream.state_manager import StateManager
+
+logger = logging.getLogger(__name__)
+
+
+class ReadOnlyCancellationSignal:
+    """Read-only view over an asyncio.Event used for cancellation."""
+
+    def __init__(self, event: asyncio.Event):
+        self._event = event
+
+    def is_set(self) -> bool:
+        return self._event.is_set()
+
+    async def wait(self) -> bool:
+        return await self._event.wait()
+
+
+class RunController:
+    def __init__(self, queue, state_data, parent_id: Optional[str] = None):
+        self._queue = queue
+        self._loop = asyncio.get_running_loop()
+        self._dispose_callbacks = []
+        self._stream_tasks = []
+        self._state_manager = StateManager(self._put_chunk_nowait, state_data)
+        self._parent_id = parent_id
+        self._cancelled_event = asyncio.Event()
+        self._cancelled_signal = ReadOnlyCancellationSignal(self._cancelled_event)
+
+    def with_parent_id(self, parent_id: str) -> 'RunController':
+        """Create a new RunController instance with the specified parent_id."""
+        controller = RunController(self._queue, self._state_manager._state_data, parent_id)
+        controller._loop = self._loop
+        controller._dispose_callbacks = self._dispose_callbacks
+        controller._stream_tasks = self._stream_tasks
+        controller._state_manager = self._state_manager
+        controller._cancelled_event = self._cancelled_event
+        controller._cancelled_signal = self._cancelled_signal
+        return controller
+
+    def append_text(self, text_delta: str) -> None:
+        """Append a text delta to the stream."""
+        chunk = TextDeltaChunk(text_delta=text_delta, parent_id=self._parent_id)
+        self._flush_and_put_chunk(chunk)
+
+    def append_reasoning(self, reasoning_delta: str) -> None:
+        """Append a reasoning delta to the stream."""
+        chunk = ReasoningDeltaChunk(reasoning_delta=reasoning_delta, parent_id=self._parent_id)
+        self._flush_and_put_chunk(chunk)
+
+    def append_state_text(
+        self, path: Sequence[Union[str, int]], text_delta: str
+    ) -> None:
+        """Append a text delta at a state path using an append-text operation."""
+        self._state_manager.append_text(path, text_delta)
+
+    async def add_tool_call(
+        self, tool_name: str, tool_call_id: str = None
+    ) -> ToolCallController:
+        """Add a tool call to the stream."""
+        if tool_call_id is None:
+            tool_call_id = generate_openai_style_tool_call_id()
+
+        stream, controller = await create_tool_call(tool_name, tool_call_id, self._parent_id)
+        self._dispose_callbacks.append(controller.close)
+
+        self.add_stream(stream)
+        return controller
+
+    def add_tool_result(self, tool_call_id: str, result: Any) -> None:
+        """Add a tool result to the stream."""
+        chunk = ToolResultChunk(
+            tool_call_id=tool_call_id,
+            result=result,
+        )
+        self._flush_and_put_chunk(chunk)
+
+    def add_stream(self, stream: AsyncGenerator[AssistantStreamChunk, None]) -> None:
+        """Append a substream to the main stream."""
+
+        async def reader():
+            async for chunk in stream:
+                self._flush_and_put_chunk(chunk)
+
+        task = asyncio.create_task(reader())
+        self._stream_tasks.append(task)
+
+    def add_data(self, data: Any) -> None:
+        """Emit an event to the main stream."""
+        chunk = DataChunk(data=data)
+        self._flush_and_put_chunk(chunk)
+
+    def add_error(self, error: str) -> None:
+        """Emit an error to the main stream."""
+        chunk = ErrorChunk(error=error)
+        self._flush_and_put_chunk(chunk)
+    
+    def add_source(self, id: str, url: str, title: Optional[str] = None) -> None:
+        """Add a source to the stream."""
+        chunk = SourceChunk(
+            id=id,
+            url=url,
+            title=title,
+            parent_id=self._parent_id
+        )
+        self._flush_and_put_chunk(chunk)
+
+    def _put_chunk_nowait(self, chunk):
+        """Helper method to put a chunk in the queue without waiting.
+
+        This is used as a callback for the StateManager.
+        """
+        self._loop.call_soon_threadsafe(self._queue.put_nowait, chunk)
+
+    def _flush_and_put_chunk(self, chunk):
+        """Helper method to flush state operations and put a chunk in the queue.
+
+        This ensures state operations are sent before other operations.
+        """
+        # Flush any pending state operations first
+        self._state_manager.flush()
+        # Add the chunk to the queue
+        self._loop.call_soon_threadsafe(self._queue.put_nowait, chunk)
+
+    @property
+    def state(self):
+        """Access the state proxy object for making state updates.
+
+        This property provides a proxy object that allows navigating to any path
+        in the state, reading values, and setting values, which will trigger the
+        appropriate state update operation.
+
+        If the state is None, this property returns None directly.
+        You can set the root state directly by assigning to this property.
+
+        Example:
+            controller.state = {"user": {"name": "John"}, "messages": [{"text": "Hi"}]}
+            controller.state["user"]["name"] = "Bob"
+            name = controller.state["user"]["name"]
+            controller.state["messages"][0]["text"] += " chunk"  # emits append-text once non-empty
+
+            # Explicit equivalent:
+            controller.append_state_text(["messages", 0, "text"], " chunk")
+        """
+        return self._state_manager.state
+
+    @state.setter
+    def state(self, value):
+        """Set the entire state object.
+
+        Args:
+            value: The new state value to set
+        """
+        self._state_manager.add_operations(
+            [{"type": "set", "path": [], "value": value}]
+        )
+
+    @property
+    def cancelled_event(self) -> ReadOnlyCancellationSignal:
+        """Expose cancellation signal for cooperative cancellation."""
+        return self._cancelled_signal
+
+    @property
+    def is_cancelled(self) -> bool:
+        """Return whether this run has been cancelled."""
+        return self._cancelled_event.is_set()
+
+    def _mark_cancelled(self) -> None:
+        """Set cancellation signal once."""
+        if not self._cancelled_event.is_set():
+            self._cancelled_event.set()
+
+
+async def create_run(
+    callback: Callable[[RunController], Coroutine[Any, Any, None]],
+    *,
+    state: Any | None = None,
+) -> AsyncGenerator[AssistantStreamChunk, None]:
+    queue = asyncio.Queue()
+    controller = RunController(queue, state_data=state)
+
+    async def background_task():
+        try:
+            await callback(controller)
+        except Exception as e:
+            controller.add_error(str(e))
+            raise
+        finally:
+            # Flush any pending state updates before disposing
+            controller._state_manager.flush()
+
+            for dispose in controller._dispose_callbacks:
+                dispose()
+            try:
+                for task in controller._stream_tasks:
+                    await task
+            finally:
+                asyncio.get_running_loop().call_soon_threadsafe(queue.put_nowait, None)
+
+    task = asyncio.create_task(background_task())
+    ended_normally = False
+
+    try:
+        while True:
+            chunk = await controller._queue.get()
+            if chunk is None:
+                ended_normally = True
+                break
+            yield chunk
+            controller._queue.task_done()
+    finally:
+        if ended_normally:
+            # The `None` sentinel is queued at the end of `background_task`, so
+            # normal stream completion implies `task` is already done here.
+            # `result()` preserves normal-path error propagation.
+            task.result()
+        else:
+            controller._mark_cancelled()
+            # Yield to the event loop to allow the cancel signal to propagate.
+            await asyncio.sleep(0)
+            if not task.done():
+                # Give callbacks a brief chance to observe `is_cancelled`
+                # and exit cooperatively before forcing cancellation.
+                # 50ms keeps disconnect cleanup responsive without immediately
+                # interrupting callbacks that can stop themselves quickly.
+                try:
+                    await asyncio.wait_for(asyncio.shield(task), timeout=0.05)
+                except asyncio.TimeoutError:
+                    # Timeout means cooperative shutdown did not finish in time.
+                    pass
+                except Exception:
+                    # The stream consumer already disconnected, so suppress callback errors
+                    # but keep a log signal for postmortem debugging.
+                    logger.warning(
+                        "Suppressed callback exception during early-close grace period",
+                        exc_info=True,
+                    )
+            if not task.done():
+                task.cancel()
+            try:
+                # `shield()` lets caller-initiated cancellation interrupt `aclose()`
+                # without conflating it with our own forced `task.cancel()`.
+                await asyncio.shield(task)
+            except asyncio.CancelledError:
+                if task.cancelled():
+                    # Expected forced cancellation for early-close cleanup.
+                    pass
+                else:
+                    # Preserve caller-initiated cancellation (e.g. wait_for timeout).
+                    raise
+            except Exception:
+                # The stream consumer already disconnected, so suppress callback errors
+                # but keep a log signal for postmortem debugging.
+                logger.warning(
+                    "Suppressed callback exception after forced early-close cancellation",
+                    exc_info=True,
+                )

python/assistant-stream/src/assistant_stream/modules/langgraph.py

@@ -0,0 +1,276 @@
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+
+from assistant_stream.create_run import RunController
+from langchain_core.messages.ai import AIMessageChunk, add_ai_message_chunks
+from langchain_core.messages.tool import ToolMessage
+
+
+def _plain(value: Any) -> Any:
+    if hasattr(value, "_get_value"):
+        return value._get_value()
+    return value
+
+
+def _message_matches(existing_message: Any, message_dict: Dict[str, Any]) -> bool:
+    existing_message = _plain(existing_message)
+    if not isinstance(existing_message, dict):
+        return False
+
+    message_id = message_dict.get("id")
+    if message_id is not None and existing_message.get("id") == message_id:
+        return True
+
+    tool_call_id = message_dict.get("tool_call_id")
+    return (
+        tool_call_id is not None
+        and existing_message.get("tool_call_id") == tool_call_id
+    )
+
+
+def _find_existing_message_index(
+    messages: Any, message_dict: Dict[str, Any]
+) -> int | None:
+    for i, existing_message in enumerate(messages):
+        if _message_matches(existing_message, message_dict):
+            return i
+    return None
+
+
+def _can_patch_value(current_value: Any, next_value: Any) -> bool:
+    # Keep these guards in sync with _patch_child so the later mutation cannot
+    # partially apply before discovering an unsupported deletion/truncation.
+    current_value = _plain(current_value)
+
+    if current_value == next_value:
+        return True
+
+    if isinstance(current_value, dict) and isinstance(next_value, dict):
+        current_keys = set(current_value.keys())
+        next_keys = set(next_value.keys())
+        return current_keys.issubset(next_keys) and all(
+            _can_patch_value(current_value[key], next_value[key])
+            for key in current_keys
+        )
+
+    if isinstance(current_value, list) and isinstance(next_value, list):
+        return len(current_value) <= len(next_value) and all(
+            _can_patch_value(item, next_value[index])
+            for index, item in enumerate(current_value)
+        )
+
+    return True
+
+
+def _patch_child(parent: Any, key: str | int, next_value: Any) -> None:
+    """Patch through StateProxy containers so writes emit granular object ops."""
+    current_value = _plain(parent[key])
+
+    if current_value == next_value:
+        return
+
+    if isinstance(current_value, dict) and isinstance(next_value, dict):
+        current_keys = set(current_value.keys())
+        next_keys = set(next_value.keys())
+        if not current_keys.issubset(next_keys):
+            raise ValueError("Cannot represent deleted dictionary keys as object ops")
+
+        target = parent[key]
+        for child_key in next_keys:
+            if child_key not in current_value:
+                target[child_key] = next_value[child_key]
+            else:
+                _patch_child(target, child_key, next_value[child_key])
+        return
+
+    if isinstance(current_value, list) and isinstance(next_value, list):
+        if len(next_value) < len(current_value):
+            raise ValueError("Cannot represent list truncation as object ops")
+
+        target = parent[key]
+        for index, item in enumerate(next_value):
+            if index >= len(current_value):
+                target.append(item)
+            else:
+                _patch_child(target, index, item)
+        return
+
+    parent[key] = next_value
+
+
+def _patch_message(messages: Any, index: int, next_message: Dict[str, Any]) -> None:
+    current_message = _plain(messages[index])
+    if not isinstance(current_message, dict) or not _can_patch_value(
+        current_message, next_message
+    ):
+        messages[index] = next_message
+        return
+
+    _patch_child(messages, index, next_message)
+
+
+def append_langgraph_event(
+    state: Dict[str, Any], _namespace: str, type: str, payload: Any
+) -> None:
+    """
+    Append a LangGraph event to the state object.
+
+    Args:
+        state: The state dictionary to update
+        _namespace: Event namespace (currently unused)
+        type: Event type ('messages' or 'updates')
+        payload: Event payload containing the data to append
+    """
+
+    if type == "messages":
+        if "messages" not in state:
+            state["messages"] = []
+
+        message = payload[0]
+        message_dict = message.model_dump()
+
+        # Check if this is an AIMessageChunk
+        is_ai_message_chunk = message_dict.get("type") == "AIMessageChunk"
+        if is_ai_message_chunk:
+            message_dict["type"] = "ai"
+        existing_message_index = _find_existing_message_index(
+            state["messages"], message_dict
+        )
+
+        if existing_message_index is not None:
+            if is_ai_message_chunk:
+                existing_message = _plain(state["messages"][existing_message_index])
+                new_message_dict = add_ai_message_chunks(
+                    AIMessageChunk(**{**existing_message, "type": "AIMessageChunk"}),
+                    AIMessageChunk(**{**message_dict, "type": "AIMessageChunk"}),
+                ).model_dump()
+                new_message_dict["type"] = "ai"
+                _patch_message(
+                    state["messages"],
+                    existing_message_index,
+                    new_message_dict,
+                )
+
+            else:
+                state["messages"][existing_message_index] = message_dict
+        else:
+            state["messages"].append(message_dict)
+
+    elif type == "updates":
+        for _node_name, channels in payload.items():
+            if not isinstance(channels, dict):
+                continue
+            for channel_name, channel_value in channels.items():
+                if channel_name == "messages":
+                    continue
+
+                state[channel_name] = channel_value
+
+
+def get_tool_call_subgraph_state(
+    controller: RunController,
+    namespace: Tuple[str, ...],
+    subgraph_node: Union[str, List[str], Callable[[List[str]], bool]],
+    default_state: Dict[str, Any],
+    *,
+    artifact_field_name: Optional[str] = None,
+    tool_name: Union[str, List[str]] | None = None,
+) -> Dict[str, Any]:
+    """
+    Get the state for a tool call subgraph by traversing the namespace and checking for subgraph nodes.
+    Ensures there's a ToolMessage as the last message and returns its artifact field value.
+
+    Args:
+        controller: The run controller managing the state
+        subgraph_node: Node name(s) to check against, or a function that checks node names
+        namespace: Tuple of strings in format 'node_name:task_id'
+        artifact_field_name: Optional field name to extract from artifact
+        default_state: Default state to use if artifact field is None
+
+    Returns:
+        The artifact field value from the ToolMessage. If the last message is already a ToolMessage,
+        returns its artifact field. If it's an AI message with tool calls, creates a ToolMessage
+        and returns the appropriate artifact field value.
+    """
+    # Helper function to check if a node is a subgraph node
+    def is_subgraph_node(node_name: str) -> bool:
+        if isinstance(subgraph_node, str):
+            return node_name == subgraph_node
+        elif isinstance(subgraph_node, list):
+            return node_name in subgraph_node
+        elif callable(subgraph_node):
+            return subgraph_node([node_name])
+        return False
+
+    def is_subgraph_tool(tool: str) -> bool:
+        if isinstance(tool_name, str):
+            return tool == tool_name
+        elif isinstance(tool_name, list):
+            return tool in tool_name
+        return True
+
+    # Start with the controller's state
+    if controller.state is None:
+        controller.state = default_state
+    current_state = controller.state
+
+    # Traverse each level of the namespace
+    for namespace_part in namespace:
+        # Split the namespace part to get node_name
+        node_name = namespace_part.split(':')[0]
+
+        # Check if this node is a subgraph node
+        if is_subgraph_node(node_name):
+            # Check for messages in the current state
+            if "messages" not in current_state:
+                return current_state
+
+            messages = current_state["messages"]
+            if not messages or len(messages) == 0:
+                return current_state
+
+            # Get the last message
+            last_message = messages[-1]
+
+            # Check if it's an AI message
+            if last_message["type"] == "ai":
+                # Check if the AI message has tool calls
+                tool_calls = last_message.get("tool_calls", [])
+                if not tool_calls:
+                    # No tool calls, return current state
+                    return current_state
+
+                # Get the last tool call
+                last_tool_call = tool_calls[-1]
+                if not is_subgraph_tool(last_tool_call["name"]):
+                    return current_state
+
+
+                # Create a new tool message for this tool call
+                tool_message = ToolMessage(
+                    tool_call_id=last_tool_call["id"],
+                    name=last_tool_call["name"],
+                    artifact={} if artifact_field_name else default_state,
+                    content="",
+                    additional_kwargs={
+                        "streaming": True
+                    }
+                ).model_dump()
+
+                messages.append(tool_message)
+                last_message = tool_message
+
+            # Check if last message is already a ToolMessage
+            if last_message["type"] == "tool":
+                # Last message is already a ToolMessage, extract and return artifact field
+                if "artifact" not in last_message:
+                    last_message["artifact"] = {} if artifact_field_name else default_state
+                artifact = last_message["artifact"]
+
+                if artifact_field_name:
+                    if artifact_field_name not in artifact:
+                        artifact[artifact_field_name] = default_state
+                    return artifact[artifact_field_name]
+                else:
+                    return artifact
+
+    return current_state

python/assistant-stream/src/assistant_stream/modules/langgraph_chunk_merge.py

@@ -0,0 +1,305 @@
+from typing import Any, Dict, List, Optional, Union
+
+from assistant_stream.create_run import RunController
+from langchain_core.messages.base import BaseMessage, message_to_dict
+
+
+# def merge_ai_message_chunk_dicts(
+#     chunk1_dict: Dict[str, Any], chunk2_dict: Dict[str, Any]
+# ) -> None:
+#     """
+#     Merge two AIMessageChunk dictionaries by modifying chunk1_dict in-place.
+
+#     This function is designed to work with proxy objects that send updates
+#     over the network, so it only sets complete values at each field path.
+
+#     Args:
+#         chunk1_dict: Dictionary representation of first AIMessageChunk (will be modified)
+#         chunk2_dict: Dictionary representation of second AIMessageChunk
+
+#     Returns:
+#         None (modifies chunk1_dict in-place)
+#     """
+#     # Merge content following merge_content logic
+#     _merge_content_inplace(chunk1_dict, chunk2_dict)
+
+#     # Merge additional_kwargs
+#     _merge_dict_field_inplace(chunk1_dict, chunk2_dict, "additional_kwargs")
+
+#     # Merge response_metadata
+#     _merge_dict_field_inplace(chunk1_dict, chunk2_dict, "response_metadata")
+
+#     # Merge tool call chunks
+#     _merge_tool_call_chunks_inplace(chunk1_dict, chunk2_dict)
+
+#     # Merge usage metadata
+#     _merge_usage_metadata_inplace(chunk1_dict, chunk2_dict)
+
+#     # Handle ID selection (prefer non-run IDs)
+#     _merge_id_inplace(chunk1_dict, chunk2_dict)
+
+#     # Ensure type is set
+#     if "type" not in chunk1_dict:
+#         chunk1_dict["type"] = "AIMessageChunk"
+
+#     # Handle example field - check they match if both present
+#     _merge_example_inplace(chunk1_dict, chunk2_dict)
+
+
+# def _merge_content_inplace(
+#     chunk1_dict: Dict[str, Any], chunk2_dict: Dict[str, Any]
+# ) -> None:
+#     """Merge content fields following merge_content logic."""
+#     content1 = chunk1_dict.get("content", "")
+#     content2 = chunk2_dict.get("content", "")
+
+#     if isinstance(content1, str):
+#         if isinstance(content2, str):
+#             # Both are strings - concatenate
+#             chunk1_dict["content"] += content2
+#         elif isinstance(content2, list):
+#             # First is string, second is list - prepend string to list
+#             chunk1_dict["content"] = [content1] + content2
+#         # If content2 is empty string, it's a no-op
+#     elif isinstance(content1, list):
+#         if isinstance(content2, str):
+#             # First is list, second is string
+#             if content1 and isinstance(content1[-1], str):
+#                 # If last element of list is string, append to it
+#                 content1[-1] += content2
+#             elif content2 == "":
+#                 # Empty string is a no-op
+#                 pass
+#             else:
+#                 # Otherwise append as new element
+#                 content1.append(content2)
+#         elif isinstance(content2, list):
+#             # Both are lists - merge them
+#             chunk1_dict["content"] = _merge_lists(content1, content2)
+
+
+# def _merge_dict_field_inplace(
+#     chunk1_dict: Dict[str, Any],
+#     chunk2_dict: Dict[str, Any],
+#     field_name: str
+# ) -> None:
+#     """Merge a dictionary field from chunk2_dict into chunk1_dict."""
+#     if field_name not in chunk2_dict:
+#         return
+
+#     if field_name not in chunk1_dict:
+#         chunk1_dict[field_name] = {}
+
+#     dict1 = chunk1_dict[field_name]
+#     dict2 = chunk2_dict[field_name]
+
+#     for key, value in dict2.items():
+#         if key not in dict1 or (value is not None and dict1[key] is None):
+#             dict1[key] = value
+#         elif value is None:
+#             continue
+#         elif type(dict1[key]) is not type(value):
+#             raise TypeError(
+#                 f'{field_name}["{key}"] already exists in this message, '
+#                 "but with a different type."
+#             )
+#         elif isinstance(dict1[key], str):
+#             dict1[key] += value
+#         elif isinstance(dict1[key], dict):
+#             # Recursively merge nested dictionaries
+#             _merge_dicts_inplace(dict1[key], value)
+#         elif isinstance(dict1[key], list):
+#             # Merge lists following merge_lists logic
+#             dict1[key] = _merge_lists(dict1[key], value)
+#         elif dict1[key] == value:
+#             continue
+#         else:
+#             raise TypeError(
+#                 f"{field_name} key {key} already exists in left dict and "
+#                 f"value has unsupported type {type(dict1[key])}."
+#             )
+
+
+# def _merge_tool_call_chunks_inplace(
+#     chunk1_dict: Dict[str, Any], chunk2_dict: Dict[str, Any]
+# ) -> None:
+#     """Merge tool call chunks from chunk2_dict into chunk1_dict."""
+#     tool_call_chunks1 = chunk1_dict.get("tool_call_chunks", [])
+#     tool_call_chunks2 = chunk2_dict.get("tool_call_chunks", [])
+
+#     if tool_call_chunks2:
+#         chunk1_dict["tool_call_chunks"] = _merge_lists(
+#             tool_call_chunks1, tool_call_chunks2
+#         )
+
+
+# def _merge_usage_metadata_inplace(
+#     chunk1_dict: Dict[str, Any], chunk2_dict: Dict[str, Any]
+# ) -> None:
+#     """Merge usage metadata following add_usage logic."""
+#     usage1 = chunk1_dict.get("usage_metadata")
+#     usage2 = chunk2_dict.get("usage_metadata")
+
+#     if usage1 or usage2:
+#         if not usage1:
+#             # Only usage2 exists
+#             chunk1_dict["usage_metadata"] = dict(usage2)
+#         elif not usage2:
+#             # Only usage1 exists, already in place
+#             pass
+#         else:
+#             # Both exist - merge them by adding numeric fields
+#             merged_usage = {}
+
+#             # Start with all keys from usage1
+#             for key, value in usage1.items():
+#                 merged_usage[key] = value
+
+#             # Add or merge keys from usage2
+#             for key, value in usage2.items():
+#                 if key in merged_usage:
+#                     # For numeric values, add them
+#                     if isinstance(value, (int, float)) and isinstance(
+#                         merged_usage[key], (int, float)
+#                     ):
+#                         merged_usage[key] += value
+#                     # For dicts (like input_token_details), recursively merge
+#                     elif isinstance(value, dict) and isinstance(
+#                         merged_usage[key], dict
+#                     ):
+#                         merged_dict = dict(merged_usage[key])
+#                         for sub_key, sub_value in value.items():
+#                             if sub_key in merged_dict and isinstance(
+#                                 sub_value, (int, float)
+#                             ):
+#                                 merged_dict[sub_key] += sub_value
+#                             else:
+#                                 merged_dict[sub_key] = sub_value
+#                         merged_usage[key] = merged_dict
+#                     else:
+#                         # For other types, take the new value
+#                         merged_usage[key] = value
+#                 else:
+#                     merged_usage[key] = value
+
+#             chunk1_dict["usage_metadata"] = merged_usage
+
+
+# def _merge_id_inplace(
+#     chunk1_dict: Dict[str, Any], chunk2_dict: Dict[str, Any]
+# ) -> None:
+#     """Handle ID selection, preferring non-run IDs."""
+#     if "id" not in chunk2_dict:
+#         return
+
+#     chunk2_id = chunk2_dict.get("id")
+#     chunk1_id = chunk1_dict.get("id")
+
+#     # Only update ID if chunk2 has a better ID
+#     if chunk2_id:
+#         if not chunk1_id:
+#             # No ID in chunk1, use chunk2's ID
+#             chunk1_dict["id"] = chunk2_id
+#         elif chunk1_id.startswith("run-") and not chunk2_id.startswith("run-"):
+#             # Prefer non-run IDs
+#             chunk1_dict["id"] = chunk2_id
+
+
+# def _merge_example_inplace(
+#     chunk1_dict: Dict[str, Any], chunk2_dict: Dict[str, Any]
+# ) -> None:
+#     """Handle example field merging with validation."""
+#     if "example" in chunk1_dict and "example" in chunk2_dict:
+#         if chunk1_dict["example"] != chunk2_dict["example"]:
+#             raise ValueError(
+#                 "Cannot concatenate AIMessageChunks with different example values."
+#             )
+#     elif "example" in chunk2_dict:
+#         chunk1_dict["example"] = chunk2_dict["example"]
+
+
+# def _merge_dicts_inplace(dict1: Dict[str, Any], dict2: Dict[str, Any]) -> None:
+#     """
+#     Recursively merge dict2 into dict1 in-place.
+
+#     Args:
+#         dict1: Target dictionary (will be modified)
+#         dict2: Source dictionary
+#     """
+#     for key, value in dict2.items():
+#         if key not in dict1 or (value is not None and dict1[key] is None):
+#             dict1[key] = value
+#         elif value is None:
+#             continue
+#         elif type(dict1[key]) is not type(value):
+#             raise TypeError(
+#                 f'Key "{key}" already exists in dictionary, '
+#                 "but with a different type."
+#             )
+#         elif isinstance(dict1[key], str):
+#             dict1[key] += value
+#         elif isinstance(dict1[key], dict):
+#             _merge_dicts_inplace(dict1[key], value)
+#         elif isinstance(dict1[key], list):
+#             dict1[key] = _merge_lists(dict1[key], value)
+#         elif dict1[key] == value:
+#             continue
+#         else:
+#             raise TypeError(
+#                 f"Key {key} already exists in left dict and "
+#                 f"value has unsupported type {type(dict1[key])}."
+#             )
+
+
+# def _merge_lists(
+#     left: Optional[List[Any]], right: Optional[List[Any]]
+# ) -> Optional[List[Any]]:
+#     """
+#     Merge two lists following merge_lists logic.
+
+#     Handles special merging for elements with 'index' fields.
+
+#     Args:
+#         left: First list (or None)
+#         right: Second list (or None)
+
+#     Returns:
+#         Merged list (or None if both inputs are None)
+#     """
+#     if right is None:
+#         return left
+#     if left is None:
+#         return right.copy() if isinstance(right, list) else right
+
+#     merged = left.copy() if isinstance(left, list) else list(left)
+
+#     for e in right:
+#         if isinstance(e, dict) and "index" in e and isinstance(e["index"], int):
+#             # Find existing element with same index
+#             to_merge = [
+#                 i
+#                 for i, e_left in enumerate(merged)
+#                 if isinstance(e_left, dict) and e_left.get("index") == e["index"]
+#             ]
+#             if to_merge:
+#                 # Merge with existing element at same index
+#                 # Remove 'type' key if present (following TODO in merge_lists)
+#                 new_e = (
+#                     {k: v for k, v in e.items() if k != "type"} if "type" in e else e
+#                 )
+#                 # Merge dictionaries
+#                 for key, value in new_e.items():
+#                     if key not in merged[to_merge[0]]:
+#                         merged[to_merge[0]][key] = value
+#                     elif isinstance(merged[to_merge[0]][key], str) and isinstance(
+#                         value, str
+#                     ):
+#                         merged[to_merge[0]][key] += value
+#                     else:
+#                         merged[to_merge[0]][key] = value
+#             else:
+#                 merged.append(e)
+#         else:
+#             merged.append(e)
+
+#     return merged

python/assistant-stream/src/assistant_stream/modules/tool_call.py

@@ -0,0 +1,92 @@
+import asyncio
+from typing import Any, AsyncGenerator
+from assistant_stream.assistant_stream_chunk import (
+    AssistantStreamChunk,
+    ToolCallBeginChunk,
+    ToolCallDeltaChunk,
+    ToolResultChunk,
+)
+import string
+import random
+
+
+def generate_openai_style_tool_call_id():
+    prefix = "call_"
+    characters = string.ascii_letters + string.digits
+    random_id = "".join(random.choices(characters, k=24))
+    return prefix + random_id
+
+
+class ToolCallController:
+    def __init__(self, queue, tool_name: str, tool_call_id: str, parent_id: str = None):
+        self.tool_name = tool_name
+        self.tool_call_id = tool_call_id
+        self.queue = queue
+        self.loop = asyncio.get_running_loop()
+
+        begin_chunk = ToolCallBeginChunk(
+            tool_call_id=self.tool_call_id,
+            tool_name=self.tool_name,
+            parent_id=parent_id,
+        )
+        self.queue.put_nowait(begin_chunk)
+
+    def append_args_text(self, args_text_delta: str) -> None:
+        """Append an args text delta to the stream."""
+        chunk = ToolCallDeltaChunk(
+            tool_call_id=self.tool_call_id,
+            args_text_delta=args_text_delta,
+        )
+        self.loop.call_soon_threadsafe(self.queue.put_nowait, chunk)
+
+    def set_result(self, result: Any) -> None:
+        """
+        Set the result of the tool call.
+
+        Deprecated: Use set_response() instead.
+        """
+        import warnings
+
+        warnings.warn(
+            "set_result() is deprecated. Use set_response() instead.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        return self.set_response(result)
+
+    def set_response(
+        self, result: Any, *, artifact: Any | None = None, is_error: bool = False
+    ) -> None:
+        """Set the result of the tool call."""
+
+        chunk = ToolResultChunk(
+            tool_call_id=self.tool_call_id,
+            result=result,
+            artifact=artifact,
+            is_error=is_error,
+        )
+        self.loop.call_soon_threadsafe(self.queue.put_nowait, chunk)
+        self.close()
+
+    def close(self) -> None:
+        """Close the stream."""
+        self.loop.call_soon_threadsafe(self.queue.put_nowait, None)
+
+
+async def create_tool_call(
+    tool_name: str,
+    tool_call_id: str,
+    parent_id: str = None,
+) -> tuple[AsyncGenerator[AssistantStreamChunk, None], ToolCallController]:
+    queue = asyncio.Queue()
+    controller = ToolCallController(queue, tool_name, tool_call_id, parent_id)
+
+    async def stream():
+        while True:
+            chunk = await controller.queue.get()
+            if chunk is None:
+                break
+            yield chunk
+            controller.queue.task_done()
+
+    return stream(), controller

python/assistant-stream/src/assistant_stream/resumable/__init__.py

@@ -0,0 +1,58 @@
+from contextlib import suppress
+
+from assistant_stream.resumable.context import (
+    ResumableStreamContext,
+    create_resumable_stream_context,
+)
+from assistant_stream.resumable.errors import (
+    DEFAULT_TTL_MS,
+    ResumableStreamError,
+    ResumableStreamErrorCode,
+    validate_stream_id,
+)
+from assistant_stream.resumable.response import (
+    RESUMABLE_STREAM_ID_HEADER,
+    create_resumable_assistant_stream_response,
+    create_resume_assistant_stream_response,
+)
+from assistant_stream.resumable.stores.in_memory import (
+    create_in_memory_resumable_stream_store,
+)
+from assistant_stream.resumable.types import (
+    CancellationSignal,
+    ResumableStreamEntry,
+    ResumableStreamRole,
+    ResumableStreamStatus,
+    ResumableStreamStore,
+)
+
+__all__ = [
+    "CancellationSignal",
+    "DEFAULT_TTL_MS",
+    "RESUMABLE_STREAM_ID_HEADER",
+    "ResumableStreamContext",
+    "ResumableStreamEntry",
+    "ResumableStreamError",
+    "ResumableStreamErrorCode",
+    "ResumableStreamRole",
+    "ResumableStreamStatus",
+    "ResumableStreamStore",
+    "create_in_memory_resumable_stream_store",
+    "create_resumable_assistant_stream_response",
+    "create_resumable_stream_context",
+    "create_resume_assistant_stream_response",
+    "validate_stream_id",
+]
+
+with suppress(ImportError):
+    from assistant_stream.resumable.stores.redis import (
+        RedisLikeClient,
+        RedisResumableStreamStore,
+        create_redis_resumable_stream_store,
+    )
+
+    __all__ += [
+        "RedisLikeClient",
+        "RedisResumableStreamStore",
+        "create_redis_resumable_stream_store",
+    ]

python/assistant-stream/src/assistant_stream/resumable/context.py

@@ -0,0 +1,165 @@
+from __future__ import annotations
+
+import asyncio
+import logging
+from collections.abc import AsyncIterator, Callable
+from dataclasses import dataclass, field
+from typing import Any, Literal
+
+from assistant_stream.resumable.errors import ResumableStreamError
+from assistant_stream.resumable.types import (
+    ResumableStreamRole,
+    ResumableStreamStatus,
+    ResumableStreamStore,
+)
+
+logger = logging.getLogger(__name__)
+_hook_logger = logging.getLogger("assistant_stream.resumable")
+
+MakeStream = Callable[[], AsyncIterator[bytes]]
+WaitUntil = Callable[[asyncio.Task[None]], None]
+OnAcquire = Callable[[str, ResumableStreamRole], None]
+OnAppend = Callable[[str, int], None]
+OnFinalize = Callable[[str, Literal["done", "error"], str | None], None]
+OnError = Callable[[str, object], None]
+
+
+def _call_hook(hook: Callable[..., Any] | None, /, *args: Any) -> None:
+    if hook is None:
+        return
+    try:
+        hook(*args)
+    except Exception:
+        _hook_logger.debug("resumable stream hook failed: %r", hook, exc_info=True)
+
+
+@dataclass
+class ResumableStreamContext:
+    _store: ResumableStreamStore
+    _ttl_ms: int | None
+    _wait_until: WaitUntil | None
+    _on_acquire: OnAcquire | None
+    _on_append: OnAppend | None
+    _on_finalize: OnFinalize | None
+    _on_error: OnError | None
+    _tasks: set[asyncio.Task[None]] = field(default_factory=set, repr=False)
+
+    async def run(
+        self, stream_id: str, make_stream: MakeStream
+    ) -> AsyncIterator[bytes]:
+        role = await self._store.acquire(
+            stream_id,
+            ttl_ms=self._ttl_ms,
+        )
+        _call_hook(self._on_acquire, stream_id, role)
+        if role == "producer":
+            _start_producer_task(
+                self._store,
+                stream_id,
+                make_stream,
+                tasks=self._tasks,
+                wait_until=self._wait_until,
+                on_append=self._on_append,
+                on_finalize=self._on_finalize,
+                on_error=self._on_error,
+            )
+        return _read_from_store(self._store, stream_id)
+
+    async def resume(self, stream_id: str) -> AsyncIterator[bytes] | None:
+        status = await self._store.status(stream_id)
+        if status == "missing":
+            return None
+        return _read_from_store(self._store, stream_id)
+
+    async def require_resume(self, stream_id: str) -> AsyncIterator[bytes]:
+        status = await self._store.status(stream_id)
+        if status == "missing":
+            raise ResumableStreamError(
+                "missing",
+                f"resumable stream not found: {stream_id}",
+            )
+        return _read_from_store(self._store, stream_id)
+
+    async def status(self, stream_id: str) -> ResumableStreamStatus:
+        return await self._store.status(stream_id)
+
+    async def delete(self, stream_id: str) -> None:
+        await self._store.delete(stream_id)
+
+
+def create_resumable_stream_context(
+    *,
+    store: ResumableStreamStore,
+    ttl_ms: int | None = None,
+    wait_until: WaitUntil | None = None,
+    on_acquire: OnAcquire | None = None,
+    on_append: OnAppend | None = None,
+    on_finalize: OnFinalize | None = None,
+    on_error: OnError | None = None,
+) -> ResumableStreamContext:
+    return ResumableStreamContext(
+        _store=store,
+        _ttl_ms=ttl_ms,
+        _wait_until=wait_until,
+        _on_acquire=on_acquire,
+        _on_append=on_append,
+        _on_finalize=on_finalize,
+        _on_error=on_error,
+    )
+
+
+def _start_producer_task(
+    store: ResumableStreamStore,
+    stream_id: str,
+    make_stream: MakeStream,
+    *,
+    tasks: set[asyncio.Task[None]],
+    wait_until: WaitUntil | None,
+    on_append: OnAppend | None,
+    on_finalize: OnFinalize | None,
+    on_error: OnError | None,
+) -> None:
+    async def _pump() -> None:
+        try:
+            async for chunk in make_stream():
+                await store.append(stream_id, chunk)
+                _call_hook(on_append, stream_id, len(chunk))
+            await store.finalize(stream_id, "done")
+            _call_hook(on_finalize, stream_id, "done", None)
+        except Exception as err:
+            _call_hook(on_error, stream_id, err)
+            message = str(err) if str(err) else repr(err)
+            try:
+                await store.finalize(stream_id, "error", message)
+                _call_hook(on_finalize, stream_id, "error", message)
+            except Exception as finalize_err:
+                logger.error(
+                    "resumable stream finalize failed: %s", finalize_err
+                )
+
+    task = asyncio.create_task(_pump())
+    tasks.add(task)
+    task.add_done_callback(tasks.discard)
+    if wait_until is not None:
+        wait_until(task)
+
+    def _done_callback(done_task: asyncio.Task[None]) -> None:
+        try:
+            exc = done_task.exception()
+        except asyncio.CancelledError:
+            return
+        if exc is not None:
+            logger.error("resumable producer task failed: %s", exc)
+
+    task.add_done_callback(_done_callback)
+
+
+async def _read_from_store(
+    store: ResumableStreamStore, stream_id: str
+) -> AsyncIterator[bytes]:
+    signal = asyncio.Event()
+    try:
+        async for entry in store.read(stream_id, "", signal):
+            yield entry.chunk
+    finally:
+        signal.set()

python/assistant-stream/src/assistant_stream/resumable/errors.py

@@ -0,0 +1,25 @@
+from __future__ import annotations
+
+import re
+from typing import Literal
+
+
+ResumableStreamErrorCode = Literal["missing", "exists", "finalized", "invalid-id"]
+
+DEFAULT_TTL_MS = 24 * 60 * 60 * 1000
+
+_STREAM_ID_PATTERN = re.compile(r"^[A-Za-z0-9_.:-]{1,256}$")
+
+
+class ResumableStreamError(Exception):
+    def __init__(self, code: ResumableStreamErrorCode, message: str) -> None:
+        super().__init__(message)
+        self.code = code
+
+
+def validate_stream_id(stream_id: str) -> None:
+    if not _STREAM_ID_PATTERN.fullmatch(stream_id):
+        raise ResumableStreamError(
+            "invalid-id",
+            f"Invalid streamId: {stream_id} (must match {_STREAM_ID_PATTERN.pattern})",
+        )

python/assistant-stream/src/assistant_stream/resumable/response.py

@@ -0,0 +1,75 @@
+from __future__ import annotations
+
+from collections.abc import AsyncIterator, Callable, Coroutine, Mapping
+from typing import Any
+
+from starlette.responses import JSONResponse, Response, StreamingResponse
+
+from assistant_stream.create_run import RunController, create_run
+from assistant_stream.resumable.context import ResumableStreamContext
+from assistant_stream.resumable.errors import ResumableStreamError
+from assistant_stream.serialization.data_stream import DataStreamEncoder
+from assistant_stream.serialization.stream_encoder import StreamEncoder
+
+RESUMABLE_STREAM_ID_HEADER = "x-resumable-stream-id"
+
+
+async def create_resumable_assistant_stream_response(
+    *,
+    context: ResumableStreamContext,
+    stream_id: str,
+    callback: Callable[[RunController], Coroutine[Any, Any, None]],
+    encoder: StreamEncoder | None = None,
+    headers: Mapping[str, str] | None = None,
+) -> Response:
+    resolved_encoder = encoder if encoder is not None else DataStreamEncoder()
+
+    async def make_stream() -> AsyncIterator[bytes]:
+        async for frame in resolved_encoder.encode_stream(create_run(callback)):
+            yield frame.encode("utf-8")
+
+    body = await context.run(stream_id, make_stream)
+    return StreamingResponse(
+        body,
+        media_type=resolved_encoder.get_media_type(),
+        headers=_merge_headers(headers, stream_id),
+    )
+
+
+async def create_resume_assistant_stream_response(
+    *,
+    context: ResumableStreamContext,
+    stream_id: str,
+    encoder: StreamEncoder | None = None,
+    headers: Mapping[str, str] | None = None,
+    missing_response: Callable[[], Response] | None = None,
+) -> Response:
+    try:
+        body = await context.resume(stream_id)
+    except ResumableStreamError as err:
+        if err.code != "invalid-id":
+            raise
+        body = None
+    if body is None:
+        if missing_response is not None:
+            return missing_response()
+        return JSONResponse({"error": "stream not found"}, status_code=404)
+
+    resolved_encoder = encoder if encoder is not None else DataStreamEncoder()
+    return StreamingResponse(
+        body,
+        media_type=resolved_encoder.get_media_type(),
+        headers=_merge_headers(headers, stream_id),
+    )
+
+
+def _merge_headers(
+    extra: Mapping[str, str] | None, stream_id: str
+) -> dict[str, str]:
+    merged: dict[str, str] = {}
+    if extra is not None:
+        for key, value in extra.items():
+            if key.lower() != RESUMABLE_STREAM_ID_HEADER:
+                merged[key] = value
+    merged[RESUMABLE_STREAM_ID_HEADER] = stream_id
+    return merged

python/assistant-stream/src/assistant_stream/resumable/stores/__init__.py

@@ -0,0 +1,18 @@
+from contextlib import suppress
+
+from assistant_stream.resumable.stores.in_memory import (
+    create_in_memory_resumable_stream_store,
+)
+
+__all__ = [
+    "create_in_memory_resumable_stream_store",
+]
+
+with suppress(ImportError):
+    from assistant_stream.resumable.stores.redis import (
+        create_redis_resumable_stream_store,
+    )
+
+    __all__ += [
+        "create_redis_resumable_stream_store",
+    ]

python/assistant-stream/src/assistant_stream/resumable/stores/in_memory.py

@@ -0,0 +1,321 @@
+from __future__ import annotations
+
+import asyncio
+import time
+from collections.abc import AsyncIterator, Callable
+from contextlib import suppress
+from dataclasses import dataclass, field
+from typing import Literal
+
+from assistant_stream.resumable.errors import (
+    DEFAULT_TTL_MS,
+    ResumableStreamError,
+    validate_stream_id,
+)
+from assistant_stream.resumable.types import (
+    CancellationSignal,
+    ResumableStreamEntry,
+    ResumableStreamRole,
+    ResumableStreamStatus,
+)
+
+
+@dataclass
+class _FinalizeMarker:
+    kind: Literal["done", "error"]
+    error: str | None = None
+
+
+@dataclass
+class _StreamState:
+    entries: list[ResumableStreamEntry] = field(default_factory=list)
+    next_seq: int = 1
+    expires_at: float = 0.0
+    ttl_ms: int = 0
+    final: _FinalizeMarker | None = None
+    waiters: list[asyncio.Event] = field(default_factory=list)
+
+
+def _cursor_of(seq: int) -> str:
+    if seq == 0:
+        return "0"
+    alphabet = "0123456789abcdefghijklmnopqrstuvwxyz"
+    n = seq
+    chars: list[str] = []
+    while n:
+        n, rem = divmod(n, 36)
+        chars.append(alphabet[rem])
+    return "".join(reversed(chars))
+
+
+def _seq_from_cursor(cursor: str) -> int:
+    if cursor == "":
+        return 0
+    try:
+        return int(cursor, 36)
+    except ValueError:
+        return 0
+
+
+class _InMemoryResumableStreamStore:
+    def __init__(
+        self,
+        *,
+        default_ttl_ms: int,
+        now: Callable[[], float],
+        max_chunk_bytes: int | None,
+        max_entries_per_stream: int | None,
+        max_streams: int | None,
+        gc_interval_ms: int | None,
+    ) -> None:
+        self._streams: dict[str, _StreamState] = {}
+        self._default_ttl_ms = default_ttl_ms
+        self._now = now
+        self._max_chunk_bytes = max_chunk_bytes
+        self._max_entries_per_stream = max_entries_per_stream
+        self._max_streams = max_streams
+        self._gc_interval_ms = gc_interval_ms
+        self._gc_task: asyncio.Task[None] | None = None
+        self._disposed = False
+
+    def _ensure_gc(self) -> None:
+        if (
+            self._gc_interval_ms is None
+            or self._disposed
+            or self._gc_task is not None
+        ):
+            return
+        try:
+            loop = asyncio.get_running_loop()
+        except RuntimeError:
+            return
+
+        interval_s = self._gc_interval_ms / 1000.0
+
+        async def _gc_loop() -> None:
+            while True:
+                await asyncio.sleep(interval_s)
+                self._evict_expired()
+
+        self._gc_task = loop.create_task(_gc_loop())
+
+    def _evict_expired(self) -> None:
+        t = self._now()
+        expired: list[tuple[str, _StreamState]] = []
+        for stream_id, state in self._streams.items():
+            if state.expires_at <= t:
+                expired.append((stream_id, state))
+        for stream_id, state in expired:
+            del self._streams[stream_id]
+            if state.final is None:
+                state.final = _FinalizeMarker(kind="error", error="Stream expired")
+            self._notify(state)
+
+    def _notify(self, state: _StreamState) -> None:
+        waiters = state.waiters
+        state.waiters = []
+        for event in waiters:
+            event.set()
+
+    def _find_start_index(self, state: _StreamState, cursor: str) -> int:
+        if cursor == "":
+            return 0
+        after = _seq_from_cursor(cursor)
+        lo = 0
+        hi = len(state.entries)
+        while lo < hi:
+            mid = (lo + hi) // 2
+            seq = _seq_from_cursor(state.entries[mid].cursor)
+            if seq <= after:
+                lo = mid + 1
+            else:
+                hi = mid
+        return lo
+
+    async def _wait_for_update(
+        self,
+        state: _StreamState,
+        signal: CancellationSignal,
+        wake_by: float | None,
+    ) -> None:
+        if signal.is_set():
+            return
+        if wake_by is not None and wake_by <= 0:
+            return
+
+        event = asyncio.Event()
+        state.waiters.append(event)
+
+        notify_task = asyncio.create_task(event.wait())
+        signal_task = asyncio.create_task(signal.wait())
+        tasks: set[asyncio.Task[object]] = {notify_task, signal_task}
+        if wake_by is not None:
+            tasks.add(asyncio.create_task(asyncio.sleep(wake_by / 1000.0)))
+
+        try:
+            await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
+        finally:
+            for task in tasks:
+                if not task.done():
+                    task.cancel()
+            for task in tasks:
+                with suppress(asyncio.CancelledError):
+                    await task
+            with suppress(ValueError):
+                state.waiters.remove(event)
+
+    def _require_active(self, stream_id: str) -> _StreamState:
+        self._evict_expired()
+        state = self._streams.get(stream_id)
+        if state is None:
+            raise RuntimeError(f"Stream not found: {stream_id}")
+        if state.final is not None:
+            raise ResumableStreamError(
+                "finalized",
+                f"Stream already finalized: {stream_id}",
+            )
+        return state
+
+    async def acquire(
+        self, stream_id: str, *, ttl_ms: int | None = None
+    ) -> ResumableStreamRole:
+        self._ensure_gc()
+        validate_stream_id(stream_id)
+        self._evict_expired()
+        if stream_id in self._streams:
+            return "consumer"
+
+        if self._max_streams is not None and len(self._streams) >= self._max_streams:
+            raise RuntimeError("maxStreams exceeded")
+
+        resolved_ttl = ttl_ms if ttl_ms is not None else self._default_ttl_ms
+        self._streams[stream_id] = _StreamState(
+            entries=[],
+            next_seq=1,
+            expires_at=self._now() + resolved_ttl,
+            ttl_ms=resolved_ttl,
+            final=None,
+            waiters=[],
+        )
+        return "producer"
+
+    async def append(self, stream_id: str, chunk: bytes) -> None:
+        self._ensure_gc()
+        validate_stream_id(stream_id)
+        if (
+            self._max_chunk_bytes is not None
+            and len(chunk) > self._max_chunk_bytes
+        ):
+            raise RuntimeError(f"Chunk exceeds maxChunkBytes: {len(chunk)}")
+        state = self._require_active(stream_id)
+        if (
+            self._max_entries_per_stream is not None
+            and len(state.entries) >= self._max_entries_per_stream
+        ):
+            raise RuntimeError(f"Stream exceeded maxEntriesPerStream: {stream_id}")
+        seq = state.next_seq
+        state.next_seq += 1
+        state.entries.append(ResumableStreamEntry(cursor=_cursor_of(seq), chunk=chunk))
+        state.expires_at = self._now() + state.ttl_ms
+        self._notify(state)
+
+    async def finalize(
+        self,
+        stream_id: str,
+        status: Literal["done", "error"],
+        error: str | None = None,
+    ) -> None:
+        self._ensure_gc()
+        validate_stream_id(stream_id)
+        self._evict_expired()
+        state = self._streams.get(stream_id)
+        if state is None:
+            raise RuntimeError(f"Stream not found: {stream_id}")
+        if state.final is not None:
+            return
+        if status == "done":
+            state.final = _FinalizeMarker(kind="done")
+        else:
+            state.final = _FinalizeMarker(
+                kind="error", error=error if error is not None else "Stream errored"
+            )
+        state.expires_at = self._now() + state.ttl_ms
+        self._notify(state)
+
+    async def read(
+        self, stream_id: str, cursor: str, signal: CancellationSignal
+    ) -> AsyncIterator[ResumableStreamEntry]:
+        self._ensure_gc()
+        validate_stream_id(stream_id)
+        self._evict_expired()
+        state = self._streams.get(stream_id)
+        if state is None:
+            raise RuntimeError(f"Stream not found: {stream_id}")
+
+        idx = self._find_start_index(state, cursor)
+
+        while True:
+            if signal.is_set():
+                return
+
+            while idx < len(state.entries):
+                if signal.is_set():
+                    return
+                yield state.entries[idx]
+                idx += 1
+
+            if state.final is not None:
+                if state.final.kind == "error":
+                    raise RuntimeError(state.final.error or "Stream errored")
+                return
+
+            wake_by = state.expires_at - self._now()
+            await self._wait_for_update(state, signal, wake_by)
+            self._evict_expired()
+
+    async def status(self, stream_id: str) -> ResumableStreamStatus:
+        self._ensure_gc()
+        validate_stream_id(stream_id)
+        self._evict_expired()
+        state = self._streams.get(stream_id)
+        if state is None:
+            return "missing"
+        if state.final is None:
+            return "streaming"
+        return "error" if state.final.kind == "error" else "done"
+
+    async def delete(self, stream_id: str) -> None:
+        self._ensure_gc()
+        validate_stream_id(stream_id)
+        state = self._streams.get(stream_id)
+        if state is None:
+            return
+        del self._streams[stream_id]
+        if state.final is None:
+            state.final = _FinalizeMarker(kind="done")
+        self._notify(state)
+
+    def dispose(self) -> None:
+        self._disposed = True
+        if self._gc_task is not None:
+            self._gc_task.cancel()
+            self._gc_task = None
+
+
+def create_in_memory_resumable_stream_store(
+    *,
+    default_ttl_ms: int = DEFAULT_TTL_MS,
+    now: Callable[[], float] | None = None,
+    max_chunk_bytes: int | None = None,
+    max_entries_per_stream: int | None = None,
+    max_streams: int | None = None,
+    gc_interval_ms: int | None = None,
+) -> _InMemoryResumableStreamStore:
+    return _InMemoryResumableStreamStore(
+        default_ttl_ms=default_ttl_ms,
+        now=now if now is not None else (lambda: time.time() * 1000),
+        max_chunk_bytes=max_chunk_bytes,
+        max_entries_per_stream=max_entries_per_stream,
+        max_streams=max_streams,
+        gc_interval_ms=gc_interval_ms,
+    )

python/assistant-stream/src/assistant_stream/resumable/stores/redis.py

@@ -0,0 +1,358 @@
+from __future__ import annotations
+
+import asyncio
+import json
+import math
+from collections.abc import AsyncIterator
+from contextlib import suppress
+from typing import Any, Literal, Protocol
+
+from assistant_stream.resumable.errors import (
+    DEFAULT_TTL_MS,
+    ResumableStreamError,
+    validate_stream_id,
+)
+from assistant_stream.resumable.types import (
+    CancellationSignal,
+    ResumableStreamEntry,
+    ResumableStreamRole,
+    ResumableStreamStatus,
+)
+
+DEFAULT_POLL_INTERVAL_MS = 100
+DEFAULT_KEY_PREFIX = "aui:resumable"
+
+FIELD_CHUNK = "c"
+FIELD_FIN = "fin"
+FIELD_ERROR = "error"
+
+FIN_DONE = "done"
+FIN_ERROR = "error"
+
+STREAM_START_ID = "0-0"
+
+
+class RedisLikeClient(Protocol):
+    async def set_nx(self, key: str, value: str, ttl_sec: int) -> bool: ...
+
+    async def get(self, key: str) -> str | None: ...
+
+    async def exists(self, key: str) -> bool: ...
+
+    async def delete(self, keys: list[str]) -> None: ...
+
+    async def xrange(
+        self, key: str, start: str, end: str
+    ) -> list[dict[str, Any]]: ...
+
+    async def pipeline(self, commands: list[dict[str, Any]]) -> None: ...
+
+
+class RedisResumableStreamStore:
+    def __init__(
+        self,
+        client: RedisLikeClient,
+        *,
+        key_prefix: str = DEFAULT_KEY_PREFIX,
+        default_ttl_ms: int = DEFAULT_TTL_MS,
+        poll_interval_ms: int = DEFAULT_POLL_INTERVAL_MS,
+        max_chunk_bytes: int | None = None,
+    ) -> None:
+        self._client = client
+        self._key_prefix = key_prefix
+        self._default_ttl_ms = default_ttl_ms
+        self._poll_interval_ms = poll_interval_ms
+        self._max_chunk_bytes = max_chunk_bytes
+
+    def _meta_key(self, stream_id: str) -> str:
+        return f"{self._key_prefix}:{{{stream_id}}}:meta"
+
+    def _data_key(self, stream_id: str) -> str:
+        return f"{self._key_prefix}:{{{stream_id}}}:data"
+
+    async def _read_meta(self, stream_id: str) -> dict[str, Any] | None:
+        raw = await self._client.get(self._meta_key(stream_id))
+        if raw is None:
+            return None
+        return _parse_meta(raw)
+
+    async def acquire(
+        self, stream_id: str, *, ttl_ms: int | None = None
+    ) -> ResumableStreamRole:
+        validate_stream_id(stream_id)
+        ttl_sec = _ms_to_sec(ttl_ms if ttl_ms is not None else self._default_ttl_ms)
+        meta = json.dumps({"status": "streaming", "ttlSec": ttl_sec})
+        acquired = await self._client.set_nx(self._meta_key(stream_id), meta, ttl_sec)
+        if acquired:
+            await self._client.delete([self._data_key(stream_id)])
+            return "producer"
+        return "consumer"
+
+    async def append(self, stream_id: str, chunk: bytes) -> None:
+        validate_stream_id(stream_id)
+        if (
+            self._max_chunk_bytes is not None
+            and len(chunk) > self._max_chunk_bytes
+        ):
+            raise RuntimeError(
+                f"Chunk exceeds maxChunkBytes ({len(chunk)} > {self._max_chunk_bytes})"
+            )
+        data_key = self._data_key(stream_id)
+        meta_key = self._meta_key(stream_id)
+        meta = await self._read_meta(stream_id)
+        if meta is None:
+            raise RuntimeError(f"Stream not found: {stream_id}")
+        if meta.get("status") != "streaming":
+            raise ResumableStreamError(
+                "finalized",
+                f"Stream already finalized: {stream_id}",
+            )
+        ttl_sec = meta.get("ttlSec")
+        if not isinstance(ttl_sec, int):
+            ttl_sec = _ms_to_sec(self._default_ttl_ms)
+        await self._client.pipeline(
+            [
+                {"type": "xAdd", "key": data_key, "fields": {FIELD_CHUNK: chunk}},
+                {"type": "expire", "key": data_key, "ttlSec": ttl_sec},
+                {"type": "expire", "key": meta_key, "ttlSec": ttl_sec},
+            ]
+        )
+
+    async def finalize(
+        self,
+        stream_id: str,
+        status: Literal["done", "error"],
+        error: str | None = None,
+    ) -> None:
+        validate_stream_id(stream_id)
+        data_key = self._data_key(stream_id)
+        meta_key = self._meta_key(stream_id)
+        existing = await self._read_meta(stream_id)
+        if existing is None:
+            raise RuntimeError(f"Stream not found: {stream_id}")
+        if existing.get("status") != "streaming":
+            return
+        ttl_sec = existing.get("ttlSec")
+        if not isinstance(ttl_sec, int):
+            ttl_sec = _ms_to_sec(self._default_ttl_ms)
+        if status == "error":
+            meta = json.dumps(
+                {
+                    "status": "error",
+                    "error": error if error is not None else "Stream errored",
+                    "ttlSec": ttl_sec,
+                }
+            )
+        else:
+            meta = json.dumps({"status": "done", "ttlSec": ttl_sec})
+        fields: dict[str, str] = {
+            FIELD_FIN: FIN_ERROR if status == "error" else FIN_DONE,
+        }
+        if status == "error":
+            fields[FIELD_ERROR] = error if error is not None else "Stream errored"
+        await self._client.pipeline(
+            [
+                {"type": "set", "key": meta_key, "value": meta, "ttlSec": ttl_sec},
+                {"type": "xAdd", "key": data_key, "fields": fields},
+                {"type": "expire", "key": data_key, "ttlSec": ttl_sec},
+            ]
+        )
+
+    async def read(
+        self, stream_id: str, cursor: str, signal: CancellationSignal
+    ) -> AsyncIterator[ResumableStreamEntry]:
+        validate_stream_id(stream_id)
+        data_key = self._data_key(stream_id)
+        meta_key = self._meta_key(stream_id)
+        initial_meta = await self._client.get(meta_key)
+        if initial_meta is None:
+            raise RuntimeError(f"Stream not found: {stream_id}")
+
+        last_id = STREAM_START_ID if cursor == "" else cursor
+
+        while True:
+            if signal.is_set():
+                return
+
+            start = "-" if last_id == STREAM_START_ID else f"({last_id}"
+            entries = await self._client.xrange(data_key, start, "+")
+
+            for entry in entries:
+                if signal.is_set():
+                    return
+                entry_id = entry["id"]
+                fields = entry["fields"]
+                last_id = entry_id
+
+                fin = _read_string(fields.get(FIELD_FIN))
+                if fin == FIN_DONE:
+                    return
+                if fin == FIN_ERROR:
+                    raise RuntimeError(
+                        _read_string(fields.get(FIELD_ERROR)) or "Stream errored"
+                    )
+
+                raw = fields.get(FIELD_CHUNK)
+                if raw is None:
+                    continue
+                yield ResumableStreamEntry(cursor=entry_id, chunk=_to_bytes(raw))
+
+            if len(entries) > 0:
+                continue
+
+            still_exists = await self._client.exists(meta_key)
+            if not still_exists:
+                return
+
+            await _sleep(self._poll_interval_ms, signal)
+
+    async def status(self, stream_id: str) -> ResumableStreamStatus:
+        validate_stream_id(stream_id)
+        meta = await self._client.get(self._meta_key(stream_id))
+        if meta is None:
+            return "missing"
+        parsed = _parse_meta(meta)
+        if parsed is None:
+            return "missing"
+        status = parsed.get("status")
+        if status == "streaming":
+            return "streaming"
+        if status == "done":
+            return "done"
+        if status == "error":
+            return "error"
+        return "missing"
+
+    async def delete(self, stream_id: str) -> None:
+        validate_stream_id(stream_id)
+        await self._client.delete(
+            [self._meta_key(stream_id), self._data_key(stream_id)]
+        )
+
+
+def _ms_to_sec(ms: int) -> int:
+    return max(1, math.ceil(ms / 1000))
+
+
+def _parse_meta(value: str) -> dict[str, Any] | None:
+    try:
+        parsed = json.loads(value)
+        if isinstance(parsed, dict):
+            return parsed
+        return None
+    except (json.JSONDecodeError, TypeError):
+        return None
+
+
+async def _sleep(ms: int, signal: CancellationSignal) -> None:
+    if signal.is_set():
+        return
+    with suppress(asyncio.TimeoutError):
+        await asyncio.wait_for(signal.wait(), timeout=ms / 1000.0)
+
+
+def _read_string(value: str | bytes | None) -> str | None:
+    if value is None:
+        return None
+    if isinstance(value, str):
+        return value
+    return value.decode("utf-8")
+
+
+def _to_bytes(value: str | bytes) -> bytes:
+    if isinstance(value, bytes):
+        return value
+    return value.encode("utf-8")
+
+
+class _RedisAsyncioAdapter:
+    def __init__(self, client: Any) -> None:
+        self._client = client
+
+    async def set_nx(self, key: str, value: str, ttl_sec: int) -> bool:
+        result = await self._client.set(key, value, nx=True, ex=ttl_sec)
+        return result is True or result == b"OK" or result == "OK"
+
+    async def get(self, key: str) -> str | None:
+        value = await self._client.get(key)
+        if value is None:
+            return None
+        if isinstance(value, bytes):
+            return value.decode("utf-8")
+        return str(value)
+
+    async def exists(self, key: str) -> bool:
+        result = await self._client.exists(key)
+        return int(result) > 0
+
+    async def delete(self, keys: list[str]) -> None:
+        if not keys:
+            return
+        await self._client.delete(*keys)
+
+    async def xrange(
+        self, key: str, start: str, end: str
+    ) -> list[dict[str, Any]]:
+        reply = await self._client.xrange(key, min=start, max=end)
+        out: list[dict[str, Any]] = []
+        for entry_id, fields in reply:
+            if isinstance(entry_id, bytes):
+                entry_id = entry_id.decode("utf-8")
+            parsed_fields: dict[str, str | bytes] = {}
+            for field_key, field_value in fields.items():
+                if isinstance(field_key, bytes):
+                    field_key = field_key.decode("utf-8")
+                parsed_fields[field_key] = field_value
+            out.append({"id": entry_id, "fields": parsed_fields})
+        return out
+
+    async def pipeline(self, commands: list[dict[str, Any]]) -> None:
+        if not commands:
+            return
+        pipe = self._client.pipeline()
+        for cmd in commands:
+            cmd_type = cmd["type"]
+            if cmd_type == "xAdd":
+                pipe.xadd(cmd["key"], dict(cmd["fields"]))
+            elif cmd_type == "expire":
+                pipe.expire(cmd["key"], cmd["ttlSec"])
+            elif cmd_type == "set":
+                pipe.set(cmd["key"], cmd["value"], ex=cmd["ttlSec"])
+        await pipe.execute()
+
+
+def create_redis_resumable_stream_store(
+    client: Any,
+    *,
+    key_prefix: str = DEFAULT_KEY_PREFIX,
+    default_ttl_ms: int = DEFAULT_TTL_MS,
+    poll_interval_ms: int = DEFAULT_POLL_INTERVAL_MS,
+    max_chunk_bytes: int | None = None,
+) -> RedisResumableStreamStore:
+    """Create a store backed by a ``redis.asyncio.Redis`` client."""
+    try:
+        import redis.asyncio  # noqa: F401
+    except ImportError as exc:
+        raise ImportError(
+            "redis is required for create_redis_resumable_stream_store; "
+            "install with: pip install assistant-stream[redis]"
+        ) from exc
+
+    connection_pool = getattr(client, "connection_pool", None)
+    if connection_pool is not None:
+        connection_kwargs = getattr(connection_pool, "connection_kwargs", None)
+        if isinstance(connection_kwargs, dict) and connection_kwargs.get(
+            "decode_responses"
+        ):
+            raise ValueError(
+                "redis client must be constructed without decode_responses "
+                "(decode_responses=True cannot round-trip binary stream chunks)"
+            )
+
+    return RedisResumableStreamStore(
+        _RedisAsyncioAdapter(client),
+        key_prefix=key_prefix,
+        default_ttl_ms=default_ttl_ms,
+        poll_interval_ms=poll_interval_ms,
+        max_chunk_bytes=max_chunk_bytes,
+    )

python/assistant-stream/src/assistant_stream/resumable/types.py

@@ -0,0 +1,44 @@
+from __future__ import annotations
+
+from collections.abc import AsyncIterator
+from dataclasses import dataclass
+from typing import Literal, Protocol
+
+
+ResumableStreamRole = Literal["producer", "consumer"]
+ResumableStreamStatus = Literal["streaming", "done", "error", "missing"]
+
+
+@dataclass(frozen=True)
+class ResumableStreamEntry:
+    cursor: str
+    chunk: bytes
+
+
+class CancellationSignal(Protocol):
+    def is_set(self) -> bool: ...
+
+    async def wait(self) -> bool: ...
+
+
+class ResumableStreamStore(Protocol):
+    async def acquire(
+        self, stream_id: str, *, ttl_ms: int | None = None
+    ) -> ResumableStreamRole: ...
+
+    async def append(self, stream_id: str, chunk: bytes) -> None: ...
+
+    async def finalize(
+        self,
+        stream_id: str,
+        status: Literal["done", "error"],
+        error: str | None = None,
+    ) -> None: ...
+
+    def read(
+        self, stream_id: str, cursor: str, signal: CancellationSignal
+    ) -> AsyncIterator[ResumableStreamEntry]: ...
+
+    async def status(self, stream_id: str) -> ResumableStreamStatus: ...
+
+    async def delete(self, stream_id: str) -> None: ...

python/assistant-stream/src/assistant_stream/serialization/__init__.py

@@ -0,0 +1,21 @@
+from assistant_stream.serialization.data_stream import (
+    DataStreamEncoder,
+    DataStreamResponse,
+)
+from assistant_stream.serialization.openai_stream import (
+    OpenAIStreamEncoder,
+    OpenAIStreamResponse,
+)
+from assistant_stream.serialization.assistant_transport import (
+    AssistantTransportEncoder,
+    AssistantTransportResponse,
+)
+
+__all__ = [
+    "DataStreamEncoder",
+    "DataStreamResponse",
+    "OpenAIStreamEncoder",
+    "OpenAIStreamResponse",
+    "AssistantTransportEncoder",
+    "AssistantTransportResponse",
+]

python/assistant-stream/src/assistant_stream/serialization/assistant_stream_response.py

@@ -0,0 +1,17 @@
+from assistant_stream.assistant_stream_chunk import AssistantStreamChunk
+from assistant_stream.serialization.stream_encoder import StreamEncoder
+from typing import AsyncGenerator
+
+from starlette.responses import StreamingResponse
+
+
+class AssistantStreamResponse(StreamingResponse):
+    def __init__(
+        self,
+        stream: AsyncGenerator[AssistantStreamChunk, None],
+        stream_encoder: StreamEncoder,
+    ):
+        super().__init__(
+            stream_encoder.encode_stream(stream),
+            media_type=stream_encoder.get_media_type(),
+        )

python/assistant-stream/src/assistant_stream/serialization/assistant_transport.py

@@ -0,0 +1,62 @@
+from assistant_stream.assistant_stream_chunk import AssistantStreamChunk
+from assistant_stream.serialization.assistant_stream_response import (
+    AssistantStreamResponse,
+)
+from assistant_stream.serialization.stream_encoder import StreamEncoder
+from assistant_stream.state_proxy import StateProxy
+from typing import AsyncGenerator, Any
+import json
+
+
+class StateProxyJSONEncoder(json.JSONEncoder):
+    """Custom JSON encoder that can handle StateProxy objects."""
+
+    def default(self, obj: Any) -> Any:
+        if isinstance(obj, StateProxy):
+            return obj._get_value()
+        return super().default(obj)
+
+
+class AssistantTransportEncoder(StreamEncoder):
+    """
+    AssistantTransportEncoder encodes AssistantStreamChunks into SSE format
+    and emits [DONE] when the stream completes.
+    """
+
+    def get_media_type(self) -> str:
+        return "text/event-stream"
+
+    def _chunk_to_dict(self, chunk: AssistantStreamChunk) -> dict[str, Any]:
+        """Convert a chunk to a JSON-serializable dictionary."""
+        chunk_dict = {"type": chunk.type}
+
+        # Add all attributes from the chunk
+        for key, value in vars(chunk).items():
+            if key != "type":  # Already added
+                chunk_dict[self._snake_to_camel(key)] = value
+
+        return chunk_dict
+
+    def _snake_to_camel(self, snake_str: str) -> str:
+        """Convert snake_case to camelCase."""
+        components = snake_str.split("_")
+        return components[0] + "".join(x.title() for x in components[1:])
+
+    async def encode_stream(
+        self, stream: AsyncGenerator[AssistantStreamChunk, None]
+    ) -> AsyncGenerator[str, None]:
+        async for chunk in stream:
+            chunk_dict = self._chunk_to_dict(chunk)
+            chunk_json = json.dumps(chunk_dict, cls=StateProxyJSONEncoder)
+            yield f"data: {chunk_json}\n\n"
+
+        # Emit [DONE] marker when stream completes
+        yield "data: [DONE]\n\n"
+
+
+class AssistantTransportResponse(AssistantStreamResponse):
+    def __init__(
+        self,
+        stream: AsyncGenerator[AssistantStreamChunk, None],
+    ):
+        super().__init__(stream, AssistantTransportEncoder())

python/assistant-stream/src/assistant_stream/serialization/data_stream.py

@@ -0,0 +1,86 @@
+from assistant_stream.assistant_stream_chunk import (
+    AssistantStreamChunk,
+)
+import json
+from typing import AsyncGenerator, Any
+from assistant_stream.serialization.assistant_stream_response import (
+    AssistantStreamResponse,
+)
+from assistant_stream.serialization.stream_encoder import StreamEncoder
+from assistant_stream.state_proxy import StateProxy
+
+
+class StateProxyJSONEncoder(json.JSONEncoder):
+    """Custom JSON encoder that can handle StateProxy objects."""
+    def default(self, obj: Any) -> Any:
+        if isinstance(obj, StateProxy):
+            return obj._get_value()
+        return super().default(obj)
+
+
+class DataStreamEncoder(StreamEncoder):
+    def __init__(self):
+        pass
+
+    def encode_chunk(self, chunk: AssistantStreamChunk) -> str:
+        if chunk.type == "text-delta":
+            if hasattr(chunk, 'parent_id') and chunk.parent_id:
+                return f"aui-text-delta:{json.dumps({'textDelta': chunk.text_delta, 'parentId': chunk.parent_id}, cls=StateProxyJSONEncoder)}\n"
+            else:
+                return f"0:{json.dumps(chunk.text_delta, cls=StateProxyJSONEncoder)}\n"
+        elif chunk.type == "reasoning-delta":
+            if hasattr(chunk, 'parent_id') and chunk.parent_id:
+                return f"aui-reasoning-delta:{json.dumps({'reasoningDelta': chunk.reasoning_delta, 'parentId': chunk.parent_id}, cls=StateProxyJSONEncoder)}\n"
+            else:
+                return f"g:{json.dumps(chunk.reasoning_delta, cls=StateProxyJSONEncoder)}\n"
+        elif chunk.type == "tool-call-begin":
+            data = {"toolCallId": chunk.tool_call_id, "toolName": chunk.tool_name}
+            if hasattr(chunk, 'parent_id') and chunk.parent_id:
+                data["parentId"] = chunk.parent_id
+            return f'b:{json.dumps(data, cls=StateProxyJSONEncoder)}\n'
+        elif chunk.type == "tool-call-delta":
+            return f'c:{json.dumps({ "toolCallId": chunk.tool_call_id, "argsTextDelta": chunk.args_text_delta }, cls=StateProxyJSONEncoder)}\n'
+        elif chunk.type == "tool-result":
+            res = {"toolCallId": chunk.tool_call_id, "result": chunk.result}
+            if chunk.artifact is not None:
+                res["artifact"] = chunk.artifact
+            if chunk.is_error:
+                res["isError"] = chunk.is_error
+            return f"a:{json.dumps(res, cls=StateProxyJSONEncoder)}\n"
+        elif chunk.type == "data":
+            return f"2:{json.dumps([chunk.data], cls=StateProxyJSONEncoder)}\n"
+        elif chunk.type == "error":
+            return f"3:{json.dumps(chunk.error, cls=StateProxyJSONEncoder)}\n"
+        elif chunk.type == "source":
+            source_data = {
+                "sourceType": chunk.source_type,
+                "id": chunk.id,
+                "url": chunk.url
+            }
+            if chunk.title is not None:
+                source_data["title"] = chunk.title
+            if hasattr(chunk, 'parent_id') and chunk.parent_id:
+                source_data["parentId"] = chunk.parent_id
+            return f"h:{json.dumps(source_data, cls=StateProxyJSONEncoder)}\n"
+        elif chunk.type == "update-state":
+            return f"aui-state:{json.dumps(chunk.operations, cls=StateProxyJSONEncoder)}\n"
+
+    def get_media_type(self) -> str:
+        return "text/plain"
+
+    async def encode_stream(
+        self, stream: AsyncGenerator[AssistantStreamChunk, None]
+    ) -> AsyncGenerator[str, None]:
+        async for chunk in stream:
+            encoded = self.encode_chunk(chunk)
+            if encoded is None:
+                continue
+            yield encoded
+
+
+class DataStreamResponse(AssistantStreamResponse):
+    def __init__(
+        self,
+        stream: AsyncGenerator[AssistantStreamChunk, None],
+    ):
+        super().__init__(stream, DataStreamEncoder())

python/assistant-stream/src/assistant_stream/serialization/openai_stream.py

@@ -0,0 +1,82 @@
+from assistant_stream.assistant_stream_chunk import AssistantStreamChunk
+import json
+import time
+import string
+import random
+from typing import AsyncGenerator
+from assistant_stream.serialization.assistant_stream_response import (
+    AssistantStreamResponse,
+)
+from assistant_stream.serialization.stream_encoder import StreamEncoder
+
+
+def generate_openai_style_id():
+    prefix = "chatcmpl-"
+    characters = string.ascii_letters + string.digits
+    random_id = "".join(random.choices(characters, k=24))
+    return prefix + random_id
+
+
+class OpenAIStreamEncoder(StreamEncoder):
+    def __init__(self, model="assistant_stream", system_fingerprint="fp_0000000000"):
+        self.id = generate_openai_style_id()
+        self.model = model
+        self.system_fingerprint = system_fingerprint
+        pass
+
+    def get_media_type(self) -> str:
+        return "text/event-stream"
+
+    def _create_chunk(self, delta={}, finish_reason=None):
+        response = {
+            "id": self.id,
+            "object": "chat.completion.chunk",
+            "created": int(time.time()),
+            "model": self.model,
+            "system_fingerprint": self.system_fingerprint,
+            "choices": [
+                {
+                    "index": 0,
+                    "delta": delta,
+                    "logprobs": None,
+                    "finish_reason": finish_reason,
+                }
+            ],
+        }
+        return f"data: {json.dumps(response, ensure_ascii=False)}\n\n"
+
+    def encode_chunk(self, chunk: AssistantStreamChunk) -> str:
+        """
+        Encodes the chunk into OpenAI's SSE format.
+        """
+        if chunk.type == "text-delta":
+            # Construct the delta for text content
+            return self._create_chunk({"content": chunk.text_delta})
+        else:
+            # Handle unknown chunk types gracefully
+            return ""
+
+    async def encode_stream(
+        self, stream: AsyncGenerator[AssistantStreamChunk, None]
+    ) -> AsyncGenerator[str, None]:
+        """
+        Asynchronously encodes chunks into SSE-formatted strings.
+        """
+        async for chunk in stream:
+            encoded_chunk = self.encode_chunk(chunk)
+            if encoded_chunk:
+                yield encoded_chunk
+
+        yield self._create_chunk(finish_reason="stop")
+        yield "data: [DONE]\n\n"
+
+
+class OpenAIStreamResponse(AssistantStreamResponse):
+    def __init__(
+        self,
+        stream: AsyncGenerator[AssistantStreamChunk, None],
+    ):
+        """
+        Initializes the response with the OpenAI SSE encoder.
+        """
+        super().__init__(stream, OpenAIStreamEncoder())

python/assistant-stream/src/assistant_stream/serialization/stream_encoder.py

@@ -0,0 +1,26 @@
+from abc import ABC, abstractmethod
+from typing import AsyncGenerator
+from assistant_stream.assistant_stream_chunk import AssistantStreamChunk
+
+
+class StreamEncoder(ABC):
+    """
+    Abstract base class for stream encoders, requiring an implementation of `encode_stream`.
+    """
+
+    @abstractmethod
+    def get_media_type(self) -> str:
+        """
+        Returns the MIME type of the stream.
+        """
+        pass
+
+    @abstractmethod
+    async def encode_stream(
+        self, stream: AsyncGenerator[AssistantStreamChunk, None]
+    ) -> AsyncGenerator[str, None]:
+        """
+        Encode the stream of AssistantStreamChunk into a specific format.
+        This method must be implemented by subclasses.
+        """
+        pass

python/assistant-stream/src/assistant_stream/state_manager.py

@@ -0,0 +1,208 @@
+import asyncio
+from typing import Any, Callable, Dict, List, Sequence, Union
+
+from assistant_stream.assistant_stream_chunk import (
+    ObjectStreamOperation,
+    UpdateStateChunk,
+)
+from assistant_stream.state_proxy import StateProxy
+
+
+class StateManager:
+    """Manages state operations with efficient batching and local updates."""
+
+    def __init__(
+        self,
+        put_chunk_callback: Callable[[UpdateStateChunk], None],
+        state_data: Any | None = None,
+    ):
+        """Initialize with callback for sending state updates."""
+        self._state_data = state_data
+        self._pending_operations = []
+        self._update_scheduled = False
+        self._put_chunk_callback = put_chunk_callback
+        self._loop = asyncio.get_running_loop()
+        self._state_proxy = StateProxy(self, [])
+
+    @property
+    def state(self) -> Any:
+        """Access the state proxy object for making state updates.
+
+        If state is None, returns None directly instead of a proxy.
+        Otherwise returns a proxy object for the state.
+        """
+        if self._state_data is None:
+            return None
+        return self._state_proxy
+
+    @property
+    def state_data(self) -> Dict[str, Any]:
+        """Current state data."""
+        return self._state_data
+
+    def add_operations(self, operations: List[ObjectStreamOperation]) -> None:
+        """Add operations to pending batch and apply locally."""
+        # Apply to local state immediately
+        for operation in operations:
+            self._apply_operation_to_local_state(operation)
+
+        # Add to pending operations
+        self._pending_operations.extend(operations)
+
+        # Schedule batch update if needed
+        if not self._update_scheduled:
+            self._update_scheduled = True
+            self._loop.call_soon_threadsafe(self._flush_updates)
+
+    def append_text(self, path: Sequence[Union[str, int]], value: str) -> None:
+        """Append text at a path using an explicit append-text delta operation."""
+        if not isinstance(value, str):
+            raise TypeError(
+                f"Can only append str (not '{type(value).__name__}') as text delta"
+            )
+
+        self.add_operations(
+            [
+                {
+                    "type": "append-text",
+                    "path": [str(segment) for segment in path],
+                    "value": value,
+                }
+            ]
+        )
+
+    def _flush_updates(self) -> None:
+        """Send pending operations as a batch."""
+        if self._pending_operations:
+            operations_to_send = self._pending_operations.copy()
+            self._pending_operations.clear()
+            self._put_chunk_callback(UpdateStateChunk(operations=operations_to_send))
+
+        self._update_scheduled = False
+
+    def flush(self) -> None:
+        """Explicitly flush any pending operations.
+
+        This should be called before the run completes to ensure all state updates are sent.
+        """
+        if self._pending_operations:
+            self._flush_updates()
+
+    def _apply_operation_to_local_state(self, operation: ObjectStreamOperation) -> None:
+        """Apply operation to local state."""
+        op_type = operation["type"]
+
+        if op_type == "set":
+            self._update_path(operation["path"], lambda _: operation["value"])
+
+        elif op_type == "append-text":
+
+            def append_text(current):
+                if not isinstance(current, str):
+                    path_str = ", ".join(operation["path"])
+                    raise TypeError(f"Expected string at path [{path_str}]")
+                return current + operation["value"]
+
+            self._update_path(operation["path"], append_text)
+
+        else:
+            raise TypeError(f"Invalid operation type: {op_type}")
+
+    def get_value_at_path(self, path: List[str]) -> Any:
+        """Get value at path, raising KeyError for invalid paths."""
+        if not path:
+            return self._state_data
+
+        # If state is None, we can't navigate further
+        if self._state_data is None:
+            raise KeyError(path[0] if path else "")
+
+        current = self._state_data
+
+        for key in path:
+            try:
+                if isinstance(current, list):
+                    idx = int(key)
+                    if idx < 0 or idx >= len(current):
+                        raise KeyError(key)
+                    current = current[idx]
+                elif isinstance(current, dict):
+                    current = current[key]
+                else:
+                    raise KeyError(key)
+            except (ValueError, KeyError, IndexError):
+                raise KeyError(key)
+
+        return current
+
+    def _update_path(self, path: List[str], updater: Callable[[Any], Any]) -> None:
+        """Update value at path without creating parent objects."""
+        # Handle empty path (update root state)
+        if not path:
+            self._state_data = updater(self._state_data)
+            return
+
+        # Initialize state as empty object if it's null
+        if self._state_data is None:
+            self._state_data = {}
+
+        if not isinstance(self._state_data, (dict, list)):
+            raise KeyError(f"Invalid path: [{', '.join(path)}]")
+
+        key, *rest = path
+
+        # Handle list access
+        if isinstance(self._state_data, list):
+            try:
+                idx = int(key)
+                if idx < 0 or idx > len(self._state_data):
+                    raise KeyError(key)
+
+                if not rest:
+                    # For direct update
+                    if idx == len(self._state_data):  # Append case
+                        value = updater(None)
+                        if value is not None:
+                            self._state_data.append(value)
+                    else:  # Update existing element
+                        self._state_data[idx] = updater(self._state_data[idx])
+                else:
+                    # For nested update
+                    if idx == len(self._state_data):
+                        raise KeyError(key)
+
+                    # Create a copy for the nested update
+                    next_state = self._state_data.copy()
+
+                    # Create a temporary manager for the nested path
+                    temp_manager = type(self)(lambda _: None)
+                    temp_manager._state_data = next_state[idx]
+
+                    # Update nested path
+                    temp_manager._update_path(rest, updater)
+                    next_state[idx] = temp_manager._state_data
+                    self._state_data = next_state
+            except ValueError:
+                raise KeyError(key)
+        else:  # Handle dict access
+            if not rest:
+                # For direct update
+                if key not in self._state_data and updater(None) is None:
+                    return
+                self._state_data[key] = updater(self._state_data.get(key))
+            else:
+                # For nested update
+                if key not in self._state_data:
+                    raise KeyError(key)
+
+                # Create a copy for the nested update
+                next_state = dict(self._state_data)
+
+                # Create a temporary manager for the nested path
+                temp_manager = type(self)(lambda _: None)
+                temp_manager._state_data = next_state[key]
+
+                # Update nested path
+                temp_manager._update_path(rest, updater)
+                next_state[key] = temp_manager._state_data
+                self._state_data = next_state

python/assistant-stream/src/assistant_stream/state_proxy.py

@@ -0,0 +1,358 @@
+from typing import Any, List, Optional, Union, TYPE_CHECKING
+
+
+# Avoid circular import
+if TYPE_CHECKING:
+    from assistant_stream.state_manager import StateManager
+
+
+class StateProxy:
+    """Proxy object for state access and updates using dictionary-style access.
+
+    Example:
+        state_proxy["user"]["name"] = "John"
+        name = state_proxy["user"]["name"]
+        state_proxy["messages"] += "Hello"
+        state_proxy["items"].append("item")
+    """
+
+    def _get_value(self):
+        return self._manager.get_value_at_path(self._path)
+
+    def __init__(
+        self,
+        state_manager: "StateManager",
+        path: Optional[List[str]] | None = None,
+    ) -> None:
+        """Initialize with state manager and current path."""
+        self._manager = state_manager
+        self._path = path or []
+
+    def __getitem__(self, key: Union[str, int]) -> Union["StateProxy", Any]:
+        """Access nested values with dict-style syntax. Returns primitives directly."""
+        current_value = self._manager.get_value_at_path(self._path)
+
+        # Handle list indexing
+        if isinstance(current_value, list):
+            try:
+                index = int(key)
+                list_len = len(current_value)
+
+                # Handle negative indices
+                if index < 0:
+                    index = list_len + index
+
+                # Validate index is in bounds
+                if index < 0 or index >= list_len:
+                    raise KeyError(key)
+
+                # Use the normalized index as string key
+                str_key = str(index)
+            except (ValueError, TypeError):
+                raise KeyError(key)
+        else:
+            # For dicts, use string representation of key
+            str_key = str(key)
+
+            # Validate key exists
+            if isinstance(current_value, dict):
+                if str_key not in current_value:
+                    raise KeyError(key)
+            elif not isinstance(current_value, list):
+                raise KeyError(key)
+
+        # Get value at path
+        value = self._manager.get_value_at_path(self._path + [str_key])
+
+        # Return primitives directly (including strings)
+        if value is None or isinstance(value, (int, float, bool, str)):
+            return value
+
+        # Return proxy only for collections
+        return StateProxy(self._manager, self._path + [str_key])
+
+    def __setitem__(self, key: Union[str, int], value: Any) -> None:
+        """Set value with dict-style syntax."""
+        current_value = self._manager.get_value_at_path(self._path)
+
+        # Handle list indexing
+        if isinstance(current_value, list):
+            try:
+                index = int(key)
+                list_len = len(current_value)
+
+                # Handle negative indices
+                if index < 0:
+                    index = list_len + index
+
+                # Validate index is in bounds
+                if index < 0 or index >= list_len:
+                    raise KeyError(key)
+
+                # Use the normalized index as string key
+                str_key = str(index)
+            except (ValueError, TypeError):
+                raise KeyError(key)
+        else:
+            # For dicts and other types, use string representation of key
+            str_key = str(key)
+
+        target_path = self._path + [str_key]
+
+        # Encode string extensions as append-text. Skip empty current values:
+        # any.startswith("") matches all strings and would convert first writes too.
+        try:
+            current_target_value = self._manager.get_value_at_path(target_path)
+            if (
+                isinstance(current_target_value, str)
+                and isinstance(value, str)
+                and current_target_value
+                and value.startswith(current_target_value)
+            ):
+                delta = value[len(current_target_value) :]
+                if delta:
+                    self._manager.append_text(target_path, delta)
+                    return
+        except KeyError:
+            pass
+
+        self._manager.add_operations(
+            [{"type": "set", "path": target_path, "value": value}]
+        )
+
+    def __iadd__(self, other: Any) -> "StateProxy":
+        """Support += on list-valued proxies.
+
+        String += on a leaf goes through __setitem__ instead, since
+        __getitem__ returns the raw str rather than a proxy.
+        """
+        current_value = self._manager.get_value_at_path(self._path)
+
+        # String concatenation
+        if isinstance(current_value, str):
+            if not isinstance(other, str):
+                raise TypeError(
+                    f"Can only concatenate str (not '{type(other).__name__}') to str"
+                )
+
+            self._manager.append_text(self._path, other)
+            return self
+
+        # List extension
+        if isinstance(current_value, list):
+            try:
+                # Ensure other is iterable
+                iterator = iter(other)
+
+                # Add each item at the end of the list
+                operations = []
+                current_len = len(current_value)
+
+                for i, item in enumerate(iterator):
+                    operations.append(
+                        {
+                            "type": "set",
+                            "path": self._path + [str(current_len + i)],
+                            "value": item,
+                        }
+                    )
+
+                if operations:
+                    self._manager.add_operations(operations)
+                return self
+            except TypeError:
+                raise TypeError(
+                    f"can only concatenate list (not '{type(other).__name__}') to list"
+                )
+
+        raise TypeError(
+            f"unsupported operand type(s) for +=: '{type(current_value).__name__}' and '{type(other).__name__}'"
+        )
+
+    def __repr__(self) -> str:
+        """String representation of the value."""
+        return repr(self._manager.get_value_at_path(self._path))
+
+    def __str__(self) -> str:
+        """String representation of the value."""
+        return str(self._manager.get_value_at_path(self._path))
+
+    def __len__(self) -> int:
+        """Length of the value."""
+        return len(self._manager.get_value_at_path(self._path))
+
+    def __contains__(self, item: Any) -> bool:
+        """Check if item is in the value."""
+        return item in self._manager.get_value_at_path(self._path)
+
+    def __eq__(self, other: Any) -> bool:
+        """Compare equality with another value."""
+        return self._manager.get_value_at_path(self._path) == other
+
+    def __ne__(self, other: Any) -> bool:
+        """Compare inequality with another value."""
+        return self._manager.get_value_at_path(self._path) != other
+
+    def __hash__(self) -> int:
+        """Hash the underlying value if hashable."""
+        value = self._manager.get_value_at_path(self._path)
+        if isinstance(value, (str, int, float, bool, tuple)):
+            return hash(value)
+        raise TypeError(f"unhashable type: '{type(value).__name__}'")
+
+    def __bool__(self) -> bool:
+        """Truth value of the underlying value."""
+        return bool(self._manager.get_value_at_path(self._path))
+
+    def __int__(self) -> int:
+        """Convert to int if possible."""
+        return int(self._manager.get_value_at_path(self._path))
+
+    def __float__(self) -> float:
+        """Convert to float if possible."""
+        return float(self._manager.get_value_at_path(self._path))
+
+    def __add__(self, other: Any) -> Any:
+        """Add operation for strings and lists."""
+        value = self._manager.get_value_at_path(self._path)
+        if isinstance(value, str) and isinstance(other, str):
+            return value + other
+        if isinstance(value, list) and hasattr(other, "__iter__"):
+            return value + list(other)
+        return NotImplemented
+
+    def __getattr__(self, name: str) -> Any:
+        """Forward attribute access to the underlying value."""
+        value = self._manager.get_value_at_path(self._path)
+
+        # Handle string methods
+        if isinstance(value, str):
+            attr = getattr(value, name)
+            if callable(attr):
+
+                def method_wrapper(*args, **kwargs):
+                    result = attr(*args, **kwargs)
+                    return self if result is value else result
+
+                return method_wrapper
+            return attr
+
+        # Forward non-modifying methods for lists and dicts
+        try:
+            attr = getattr(value, name)
+            if callable(attr):
+                return lambda *args, **kwargs: attr(*args, **kwargs)
+            return attr
+        except (AttributeError, TypeError):
+            pass
+
+        raise AttributeError(
+            f"'{type(value).__name__}' object has no attribute '{name}'"
+        )
+
+    def __iter__(self):
+        """Make the proxy iterable."""
+        return iter(self._manager.get_value_at_path(self._path))
+
+    # Efficient list operations
+    def append(self, item: Any) -> None:
+        """Append an item to a list."""
+        value = self._manager.get_value_at_path(self._path)
+        if not isinstance(value, list):
+            raise TypeError(f"'append' not supported for type {type(value).__name__}")
+
+        self._manager.add_operations(
+            [{"type": "set", "path": self._path + [str(len(value))], "value": item}]
+        )
+
+    def extend(self, iterable: Any) -> None:
+        """Extend a list with items from an iterable."""
+        if isinstance(iterable, StateProxy):
+            iterable = iterable._manager.get_value_at_path(iterable._path)
+        self.__iadd__(iterable)
+
+    def clear(self) -> None:
+        """Clear a list or dictionary."""
+        value = self._manager.get_value_at_path(self._path)
+
+        if isinstance(value, (list, dict)):
+            empty_value = [] if isinstance(value, list) else {}
+            self._manager.add_operations(
+                [{"type": "set", "path": self._path, "value": empty_value}]
+            )
+        else:
+            raise TypeError(f"'clear' not supported for type {type(value).__name__}")
+
+    # Dictionary operations
+    def get(self, key: Any, default: Any = None) -> Any:
+        """Get dictionary value with default."""
+        value = self._manager.get_value_at_path(self._path)
+        if not isinstance(value, dict):
+            raise TypeError(f"'get' not supported for type {type(value).__name__}")
+
+        try:
+            return self[key]
+        except KeyError:
+            return default
+
+    def keys(self):
+        """Dictionary keys view."""
+        value = self._manager.get_value_at_path(self._path)
+        if not isinstance(value, dict):
+            raise TypeError(f"'keys' not supported for type {type(value).__name__}")
+        return value.keys()
+
+    def values(self):
+        """Dictionary values view."""
+        value = self._manager.get_value_at_path(self._path)
+        if not isinstance(value, dict):
+            raise TypeError(f"'values' not supported for type {type(value).__name__}")
+        return value.values()
+
+    def items(self):
+        """Dictionary items view."""
+        value = self._manager.get_value_at_path(self._path)
+        if not isinstance(value, dict):
+            raise TypeError(f"'items' not supported for type {type(value).__name__}")
+        return value.items()
+
+    def setdefault(self, key, default=None):
+        """Set default value if key doesn't exist."""
+        value = self._manager.get_value_at_path(self._path)
+        if not isinstance(value, dict):
+            raise TypeError(
+                f"'setdefault' not supported for type {type(value).__name__}"
+            )
+
+        if key in value:
+            return self[key]
+
+        self[key] = default
+        return default
+
+    # Unsupported operations that would be inefficient
+    def insert(self, index: int, item: Any) -> None:
+        """Not supported - would require sending entire list."""
+        raise NotImplementedError("Use indexing or append() instead")
+
+    def pop(self, *args):
+        """Not supported - would require sending entire collection."""
+        raise NotImplementedError(
+            "Would require sending the entire collection over the network"
+        )
+
+    def remove(self, item: Any) -> None:
+        """Not supported - would require sending entire list."""
+        raise NotImplementedError(
+            "Would require sending the entire list over the network"
+        )
+
+    def update(self, *args, **kwargs):
+        """Not supported - would require sending entire dictionary."""
+        raise NotImplementedError("Use individual assignments instead")
+
+    def popitem(self):
+        """Not supported - would require sending entire dictionary."""
+        raise NotImplementedError(
+            "Would require sending the entire dictionary over the network"
+        )

python/assistant-stream/tests/__init__.py

@@ -0,0 +1 @@
+

python/assistant-stream/tests/test_assistant_transport.py

@@ -0,0 +1,150 @@
+import pytest
+from assistant_stream import create_run, RunController
+from assistant_stream.assistant_stream_chunk import UpdateStateChunk
+from assistant_stream.serialization.assistant_transport import AssistantTransportEncoder
+import json
+
+
+@pytest.mark.anyio
+async def test_assistant_transport_encoder_format():
+    """Test that AssistantTransportEncoder produces SSE format."""
+    encoder = AssistantTransportEncoder()
+    collected_output = []
+
+    async def run_callback(controller: RunController):
+        controller.append_text("Hello")
+        controller.append_text(" world")
+
+    # Create the run and encode it
+    chunks = create_run(run_callback)
+    encoded = encoder.encode_stream(chunks)
+
+    async for line in encoded:
+        collected_output.append(line)
+
+    # Verify SSE format
+    assert len(collected_output) > 0
+
+    # All lines except the last should be SSE formatted chunks
+    for line in collected_output[:-1]:
+        assert line.startswith("data: ")
+        assert line.endswith("\n\n")
+        # Verify it's valid JSON (excluding the "data: " prefix and newlines)
+        json_str = line[6:-2]  # Remove "data: " and "\n\n"
+        chunk_data = json.loads(json_str)
+        assert "type" in chunk_data
+
+    # Last line should be [DONE]
+    assert collected_output[-1] == "data: [DONE]\n\n"
+
+
+@pytest.mark.anyio
+async def test_assistant_transport_encoder_text_chunks():
+    """Test that text chunks are properly encoded."""
+    encoder = AssistantTransportEncoder()
+    collected_chunks = []
+
+    async def run_callback(controller: RunController):
+        controller.append_text("Hello")
+        controller.append_text(" world")
+
+    # Create the run and encode it
+    chunks = create_run(run_callback)
+    encoded = encoder.encode_stream(chunks)
+
+    async for line in encoded:
+        if line != "data: [DONE]\n\n":
+            json_str = line[6:-2]  # Remove "data: " and "\n\n"
+            chunk_data = json.loads(json_str)
+            collected_chunks.append(chunk_data)
+
+    # Verify we got text-delta chunks with camelCase
+    text_chunks = [c for c in collected_chunks if c["type"] == "text-delta"]
+    assert len(text_chunks) == 2
+    assert text_chunks[0]["textDelta"] == "Hello"
+    assert text_chunks[1]["textDelta"] == " world"
+
+
+@pytest.mark.anyio
+async def test_assistant_transport_encoder_reasoning():
+    """Test that reasoning chunks are properly encoded."""
+    encoder = AssistantTransportEncoder()
+    collected_chunks = []
+
+    async def run_callback(controller: RunController):
+        controller.append_reasoning("Thinking...")
+
+    # Create the run and encode it
+    chunks = create_run(run_callback)
+    encoded = encoder.encode_stream(chunks)
+
+    async for line in encoded:
+        if line != "data: [DONE]\n\n":
+            json_str = line[6:-2]
+            chunk_data = json.loads(json_str)
+            collected_chunks.append(chunk_data)
+
+    # Verify we got reasoning-delta chunks with camelCase
+    reasoning_chunks = [c for c in collected_chunks if c["type"] == "reasoning-delta"]
+    assert len(reasoning_chunks) == 1
+    assert reasoning_chunks[0]["reasoningDelta"] == "Thinking..."
+
+
+@pytest.mark.anyio
+async def test_assistant_transport_encoder_media_type():
+    """Test that the encoder returns the correct media type."""
+    encoder = AssistantTransportEncoder()
+    assert encoder.get_media_type() == "text/event-stream"
+
+
+@pytest.mark.anyio
+async def test_assistant_transport_encoder_tool_calls():
+    """Test that tool call chunks are properly encoded."""
+    encoder = AssistantTransportEncoder()
+    collected_chunks = []
+
+    async def run_callback(controller: RunController):
+        tool_controller = await controller.add_tool_call("get_weather", "tool_1")
+        tool_controller.append_args_text('{"location": "NYC"}')
+        tool_controller.set_response({"temp": 70})
+
+    # Create the run and encode it
+    chunks = create_run(run_callback)
+    encoded = encoder.encode_stream(chunks)
+
+    async for line in encoded:
+        if line != "data: [DONE]\n\n":
+            json_str = line[6:-2]
+            chunk_data = json.loads(json_str)
+            collected_chunks.append(chunk_data)
+
+    # Verify we got tool-call chunks with camelCase
+    tool_begin_chunks = [c for c in collected_chunks if c["type"] == "tool-call-begin"]
+    assert len(tool_begin_chunks) == 1
+    assert tool_begin_chunks[0]["toolCallId"] == "tool_1"
+    assert tool_begin_chunks[0]["toolName"] == "get_weather"
+
+    tool_delta_chunks = [c for c in collected_chunks if c["type"] == "tool-call-delta"]
+    assert len(tool_delta_chunks) > 0
+
+    tool_result_chunks = [c for c in collected_chunks if c["type"] == "tool-result"]
+    assert len(tool_result_chunks) == 1
+    assert tool_result_chunks[0]["toolCallId"] == "tool_1"
+
+
+@pytest.mark.anyio
+async def test_assistant_transport_encoder_update_state_shape():
+    """Test that update-state chunks preserve operation payload shape."""
+    encoder = AssistantTransportEncoder()
+    operations = [
+        {"type": "append-text", "path": ["messages", "0", "text"], "value": "hi"}
+    ]
+
+    async def stream():
+        yield UpdateStateChunk(operations=operations)
+
+    collected_output = [line async for line in encoder.encode_stream(stream())]
+
+    assert collected_output[-1] == "data: [DONE]\n\n"
+    update_state_payload = json.loads(collected_output[0][6:-2])
+    assert update_state_payload == {"type": "update-state", "operations": operations}

python/assistant-stream/tests/test_cancellation.py

@@ -0,0 +1,178 @@
+import asyncio
+
+import pytest
+
+from assistant_stream import RunController, create_run
+
+
+@pytest.mark.anyio
+async def test_controller_exposes_cancel_signal():
+    observed: dict[str, object] = {}
+
+    async def run_callback(controller: RunController):
+        observed["cancel_signal"] = controller.cancelled_event
+        observed["initial_is_cancelled"] = controller.is_cancelled
+        controller.append_text("hello")
+
+    chunks = [chunk async for chunk in create_run(run_callback)]
+
+    assert len(chunks) == 1
+    assert chunks[0].type == "text-delta"
+    cancel_signal = observed["cancel_signal"]
+    assert callable(getattr(cancel_signal, "wait", None))
+    assert callable(getattr(cancel_signal, "is_set", None))
+    assert getattr(cancel_signal, "set", None) is None
+    assert observed["initial_is_cancelled"] is False
+
+
+@pytest.mark.anyio
+async def test_early_stream_close_sets_cancel_signal():
+    callback_done = asyncio.Event()
+    observed: dict[str, object] = {}
+
+    async def run_callback(controller: RunController):
+        controller.append_text("start")
+        for _ in range(100):
+            if controller.is_cancelled:
+                break
+            await asyncio.sleep(0.01)
+
+        observed["is_cancelled_after_close"] = controller.is_cancelled
+        callback_done.set()
+
+    stream = create_run(run_callback)
+    first_chunk = await anext(stream)
+    assert first_chunk.type == "text-delta"
+
+    await stream.aclose()
+    await asyncio.wait_for(callback_done.wait(), timeout=2)
+
+    assert observed["is_cancelled_after_close"] is True
+
+
+@pytest.mark.anyio
+async def test_early_stream_close_stops_background_task():
+    callback_done = asyncio.Event()
+
+    async def run_callback(controller: RunController):
+        controller.append_text("start")
+        try:
+            for _ in range(100):
+                if controller.is_cancelled:
+                    break
+                await asyncio.sleep(0.01)
+        finally:
+            callback_done.set()
+
+    stream = create_run(run_callback)
+    first_chunk = await anext(stream)
+    assert first_chunk.type == "text-delta"
+
+    await stream.aclose()
+    await asyncio.wait_for(callback_done.wait(), timeout=0.2)
+
+
+@pytest.mark.anyio
+async def test_normal_completion_does_not_set_cancel_signal():
+    observed: dict[str, object] = {}
+
+    async def run_callback(controller: RunController):
+        controller.append_text("done")
+        observed["is_cancelled_on_complete"] = controller.is_cancelled
+
+    chunks = [chunk async for chunk in create_run(run_callback)]
+
+    assert len(chunks) == 1
+    assert chunks[0].type == "text-delta"
+    assert observed["is_cancelled_on_complete"] is False
+
+
+@pytest.mark.anyio
+async def test_normal_completion_surfaces_callback_exception():
+    chunk_types: list[str] = []
+
+    async def run_callback(controller: RunController):
+        controller.append_text("start")
+        raise RuntimeError("boom")
+
+    with pytest.raises(RuntimeError, match="boom"):
+        async for chunk in create_run(run_callback):
+            chunk_types.append(chunk.type)
+
+    assert chunk_types == ["text-delta", "error"]
+
+
+@pytest.mark.anyio
+async def test_early_stream_close_forces_background_task_cancellation():
+    callback_cancelled = asyncio.Event()
+    callback_finished = asyncio.Event()
+
+    async def run_callback(controller: RunController):
+        controller.append_text("start")
+        try:
+            # Intentionally ignore `is_cancelled` to exercise forced cancellation.
+            while True:
+                await asyncio.sleep(1)
+        except asyncio.CancelledError:
+            callback_cancelled.set()
+            raise
+        finally:
+            callback_finished.set()
+
+    stream = create_run(run_callback)
+    first_chunk = await anext(stream)
+    assert first_chunk.type == "text-delta"
+
+    await stream.aclose()
+    await asyncio.wait_for(callback_cancelled.wait(), timeout=2)
+    await asyncio.wait_for(callback_finished.wait(), timeout=2)
+
+
+@pytest.mark.anyio
+async def test_early_stream_close_does_not_raise_callback_exception():
+    async def run_callback(controller: RunController):
+        controller.append_text("start")
+        await asyncio.sleep(0.01)
+        raise RuntimeError("boom")
+
+    stream = create_run(run_callback)
+    first_chunk = await anext(stream)
+    assert first_chunk.type == "text-delta"
+
+    await stream.aclose()
+
+
+@pytest.mark.anyio
+async def test_early_stream_close_does_not_swallow_close_task_cancellation():
+    callback_finished = asyncio.Event()
+    loop = asyncio.get_running_loop()
+
+    async def run_callback(controller: RunController):
+        controller.append_text("start")
+        deadline = loop.time() + 0.5
+        try:
+            while loop.time() < deadline:
+                try:
+                    await asyncio.sleep(0.01)
+                except asyncio.CancelledError:
+                    # Simulate non-cooperative callback behavior: ignore cancellation.
+                    continue
+        finally:
+            callback_finished.set()
+
+    stream = create_run(run_callback)
+    first_chunk = await anext(stream)
+    assert first_chunk.type == "text-delta"
+
+    close_task = asyncio.create_task(stream.aclose())
+    await asyncio.sleep(0)
+    close_task.cancel()
+
+    try:
+        done, _ = await asyncio.wait({close_task}, timeout=0.2)
+        assert close_task in done
+        assert close_task.cancelled()
+    finally:
+        await asyncio.wait_for(callback_finished.wait(), timeout=2)
+        if not close_task.done():
+            await asyncio.wait({close_task}, timeout=1)

python/assistant-stream/tests/test_data_stream.py

@@ -0,0 +1,17 @@
+import json
+
+from assistant_stream.assistant_stream_chunk import UpdateStateChunk
+from assistant_stream.serialization.data_stream import DataStreamEncoder
+
+
+def test_data_stream_encoder_update_state_shape() -> None:
+    encoder = DataStreamEncoder()
+    operations = [
+        {"type": "append-text", "path": ["messages", "0", "text"], "value": "hi"}
+    ]
+
+    encoded = encoder.encode_chunk(UpdateStateChunk(operations=operations))
+
+    assert encoded.startswith("aui-state:")
+    assert encoded.endswith("\n")
+    assert json.loads(encoded[len("aui-state:") :].strip()) == operations

python/assistant-stream/tests/test_langgraph.py

@@ -0,0 +1,223 @@
+"""Tests for the LangGraph integration (assistant_stream.modules.langgraph).
+
+These exercise append_langgraph_event against a real StateManager proxy, mirroring
+how the assistant-transport langgraph backend feeds langgraph's native
+``stream_mode=["messages", "updates"]`` output into ``controller.state``: the
+first argument is the state proxy, the event type is langgraph's stream mode name
+("messages" or "updates"), and a "messages" payload is a ``(message, metadata)``
+tuple carrying a single message or message chunk.
+"""
+
+from typing import Any
+
+import pytest
+from langchain_core.messages import AIMessage, AIMessageChunk, HumanMessage
+
+from assistant_stream.modules.langgraph import append_langgraph_event
+from assistant_stream.state_manager import StateManager
+
+
+def _manager(initial: Any) -> StateManager:
+    return StateManager(lambda _chunk: None, initial)
+
+
+def _manager_with_ops(initial: Any) -> tuple[StateManager, list[dict[str, Any]]]:
+    ops: list[dict[str, Any]] = []
+    return StateManager(lambda chunk: ops.extend(chunk.operations), initial), ops
+
+
+@pytest.mark.anyio
+async def test_appends_single_message_to_empty_state() -> None:
+    manager = _manager({})
+
+    append_langgraph_event(
+        manager.state, (), "messages", (HumanMessage(content="Hello", id="m1"), {})
+    )
+
+    messages = manager.state_data["messages"]
+    assert len(messages) == 1
+    assert messages[0]["content"] == "Hello"
+    assert messages[0]["id"] == "m1"
+    assert messages[0]["type"] == "human"
+
+
+@pytest.mark.anyio
+async def test_appends_messages_in_order() -> None:
+    manager = _manager({"messages": []})
+
+    append_langgraph_event(
+        manager.state, (), "messages", (HumanMessage(content="A", id="a"), {})
+    )
+    append_langgraph_event(
+        manager.state, (), "messages", (AIMessage(content="B", id="b"), {})
+    )
+
+    assert [m["content"] for m in manager.state_data["messages"]] == ["A", "B"]
+
+
+@pytest.mark.anyio
+async def test_merges_ai_message_chunks_by_id() -> None:
+    manager = _manager({"messages": []})
+
+    append_langgraph_event(
+        manager.state, (), "messages", (AIMessageChunk(content="Hello", id="m1"), {})
+    )
+    append_langgraph_event(
+        manager.state, (), "messages", (AIMessageChunk(content=" world", id="m1"), {})
+    )
+
+    messages = manager.state_data["messages"]
+    assert len(messages) == 1
+    assert messages[0]["content"] == "Hello world"
+    assert messages[0]["id"] == "m1"
+    assert messages[0]["type"] == "ai"
+
+
+@pytest.mark.anyio
+async def test_merging_ai_message_chunk_emits_content_append_text_delta() -> None:
+    manager, ops = _manager_with_ops({"messages": []})
+
+    append_langgraph_event(
+        manager.state, (), "messages", (AIMessageChunk(content="Hello", id="m1"), {})
+    )
+    manager.flush()
+    ops.clear()
+
+    append_langgraph_event(
+        manager.state, (), "messages", (AIMessageChunk(content=" world", id="m1"), {})
+    )
+    manager.flush()
+
+    assert manager.state_data["messages"][0]["content"] == "Hello world"
+    assert {
+        "type": "append-text",
+        "path": ["messages", "0", "content"],
+        "value": " world",
+    } in ops
+    assert not any(
+        op["type"] == "set" and op["path"] == ["messages", "0"] for op in ops
+    )
+
+
+@pytest.mark.anyio
+async def test_merging_ai_message_chunk_handles_plain_dict_messages() -> None:
+    state: dict[str, Any] = {"messages": []}
+
+    append_langgraph_event(
+        state, (), "messages", (AIMessageChunk(content="Hello", id="m1"), {})
+    )
+    append_langgraph_event(
+        state, (), "messages", (AIMessageChunk(content=" world", id="m1"), {})
+    )
+
+    assert state["messages"][0]["content"] == "Hello world"
+
+
+@pytest.mark.anyio
+async def test_merging_ai_message_chunk_patches_nested_tool_call_args() -> None:
+    manager, ops = _manager_with_ops({"messages": []})
+
+    append_langgraph_event(
+        manager.state,
+        (),
+        "messages",
+        (
+            AIMessageChunk(
+                content="",
+                id="m1",
+                tool_call_chunks=[
+                    {
+                        "name": "search",
+                        "args": '{"query"',
+                        "id": "call_1",
+                        "index": 0,
+                    }
+                ],
+            ),
+            {},
+        ),
+    )
+    manager.flush()
+    ops.clear()
+
+    append_langgraph_event(
+        manager.state,
+        (),
+        "messages",
+        (
+            AIMessageChunk(
+                content="",
+                id="m1",
+                tool_call_chunks=[
+                    {
+                        "name": None,
+                        "args": ':"docs"}',
+                        "id": None,
+                        "index": 0,
+                    }
+                ],
+            ),
+            {},
+        ),
+    )
+    manager.flush()
+
+    assert (
+        manager.state_data["messages"][0]["tool_call_chunks"][0]["args"]
+        == '{"query":"docs"}'
+    )
+    assert {
+        "type": "append-text",
+        "path": ["messages", "0", "tool_call_chunks", "0", "args"],
+        "value": ':"docs"}',
+    } in ops
+    assert not any(
+        op["type"] == "set" and op["path"] == ["messages", "0"] for op in ops
+    )
+
+
+@pytest.mark.anyio
+async def test_replaces_existing_message_with_same_id() -> None:
+    manager = _manager({"messages": [{"type": "human", "id": "m1", "content": "old"}]})
+
+    append_langgraph_event(
+        manager.state, (), "messages", (HumanMessage(content="new", id="m1"), {})
+    )
+
+    messages = manager.state_data["messages"]
+    assert len(messages) == 1
+    assert messages[0]["content"] == "new"
+
+
+@pytest.mark.anyio
+async def test_updates_event_writes_channels_onto_state() -> None:
+    manager = _manager({})
+
+    append_langgraph_event(
+        manager.state, (), "updates", {"agent": {"answer": "42", "messages": "ignored"}}
+    )
+
+    assert manager.state_data["answer"] == "42"
+    assert "messages" not in manager.state_data
+    assert "agent" not in manager.state_data
+
+
+@pytest.mark.anyio
+async def test_updates_event_skips_non_dict_nodes() -> None:
+    manager = _manager({})
+
+    append_langgraph_event(
+        manager.state, (), "updates", {"bad": "not-a-dict", "agent": {"answer": "42"}}
+    )
+
+    assert manager.state_data["answer"] == "42"
+    assert "bad" not in manager.state_data
+
+
+@pytest.mark.anyio
+async def test_unknown_event_type_is_ignored() -> None:
+    manager = _manager({"existing": "value"})
+
+    append_langgraph_event(manager.state, (), "custom", "anything")
+
+    assert manager.state_data == {"existing": "value"}

python/assistant-stream/tests/test_reasoning.py

@@ -0,0 +1,59 @@
+import asyncio
+import pytest
+from assistant_stream import create_run, RunController
+
+
+@pytest.mark.anyio
+async def test_append_reasoning():
+    """Test that append_reasoning works correctly."""
+    reasoning_chunks = []
+
+    async def collect_chunks(chunks):
+        async for chunk in chunks:
+            if chunk.type == "reasoning-delta":
+                reasoning_chunks.append(chunk)
+
+    async def run_callback(controller: RunController):
+        controller.append_reasoning("This is ")
+        controller.append_reasoning("reasoning content")
+
+    # Create the run and collect chunks
+    chunks = create_run(run_callback)
+    await collect_chunks(chunks)
+
+    # Verify the reasoning chunks
+    assert len(reasoning_chunks) == 2
+    assert reasoning_chunks[0].reasoning_delta == "This is "
+    assert reasoning_chunks[1].reasoning_delta == "reasoning content"
+
+
+@pytest.mark.anyio
+async def test_mixed_text_and_reasoning():
+    """Test that append_text and append_reasoning can be used together."""
+    collected_chunks = []
+
+    async def collect_chunks(chunks):
+        async for chunk in chunks:
+            if chunk.type in ["text-delta", "reasoning-delta"]:
+                collected_chunks.append(chunk)
+
+    async def run_callback(controller: RunController):
+        controller.append_text("This is text")
+        controller.append_reasoning("This is reasoning")
+        controller.append_text("More text")
+        controller.append_reasoning("More reasoning")
+
+    # Create the run and collect chunks
+    chunks = create_run(run_callback)
+    await collect_chunks(chunks)
+
+    # Verify the chunks
+    assert len(collected_chunks) == 4
+    assert collected_chunks[0].type == "text-delta"
+    assert collected_chunks[0].text_delta == "This is text"
+    assert collected_chunks[1].type == "reasoning-delta"
+    assert collected_chunks[1].reasoning_delta == "This is reasoning"
+    assert collected_chunks[2].type == "text-delta"
+    assert collected_chunks[2].text_delta == "More text"
+    assert collected_chunks[3].type == "reasoning-delta"
+    assert collected_chunks[3].reasoning_delta == "More reasoning"

python/assistant-stream/tests/test_resumable_context.py

@@ -0,0 +1,298 @@
+from __future__ import annotations
+
+import asyncio
+from collections.abc import AsyncIterator
+
+import pytest
+
+from assistant_stream.resumable import (
+    ResumableStreamError,
+    create_in_memory_resumable_stream_store,
+    create_resumable_stream_context,
+)
+
+
+def _bytes(s: str) -> bytes:
+    return s.encode("utf-8")
+
+
+async def _collect(stream: AsyncIterator[bytes]) -> str:
+    out = bytearray()
+    async for chunk in stream:
+        out.extend(chunk)
+    return out.decode("utf-8")
+
+
+def _make_string_stream(parts: list[str]) -> AsyncIterator[bytes]:
+    async def gen() -> AsyncIterator[bytes]:
+        for part in parts:
+            yield _bytes(part)
+            await asyncio.sleep(0)
+
+    return gen()
+
+
+@pytest.mark.anyio
+async def test_producer_receives_full_byte_stream() -> None:
+    ctx = create_resumable_stream_context(
+        store=create_in_memory_resumable_stream_store()
+    )
+    stream = await ctx.run("a", lambda: _make_string_stream(["hello ", "world"]))
+    assert await _collect(stream) == "hello world"
+
+
+@pytest.mark.anyio
+async def test_second_caller_is_consumer_with_identical_bytes() -> None:
+    store = create_in_memory_resumable_stream_store()
+    ctx = create_resumable_stream_context(store=store)
+    make_stream_calls = 0
+
+    def make_producer() -> AsyncIterator[bytes]:
+        nonlocal make_stream_calls
+        make_stream_calls += 1
+        return _make_string_stream(["one ", "two ", "three"])
+
+    def make_unused() -> AsyncIterator[bytes]:
+        nonlocal make_stream_calls
+        make_stream_calls += 1
+        return _make_string_stream(["should-not-run"])
+
+    producer_stream = await ctx.run("a", make_producer)
+    consumer_stream = await ctx.run("a", make_unused)
+
+    a, b = await asyncio.gather(
+        _collect(producer_stream), _collect(consumer_stream)
+    )
+    assert a == "one two three"
+    assert b == "one two three"
+    assert make_stream_calls == 1
+
+
+@pytest.mark.anyio
+async def test_late_consumer_after_done_replays_via_resume() -> None:
+    ctx = create_resumable_stream_context(
+        store=create_in_memory_resumable_stream_store()
+    )
+    producer = await ctx.run(
+        "a", lambda: _make_string_stream(["alpha", "beta", "gamma"])
+    )
+    assert await _collect(producer) == "alphabetagamma"
+
+    replay = await ctx.resume("a")
+    assert replay is not None
+    assert await _collect(replay) == "alphabetagamma"
+
+
+@pytest.mark.anyio
+async def test_resume_returns_none_for_missing() -> None:
+    ctx = create_resumable_stream_context(
+        store=create_in_memory_resumable_stream_store()
+    )
+    assert await ctx.resume("nope") is None
+
+
+@pytest.mark.anyio
+async def test_require_resume_raises_missing() -> None:
+    ctx = create_resumable_stream_context(
+        store=create_in_memory_resumable_stream_store()
+    )
+    with pytest.raises(ResumableStreamError) as exc:
+        await ctx.require_resume("nope")
+    assert exc.value.code == "missing"
+
+
+@pytest.mark.anyio
+async def test_require_resume_returns_replay_when_exists() -> None:
+    ctx = create_resumable_stream_context(
+        store=create_in_memory_resumable_stream_store()
+    )
+    producer = await ctx.run("a", lambda: _make_string_stream(["hi"]))
+    assert await _collect(producer) == "hi"
+
+    replay = await ctx.require_resume("a")
+    assert await _collect(replay) == "hi"
+
+
+@pytest.mark.anyio
+async def test_status_tracks_lifecycle() -> None:
+    ctx = create_resumable_stream_context(
+        store=create_in_memory_resumable_stream_store()
+    )
+    assert await ctx.status("a") == "missing"
+    stream = await ctx.run("a", lambda: _make_string_stream(["x"]))
+    await _collect(stream)
+    assert await ctx.status("a") == "done"
+
+
+@pytest.mark.anyio
+async def test_delete_removes_stream_state() -> None:
+    ctx = create_resumable_stream_context(
+        store=create_in_memory_resumable_stream_store()
+    )
+    stream = await ctx.run("a", lambda: _make_string_stream(["x"]))
+    await _collect(stream)
+    assert await ctx.status("a") == "done"
+    await ctx.delete("a")
+    assert await ctx.status("a") == "missing"
+
+
+@pytest.mark.anyio
+async def test_producer_keeps_writing_after_consumer_closes_early() -> None:
+    store = create_in_memory_resumable_stream_store()
+    ctx = create_resumable_stream_context(store=store)
+
+    producer_emitted = 0
+
+    async def slow_stream() -> AsyncIterator[bytes]:
+        nonlocal producer_emitted
+        for i in range(5):
+            await asyncio.sleep(0.01)
+            yield _bytes(f"chunk{i};")
+            producer_emitted += 1
+
+    stream = await ctx.run("a", lambda: slow_stream())
+    first = await anext(stream)
+    assert first == _bytes("chunk0;")
+    await stream.aclose()
+
+    while producer_emitted < 5:
+        await asyncio.sleep(0.01)
+    while await ctx.status("a") == "streaming":
+        await asyncio.sleep(0.01)
+
+    replay = await ctx.resume("a")
+    assert replay is not None
+    assert await _collect(replay) == "chunk0;chunk1;chunk2;chunk3;chunk4;"
+
+
+@pytest.mark.anyio
+async def test_propagates_producer_errors_to_consumers() -> None:
+    ctx = create_resumable_stream_context(
+        store=create_in_memory_resumable_stream_store()
+    )
+
+    async def failing() -> AsyncIterator[bytes]:
+        yield _bytes("partial;")
+        raise Exception("oops")
+
+    stream = await ctx.run("a", lambda: failing())
+    with pytest.raises(Exception, match="oops"):
+        await _collect(stream)
+    assert await ctx.status("a") == "error"
+
+
+@pytest.mark.anyio
+async def test_wait_until_receives_the_task() -> None:
+    tasks: list[asyncio.Task[None]] = []
+    ctx = create_resumable_stream_context(
+        store=create_in_memory_resumable_stream_store(),
+        wait_until=lambda t: tasks.append(t),
+    )
+    stream = await ctx.run("a", lambda: _make_string_stream(["x"]))
+    assert await _collect(stream) == "x"
+    assert len(tasks) == 1
+    await asyncio.gather(*tasks)
+
+
+@pytest.mark.anyio
+async def test_on_acquire_fires_for_producer_and_consumer() -> None:
+    calls: list[dict[str, str]] = []
+    ctx = create_resumable_stream_context(
+        store=create_in_memory_resumable_stream_store(),
+        on_acquire=lambda id, role: calls.append({"id": id, "role": role}),
+    )
+    producer = await ctx.run("a", lambda: _make_string_stream(["x"]))
+    consumer = await ctx.run("a", lambda: _make_string_stream(["unused"]))
+    await asyncio.gather(_collect(producer), _collect(consumer))
+    assert calls == [
+        {"id": "a", "role": "producer"},
+        {"id": "a", "role": "consumer"},
+    ]
+
+
+@pytest.mark.anyio
+async def test_on_append_fires_per_chunk_with_byte_length() -> None:
+    calls: list[dict[str, object]] = []
+    ctx = create_resumable_stream_context(
+        store=create_in_memory_resumable_stream_store(),
+        on_append=lambda id, n: calls.append({"id": id, "byteLength": n}),
+    )
+    stream = await ctx.run("a", lambda: _make_string_stream(["ab", "cde"]))
+    assert await _collect(stream) == "abcde"
+    assert calls == [
+        {"id": "a", "byteLength": 2},
+        {"id": "a", "byteLength": 3},
+    ]
+
+
+@pytest.mark.anyio
+async def test_on_finalize_fires_on_success() -> None:
+    calls: list[dict[str, object]] = []
+    ctx = create_resumable_stream_context(
+        store=create_in_memory_resumable_stream_store(),
+        on_finalize=lambda id, status, error: calls.append(
+            {"id": id, "status": status, "error": error}
+        ),
+    )
+    stream = await ctx.run("a", lambda: _make_string_stream(["x"]))
+    await _collect(stream)
+    assert calls == [{"id": "a", "status": "done", "error": None}]
+
+
+@pytest.mark.anyio
+async def test_on_finalize_and_on_error_fire_on_failure() -> None:
+    finalize_calls: list[dict[str, object]] = []
+    errors: list[dict[str, object]] = []
+    ctx = create_resumable_stream_context(
+        store=create_in_memory_resumable_stream_store(),
+        on_finalize=lambda id, status, error: finalize_calls.append(
+            {"id": id, "status": status, "error": error}
+        ),
+        on_error=lambda id, error: errors.append({"id": id, "error": error}),
+    )
+
+    async def failing() -> AsyncIterator[bytes]:
+        raise Exception("boom")
+        yield b""  # pragma: no cover
+
+    stream = await ctx.run("a", lambda: failing())
+    with pytest.raises(Exception, match="boom"):
+        await _collect(stream)
+    assert finalize_calls == [{"id": "a", "status": "error", "error": "boom"}]
+    assert len(errors) == 1
+    assert errors[0]["id"] == "a"
+    assert str(errors[0]["error"]) == "boom"
+
+
+@pytest.mark.anyio
+async def test_raising_on_acquire_does_not_break_producer_pipeline() -> None:
+    def boom_acquire(_id: str, _role: str) -> None:
+        raise RuntimeError("hook-acquire-boom")
+
+    ctx = create_resumable_stream_context(
+        store=create_in_memory_resumable_stream_store(),
+        on_acquire=boom_acquire,
+    )
+    stream = await ctx.run("a", lambda: _make_string_stream(["hello ", "world"]))
+    assert await _collect(stream) == "hello world"
+    assert await ctx.status("a") == "done"
+
+
+@pytest.mark.anyio
+async def test_raising_on_error_still_finalizes_error_status() -> None:
+    def boom_error(_id: str, _err: object) -> None:
+        raise RuntimeError("hook-error-boom")
+
+    ctx = create_resumable_stream_context(
+        store=create_in_memory_resumable_stream_store(),
+        on_error=boom_error,
+    )
+
+    async def failing() -> AsyncIterator[bytes]:
+        yield _bytes("partial;")
+        raise Exception("producer-failed")
+
+    stream = await ctx.run("a", lambda: failing())
+    with pytest.raises(Exception, match="producer-failed"):
+        await _collect(stream)
+    assert await ctx.status("a") == "error"

python/assistant-stream/tests/test_resumable_in_memory.py

@@ -0,0 +1,308 @@
+from __future__ import annotations
+
+import asyncio
+
+import pytest
+
+from assistant_stream.resumable import (
+    ResumableStreamError,
+    create_in_memory_resumable_stream_store,
+)
+from assistant_stream.resumable.types import ResumableStreamEntry
+
+
+def _bytes(s: str) -> bytes:
+    return s.encode("utf-8")
+
+
+def _decode(chunk: bytes) -> str:
+    return chunk.decode("utf-8")
+
+
+async def _drain(iter_) -> list[str]:
+    out: list[str] = []
+    async for entry in iter_:
+        out.append(_decode(entry.chunk))
+    return out
+
+
+@pytest.mark.anyio
+async def test_elects_exactly_one_producer_per_stream_id() -> None:
+    store = create_in_memory_resumable_stream_store()
+    first = await store.acquire("a")
+    second = await store.acquire("a")
+    third = await store.acquire("a")
+    assert first == "producer"
+    assert second == "consumer"
+    assert third == "consumer"
+
+
+@pytest.mark.anyio
+async def test_stream_id_with_trailing_newline_is_invalid() -> None:
+    store = create_in_memory_resumable_stream_store()
+    with pytest.raises(ResumableStreamError) as exc:
+        await store.acquire("valid-id\n")
+    assert exc.value.code == "invalid-id"
+
+
+@pytest.mark.anyio
+async def test_post_finalize_acquire_is_consumer() -> None:
+    store = create_in_memory_resumable_stream_store()
+    assert await store.acquire("a") == "producer"
+    await store.finalize("a", "done")
+    assert await store.acquire("a") == "consumer"
+
+
+@pytest.mark.anyio
+async def test_isolates_streams_by_id() -> None:
+    store = create_in_memory_resumable_stream_store()
+    assert await store.acquire("a") == "producer"
+    assert await store.acquire("b") == "producer"
+
+
+@pytest.mark.anyio
+async def test_replays_buffered_entries_and_tails_until_finalize() -> None:
+    store = create_in_memory_resumable_stream_store()
+    await store.acquire("a")
+    await store.append("a", _bytes("hello "))
+    await store.append("a", _bytes("world"))
+
+    signal = asyncio.Event()
+    collected: list[str] = []
+
+    async def reading() -> None:
+        async for entry in store.read("a", "", signal):
+            collected.append(_decode(entry.chunk))
+            if len(collected) == 3:
+                await store.finalize("a", "done")
+
+    task = asyncio.create_task(reading())
+    await store.append("a", _bytes("!"))
+    await task
+    assert collected == ["hello ", "world", "!"]
+
+
+@pytest.mark.anyio
+async def test_status_transitions_missing_streaming_done() -> None:
+    store = create_in_memory_resumable_stream_store()
+    assert await store.status("a") == "missing"
+    await store.acquire("a")
+    assert await store.status("a") == "streaming"
+    await store.finalize("a", "done")
+    assert await store.status("a") == "done"
+
+
+@pytest.mark.anyio
+async def test_status_reports_error_after_error_finalize() -> None:
+    store = create_in_memory_resumable_stream_store()
+    await store.acquire("a")
+    await store.finalize("a", "error", "boom")
+    assert await store.status("a") == "error"
+
+
+@pytest.mark.anyio
+async def test_read_throws_after_error_finalize_after_draining() -> None:
+    store = create_in_memory_resumable_stream_store()
+    await store.acquire("a")
+    await store.append("a", _bytes("partial"))
+    await store.finalize("a", "error", "boom")
+
+    signal = asyncio.Event()
+    seen: list[str] = []
+    with pytest.raises(Exception, match="boom"):
+        async for entry in store.read("a", "", signal):
+            seen.append(_decode(entry.chunk))
+    assert seen == ["partial"]
+
+
+@pytest.mark.anyio
+async def test_late_consumer_after_done_replays_everything() -> None:
+    store = create_in_memory_resumable_stream_store()
+    await store.acquire("a")
+    await store.append("a", _bytes("a"))
+    await store.append("a", _bytes("b"))
+    await store.append("a", _bytes("c"))
+    await store.finalize("a", "done")
+
+    signal = asyncio.Event()
+    assert await _drain(store.read("a", "", signal)) == ["a", "b", "c"]
+
+
+@pytest.mark.anyio
+async def test_cursor_advances_and_skips_already_seen() -> None:
+    store = create_in_memory_resumable_stream_store()
+    await store.acquire("a")
+    await store.append("a", _bytes("1"))
+    await store.append("a", _bytes("2"))
+    await store.append("a", _bytes("3"))
+    await store.finalize("a", "done")
+
+    signal = asyncio.Event()
+    seen: list[ResumableStreamEntry] = []
+    async for entry in store.read("a", "", signal):
+        seen.append(entry)
+    assert [_decode(s.chunk) for s in seen] == ["1", "2", "3"]
+
+    after_first = seen[0].cursor
+    assert await _drain(store.read("a", after_first, signal)) == ["2", "3"]
+
+
+@pytest.mark.anyio
+async def test_signal_set_terminates_read_without_raising() -> None:
+    store = create_in_memory_resumable_stream_store()
+    await store.acquire("a")
+    signal = asyncio.Event()
+    collected: list[str] = []
+
+    async def reading() -> None:
+        async for entry in store.read("a", "", signal):
+            collected.append(_decode(entry.chunk))
+
+    task = asyncio.create_task(reading())
+    await store.append("a", _bytes("x"))
+    await asyncio.sleep(0.01)
+    signal.set()
+    await task
+    assert collected == ["x"]
+
+
+@pytest.mark.anyio
+async def test_multiple_consumers_read_concurrently() -> None:
+    store = create_in_memory_resumable_stream_store()
+    await store.acquire("a")
+    signal = asyncio.Event()
+    a = asyncio.create_task(_drain(store.read("a", "", signal)))
+    b = asyncio.create_task(_drain(store.read("a", "", signal)))
+    await store.append("a", _bytes("x"))
+    await store.append("a", _bytes("y"))
+    await store.finalize("a", "done")
+    assert await a == ["x", "y"]
+    assert await b == ["x", "y"]
+
+
+@pytest.mark.anyio
+async def test_delete_ends_in_flight_reads_and_status_missing() -> None:
+    store = create_in_memory_resumable_stream_store()
+    await store.acquire("a")
+    signal = asyncio.Event()
+    reading = asyncio.create_task(_drain(store.read("a", "", signal)))
+    await store.append("a", _bytes("x"))
+    await asyncio.sleep(0)
+    await store.delete("a")
+    assert await reading == ["x"]
+    assert await store.status("a") == "missing"
+
+
+@pytest.mark.anyio
+async def test_expired_streams_evicted_on_next_access() -> None:
+    now = 1_000.0
+
+    def clock() -> float:
+        return now
+
+    store = create_in_memory_resumable_stream_store(default_ttl_ms=100, now=clock)
+    await store.acquire("a")
+    await store.append("a", _bytes("hi"))
+    assert await store.status("a") == "streaming"
+    now += 200
+    assert await store.status("a") == "missing"
+
+
+@pytest.mark.anyio
+async def test_appending_refreshes_ttl() -> None:
+    now = 1_000.0
+
+    def clock() -> float:
+        return now
+
+    store = create_in_memory_resumable_stream_store(default_ttl_ms=100, now=clock)
+    await store.acquire("a")
+    now += 80
+    await store.append("a", _bytes("x"))
+    now += 80
+    assert await store.status("a") == "streaming"
+
+
+@pytest.mark.anyio
+async def test_rejects_append_on_finalized_stream() -> None:
+    store = create_in_memory_resumable_stream_store()
+    await store.acquire("a")
+    await store.finalize("a", "done")
+    with pytest.raises(ResumableStreamError, match="already finalized") as exc:
+        await store.append("a", _bytes("late"))
+    assert exc.value.code == "finalized"
+
+
+@pytest.mark.anyio
+async def test_rejects_append_on_missing_stream() -> None:
+    store = create_in_memory_resumable_stream_store()
+    with pytest.raises(Exception, match="Stream not found"):
+        await store.append("a", _bytes("x"))
+
+
+@pytest.mark.anyio
+async def test_finalize_is_idempotent() -> None:
+    store = create_in_memory_resumable_stream_store()
+    await store.acquire("a")
+    await store.finalize("a", "done")
+    await store.finalize("a", "done")
+    assert await store.status("a") == "done"
+
+
+@pytest.mark.anyio
+async def test_rejects_append_when_chunk_exceeds_max_chunk_bytes() -> None:
+    store = create_in_memory_resumable_stream_store(max_chunk_bytes=4)
+    await store.acquire("a")
+    with pytest.raises(Exception, match="Chunk exceeds maxChunkBytes: 5"):
+        await store.append("a", _bytes("hello"))
+    await store.append("a", _bytes("ok"))
+    assert await store.status("a") == "streaming"
+
+
+@pytest.mark.anyio
+async def test_rejects_append_when_stream_reaches_max_entries() -> None:
+    store = create_in_memory_resumable_stream_store(max_entries_per_stream=2)
+    await store.acquire("a")
+    await store.append("a", _bytes("1"))
+    await store.append("a", _bytes("2"))
+    with pytest.raises(Exception, match="Stream exceeded maxEntriesPerStream: a"):
+        await store.append("a", _bytes("3"))
+
+
+@pytest.mark.anyio
+async def test_rejects_acquire_when_max_streams_exceeded() -> None:
+    store = create_in_memory_resumable_stream_store(max_streams=2)
+    await store.acquire("a")
+    await store.acquire("b")
+    with pytest.raises(Exception, match="maxStreams exceeded"):
+        await store.acquire("c")
+    assert await store.acquire("a") == "consumer"
+
+
+@pytest.mark.anyio
+async def test_gc_sweeper_evicts_expired_streams() -> None:
+    now = 1_000.0
+
+    def clock() -> float:
+        return now
+
+    store = create_in_memory_resumable_stream_store(
+        default_ttl_ms=100,
+        gc_interval_ms=50,
+        now=clock,
+    )
+    await store.acquire("a")
+    await store.append("a", _bytes("hi"))
+    now += 200
+    await asyncio.sleep(0.08)
+    assert await store.status("a") == "missing"
+    store.dispose()
+
+
+@pytest.mark.anyio
+async def test_dispose_clears_gc_and_is_noop_without_gc() -> None:
+    store = create_in_memory_resumable_stream_store(gc_interval_ms=50)
+    await store.acquire("a")
+    store.dispose()
+    store2 = create_in_memory_resumable_stream_store()
+    store2.dispose()

python/assistant-stream/tests/test_resumable_redis.py

@@ -0,0 +1,164 @@
+from __future__ import annotations
+
+import asyncio
+import os
+import time
+import uuid
+
+import pytest
+
+REDIS_URL = os.environ.get("REDIS_URL", "redis://127.0.0.1:6379")
+REDIS_TESTS_DISABLED = (
+    os.environ.get("REDIS_URL") is None and os.environ.get("REDIS_TESTS") != "1"
+)
+
+pytestmark = pytest.mark.skipif(
+    REDIS_TESTS_DISABLED,
+    reason="set REDIS_URL or REDIS_TESTS=1 to run redis adapter tests",
+)
+
+
+def _bytes(s: str) -> bytes:
+    return s.encode("utf-8")
+
+
+def _decode(chunk: bytes) -> str:
+    return chunk.decode("utf-8")
+
+
+@pytest.fixture
+async def redis_store():
+    import redis.asyncio as redis
+
+    from assistant_stream.resumable import create_redis_resumable_stream_store
+
+    client = redis.from_url(REDIS_URL)
+    key_prefix = f"aui:resumable:test:{int(time.time() * 1000)}:{uuid.uuid4().hex[:8]}"
+    store = create_redis_resumable_stream_store(
+        client,
+        key_prefix=key_prefix,
+        poll_interval_ms=25,
+    )
+    try:
+        yield store, client, key_prefix
+    finally:
+        keys = [key async for key in client.scan_iter(match=f"{key_prefix}:*")]
+        if keys:
+            await client.delete(*keys)
+        await client.aclose()
+
+
+@pytest.mark.anyio
+async def test_elects_exactly_one_producer(redis_store) -> None:
+    store, _, _ = redis_store
+    stream_id = f"id-{uuid.uuid4().hex}"
+    first = await store.acquire(stream_id)
+    second = await store.acquire(stream_id)
+    assert first == "producer"
+    assert second == "consumer"
+
+
+@pytest.mark.anyio
+async def test_replays_buffered_and_tails_until_done(redis_store) -> None:
+    store, _, _ = redis_store
+    stream_id = f"id-{uuid.uuid4().hex}"
+    await store.acquire(stream_id)
+    await store.append(stream_id, _bytes("hello"))
+    await store.append(stream_id, _bytes(" world"))
+
+    signal = asyncio.Event()
+    collected: list[str] = []
+
+    async def reading() -> None:
+        async for entry in store.read(stream_id, "", signal):
+            collected.append(_decode(entry.chunk))
+            if len(collected) == 3:
+                await store.finalize(stream_id, "done")
+
+    task = asyncio.create_task(reading())
+    await asyncio.sleep(0.05)
+    await store.append(stream_id, _bytes("!"))
+    await task
+    assert "".join(collected) == "hello world!"
+
+
+@pytest.mark.anyio
+async def test_error_finalize_raises_after_draining(redis_store) -> None:
+    store, _, _ = redis_store
+    stream_id = f"id-{uuid.uuid4().hex}"
+    await store.acquire(stream_id)
+    await store.append(stream_id, _bytes("partial"))
+    await store.finalize(stream_id, "error", "boom")
+
+    signal = asyncio.Event()
+    seen: list[str] = []
+    with pytest.raises(Exception, match="boom"):
+        async for entry in store.read(stream_id, "", signal):
+            seen.append(_decode(entry.chunk))
+    assert seen == ["partial"]
+
+
+@pytest.mark.anyio
+async def test_cursor_skip(redis_store) -> None:
+    store, _, _ = redis_store
+    stream_id = f"id-{uuid.uuid4().hex}"
+    await store.acquire(stream_id)
+    await store.append(stream_id, _bytes("1"))
+    await store.append(stream_id, _bytes("2"))
+    await store.append(stream_id, _bytes("3"))
+    await store.finalize(stream_id, "done")
+
+    signal = asyncio.Event()
+    seen: list[tuple[str, str]] = []
+    async for entry in store.read(stream_id, "", signal):
+        seen.append((entry.cursor, _decode(entry.chunk)))
+    assert [t for _, t in seen] == ["1", "2", "3"]
+
+    out: list[str] = []
+    async for entry in store.read(stream_id, seen[0][0], signal):
+        out.append(_decode(entry.chunk))
+    assert out == ["2", "3"]
+
+
+@pytest.mark.anyio
+async def test_delete_removes_stream(redis_store) -> None:
+    store, _, _ = redis_store
+    stream_id = f"id-{uuid.uuid4().hex}"
+    await store.acquire(stream_id)
+    await store.append(stream_id, _bytes("x"))
+    assert await store.status(stream_id) == "streaming"
+    await store.delete(stream_id)
+    assert await store.status(stream_id) == "missing"
+
+
+@pytest.mark.anyio
+async def test_binary_round_trip_0_to_255(redis_store) -> None:
+    store, _, _ = redis_store
+    stream_id = f"id-{uuid.uuid4().hex}"
+    await store.acquire(stream_id)
+
+    producer = bytes(i & 0xFF for i in range(512))
+    half = len(producer) // 2
+    await store.append(stream_id, producer[:half])
+    await store.append(stream_id, producer[half:])
+    await store.finalize(stream_id, "done")
+
+    signal = asyncio.Event()
+    replayed = bytearray()
+    async for entry in store.read(stream_id, "", signal):
+        replayed.extend(entry.chunk)
+    assert bytes(replayed) == producer
+
+
+@pytest.mark.anyio
+async def test_decode_responses_true_is_rejected() -> None:
+    import redis.asyncio as redis
+
+    from assistant_stream.resumable import create_redis_resumable_stream_store
+
+    client = redis.from_url(REDIS_URL, decode_responses=True)
+    try:
+        with pytest.raises(ValueError, match="decode_responses"):
+            create_redis_resumable_stream_store(client)
+    finally:
+        await client.aclose()

python/assistant-stream/tests/test_resumable_response.py

@@ -0,0 +1,177 @@
+from __future__ import annotations
+
+import pytest
+from starlette.responses import StreamingResponse
+
+from assistant_stream.create_run import RunController
+from assistant_stream.resumable import (
+    RESUMABLE_STREAM_ID_HEADER,
+    create_in_memory_resumable_stream_store,
+    create_resumable_assistant_stream_response,
+    create_resumable_stream_context,
+    create_resume_assistant_stream_response,
+)
+from assistant_stream.serialization.assistant_transport import (
+    AssistantTransportEncoder,
+)
+
+
+async def _collect_body(response: StreamingResponse) -> bytes:
+    parts: list[bytes] = []
+    async for chunk in response.body_iterator:
+        if isinstance(chunk, str):
+            parts.append(chunk.encode("utf-8"))
+        else:
+            parts.append(bytes(chunk))
+    return b"".join(parts)
+
+
+@pytest.mark.anyio
+async def test_producer_response_carries_id_header_and_media_type() -> None:
+    ctx = create_resumable_stream_context(
+        store=create_in_memory_resumable_stream_store()
+    )
+
+    async def callback(controller: RunController) -> None:
+        controller.append_text("hello ")
+        controller.append_text("world")
+
+    response = await create_resumable_assistant_stream_response(
+        context=ctx,
+        stream_id="s1",
+        callback=callback,
+    )
+    assert isinstance(response, StreamingResponse)
+    assert response.headers[RESUMABLE_STREAM_ID_HEADER] == "s1"
+    assert "text/plain" in response.media_type
+    body = await _collect_body(response)
+    assert b"hello " in body
+    assert b"world" in body
+
+
+@pytest.mark.anyio
+async def test_resume_replays_byte_for_byte() -> None:
+    store = create_in_memory_resumable_stream_store()
+    ctx = create_resumable_stream_context(store=store)
+
+    async def callback(controller: RunController) -> None:
+        controller.append_text("alpha")
+        tool = await controller.add_tool_call("echo", "t1")
+        tool.append_args_text('{"v": 1}')
+        tool.set_response({"ok": True})
+
+    first = await create_resumable_assistant_stream_response(
+        context=ctx,
+        stream_id="s1",
+        callback=callback,
+    )
+    first_bytes = await _collect_body(first)
+
+    second = await create_resume_assistant_stream_response(
+        context=ctx,
+        stream_id="s1",
+    )
+    second_bytes = await _collect_body(second)
+
+    assert second_bytes == first_bytes
+    assert second.headers[RESUMABLE_STREAM_ID_HEADER] == "s1"
+
+
+@pytest.mark.anyio
+async def test_resume_returns_404_json_when_missing() -> None:
+    ctx = create_resumable_stream_context(
+        store=create_in_memory_resumable_stream_store()
+    )
+    response = await create_resume_assistant_stream_response(
+        context=ctx,
+        stream_id="missing",
+    )
+    assert response.status_code == 404
+    body = response.body
+    assert b"stream not found" in body
+
+
+@pytest.mark.anyio
+async def test_resume_returns_404_json_for_invalid_stream_id() -> None:
+    ctx = create_resumable_stream_context(
+        store=create_in_memory_resumable_stream_store()
+    )
+    response = await create_resume_assistant_stream_response(
+        context=ctx,
+        stream_id="x" * 300,
+    )
+    assert response.status_code == 404
+    body = response.body
+    assert b"stream not found" in body
+
+
+@pytest.mark.anyio
+async def test_assistant_transport_encoder_round_trips() -> None:
+    ctx = create_resumable_stream_context(
+        store=create_in_memory_resumable_stream_store()
+    )
+
+    async def callback(controller: RunController) -> None:
+        controller.append_text("hi")
+
+    response = await create_resumable_assistant_stream_response(
+        context=ctx,
+        stream_id="s1",
+        encoder=AssistantTransportEncoder(),
+        callback=callback,
+    )
+    assert response.media_type == "text/event-stream"
+    text = (await _collect_body(response)).decode("utf-8")
+    assert "text-delta" in text
+    assert "[DONE]" in text
+
+    resume = await create_resume_assistant_stream_response(
+        context=ctx,
+        stream_id="s1",
+        encoder=AssistantTransportEncoder(),
+    )
+    resume_text = (await _collect_body(resume)).decode("utf-8")
+    assert resume_text == text
+
+
+@pytest.mark.anyio
+async def test_user_headers_merge_but_cannot_override_id() -> None:
+    ctx = create_resumable_stream_context(
+        store=create_in_memory_resumable_stream_store()
+    )
+
+    async def callback(controller: RunController) -> None:
+        controller.append_text("hi")
+
+    response = await create_resumable_assistant_stream_response(
+        context=ctx,
+        stream_id="s1",
+        headers={
+            "Cache-Control": "private, max-age=0",
+            RESUMABLE_STREAM_ID_HEADER: "spoofed",
+        },
+        callback=callback,
+    )
+    assert response.headers["cache-control"] == "private, max-age=0"
+    assert response.headers[RESUMABLE_STREAM_ID_HEADER] == "s1"
+    await response.body_iterator.aclose()
+
+
+@pytest.mark.anyio
+async def test_differently_cased_user_id_header_is_filtered() -> None:
+    ctx = create_resumable_stream_context(
+        store=create_in_memory_resumable_stream_store()
+    )
+
+    async def callback(controller: RunController) -> None:
+        controller.append_text("hi")
+
+    response = await create_resumable_assistant_stream_response(
+        context=ctx,
+        stream_id="s1",
+        headers={"X-RESUMABLE-STREAM-ID": "spoof"},
+        callback=callback,
+    )
+    assert response.headers[RESUMABLE_STREAM_ID_HEADER] == "s1"
+    assert "spoof" not in response.headers.values()
+    await response.body_iterator.aclose()

python/assistant-stream/tests/test_state_updates.py

@@ -0,0 +1,164 @@
+from typing import Any
+
+import pytest
+
+from assistant_stream import RunController, create_run
+from assistant_stream.state_manager import StateManager
+
+
+@pytest.mark.anyio
+async def test_append_text_helper_emits_append_text_op() -> None:
+    ops: list[dict[str, Any]] = []
+    manager = StateManager(
+        lambda chunk: ops.extend(chunk.operations),
+        {"messages": [{"text": ""}]},
+    )
+
+    manager.append_text(["messages", "0", "text"], "hi")
+    manager.flush()
+
+    assert ops == [{"type": "append-text", "path": ["messages", "0", "text"], "value": "hi"}]
+
+
+@pytest.mark.anyio
+async def test_nested_plus_equals_on_string_emits_append_text() -> None:
+    ops: list[dict[str, Any]] = []
+    manager = StateManager(
+        lambda chunk: ops.extend(chunk.operations),
+        {"messages": [{"text": "h"}]},
+    )
+
+    manager.state["messages"][0]["text"] += "i"
+    manager.flush()
+
+    assert ops == [
+        {"type": "append-text", "path": ["messages", "0", "text"], "value": "i"}
+    ]
+
+
+@pytest.mark.anyio
+async def test_nested_plus_equals_multiple_operations_accumulate() -> None:
+    ops: list[dict[str, Any]] = []
+    manager = StateManager(
+        lambda chunk: ops.extend(chunk.operations),
+        {"messages": [{"text": ""}]},
+    )
+
+    manager.state["messages"][0]["text"] += "Hel"
+    manager.state["messages"][0]["text"] += "lo"
+    manager.flush()
+
+    assert ops == [
+        {"type": "set", "path": ["messages", "0", "text"], "value": "Hel"},
+        {"type": "append-text", "path": ["messages", "0", "text"], "value": "lo"},
+    ]
+    assert manager.state_data["messages"][0]["text"] == "Hello"
+
+
+@pytest.mark.anyio
+async def test_streaming_path_emits_append_text_deltas_not_set() -> None:
+    async def run_callback(controller: RunController):
+        controller.append_state_text(["messages", 0, "text"], "Hel")
+        controller.append_state_text(["messages", 0, "text"], "lo")
+
+    chunks = [
+        chunk
+        async for chunk in create_run(
+            run_callback, state={"messages": [{"text": ""}]}
+        )
+    ]
+    operations = [
+        operation
+        for chunk in chunks
+        if chunk.type == "update-state"
+        for operation in chunk.operations
+    ]
+
+    assert operations == [
+        {"type": "append-text", "path": ["messages", "0", "text"], "value": "Hel"},
+        {"type": "append-text", "path": ["messages", "0", "text"], "value": "lo"},
+    ]
+
+
+@pytest.mark.anyio
+async def test_string_assignment_non_extension_emits_set() -> None:
+    ops: list[dict[str, Any]] = []
+    manager = StateManager(
+        lambda chunk: ops.extend(chunk.operations),
+        {"messages": [{"text": "hello"}]},
+    )
+
+    manager.state["messages"][0]["text"] = "goodbye"
+    manager.flush()
+
+    assert ops == [{"type": "set", "path": ["messages", "0", "text"], "value": "goodbye"}]
+
+
+@pytest.mark.anyio
+async def test_string_assignment_non_string_override_emits_set() -> None:
+    ops: list[dict[str, Any]] = []
+    manager = StateManager(
+        lambda chunk: ops.extend(chunk.operations),
+        {"messages": [{"text": "hello"}]},
+    )
+
+    manager.state["messages"][0]["text"] = 42
+    manager.flush()
+
+    assert ops == [{"type": "set", "path": ["messages", "0", "text"], "value": 42}]
+
+
+@pytest.mark.anyio
+async def test_string_assignment_unchanged_value_emits_set() -> None:
+    ops: list[dict[str, Any]] = []
+    manager = StateManager(
+        lambda chunk: ops.extend(chunk.operations),
+        {"messages": [{"text": "hello"}]},
+    )
+
+    manager.state["messages"][0]["text"] = "hello"
+    manager.flush()
+
+    assert ops == [{"type": "set", "path": ["messages", "0", "text"], "value": "hello"}]
+
+
+@pytest.mark.anyio
+async def test_empty_string_initial_write_emits_set() -> None:
+    """First write to a field initialized to "" should emit set, not append-text."""
+    ops: list[dict[str, Any]] = []
+    manager = StateManager(
+        lambda chunk: ops.extend(chunk.operations),
+        {"messages": [{"text": ""}]},
+    )
+
+    manager.state["messages"][0]["text"] = "Hello"
+    manager.flush()
+
+    assert ops == [{"type": "set", "path": ["messages", "0", "text"], "value": "Hello"}]
+
+
+@pytest.mark.anyio
+async def test_run_controller_append_state_text() -> None:
+    """RunController.append_state_text() should emit append-text operations."""
+
+    async def run_callback(controller: RunController):
+        controller.append_state_text(["user", "name"], "Al")
+        controller.append_state_text(["user", "name"], "ice")
+
+    chunks = [
+        chunk
+        async for chunk in create_run(
+            run_callback, state={"user": {"name": ""}}
+        )
+    ]
+    operations = [
+        operation
+        for chunk in chunks
+        if chunk.type == "update-state"
+        for operation in chunk.operations
+    ]
+
+    assert operations == [
+        {"type": "append-text", "path": ["user", "name"], "value": "Al"},
+        {"type": "append-text", "path": ["user", "name"], "value": "ice"},
+    ]

python/assistant-transport-backend-langgraph/.env.example

@@ -0,0 +1,11 @@
+# Server Configuration
+HOST=0.0.0.0
+PORT=8001
+DEBUG=false
+LOG_LEVEL=info
+
+# CORS Configuration
+CORS_ORIGINS=http://localhost:3000
+
+# OpenAI API Key (for LangChain/LangGraph)
+OPENAI_API_KEY=your-openai-api-key-here

python/assistant-transport-backend-langgraph/README.md

@@ -0,0 +1,196 @@
+# Assistant Transport Backend with LangGraph
+
+This is a LangGraph-based implementation of the assistant transport backend, providing streaming chat capabilities using FastAPI, assistant-stream, and LangGraph.
+
+## Features
+
+- Streaming responses using LangGraph's astream and astream_events
+- Synchronization of LangGraph state to the frontend
+- Support for both message streaming and state updates
+- DeltaChannel-backed LangGraph message checkpoints (`langgraph>=1.2`)
+- Optional Postgres checkpoint storage via `langgraph-checkpoint-postgres`
+- Compatible with the assistant-ui frontend
+
+## Installation
+
+### Using uv (Recommended)
+
+1. Initialize and install dependencies:
+```bash
+uv init --name assistant-transport-backend-langgraph --package
+uv add fastapi uvicorn[standard] assistant-stream pydantic python-dotenv "langgraph>=1.2.0" langgraph-checkpoint-postgres langchain langchain-core langchain-openai httpx
+# Or simply:
+uv sync
+```
+
+2. Set up environment variables:
+```bash
+cp .env.example .env
+# Edit .env to add your OpenAI API key
+```
+
+### Using pip
+
+1. Install dependencies:
+```bash
+pip install -r requirements.txt
+```
+
+2. Set up environment variables:
+```bash
+cp .env.example .env
+# Edit .env to add your OpenAI API key
+```
+
+## Configuration
+
+The server can be configured via environment variables:
+
+- `HOST`: Server host (default: 0.0.0.0)
+- `PORT`: Server port (default: 8001)
+- `DEBUG`: Enable debug mode (default: false)
+- `LOG_LEVEL`: Log level (default: info)
+- `CORS_ORIGINS`: CORS origins (default: http://localhost:3000)
+- `OPENAI_API_KEY`: Your OpenAI API key (required)
+- `LANGGRAPH_POSTGRES_URL`: Optional Postgres connection URL for LangGraph checkpoints
+- `DATABASE_URL`: Fallback Postgres connection URL when `LANGGRAPH_POSTGRES_URL` is not set
+
+## Running the Server
+
+### Using uv
+```bash
+uv run python main.py
+```
+
+Or with uvicorn directly:
+```bash
+uv run uvicorn main:app --reload --host 0.0.0.0 --port 8001
+```
+
+### Using standard Python
+```bash
+python main.py
+```
+
+Or with uvicorn directly:
+```bash
+uvicorn main:app --reload --host 0.0.0.0 --port 8001
+```
+
+## API Endpoints
+
+### POST /api/chat
+Main chat endpoint that processes commands and streams responses using LangGraph.
+
+Request body:
+```json
+{
+  "commands": [
+    {
+      "type": "add-message",
+      "message": {
+        "role": "user",
+        "parts": [
+          {
+            "type": "text",
+            "text": "Hello, how are you?"
+          }
+        ]
+      }
+    }
+  ],
+  "system": "Optional system prompt",
+  "state": {}
+}
+```
+
+### GET /health
+Health check endpoint.
+
+## How It Works
+
+1. The server receives chat requests at `/api/chat`
+2. Commands are converted to LangGraph messages (HumanMessage, AIMessage, etc.)
+3. The LangGraph processes the messages through its nodes using a per-thread checkpoint keyed by the AssistantTransport `threadId`
+4. Two streaming tasks run concurrently:
+   - `astream` provides state updates
+   - `astream_events` provides message streaming
+5. Both streams are synchronized to the frontend using `append_langgraph_event`
+6. The response is streamed back using assistant-stream's DataStreamResponse
+
+Frontend tools declared by `useAssistantTransportRuntime` are bound to the LangGraph model from the request `tools` payload, but they are not executed by this backend. For example, the `with-assistant-transport` demo keeps `get_weather` frontend-only: the backend streams the tool call, the browser runs the tool and sends an `add-tool-result` command, and LangGraph continues from that result. Server-owned smoke tools such as `calculate_sum`, `save_note`, and `task_tool` still execute inside the backend graph.
+
+## DeltaChannel Prototype Notes
+
+The graph's `messages` state uses LangGraph's `DeltaChannel` with a bulk reducer:
+
+```python
+def add_messages_delta(state, writes):
+    result = list(state)
+    for write in writes:
+        if isinstance(write, BaseMessage):
+            result = add_messages(result, [write])
+        else:
+            result = add_messages(result, list(write))
+    return result
+```
+
+This keeps the assistant-ui API unchanged. The frontend still uses `useAssistantTransportRuntime`; the backend still accepts normal AssistantTransport `add-message` and `add-tool-result` commands; and the response remains the default data-stream encoding. The only required API adjustment is inside the LangGraph state definition: a delta-backed channel reducer receives `(state, writes)` where `writes` is a batch, not the old pairwise `(state, update)` reducer shape.
+
+Postgres works through LangGraph's async checkpointer path:
+
+```bash
+docker run --rm -p 127.0.0.1:55432:5432 \
+  -e POSTGRES_PASSWORD=postgres \
+  -e POSTGRES_DB=assistant_ui \
+  postgres:16-alpine
+
+LANGGRAPH_POSTGRES_URL=postgresql://postgres:postgres@127.0.0.1:55432/assistant_ui \
+  uv run python main.py
+```
+
+Because the FastAPI route streams with `graph.astream`, the backend uses `AsyncPostgresSaver`; the synchronous `PostgresSaver` does not implement the async checkpointer methods used by this route.
+
+## Integration with Frontend
+
+This backend is designed to work with the assistant-ui frontend. Update your frontend configuration to point to this server:
+
+```typescript
+const runtime = useExternalStoreRuntime({
+  endpoint: "http://localhost:8001/api/chat"
+});
+```
+
+## Customizing the LangGraph
+
+You can customize the graph in the `create_graph()` function. Currently, it implements a simple chat node using OpenAI's GPT-5.4 Nano model. You can:
+
+- Add more nodes for different functionalities
+- Implement tool calling
+- Add conditional edges
+- Integrate with different LLMs
+- Add memory or persistence
+
+Example of adding a tool node:
+```python
+from langgraph.prebuilt import ToolExecutor
+
+def create_graph():
+    workflow = StateGraph(GraphState)
+
+    # Add nodes
+    workflow.add_node("chat", chat_node)
+    workflow.add_node("tools", tool_node)
+
+    # Add conditional routing
+    workflow.add_conditional_edges(
+        "chat",
+        should_use_tools,
+        {
+            "tools": "tools",
+            "end": END
+        }
+    )
+
+    return workflow.compile()
+```

python/assistant-transport-backend-langgraph/client_demo.py

@@ -0,0 +1,93 @@
+#!/usr/bin/env python3
+"""
+Test client for the LangGraph backend.
+"""
+
+import asyncio
+import httpx
+import json
+
+async def test_chat():
+    """Test the chat endpoint with streaming."""
+    url = "http://localhost:8001/api/chat"
+
+    payload = {
+        "commands": [
+            {
+                "type": "add-message",
+                "message": {
+                    "role": "user",
+                    "parts": [
+                        {
+                            "type": "text",
+                            "text": "Hello! Can you tell me a short joke?"
+                        }
+                    ]
+                }
+            }
+        ],
+        "system": "You are a helpful assistant.",
+        "state": {}
+    }
+
+    print("📤 Sending request to", url)
+    print(f"📝 Message: {payload['commands'][0]['message']['parts'][0]['text']}")
+    print("-" * 50)
+
+    async with httpx.AsyncClient(timeout=30.0) as client:
+        async with client.stream(
+            "POST",
+            url,
+            json=payload,
+            headers={"Content-Type": "application/json"}
+        ) as response:
+            print(f"📥 Response status: {response.status_code}")
+            print(f"📥 Response headers: {dict(response.headers)}")
+            print("-" * 50)
+            print("📥 Streaming response:")
+
+            buffer = ""
+            async for chunk in response.aiter_bytes():
+                chunk_str = chunk.decode('utf-8')
+                buffer += chunk_str
+
+                # Parse SSE events
+                lines = buffer.split('\n')
+                buffer = lines[-1]  # Keep incomplete line in buffer
+
+                for line in lines[:-1]:
+                    if line.startswith('data: '):
+                        data_str = line[6:]  # Remove 'data: ' prefix
+                        if data_str.strip():
+                            try:
+                                data = json.loads(data_str)
+                                if data.get("type") == "state-update":
+                                    print(f"  State update: {json.dumps(data.get('value', {}), indent=2)}")
+                                else:
+                                    print(f"  Event: {data.get('type', 'unknown')}")
+                            except json.JSONDecodeError:
+                                print(f"  Raw: {data_str[:100]}...")
+
+    print("-" * 50)
+    print("✅ Test completed")
+
+async def test_health():
+    """Test the health endpoint."""
+    url = "http://localhost:8001/health"
+
+    async with httpx.AsyncClient() as client:
+        response = await client.get(url)
+        print(f"Health check: {response.status_code}")
+        print(f"Response: {response.json()}")
+
+if __name__ == "__main__":
+    print("🧪 Testing LangGraph Backend")
+    print("=" * 50)
+
+    # Test health check
+    print("\n1. Testing health endpoint...")
+    asyncio.run(test_health())
+
+    # Test chat
+    print("\n2. Testing chat endpoint...")
+    asyncio.run(test_chat())

python/assistant-transport-backend-langgraph/langgraph.json

@@ -0,0 +1,7 @@
+{
+  "dependencies": ["."],
+  "graphs": {
+    "agent": "./main.py:graph"
+  },
+  "env": ".env"
+}

python/assistant-transport-backend-langgraph/main.py

@@ -0,0 +1,575 @@
+#!/usr/bin/env python3
+"""
+Assistant Transport Backend with LangGraph - FastAPI + assistant-stream + LangGraph server
+"""
+
+import json
+import os
+from collections.abc import Sequence
+from contextlib import asynccontextmanager
+from typing import Annotated, Any, TypedDict
+from uuid import uuid4
+
+import uvicorn
+from assistant_stream import RunController, create_run
+from assistant_stream.modules.langgraph import append_langgraph_event, get_tool_call_subgraph_state
+from assistant_stream.serialization import DataStreamResponse
+from dotenv import load_dotenv
+from fastapi import FastAPI
+from fastapi.middleware.cors import CORSMiddleware
+from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, SystemMessage, ToolMessage
+from langchain_core.tools import tool
+from langchain_openai import ChatOpenAI
+from langgraph.channels import DeltaChannel
+from langgraph.checkpoint.memory import InMemorySaver
+from langgraph.graph import END, StateGraph, add_messages
+from langgraph.graph.state import CompiledStateGraph
+from pydantic import BaseModel, ConfigDict, Field
+
+# Load environment variables
+load_dotenv()
+
+_postgres_checkpointer_context = None
+
+
+class MessagePart(BaseModel):
+    """A part of a user message."""
+    type: str = Field(..., description="The type of message part")
+    text: str | None = Field(None, description="Text content")
+    image: str | None = Field(None, description="Image URL or data")
+
+
+class UserMessage(BaseModel):
+    """A user message."""
+    role: str = Field(default="user", description="Message role")
+    parts: list[MessagePart] = Field(..., description="Message parts")
+
+
+class AddMessageCommand(BaseModel):
+    """Command to add a new message to the conversation."""
+    type: str = Field(default="add-message", description="Command type")
+    message: UserMessage = Field(..., description="User message")
+
+
+class AddToolResultCommand(BaseModel):
+    """Command to add a tool result to the conversation."""
+    model_config = ConfigDict(populate_by_name=True)
+
+    type: str = Field(default="add-tool-result", description="Command type")
+    tool_call_id: str = Field(..., alias="toolCallId", description="ID of the tool call")
+    tool_name: str | None = Field(None, alias="toolName", description="Name of the tool")
+    result: Any = Field(..., description="Tool execution result")
+    is_error: bool | None = Field(None, alias="isError", description="Whether the tool failed")
+    artifact: Any | None = Field(None, description="UI-only tool result artifact")
+    model_content: Any | None = Field(
+        None, alias="modelContent", description="Tool result content for the model"
+    )
+
+
+class ChatRequest(BaseModel):
+    """Request payload for the chat endpoint."""
+    model_config = ConfigDict(populate_by_name=True)
+
+    commands: list[AddMessageCommand | AddToolResultCommand] = Field(
+        ..., description="List of commands to execute"
+    )
+    system: str | None = Field(None, description="System prompt")
+    tools: dict[str, Any] | None = Field(None, description="Available tools")
+    run_config: dict[str, Any] | None = Field(
+        None, alias="runConfig", description="Run configuration"
+    )
+    thread_id: str | None = Field(None, alias="threadId", description="Assistant UI thread ID")
+    state: dict[str, Any] | None = Field(None, description="State")
+
+
+def get_thread_id(request: ChatRequest) -> str:
+    """Use persistent checkpoints only when the AssistantTransport request identifies a thread."""
+    return request.thread_id or f"anonymous-{uuid4()}"
+
+
+def add_messages_delta(
+    state: Sequence[BaseMessage],
+    writes: Sequence[BaseMessage | Sequence[BaseMessage]],
+) -> list[BaseMessage]:
+    result = list(state)
+    for write in writes:
+        if isinstance(write, BaseMessage):
+            result = add_messages(result, [write])
+        else:
+            result = add_messages(result, list(write))
+    return result
+
+
+def request_tool_schemas(tools: dict[str, Any] | None) -> list[dict[str, Any]]:
+    if not tools:
+        return []
+
+    schemas = []
+    for name, tool_definition in tools.items():
+        if (
+            name in TOOL_BY_NAME
+            or not isinstance(tool_definition, dict)
+            or tool_definition.get("disabled") is True
+        ):
+            continue
+
+        parameters = tool_definition.get("parameters") or {
+            "type": "object",
+            "properties": {},
+        }
+        schemas.append(
+            {
+                "type": "function",
+                "function": {
+                    "name": name,
+                    "description": tool_definition.get("description", ""),
+                    "parameters": parameters,
+                },
+            }
+        )
+    return schemas
+
+
+def bindable_tools(tools: dict[str, Any] | None) -> list[Any]:
+    return [*TOOLS, *request_tool_schemas(tools)]
+
+
+def tool_result_content(command: AddToolResultCommand) -> str:
+    content = command.model_content if command.model_content is not None else command.result
+    return content if isinstance(content, str) else json.dumps(content)
+
+
+# Define LangGraph state
+class GraphState(TypedDict, total=False):
+    """State for the conversation graph."""
+    messages: Annotated[
+        Sequence[BaseMessage],
+        DeltaChannel(reducer=add_messages_delta, snapshot_frequency=50),
+    ]
+    tools: dict[str, Any] | None
+
+
+# Define subagent state
+class SubagentState(TypedDict):
+    """State for the subagent."""
+    messages: Annotated[
+        Sequence[BaseMessage],
+        DeltaChannel(reducer=add_messages_delta, snapshot_frequency=50),
+    ]
+    task: str
+    result: str
+
+
+# Create the Task tool
+@tool
+def task_tool(task_description: str) -> str:
+    """
+    Execute a complex task using a subagent.
+
+    Args:
+        task_description: Description of the task to perform
+
+    Returns:
+        The result of the task execution
+    """
+    # This is a placeholder - the actual execution will be handled by the subgraph
+    return f"Task '{task_description}' will be executed by the subagent."
+
+
+@tool
+def calculate_sum(numbers: list[float]) -> dict[str, Any]:
+    """
+    Add a list of numbers.
+
+    Args:
+        numbers: Numbers to add together.
+    """
+    return {
+        "numbers": numbers,
+        "sum": sum(numbers),
+        "count": len(numbers),
+    }
+
+
+@tool
+def save_note(title: str, body: str) -> dict[str, Any]:
+    """
+    Save a sample note and return a note ID.
+
+    Args:
+        title: Short note title.
+        body: Note body.
+    """
+    note_seed = f"{title}\n{body}"
+    note_id = sum((index + 1) * ord(char) for index, char in enumerate(note_seed))
+    return {
+        "id": f"note-{note_id % 100000}",
+        "title": title,
+        "body": body,
+        "saved": True,
+    }
+
+
+TOOLS = [task_tool, calculate_sum, save_note]
+TOOL_BY_NAME = {tool.name: tool for tool in TOOLS}
+
+
+# Subagent node for executing tasks
+async def subagent_node(state: SubagentState) -> dict[str, Any]:
+    """Subagent that executes the task."""
+    task = state.get("task", "")
+
+    # Create a prompt for the subagent
+    subagent_messages = [
+        SystemMessage(content=f"You are a helpful subagent. Execute this task: {task}"),
+        HumanMessage(content=f"Please complete the following task: {task}")
+    ]
+
+    # Generate response
+    if os.getenv("OPENAI_API_KEY"):
+        # Initialize a simpler LLM for the subagent
+        llm = ChatOpenAI(
+            model="gpt-5.4-nano",
+            temperature=0.7,
+            streaming=True
+        )
+        response = await llm.ainvoke(subagent_messages)
+        result = response.content
+    else:
+        result = f"Mock subagent result for task: {task}"
+
+    return {
+        "messages": [AIMessage(content=result)],
+        "result": result
+    }
+
+
+def create_subagent_graph() -> CompiledStateGraph:
+    """Create the subagent graph."""
+    workflow = StateGraph(SubagentState)
+
+    # Add the subagent node
+    workflow.add_node("execute_task", subagent_node)
+
+    # Set entry and exit points
+    workflow.set_entry_point("execute_task")
+    workflow.add_edge("execute_task", END)
+
+    return workflow.compile()
+
+
+async def agent_node(state: GraphState) -> dict[str, Any]:
+    """Main agent node that can call tools."""
+    messages = state.get("messages", [])
+    tools = state.get("tools")
+
+    # Check if OpenAI API key is set
+    if os.getenv("OPENAI_API_KEY"):
+        # Initialize the LLM with tool binding
+        llm = ChatOpenAI(
+            model="gpt-5.4-nano",
+            temperature=0.7,
+            streaming=True,
+        )
+
+        # Bind server tools plus request-provided frontend tool declarations.
+        llm_with_tools = llm.bind_tools(bindable_tools(tools))
+        response = await llm_with_tools.ainvoke(messages)
+    elif messages and isinstance(messages[-1], ToolMessage):
+        response = AIMessage(
+            content=f"Task complete: {messages[-1].content}",
+        )
+    else:
+        # Mock response with a tool call for testing
+        print("⚠️ No OpenAI API key found - using mock response with tool call")
+        last_content = messages[-1].content if messages else ""
+        if "weather" in str(last_content).lower() and tools and "get_weather" in tools:
+            tool_call = {
+                "id": "weather_001",
+                "name": "get_weather",
+                "args": {"location": "San Francisco", "unit": "fahrenheit"},
+            }
+        elif "sum" in str(last_content).lower() or "add" in str(last_content).lower():
+            tool_call = {
+                "id": "sum_001",
+                "name": "calculate_sum",
+                "args": {"numbers": [2, 3, 5]},
+            }
+        elif "note" in str(last_content).lower():
+            tool_call = {
+                "id": "note_001",
+                "name": "save_note",
+                "args": {"title": "Smoke test", "body": "Saved from mock mode"},
+            }
+        else:
+            tool_call = {
+                "id": "task_001",
+                "name": "task_tool",
+                "args": {"task_description": "Complete the requested task"},
+            }
+
+        response = AIMessage(content="I'll call a tool for that.", tool_calls=[tool_call])
+
+    return {"messages": [response]}
+
+
+def should_call_tools(state: GraphState) -> str:
+    """Run only backend-owned tools. Frontend tools are returned to the client."""
+    messages = state.get("messages", [])
+    if not messages:
+        return "end"
+
+    last_message = messages[-1]
+    if (
+        hasattr(last_message, 'tool_calls')
+        and last_message.tool_calls
+        and any(tool_call["name"] in TOOL_BY_NAME for tool_call in last_message.tool_calls)
+    ):
+        return "tools"
+
+    return "end"
+
+
+async def tool_executor_node(state: GraphState) -> dict[str, Any]:
+    """Execute tool calls, including Task tool which spawns subagents."""
+    messages = state.get("messages", [])
+    if not messages:
+        return {"messages": []}
+
+    last_message = messages[-1]
+    if not hasattr(last_message, 'tool_calls') or not last_message.tool_calls:
+        return {"messages": []}
+
+    # Process each tool call
+    tool_messages = []
+    for tool_call in last_message.tool_calls:
+        tool_name = tool_call["name"]
+        if tool_name == "task_tool":
+            # Extract task description
+            task_description = tool_call["args"].get("task_description", "")
+
+            # Create and run the subagent graph
+
+            # Initialize subagent state
+            subagent_state = {
+                "messages": [],
+                "task": task_description,
+                "result": ""
+            }
+
+            # Run the subagent
+            final_state = await subagent_graph.ainvoke(subagent_state)
+
+            # Create tool message with the result
+            tool_message = ToolMessage(
+                content=final_state.get("result", "Task completed"),
+                tool_call_id=tool_call["id"],
+                artifact={"subgraph_state": final_state}
+            )
+            tool_messages.append(tool_message)
+        else:
+            tool = TOOL_BY_NAME.get(tool_name)
+            if tool is None:
+                result = {"error": f"Unknown tool: {tool_name}"}
+            else:
+                result = tool.invoke(tool_call.get("args", {}))
+
+            tool_message = ToolMessage(
+                content=json.dumps(result),
+                tool_call_id=tool_call["id"],
+                name=tool_name,
+                artifact=result,
+            )
+            tool_messages.append(tool_message)
+
+    return {"messages": tool_messages}
+
+
+subagent_graph = create_subagent_graph()
+
+async def configure_checkpointer_from_env() -> None:
+    """Switch the graph to the configured persistent checkpointer, when present."""
+    postgres_url = os.getenv("LANGGRAPH_POSTGRES_URL") or os.getenv("DATABASE_URL")
+    if not postgres_url:
+        return
+
+    from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
+
+    global _postgres_checkpointer_context, graph
+    if _postgres_checkpointer_context is not None:
+        return
+
+    _postgres_checkpointer_context = AsyncPostgresSaver.from_conn_string(postgres_url)
+    checkpointer = await _postgres_checkpointer_context.__aenter__()
+    await checkpointer.setup()
+    graph = create_graph(checkpointer)
+
+
+def create_graph(checkpointer=None) -> CompiledStateGraph:
+    """Create and compile the LangGraph with subgraph support."""
+    # Create the main workflow
+    workflow = StateGraph(GraphState)
+
+    # Add nodes
+    workflow.add_node("agent", agent_node)
+    workflow.add_node("tools", tool_executor_node)
+
+    # Set entry point
+    workflow.set_entry_point("agent")
+
+    # Add conditional edges
+    workflow.add_conditional_edges(
+        "agent",
+        should_call_tools,
+        {
+            "tools": "tools",
+            "end": END
+        }
+    )
+
+    # After tools, go back to agent for potential follow-up
+    workflow.add_edge("tools", "agent")
+
+    # Compile with a checkpointer so DeltaChannel is exercised across thread turns.
+    return workflow.compile(checkpointer=checkpointer or InMemorySaver())
+
+graph = create_graph()
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    """Application lifespan manager."""
+    print("🚀 Assistant Transport Backend with LangGraph starting up...")
+    await configure_checkpointer_from_env()
+    try:
+        yield
+    finally:
+        if _postgres_checkpointer_context is not None:
+            await _postgres_checkpointer_context.__aexit__(None, None, None)
+        print("🛑 Assistant Transport Backend with LangGraph shutting down...")
+
+
+# Create FastAPI app
+app = FastAPI(
+    title="Assistant Transport Backend with LangGraph",
+    description=(
+        "A server implementing the assistant-transport protocol with LangGraph and subgraphs"
+    ),
+    version="0.2.0",
+    lifespan=lifespan,
+)
+
+# Configure CORS
+cors_origins = ["*"]  # Allow all origins
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=cors_origins,
+    allow_credentials=True,
+    allow_methods=["GET", "POST", "PUT", "DELETE"],
+    allow_headers=["*"],
+)
+
+
+@app.post("/assistant")
+async def chat_endpoint(request: ChatRequest):
+    """Chat endpoint using LangGraph with streaming and subgraph support."""
+
+    async def run_callback(controller: RunController):
+        """Callback function for the run controller."""
+        # Initialize controller state if needed
+        if controller.state is None:
+            controller.state = {}
+        if "messages" not in controller.state:
+            controller.state["messages"] = []
+
+        input_messages = []
+
+        # Process commands
+        for command in request.commands:
+            if command.type == "add-message":
+                # Extract text from parts
+                text_parts = [
+                    part.text for part in command.message.parts
+                    if part.type == "text" and part.text
+                ]
+                if text_parts:
+                    input_messages.append(HumanMessage(content=" ".join(text_parts)))
+            elif command.type == "add-tool-result":
+                # Handle tool results
+                input_messages.append(ToolMessage(
+                    content=tool_result_content(command),
+                    tool_call_id=command.tool_call_id,
+                    name=command.tool_name,
+                    artifact=command.artifact if command.artifact is not None else command.result,
+                    status="error" if command.is_error else "success",
+                ))
+
+        # Add messages to controller state
+        for message in input_messages:
+            controller.state["messages"].append(message.model_dump())
+
+        # Create initial state for LangGraph
+        input_state = {"messages": input_messages, "tools": request.tools}
+
+        # Stream with subgraph support
+        config = {
+            "configurable": {
+                "thread_id": get_thread_id(request),
+            }
+        }
+
+        async for namespace, event_type, chunk in graph.astream(
+            input_state,
+            config=config,
+            stream_mode=["messages", "updates"],
+            subgraphs=True
+        ):
+            state = get_tool_call_subgraph_state(
+                controller,
+                subgraph_node="tools",
+                namespace=namespace,
+                artifact_field_name="subgraph_state",
+                default_state={}
+            )
+            # Append the event normally
+            append_langgraph_event(
+                state,
+                namespace,
+                event_type,
+                chunk
+            )
+
+    # Create streaming response using assistant-stream
+    stream = create_run(run_callback, state=request.state)
+
+    return DataStreamResponse(stream)
+
+
+@app.get("/health")
+async def health_check():
+    """Health check endpoint."""
+    return {"status": "healthy", "service": "assistant-transport-backend-langgraph"}
+
+
+def main():
+    """Main entry point for running the server."""
+    host = os.getenv("HOST", "0.0.0.0")
+    port = int(os.getenv("PORT", "8010"))
+    debug = os.getenv("DEBUG", "false").lower() == "true"
+    log_level = os.getenv("LOG_LEVEL", "info").lower()
+
+    print(f"🌟 Starting Assistant Transport Backend with LangGraph on {host}:{port}")
+    print(f"🎯 Debug mode: {debug}")
+    print(f"🌍 CORS origins: {cors_origins}")
+
+    uvicorn.run(
+        "main:app",
+        host=host,
+        port=port,
+        reload=debug,
+        log_level=log_level,
+        access_log=True,
+    )
+
+
+if __name__ == "__main__":
+    main()

python/assistant-transport-backend-langgraph/pyproject.toml

@@ -0,0 +1,50 @@
+[project]
+name = "assistant-transport-backend-langgraph"
+version = "0.1.0"
+description = "Assistant Transport Backend with LangGraph integration"
+readme = "README.md"
+requires-python = ">=3.10"
+dependencies = [
+    "fastapi>=0.100.0",
+    "uvicorn[standard]>=0.20.0",
+    "assistant-stream>=0.0.28",
+    "pydantic>=2.0.0",
+    "python-dotenv>=1.0.0",
+    "langgraph>=1.2.0",
+    "langchain>=0.2.0",
+    "langchain-core>=0.2.0",
+    "langchain-openai>=0.1.0",
+    "httpx>=0.24.0",
+    "langgraph-checkpoint-postgres>=3.0.0",
+]
+
+[project.optional-dependencies]
+dev = [
+    "pytest>=7.0.0",
+    "pytest-asyncio>=0.21.0",
+    "black>=23.0.0",
+    "isort>=5.12.0",
+    "mypy>=1.0.0",
+    "ruff>=0.1.0",
+]
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[tool.hatch.build.targets.wheel]
+packages = ["."]
+exclude = ["client_demo.py", "subgraph_demo.py"]
+
+[tool.ruff]
+target-version = "py310"
+line-length = 100
+select = ["E", "F", "I", "N", "W", "UP"]
+
+[tool.black]
+line-length = 100
+target-version = ["py310"]
+
+[tool.isort]
+profile = "black"
+line_length = 100

python/assistant-transport-backend-langgraph/requirements.txt

@@ -0,0 +1,21 @@
+# Core dependencies
+fastapi>=0.100.0
+uvicorn[standard]>=0.20.0
+assistant-stream>=0.0.28
+pydantic>=2.0.0
+python-dotenv>=1.0.0
+
+# LangGraph dependencies
+langgraph>=1.2.0
+langgraph-checkpoint-postgres>=3.0.0
+langchain>=0.2.0
+langchain-core>=0.2.0
+langchain-openai>=0.1.0
+
+# Development dependencies
+pytest>=7.0.0
+pytest-asyncio>=0.21.0
+black>=23.0.0
+isort>=5.12.0
+mypy>=1.0.0
+ruff>=0.1.0

python/assistant-transport-backend-langgraph/subgraph_demo.py

@@ -0,0 +1,181 @@
+#!/usr/bin/env python3
+"""
+Test client for the LangGraph backend with subgraph support.
+"""
+
+import asyncio
+import httpx
+import json
+
+async def test_subgraph_chat():
+    """Test the chat endpoint with task tool and subgraph streaming."""
+    url = "http://localhost:8000/assistant"
+
+    payload = {
+        "commands": [
+            {
+                "type": "add-message",
+                "message": {
+                    "role": "user",
+                    "parts": [
+                        {
+                            "type": "text",
+                            "text": "Please help me create a task that writes a haiku about Python programming."
+                        }
+                    ]
+                }
+            }
+        ],
+        "system": "You are a helpful assistant that can delegate complex tasks to subagents using the Task tool.",
+        "state": {}
+    }
+
+    print("📤 Sending request to", url)
+    print(f"📝 Message: {payload['commands'][0]['message']['parts'][0]['text']}")
+    print("-" * 50)
+
+    async with httpx.AsyncClient(timeout=30.0) as client:
+        async with client.stream(
+            "POST",
+            url,
+            json=payload,
+            headers={"Content-Type": "application/json"}
+        ) as response:
+            print(f"📥 Response status: {response.status_code}")
+            print("-" * 50)
+            print("📥 Streaming response:")
+
+            buffer = ""
+            message_count = 0
+            tool_call_count = 0
+
+            async for chunk in response.aiter_bytes():
+                chunk_str = chunk.decode('utf-8')
+                buffer += chunk_str
+
+                # Parse SSE events
+                lines = buffer.split('\n')
+                buffer = lines[-1]  # Keep incomplete line in buffer
+
+                for line in lines[:-1]:
+                    if line.startswith('data: '):
+                        data_str = line[6:]  # Remove 'data: ' prefix
+                        if data_str.strip() and data_str.strip() != '[DONE]':
+                            try:
+                                data = json.loads(data_str)
+
+                                # Log different types of events
+                                if data.get("type") == "state-update":
+                                    state = data.get("value", {})
+
+                                    # Check for messages
+                                    if "messages" in state:
+                                        messages = state["messages"]
+                                        if len(messages) > message_count:
+                                            message_count = len(messages)
+                                            last_msg = messages[-1]
+
+                                            # Log AI messages
+                                            if last_msg.get("type") == "ai" or last_msg.get("role") == "assistant":
+                                                print(f"\n🤖 AI Message:")
+                                                if "content" in last_msg:
+                                                    print(f"   Content: {last_msg['content']}")
+                                                if "tool_calls" in last_msg:
+                                                    print(f"   Tool Calls: {last_msg['tool_calls']}")
+                                                    tool_call_count = len(last_msg['tool_calls'])
+
+                                            # Log Tool messages (from subgraph)
+                                            elif last_msg.get("type") == "tool" or last_msg.get("role") == "tool":
+                                                print(f"\n🔧 Tool Message:")
+                                                print(f"   Content: {last_msg.get('content', '')}")
+                                                if "artifact" in last_msg:
+                                                    print(f"   Artifact: {json.dumps(last_msg['artifact'], indent=6)}")
+
+                                    # Check for other state updates
+                                    for key in state:
+                                        if key not in ["messages", "status"]:
+                                            print(f"\n📊 State Update - {key}: {state[key]}")
+
+                                elif data.get("type") == "message-delta":
+                                    # Handle streaming message updates
+                                    delta = data.get("value", {})
+                                    if delta.get("content"):
+                                        print(f"   Streaming: {delta['content']}", end="")
+
+                                else:
+                                    # Other event types
+                                    event_type = data.get("type", "unknown")
+                                    if event_type not in ["ping", "status"]:
+                                        print(f"\n📡 Event ({event_type}): {data.get('value', '')}")
+
+                            except json.JSONDecodeError as e:
+                                if data_str.strip():
+                                    print(f"  Parse error: {e}")
+                                    print(f"  Raw data: {data_str[:200]}...")
+
+    print("\n" + "-" * 50)
+    print(f"✅ Test completed - Processed {message_count} messages, {tool_call_count} tool calls")
+
+
+async def test_direct_tool_result():
+    """Test sending a tool result back to the assistant."""
+    url = "http://localhost:8000/assistant"
+
+    # First create a state with a tool call
+    initial_state = {
+        "messages": [
+            {
+                "type": "human",
+                "content": "Help me with a task"
+            },
+            {
+                "type": "ai",
+                "content": "I'll help you with that task.",
+                "tool_calls": [
+                    {
+                        "id": "test_tool_001",
+                        "name": "task_tool",
+                        "args": {"task_description": "Write a poem"}
+                    }
+                ]
+            }
+        ]
+    }
+
+    payload = {
+        "commands": [
+            {
+                "type": "add-tool-result",
+                "toolCallId": "test_tool_001",
+                "result": {"output": "A beautiful poem about coding"}
+            }
+        ],
+        "system": "You are a helpful assistant.",
+        "state": initial_state
+    }
+
+    print("📤 Testing tool result handling")
+    print("-" * 50)
+
+    async with httpx.AsyncClient(timeout=10.0) as client:
+        response = await client.post(url, json=payload)
+        print(f"📥 Response status: {response.status_code}")
+
+        # Read the full response for tool result test
+        content = response.text
+        print(f"📥 Response preview: {content[:500]}...")
+
+    print("✅ Tool result test completed")
+
+
+if __name__ == "__main__":
+    print("🧪 Testing LangGraph Backend with Subgraph Support")
+    print("=" * 50)
+
+    # Test 1: Subgraph chat with task tool
+    print("\n1. Testing subgraph chat with task delegation...")
+    asyncio.run(test_subgraph_chat())
+
+    # Test 2: Direct tool result handling
+    print("\n2. Testing direct tool result handling...")
+    asyncio.run(test_direct_tool_result())

python/assistant-transport-backend/.env.example

@@ -0,0 +1,19 @@
+# Server Configuration
+HOST=0.0.0.0
+PORT=8000
+DEBUG=true
+
+# AI Provider Configuration (choose one)
+# OpenAI Configuration
+OPENAI_API_KEY=your_openai_api_key_here
+OPENAI_MODEL=gpt-4o-mini
+
+# Anthropic Configuration  
+# ANTHROPIC_API_KEY=your_anthropic_api_key_here
+# ANTHROPIC_MODEL=claude-3-haiku-20240307
+
+# CORS Configuration
+CORS_ORIGINS=http://localhost:3000,https://yourapp.com
+
+# Logging
+LOG_LEVEL=INFO

python/assistant-transport-backend/README.md

@@ -0,0 +1,217 @@
+# Assistant Transport Backend
+
+A simple Python server that demonstrates the assistant-transport protocol using FastAPI and assistant-stream. This backend returns static responses to show how the streaming protocol works with assistant-ui frontend applications.
+
+## Features
+
+- 🚀 **FastAPI-based** - High-performance async server
+- 📡 **Streaming Responses** - Real-time responses using assistant-stream
+- 🔄 **State Management** - Uses assistant-stream's object-stream state utilities
+- 🔌 **Assistant-Transport Protocol** - Full compatibility with assistant-ui
+- 🌐 **CORS Enabled** - Works with any frontend origin
+- 📦 **Simple Setup** - Minimal dependencies
+- 🧪 **Static Responses** - No API keys required, perfect for testing
+
+## Prerequisites
+
+- Python 3.9 or higher
+- pip package manager
+
+## Quick Start
+
+### 1. Install Dependencies
+
+```bash
+# Install the package and dependencies
+pip install -e .
+
+# Or install dependencies directly
+pip install -r requirements.txt
+```
+
+### 2. Configure Environment (Optional)
+
+```bash
+# Copy example environment file
+cp .env.example .env
+# Edit .env file if needed
+```
+
+Default configuration:
+```env
+# Server Configuration
+HOST=0.0.0.0
+PORT=8000
+DEBUG=true
+
+# CORS Configuration  
+CORS_ORIGINS=http://localhost:3000
+```
+
+### 3. Start the Server
+
+```bash
+# Using the installed command
+assistant-transport-backend
+
+# Or run directly
+python main.py
+
+# Or using uvicorn
+uvicorn main:app --reload --host 0.0.0.0 --port 8000
+```
+
+The server will be available at:
+- **API**: http://localhost:8000
+- **Health Check**: http://localhost:8000/health
+- **Assistant Endpoint**: http://localhost:8000/assistant
+
+## API Endpoints
+
+### `POST /assistant`
+
+Main endpoint that implements the assistant-transport protocol.
+
+**Request Format:**
+```json
+{
+  "commands": [
+    {
+      "type": "add-message",
+      "message": {
+        "role": "user",
+        "parts": [
+          {"type": "text", "text": "Hello!"}
+        ]
+      }
+    }
+  ],
+  "system": "You are a helpful assistant",
+  "tools": {},
+  "runConfig": {}
+}
+```
+
+**Response:** Streaming response using assistant-stream format with static responses.
+
+### `GET /health`
+
+Health check endpoint that returns server status and current conversation state.
+
+### `POST /cancel`
+
+Cancel the current request (placeholder for request cancellation).
+
+## Static Response Patterns
+
+The backend recognizes these message patterns and returns appropriate static responses:
+
+- **Greetings** (`hello`, `hi`) → Welcome message
+- **Weather** (`weather`) → Sunny static response
+- **What/What is** → Explanation of what the backend does
+- **Help** → Available command information
+- **Other messages** → Acknowledgment with echo
+
+## Integration with Frontend
+
+This backend works with the `with-assistant-transport` frontend example:
+
+1. Start backend: `python main.py`
+2. Start frontend: `cd ../../examples/with-assistant-transport && pnpm dev`
+3. Frontend connects to `http://localhost:8000/assistant`
+
+## Project Structure
+
+```
+python/assistant-transport-backend/
+├── main.py                    # FastAPI server and main entry point
+├── pyproject.toml            # Project configuration and dependencies
+├── requirements.txt          # Pip requirements file
+├── setup.py                  # Automated setup script
+├── .env.example             # Environment variables template
+└── README.md               # This file
+```
+
+## How It Works
+
+### Assistant-Stream Integration
+
+The backend uses `assistant_stream.create_run()` to create a streaming controller that:
+
+1. **Manages State**: Updates conversation state with messages
+2. **Streams Text**: Uses `controller.append_text()` for character-by-character streaming
+3. **Updates State**: Uses `controller.state` to update the object-stream state
+4. **Handles Protocol**: Automatically formats responses for assistant-transport
+
+### State Management
+
+```python
+# Initialize state
+conversation_state = {
+    "messages": [],
+    "provider": "static"
+}
+
+# Create run controller with state
+controller = create_run(conversation_state)
+
+# Update state during processing
+controller.state.provider = "processing"
+controller.state.messages.append(new_message)
+
+# Stream text responses
+for char in response_text:
+    controller.append_text(char)
+```
+
+## Development
+
+### Running in Development Mode
+
+```bash
+# Enable debug mode and auto-reload
+DEBUG=true python main.py
+```
+
+### Adding Response Patterns
+
+Edit the static response logic in `main.py`:
+
+```python
+# Add new patterns
+if "goodbye" in user_message:
+    response_text = "Goodbye! Thanks for testing the assistant-transport backend!"
+```
+
+### Testing
+
+```bash
+# Test the health endpoint
+curl http://localhost:8000/health
+
+# Test the assistant endpoint
+curl -X POST http://localhost:8000/assistant \
+  -H "Content-Type: application/json" \
+  -d '{"commands":[{"type":"add-message","message":{"role":"user","parts":[{"type":"text","text":"Hello!"}]}}]}'
+```
+
+## Environment Variables
+
+| Variable | Description | Default |
+|----------|-------------|---------|
+| `HOST` | Server host | `0.0.0.0` |
+| `PORT` | Server port | `8000` |
+| `DEBUG` | Enable debug mode | `false` |
+| `LOG_LEVEL` | Logging level | `info` |
+| `CORS_ORIGINS` | Allowed CORS origins | `http://localhost:3000` |
+
+## License
+
+This project is part of the assistant-ui monorepo and follows the same MIT licensing terms.
+
+## Learn More
+
+- [assistant-ui Documentation](https://docs.assistant-ui.com)
+- [Assistant Transport Protocol](https://docs.assistant-ui.com/runtimes/assistant-transport)
+- [assistant-stream Package](https://github.com/assistant-ui/assistant-ui/tree/main/python/assistant-stream)
+- [FastAPI Documentation](https://fastapi.tiangolo.com)

python/assistant-transport-backend/main.py

@@ -0,0 +1,172 @@
+#!/usr/bin/env python3
+"""
+Assistant Transport Backend - Simple FastAPI + assistant-stream server
+
+This server implements the assistant-transport protocol with static responses.
+"""
+
+import os
+import asyncio
+import random
+from typing import Dict, Any, List, Optional, Union
+from contextlib import asynccontextmanager
+import uvicorn 
+
+from fastapi import FastAPI
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel, Field
+from dotenv import load_dotenv
+
+from assistant_stream.serialization import DataStreamResponse
+from assistant_stream import RunController, create_run
+
+# Load environment variables
+load_dotenv()
+
+
+class MessagePart(BaseModel):
+    """A part of a user message."""
+    type: str = Field(..., description="The type of message part")
+    text: Optional[str] = Field(None, description="Text content")
+    image: Optional[str] = Field(None, description="Image URL or data")
+
+
+class UserMessage(BaseModel):
+    """A user message."""
+    role: str = Field(default="user", description="Message role")
+    parts: List[MessagePart] = Field(..., description="Message parts")
+
+
+class AddMessageCommand(BaseModel):
+    """Command to add a new message to the conversation."""
+    type: str = Field(default="add-message", description="Command type")
+    message: UserMessage = Field(..., description="User message")
+
+
+class AddToolResultCommand(BaseModel):
+    """Command to add a tool result to the conversation."""
+    type: str = Field(default="add-tool-result", description="Command type")  
+    toolCallId: str = Field(..., description="ID of the tool call")
+    result: Dict[str, Any] = Field(..., description="Tool execution result")
+
+
+class AssistantRequest(BaseModel):
+    """Request payload for the assistant endpoint."""
+    commands: List[Union[AddMessageCommand, AddToolResultCommand]] = Field(
+        ..., description="List of commands to execute"
+    )
+    system: Optional[str] = Field(None, description="System prompt")
+    tools: Optional[Dict[str, Any]] = Field(None, description="Available tools")
+    # Additional fields that may be sent by the frontend
+    runConfig: Optional[Dict[str, Any]] = Field(None, description="Run configuration")
+    state: Optional[Dict[str, Any]] = Field(None, description="State")
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    """Application lifespan manager."""
+    print("🚀 Assistant Transport Backend starting up...")
+    yield
+    print("🛑 Assistant Transport Backend shutting down...")
+
+
+# Create FastAPI app
+app = FastAPI(
+    title="Assistant Transport Backend",
+    description="A simple server implementing the assistant-transport protocol with static responses",
+    version="0.1.0",
+    lifespan=lifespan,
+)
+
+# Configure CORS
+cors_origins = os.getenv("CORS_ORIGINS", "http://localhost:3000").split(",")
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=cors_origins,
+    allow_credentials=True,
+    allow_methods=["GET", "POST", "PUT", "DELETE"],
+    allow_headers=["*"],
+)
+
+
+@app.post("/assistant")
+async def assistant_endpoint(request: AssistantRequest): 
+    # Create streaming response using assistant-stream
+    async def run_callback(controller: RunController):
+        """Callback function for the run controller."""
+        try:
+            print("run_callback")
+            await asyncio.sleep(1)
+
+            if (request.commands[0].type == "add-message"):
+                controller.state["messages"].append(request.commands[0].message.model_dump())
+            if (request.commands[0].type == "add-tool-result"):
+                controller.state["messages"][-1]["parts"][-1]["result"] = request.commands[0].result
+
+            await asyncio.sleep(1)
+
+            controller.state["messages"].append({
+                "role": "assistant",
+                "parts": [
+                    {
+                        "type": "text",
+                        "text": "Hello22"
+                    }
+                ]
+            })
+
+            controller.state["messages"][-1]["parts"].append({
+                "type": "tool-call",
+                "toolCallId": "tool_" + str(random.randint(0, 1000000)),
+                "toolName": "get_weather",
+                "argsText": "",
+                "done": False,
+            })
+
+            await asyncio.sleep(1)
+
+            controller.state["messages"][-1]["parts"][-1]["argsText"] = "{\"location\": \"SF\""
+
+            await asyncio.sleep(1)
+
+            controller.state["messages"][-1]["parts"][-1]["argsText"] =  controller.state["messages"][-1]["parts"][-1]["argsText"] +    "}"
+            controller.state["messages"][-1]["parts"][-1]["done"] = True
+
+            await asyncio.sleep(1)
+
+            controller.state["provider"] = "completed"
+            print("run_callback3")
+        except Exception as e:
+            print(f"❌ Error in stream generation: {e}")
+            controller.state["provider"] = "error"
+            controller.append_text(f"Error: {str(e)}")
+    
+    # Create streaming response using assistant-stream
+    stream = create_run(run_callback, state=request.state)
+    
+    return DataStreamResponse(stream)
+
+
+def main():
+    """Main entry point for running the server."""
+    host = os.getenv("HOST", "0.0.0.0")
+    port = int(os.getenv("PORT", "8000"))
+    debug = os.getenv("DEBUG", "false").lower() == "true"
+    log_level = os.getenv("LOG_LEVEL", "info").lower()
+    
+    print(f"🌟 Starting Assistant Transport Backend on {host}:{port}")
+    print(f"🎯 Debug mode: {debug}")
+    print(f"🌍 CORS origins: {cors_origins}")
+    
+    uvicorn.run(
+        "main:app",
+        host=host,
+        port=port,
+        reload=debug,
+        log_level=log_level,
+        access_log=True,
+    )
+
+
+if __name__ == "__main__":
+    main()

python/assistant-transport-backend/pyproject.toml

@@ -0,0 +1,77 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "assistant-transport-backend"
+version = "0.1.0"
+description = "A simple server implementing the assistant-transport protocol with static responses"
+authors = [
+    { name = "assistant-ui Team", email = "team@assistant-ui.com" }
+]
+readme = "README.md"
+license = "MIT"
+requires-python = ">=3.10,<4.0"
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "License :: OSI Approved :: MIT License",
+    "Operating System :: OS Independent",
+    "Topic :: Software Development :: Libraries :: Python Modules",
+    "Topic :: Internet :: WWW/HTTP :: HTTP Servers",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
+dependencies = [
+    "fastapi>=0.100.0",
+    "uvicorn[standard]>=0.20.0",
+    "assistant-stream>=0.0.28",
+    "pydantic>=2.0.0",
+    "python-dotenv>=1.0.0",
+]
+
+[project.optional-dependencies]
+openai = ["openai>=1.0.0"]
+anthropic = ["anthropic>=0.18.0"]
+dev = [
+    "pytest>=7.0.0",
+    "pytest-asyncio>=0.21.0",
+    "black>=23.0.0",
+    "isort>=5.12.0",
+    "mypy>=1.0.0",
+    "ruff>=0.1.0",
+]
+
+[project.urls]
+Homepage = "https://github.com/assistant-ui/assistant-ui"
+Issues = "https://github.com/assistant-ui/assistant-ui/issues"
+Documentation = "https://docs.assistant-ui.com"
+
+[project.scripts]
+assistant-transport-backend = "main:main"
+
+[tool.hatch.metadata]
+allow-direct-references = true
+
+[tool.hatch.build.targets.wheel]
+packages = ["."]
+
+[tool.black]
+line-length = 100
+target-version = ['py310']
+
+[tool.isort]
+profile = "black"
+line_length = 100
+
+[tool.mypy]
+python_version = "3.10"
+warn_return_any = true
+warn_unused_configs = true
+disallow_untyped_defs = true
+
+[tool.ruff]
+line-length = 100
+target-version = "py310"
+
+[tool.ruff.lint]
+select = ["E", "F", "W", "C", "I"]
+ignore = ["E501"]

python/assistant-transport-backend/requirements.txt

@@ -0,0 +1,14 @@
+# Core dependencies
+fastapi>=0.100.0
+uvicorn[standard]>=0.20.0  
+assistant-stream>=0.0.28
+pydantic>=2.0.0
+python-dotenv>=1.0.0
+
+# Development dependencies
+pytest>=7.0.0
+pytest-asyncio>=0.21.0
+black>=23.0.0
+isort>=5.12.0
+mypy>=1.0.0
+ruff>=0.1.0

python/assistant-transport-backend/setup.py

@@ -0,0 +1,151 @@
+#!/usr/bin/env python3
+"""
+Setup script for Assistant Transport Backend.
+
+This script helps users get started by:
+1. Checking Python version
+2. Creating a virtual environment (optional)
+3. Installing dependencies
+4. Setting up environment variables
+"""
+
+import os
+import sys
+import subprocess
+from pathlib import Path
+
+
+def check_python_version():
+    """Check if Python version is compatible."""
+    version = sys.version_info
+    if version.major < 3 or (version.major == 3 and version.minor < 9):
+        print("❌ Python 3.9+ is required")
+        sys.exit(1)
+    print(f"✅ Python {version.major}.{version.minor}.{version.micro} detected")
+
+
+def run_command(command, description):
+    """Run a shell command and handle errors."""
+    print(f"🔄 {description}...")
+    try:
+        result = subprocess.run(command, shell=True, check=True, capture_output=True, text=True)
+        print(f"✅ {description} completed")
+        return result.stdout
+    except subprocess.CalledProcessError as e:
+        print(f"❌ {description} failed: {e}")
+        print(f"Error output: {e.stderr}")
+        return None
+
+
+def main():
+    """Main setup function."""
+    print("🚀 Assistant Transport Backend Setup Script")
+    print("=" * 40)
+    
+    # Check Python version
+    check_python_version()
+    
+    # Check if we're in the right directory
+    if not Path("pyproject.toml").exists():
+        print("❌ Please run this script from the assistant-transport-backend directory")
+        sys.exit(1)
+    
+    # Ask user about virtual environment
+    create_venv = input("\n💡 Create a virtual environment? (y/N): ").lower().strip()
+    if create_venv in ['y', 'yes']:
+        venv_name = input("Virtual environment name (default: venv): ").strip() or "venv"
+        
+        if not Path(venv_name).exists():
+            run_command(f"python -m venv {venv_name}", f"Creating virtual environment '{venv_name}'")
+        
+        # Provide activation instructions
+        if os.name == 'nt':  # Windows
+            activate_cmd = f"{venv_name}\\Scripts\\activate"
+        else:  # Unix/Linux/macOS
+            activate_cmd = f"source {venv_name}/bin/activate"
+        
+        print(f"\n💡 To activate the virtual environment, run:")
+        print(f"   {activate_cmd}")
+        print("💡 Then re-run this setup script.")
+        return
+    
+    # Install dependencies
+    print("\n📦 Installing dependencies...")
+    if Path("pyproject.toml").exists():
+        # Try pip install with pyproject.toml (editable install)
+        result = run_command("pip install -e .", "Installing package in editable mode")
+        if not result:
+            # Fallback to requirements.txt
+            run_command("pip install -r requirements.txt", "Installing from requirements.txt")
+    else:
+        run_command("pip install -r requirements.txt", "Installing from requirements.txt")
+    
+    # Set up environment variables
+    print("\n🔧 Setting up environment variables...")
+    env_file = Path(".env")
+    example_env = Path(".env.example")
+    
+    if not env_file.exists() and example_env.exists():
+        # Copy example env file
+        with open(example_env, 'r') as f:
+            example_content = f.read()
+        
+        with open(env_file, 'w') as f:
+            f.write(example_content)
+        
+        print("✅ Created .env file from .env.example")
+        print("💡 Please edit .env file to add your API keys:")
+        print("   - OPENAI_API_KEY (for OpenAI GPT models)")
+        print("   - ANTHROPIC_API_KEY (for Anthropic Claude models)")
+    elif env_file.exists():
+        print("✅ .env file already exists")
+    
+    # Final instructions
+    print("\n🎉 Setup complete!")
+    print("\n🚀 To start the server:")
+    print("   python main.py")
+    print("\n   Or:")
+    print("   uvicorn main:app --reload --host 0.0.0.0 --port 8000")
+    print("\n📖 The server will be available at http://localhost:8000")
+    print("🌐 API endpoint will be at http://localhost:8000/assistant")
+    
+    # Test server start
+    test_server = input("\n🧪 Test server startup? (y/N): ").lower().strip()
+    if test_server in ['y', 'yes']:
+        print("\n🧪 Testing server startup (will stop after 5 seconds)...")
+        try:
+            import signal
+            import time
+            from multiprocessing import Process
+            
+            # Import and start server
+            from main import app
+            import uvicorn
+            
+            def run_test_server():
+                uvicorn.run(app, host="127.0.0.1", port=8000, log_level="warning")
+            
+            process = Process(target=run_test_server)
+            process.start()
+            time.sleep(2)
+            
+            # Test if server responds
+            try:
+                import httpx
+                response = httpx.get("http://127.0.0.1:8000/health", timeout=3.0)
+                if response.status_code == 200:
+                    print("✅ Server test successful!")
+                else:
+                    print(f"⚠️ Server responded with status {response.status_code}")
+            except Exception as e:
+                print(f"⚠️ Could not test server: {e}")
+            
+            process.terminate()
+            process.join()
+            
+        except Exception as e:
+            print(f"⚠️ Server test failed: {e}")
+
+
+if __name__ == "__main__":
+    main()

python/assistant-ui-sync-server-api/.gitignore

@@ -0,0 +1,52 @@
+# Python
+__pycache__/
+*.py[cod]
+*$py.class
+*.so
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# Virtual environments
+.venv/
+venv/
+ENV/
+env/
+
+# uv
+.venv/
+uv.lock
+
+# Testing
+.pytest_cache/
+.coverage
+htmlcov/
+.tox/
+.nox/
+coverage.xml
+*.cover
+
+# IDEs
+.idea/
+.vscode/
+*.swp
+*.swo
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json

python/assistant-ui-sync-server-api/.python-version

@@ -0,0 +1 @@
+3.12.6

python/assistant-ui-sync-server-api/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 AgentbaseAI Inc.
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

python/assistant-ui-sync-server-api/README.md

@@ -0,0 +1,224 @@
+# assistant-ui Sync Server API
+
+A Python client library for interacting with assistant-ui sync server backends, providing the same API structure as the JavaScript/TypeScript `useChatRuntime`.
+
+## Installation
+
+```bash
+pip install assistant-ui-sync-server-api
+```
+
+Or using uv:
+
+```bash
+uv add assistant-ui-sync-server-api
+```
+
+## Usage
+
+### Basic Example
+
+```python
+import asyncio
+from assistant_ui import AssistantClient
+
+# Create a client
+client = AssistantClient(base_url="https://api.example.com")
+
+# Create messages
+messages = [
+    {
+        "role": "user",
+        "content": [{"type": "text", "text": "Hello, how are you?"}]
+    }
+]
+
+# Send a chat request to a specific thread
+async def main():
+    thread = client.threads("thread-123")
+    response = await thread.chat(
+        messages=messages,
+        system="You are a helpful assistant."
+    )
+    print(response.json())
+    await client.close()
+
+asyncio.run(main())
+```
+
+### Using Context Managers
+
+```python
+# Async context manager
+async with AssistantClient(base_url="https://api.example.com") as client:
+    thread = client.threads("thread-123")
+    response = await thread.chat(messages=messages)
+
+# Sync context manager
+with AssistantClient(base_url="https://api.example.com") as client:
+    thread = client.threads("thread-123")
+    response = thread.chat_sync(messages=messages)
+```
+
+### Authentication
+
+```python
+# Static headers
+client = AssistantClient(
+    base_url="https://api.example.com",
+    headers={"Authorization": "Bearer your-token"}
+)
+
+# Dynamic headers (async)
+async def get_auth_headers():
+    token = await fetch_token()
+    return {"Authorization": f"Bearer {token}"}
+
+client = AssistantClient(
+    base_url="https://api.example.com",
+    headers=get_auth_headers
+)
+```
+
+### Using Tools
+
+```python
+tools = {
+    "get_weather": {
+        "description": "Get the current weather for a location",
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "location": {"type": "string"},
+                "unit": {"type": "string", "enum": ["celsius", "fahrenheit"]}
+            },
+            "required": ["location"]
+        }
+    }
+}
+
+thread = client.threads("thread-123")
+response = await thread.chat(
+    messages=messages,
+    tools=tools
+)
+```
+
+### Canceling Operations
+
+```python
+thread = client.threads("thread-123")
+
+# Start a long-running chat
+chat_task = asyncio.create_task(thread.chat(messages=messages))
+
+# Cancel it
+await thread.cancel()
+chat_task.cancel()
+```
+
+### Message Types
+
+The package supports various message types matching the assistant-ui format:
+
+```python
+from assistant_ui.types import Message
+
+# Text message
+text_message: Message = {
+    "role": "user",
+    "content": [{"type": "text", "text": "Hello!"}]
+}
+
+# Image message
+image_message: Message = {
+    "role": "user",
+    "content": [
+        {"type": "text", "text": "What's in this image?"},
+        {"type": "image", "image": "https://example.com/image.jpg"}
+    ]
+}
+
+# File message
+file_message: Message = {
+    "role": "user",
+    "content": [
+        {"type": "file", "data": "https://example.com/file.pdf", "mimeType": "application/pdf"}
+    ]
+}
+
+# Assistant message with tool calls
+assistant_message: Message = {
+    "role": "assistant",
+    "content": [
+        {"type": "text", "text": "I'll help you with that."},
+        {
+            "type": "tool-call",
+            "toolCallId": "call-123",
+            "toolName": "get_weather",
+            "args": {"location": "San Francisco"}
+        }
+    ]
+}
+
+# Tool result message
+tool_message: Message = {
+    "role": "tool",
+    "content": [
+        {
+            "type": "tool-result",
+            "toolCallId": "call-123",
+            "toolName": "get_weather",
+            "result": {"temperature": 72, "condition": "sunny"},
+            "isError": False
+        }
+    ]
+}
+```
+
+## API Reference
+
+### `AssistantClient`
+
+Main client for interacting with assistant-ui backends.
+
+**Constructor:**
+
+```python
+AssistantClient(
+    base_url: str,
+    headers: Optional[Dict[str, str] | Callable] = None,
+    timeout: Optional[float] = None,
+    **kwargs
+)
+```
+
+**Methods:**
+
+- `threads(thread_id: str) -> ThreadClient`: Get a ThreadClient for a specific thread
+- `close()`: Close the async client
+- `close_sync()`: Close the sync client
+
+### `ThreadClient`
+
+Client for interacting with a specific thread.
+
+**Methods:**
+
+- `chat(messages, system=None, tools=None, **kwargs)`: Send an async chat request
+- `chat_sync(messages, system=None, tools=None, **kwargs)`: Send a sync chat request
+- `cancel()`: Cancel the current async operation
+- `cancel_sync()`: Cancel the current sync operation
+
+## Type Definitions
+
+The package includes TypedDict definitions for all message types and configuration options, providing full type hints for better IDE support and type checking.
+
+## Development
+
+This package is part of the assistant-ui monorepo. To contribute:
+
+1. Clone the main repository
+2. Navigate to `python/assistant-ui`
+3. Install dependencies with `uv sync`
+4. Run tests with `uv run pytest`

python/assistant-ui-sync-server-api/examples/basic_example.py

@@ -0,0 +1,215 @@
+#!/usr/bin/env python3
+"""
+Basic example of using the assistant-ui Sync Server API client.
+"""
+
+import asyncio
+from assistant_ui import AssistantClient
+from assistant_ui.types import Message, Tool
+
+
+async def basic_chat_example():
+    """Simple async chat example."""
+    print("=== Basic Async Chat Example ===")
+    
+    # Create client
+    client = AssistantClient(
+        base_url="http://localhost:3000",
+        headers={"Authorization": "Bearer your-api-key"}
+    )
+    
+    messages: list[Message] = [
+        {
+            "role": "user",
+            "content": [{"type": "text", "text": "What is the capital of France?"}],
+        }
+    ]
+    
+    try:
+        # Get thread client and send chat
+        thread = client.threads("thread-123")
+        response = await thread.chat(
+            messages=messages,
+            system="You are a helpful assistant that answers questions concisely."
+        )
+        print(f"Response status: {response.status_code}")
+        print(f"Response body: {response.json()}")
+    except Exception as e:
+        print(f"Error: {e}")
+    finally:
+        await client.close()
+
+
+def sync_chat_example():
+    """Simple synchronous chat example."""
+    print("\n=== Sync Chat Example ===")
+    
+    # Create client using context manager
+    with AssistantClient(
+        base_url="http://localhost:3000",
+        headers={"Authorization": "Bearer your-api-key"}
+    ) as client:
+        messages: list[Message] = [
+            {
+                "role": "user",
+                "content": [{"type": "text", "text": "Tell me a short joke."}],
+            }
+        ]
+        
+        try:
+            # Get thread client and send chat
+            thread = client.threads("thread-456")
+            response = thread.chat_sync(
+                messages=messages,
+                system="You are a helpful and funny assistant."
+            )
+            print(f"Response status: {response.status_code}")
+            print(f"Response body: {response.json()}")
+        except Exception as e:
+            print(f"Error: {e}")
+
+
+async def chat_with_tools_example():
+    """Example using tools."""
+    print("\n=== Chat with Tools Example ===")
+    
+    # Using async context manager
+    async with AssistantClient(base_url="http://localhost:3000") as client:
+        tools: dict[str, Tool] = {
+            "get_weather": {
+                "description": "Get the current weather for a location",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "location": {
+                            "type": "string",
+                            "description": "The city and state, e.g., San Francisco, CA"
+                        },
+                        "unit": {
+                            "type": "string",
+                            "enum": ["celsius", "fahrenheit"],
+                            "description": "The temperature unit"
+                        }
+                    },
+                    "required": ["location"],
+                }
+            },
+            "calculate": {
+                "description": "Perform mathematical calculations",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "expression": {
+                            "type": "string",
+                            "description": "The mathematical expression to evaluate"
+                        }
+                    },
+                    "required": ["expression"],
+                }
+            }
+        }
+        
+        messages: list[Message] = [
+            {
+                "role": "user",
+                "content": [{"type": "text", "text": "What's the weather in New York?"}],
+            }
+        ]
+        
+        try:
+            thread = client.threads("thread-789")
+            response = await thread.chat(
+                messages=messages,
+                tools=tools,
+                system="You are a helpful weather assistant."
+            )
+            print(f"Response status: {response.status_code}")
+            print(f"Response body: {response.json()}")
+        except Exception as e:
+            print(f"Error: {e}")
+
+
+async def cancel_example():
+    """Example of canceling a thread operation."""
+    print("\n=== Cancel Example ===")
+    
+    async with AssistantClient(base_url="http://localhost:3000") as client:
+        thread = client.threads("thread-cancel-demo")
+        
+        try:
+            # Start a chat that might take a long time
+            messages: list[Message] = [
+                {
+                    "role": "user",
+                    "content": [{"type": "text", "text": "Tell me a very long story..."}],
+                }
+            ]
+            
+            # Start chat in background
+            chat_task = asyncio.create_task(thread.chat(messages=messages))
+            
+            # Wait a bit then cancel
+            await asyncio.sleep(0.5)
+            print("Canceling operation...")
+            cancel_response = await thread.cancel()
+            print(f"Cancel response: {cancel_response.status_code}")
+            
+            # Cancel the chat task
+            chat_task.cancel()
+            
+        except asyncio.CancelledError:
+            print("Chat was cancelled")
+        except Exception as e:
+            print(f"Error: {e}")
+
+
+async def multimodal_example():
+    """Example with images and files."""
+    print("\n=== Multimodal Example ===")
+    
+    # Example with async header function
+    async def get_auth_headers():
+        # Simulate fetching auth token
+        await asyncio.sleep(0.1)
+        return {"Authorization": "Bearer dynamic-token"}
+    
+    client = AssistantClient(
+        base_url="http://localhost:3000",
+        headers=get_auth_headers
+    )
+    
+    messages: list[Message] = [
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": "Can you analyze this image and document?"},
+                {"type": "image", "image": "https://example.com/chart.png"},
+                {"type": "file", "data": "https://example.com/report.pdf", "mimeType": "application/pdf"}
+            ],
+        }
+    ]
+    
+    try:
+        thread = client.threads("thread-multimodal")
+        response = await thread.chat(
+            messages=messages,
+            system="You are an expert at analyzing visual and document content."
+        )
+        print(f"Response: {response.json()}")
+    except Exception as e:
+        print(f"Error: {e}")
+    finally:
+        await client.close()
+
+
+async def main():
+    """Run all examples."""
+    await basic_chat_example()
+    sync_chat_example()
+    await chat_with_tools_example()
+    await cancel_example()
+    await multimodal_example()
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

python/assistant-ui-sync-server-api/pyproject.toml

@@ -0,0 +1,51 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "assistant-ui-sync-server-api"
+version = "0.1.1"
+description = "Python client for assistant-ui sync server API"
+authors = [
+  { name="assistant-ui Team", email="team@assistant-ui.com" },
+]
+readme = "README.md"
+requires-python = ">=3.10"
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "License :: OSI Approved :: MIT License",
+    "Operating System :: OS Independent",
+    "Development Status :: 3 - Alpha",
+    "Intended Audience :: Developers",
+    "Topic :: Software Development :: Libraries :: Python Modules",
+]
+dependencies = [
+    "httpx>=0.24.0",
+    "typing-extensions>=4.0.0 ; python_full_version < '3.11'",
+]
+
+[project.optional-dependencies]
+dev = [
+    "pytest>=7.0.0",
+    "pytest-asyncio>=0.21.0",
+    "black>=23.0.0",
+    "isort>=5.0.0",
+    "mypy>=1.0.0",
+]
+
+[project.urls]
+Homepage = "https://github.com/assistant-ui/assistant-ui"
+Issues = "https://github.com/assistant-ui/assistant-ui/issues"
+Documentation = "https://docs.assistant-ui.com"
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/assistant_ui"]
+
+[dependency-groups]
+dev = [
+    "black>=24.8.0",
+    "isort>=5.13.2",
+    "mypy>=1.14.1",
+    "pytest>=8.3.5",
+    "pytest-asyncio>=0.24.0",
+]

python/assistant-ui-sync-server-api/scripts/deploy.sh

@@ -0,0 +1,24 @@
+#!/bin/bash
+set -e
+
+echo "🚀 Deploying assistant-ui-sync-server-api to PyPI"
+
+# Clean previous builds
+echo "🧹 Cleaning previous builds..."
+rm -rf dist/ build/ *.egg-info src/*.egg-info
+
+# Build the package
+echo "📦 Building package..."
+python -m pip install --upgrade build twine
+python -m build
+
+# Check the package
+echo "🔍 Checking package..."
+twine check dist/*
+
+# Upload to PyPI
+echo "📤 Uploading to PyPI..."
+echo "Note: You'll need to authenticate with your PyPI credentials"
+twine upload dist/*
+
+echo "✅ Deployment complete!"

python/assistant-ui-sync-server-api/src/assistant_ui/__init__.py

@@ -0,0 +1,32 @@
+from .client import AssistantClient, ThreadClient
+from .types import (
+    Message,
+    TextPart,
+    ImagePart,
+    FilePart,
+    ToolCallPart,
+    ToolResultPart,
+    Tool,
+    ToolParameters,
+    SystemMessage,
+    UserMessage,
+    AssistantMessage,
+    ToolMessage,
+)
+
+__all__ = [
+    "AssistantClient",
+    "ThreadClient",
+    "Message",
+    "SystemMessage",
+    "UserMessage",
+    "AssistantMessage",
+    "ToolMessage",
+    "TextPart",
+    "ImagePart",
+    "FilePart",
+    "ToolCallPart",
+    "ToolResultPart",
+    "Tool",
+    "ToolParameters",
+]

python/assistant-ui-sync-server-api/src/assistant_ui/client.py

@@ -0,0 +1,451 @@
+from typing import Dict, List, Any, Optional, Union, Callable, Awaitable
+import warnings
+import httpx
+from .types import (
+    Message,
+    Tool,
+    AssistantTransportCommand,
+    AddMessageCommand,
+    AddToolResultCommand
+)
+
+
+def _convert_messages_to_commands(messages: List[Message]) -> List[AssistantTransportCommand]:
+    """
+    Convert legacy messages format to commands format.
+
+    Args:
+        messages: List of messages in the legacy format
+
+    Returns:
+        List of commands in the new format
+    """
+    commands: List[AssistantTransportCommand] = []
+
+    for message in messages:
+        role = message.get("role")
+
+        if role == "system":
+            # System messages can't be directly converted to commands
+            # They should be passed via the system parameter instead
+            continue
+
+        elif role == "user":
+            content = message.get("content", [])
+            parts = []
+
+            for part in content:
+                part_type = part.get("type")
+                if part_type == "text":
+                    parts.append({"type": "text", "text": part.get("text", "")})
+                elif part_type == "image":
+                    parts.append({"type": "image", "image": part.get("image", "")})
+                # FilePart cannot be converted to command format
+
+            if parts:
+                command: AddMessageCommand = {
+                    "type": "add-message",
+                    "message": {
+                        "role": "user",
+                        "parts": parts
+                    }
+                }
+                commands.append(command)
+
+        elif role == "assistant":
+            content = message.get("content", [])
+            text_parts = []
+
+            for part in content:
+                part_type = part.get("type")
+                if part_type == "text":
+                    text_parts.append({"type": "text", "text": part.get("text", "")})
+                elif part_type == "tool-call":
+                    # Tool calls are handled separately, not in add-message
+                    continue
+
+            if text_parts:
+                command: AddMessageCommand = {
+                    "type": "add-message",
+                    "message": {
+                        "role": "assistant",
+                        "parts": text_parts
+                    }
+                }
+                commands.append(command)
+
+        elif role == "tool":
+            # Convert tool results to commands
+            content = message.get("content", [])
+            for part in content:
+                if part.get("type") == "tool-result":
+                    tool_command: AddToolResultCommand = {
+                        "type": "add-tool-result",
+                        "toolCallId": part.get("toolCallId", ""),
+                        "toolName": part.get("toolName", ""),
+                        "result": part.get("result"),
+                        "isError": part.get("isError", False),
+                        "artifact": None  # artifact not available in legacy format
+                    }
+                    commands.append(tool_command)
+
+    return commands
+
+
+class ThreadClient:
+    """Client for interacting with a specific thread."""
+    
+    def __init__(self, client: "AssistantClient", thread_id: str):
+        self._client = client
+        self._thread_id = thread_id
+    
+    async def chat(
+        self,
+        messages: Optional[List[Message]] = None,
+        commands: Optional[List[AssistantTransportCommand]] = None,
+        system: Optional[str] = None,
+        tools: Optional[Dict[str, Tool]] = None,
+        unstable_assistantMessageId: Optional[str] = None,
+        runConfig: Optional[Dict[str, Any]] = None,
+        state: Optional[Any] = None,
+        **kwargs: Any
+    ) -> httpx.Response:
+        """
+        Send a chat request for this thread.
+
+        Args:
+            messages: (Deprecated) List of messages in the conversation
+            commands: List of commands to execute
+            system: System prompt
+            tools: Dictionary of available tools
+            unstable_assistantMessageId: Optional assistant message ID
+            runConfig: Optional run configuration
+            state: Optional state data
+            **kwargs: Additional parameters to include in the request body
+
+        Returns:
+            httpx.Response object containing the backend response
+        """
+        # Build request payload
+        payload = {
+            "threadId": self._thread_id,
+        }
+
+        # Handle commands and messages
+        if messages is not None:
+            warnings.warn(
+                "The 'messages' parameter is deprecated and will be removed in a future version. "
+                "Use 'commands' parameter instead.",
+                DeprecationWarning,
+                stacklevel=2
+            )
+
+            # Convert messages to commands
+            converted_commands = _convert_messages_to_commands(messages)
+
+            # If both messages and commands provided, merge them
+            if commands is not None:
+                commands = commands + converted_commands
+            else:
+                commands = converted_commands
+
+            # Still send messages for backward compatibility
+            payload["messages"] = messages
+
+        if commands is not None:
+            payload["commands"] = commands
+
+        if system is not None:
+            payload["system"] = system
+
+        if tools is not None:
+            payload["tools"] = tools
+
+        if unstable_assistantMessageId is not None:
+            payload["unstable_assistantMessageId"] = unstable_assistantMessageId
+
+        if runConfig is not None:
+            payload["runConfig"] = runConfig
+
+        if state is not None:
+            payload["state"] = state
+
+        # Add any additional kwargs
+        payload.update(kwargs)
+
+        # Make the request
+        return await self._client._make_request("POST", "/api/chat", json=payload)
+    
+    def chat_sync(
+        self,
+        messages: Optional[List[Message]] = None,
+        commands: Optional[List[AssistantTransportCommand]] = None,
+        system: Optional[str] = None,
+        tools: Optional[Dict[str, Tool]] = None,
+        unstable_assistantMessageId: Optional[str] = None,
+        runConfig: Optional[Dict[str, Any]] = None,
+        state: Optional[Any] = None,
+        **kwargs: Any
+    ) -> httpx.Response:
+        """
+        Synchronous version of chat method.
+
+        Args:
+            messages: (Deprecated) List of messages in the conversation
+            commands: List of commands to execute
+            system: System prompt
+            tools: Dictionary of available tools
+            unstable_assistantMessageId: Optional assistant message ID
+            runConfig: Optional run configuration
+            state: Optional state data
+            **kwargs: Additional parameters to include in the request body
+
+        Returns:
+            httpx.Response object containing the backend response
+        """
+        # Build request payload
+        payload = {
+            "threadId": self._thread_id,
+        }
+
+        # Handle commands and messages
+        if messages is not None:
+            warnings.warn(
+                "The 'messages' parameter is deprecated and will be removed in a future version. "
+                "Use 'commands' parameter instead.",
+                DeprecationWarning,
+                stacklevel=2
+            )
+
+            # Convert messages to commands
+            converted_commands = _convert_messages_to_commands(messages)
+
+            # If both messages and commands provided, merge them
+            if commands is not None:
+                commands = commands + converted_commands
+            else:
+                commands = converted_commands
+
+            # Still send messages for backward compatibility
+            payload["messages"] = messages
+
+        if commands is not None:
+            payload["commands"] = commands
+
+        if system is not None:
+            payload["system"] = system
+
+        if tools is not None:
+            payload["tools"] = tools
+
+        if unstable_assistantMessageId is not None:
+            payload["unstable_assistantMessageId"] = unstable_assistantMessageId
+
+        if runConfig is not None:
+            payload["runConfig"] = runConfig
+
+        if state is not None:
+            payload["state"] = state
+
+        # Add any additional kwargs
+        payload.update(kwargs)
+
+        # Make the request
+        return self._client._make_request_sync("POST", "/api/chat", json=payload)
+    
+    async def cancel(self) -> httpx.Response:
+        """
+        Cancel the current operation for this thread.
+        
+        Returns:
+            httpx.Response object containing the backend response
+        """
+        return await self._client._make_request(
+            "POST", 
+            "/api/cancel",
+            json={"threadId": self._thread_id}
+        )
+    
+    def cancel_sync(self) -> httpx.Response:
+        """
+        Synchronous version of cancel method.
+        
+        Returns:
+            httpx.Response object containing the backend response
+        """
+        return self._client._make_request_sync(
+            "POST", 
+            "/api/cancel",
+            json={"threadId": self._thread_id}
+        )
+
+
+class AssistantClient:
+    """Main client for interacting with assistant-ui backends."""
+    
+    def __init__(
+        self,
+        base_url: str,
+        headers: Optional[Union[Dict[str, str], Callable[[], Union[Dict[str, str], Awaitable[Dict[str, str]]]]]] = None,
+        timeout: Optional[float] = None,
+        **kwargs: Any
+    ):
+        """
+        Initialize the AssistantClient.
+        
+        Args:
+            base_url: Base URL for the API (e.g., "https://api.example.com")
+            headers: Optional headers to include with requests
+            timeout: Optional timeout for requests
+            **kwargs: Additional arguments passed to httpx client
+        """
+        self.base_url = base_url.rstrip("/")
+        self._headers = headers
+        self._timeout = timeout
+        self._client_kwargs = kwargs
+        self._async_client: Optional[httpx.AsyncClient] = None
+        self._sync_client: Optional[httpx.Client] = None
+    
+    def threads(self, thread_id: str) -> ThreadClient:
+        """
+        Get a ThreadClient for a specific thread.
+        
+        Args:
+            thread_id: The ID of the thread
+            
+        Returns:
+            ThreadClient instance for the specified thread
+        """
+        return ThreadClient(self, thread_id)
+    
+    async def _get_headers(self) -> Dict[str, str]:
+        """Get headers, resolving async functions if needed."""
+        if self._headers is None:
+            return {"Content-Type": "application/json"}
+        
+        if callable(self._headers):
+            headers = self._headers()
+            if hasattr(headers, "__await__"):
+                headers = await headers
+        else:
+            headers = self._headers
+        
+        if isinstance(headers, dict):
+            headers = dict(headers)  # Make a copy
+            headers.setdefault("Content-Type", "application/json")
+        else:
+            headers = {"Content-Type": "application/json"}
+        
+        return headers
+    
+    def _get_headers_sync(self) -> Dict[str, str]:
+        """Get headers synchronously."""
+        if self._headers is None:
+            return {"Content-Type": "application/json"}
+        
+        if callable(self._headers):
+            headers = self._headers()
+            if hasattr(headers, "__await__"):
+                raise ValueError("Synchronous methods do not support async header functions")
+        else:
+            headers = self._headers
+        
+        if isinstance(headers, dict):
+            headers = dict(headers)  # Make a copy
+            headers.setdefault("Content-Type", "application/json")
+        else:
+            headers = {"Content-Type": "application/json"}
+        
+        return headers
+    
+    async def _ensure_async_client(self) -> httpx.AsyncClient:
+        """Ensure async client is initialized."""
+        if self._async_client is None:
+            self._async_client = httpx.AsyncClient(
+                base_url=self.base_url,
+                timeout=self._timeout,
+                **self._client_kwargs
+            )
+        return self._async_client
+    
+    def _ensure_sync_client(self) -> httpx.Client:
+        """Ensure sync client is initialized."""
+        if self._sync_client is None:
+            self._sync_client = httpx.Client(
+                base_url=self.base_url,
+                timeout=self._timeout,
+                **self._client_kwargs
+            )
+        return self._sync_client
+    
+    async def _make_request(
+        self,
+        method: str,
+        path: str,
+        **kwargs: Any
+    ) -> httpx.Response:
+        """Make an async HTTP request."""
+        client = await self._ensure_async_client()
+        headers = await self._get_headers()
+        
+        # Merge headers
+        if "headers" in kwargs:
+            headers.update(kwargs["headers"])
+        kwargs["headers"] = headers
+        
+        response = await client.request(method, path, **kwargs)
+        
+        if not response.is_success:
+            raise Exception(f"Request failed with status {response.status_code}: {response.text}")
+        
+        return response
+    
+    def _make_request_sync(
+        self,
+        method: str,
+        path: str,
+        **kwargs: Any
+    ) -> httpx.Response:
+        """Make a synchronous HTTP request."""
+        client = self._ensure_sync_client()
+        headers = self._get_headers_sync()
+        
+        # Merge headers
+        if "headers" in kwargs:
+            headers.update(kwargs["headers"])
+        kwargs["headers"] = headers
+        
+        response = client.request(method, path, **kwargs)
+        
+        if not response.is_success:
+            raise Exception(f"Request failed with status {response.status_code}: {response.text}")
+        
+        return response
+    
+    async def close(self):
+        """Close the async client if it's open."""
+        if self._async_client:
+            await self._async_client.aclose()
+            self._async_client = None
+    
+    def close_sync(self):
+        """Close the sync client if it's open."""
+        if self._sync_client:
+            self._sync_client.close()
+            self._sync_client = None
+    
+    async def __aenter__(self):
+        """Async context manager entry."""
+        return self
+    
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        """Async context manager exit."""
+        await self.close()
+    
+    def __enter__(self):
+        """Sync context manager entry."""
+        return self
+    
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Sync context manager exit."""
+        self.close_sync()

python/assistant-ui-sync-server-api/src/assistant_ui/types.py

@@ -0,0 +1,101 @@
+from typing import TypedDict, Literal, Union, Dict, Any, List, Optional
+
+
+class TextPart(TypedDict):
+    type: Literal["text"]
+    text: str
+
+
+class ImagePart(TypedDict):
+    type: Literal["image"]
+    image: str
+
+
+class FilePart(TypedDict):
+    type: Literal["file"]
+    data: str
+    mimeType: str
+
+
+class ToolCallPart(TypedDict):
+    type: Literal["tool-call"]
+    toolCallId: str
+    toolName: str
+    args: Dict[str, Any]
+
+
+class ToolResultPart(TypedDict):
+    type: Literal["tool-result"]
+    toolCallId: str
+    toolName: str
+    result: Any
+    isError: Optional[bool]
+
+
+ContentPart = Union[TextPart, ImagePart, FilePart, ToolCallPart, ToolResultPart]
+
+
+class SystemMessage(TypedDict):
+    role: Literal["system"]
+    content: str
+    unstable_id: Optional[str]
+
+
+class UserMessage(TypedDict):
+    role: Literal["user"]
+    content: List[Union[TextPart, ImagePart, FilePart]]
+    unstable_id: Optional[str]
+
+
+class AssistantMessage(TypedDict):
+    role: Literal["assistant"]
+    content: List[Union[TextPart, ToolCallPart]]
+    unstable_id: Optional[str]
+
+
+class ToolMessage(TypedDict):
+    role: Literal["tool"]
+    content: List[ToolResultPart]
+
+
+Message = Union[SystemMessage, UserMessage, AssistantMessage, ToolMessage]
+
+
+# Command types for AssistantTransportCommand
+class UserMessageCommand(TypedDict):
+    role: Literal["user"]
+    parts: List[Union[TextPart, ImagePart]]
+
+
+class AssistantMessageCommand(TypedDict):
+    role: Literal["assistant"]
+    parts: List[TextPart]
+
+
+class AddMessageCommand(TypedDict):
+    type: Literal["add-message"]
+    message: Union[UserMessageCommand, AssistantMessageCommand]
+
+
+class AddToolResultCommand(TypedDict):
+    type: Literal["add-tool-result"]
+    toolCallId: str
+    toolName: str
+    result: Any
+    isError: bool
+    artifact: Optional[Any]
+
+
+AssistantTransportCommand = Union[AddMessageCommand, AddToolResultCommand]
+
+
+class ToolParameters(TypedDict):
+    type: Literal["object"]
+    properties: Dict[str, Any]
+    required: Optional[List[str]]
+    additionalProperties: Optional[bool]
+
+
+class Tool(TypedDict, total=False):
+    description: str
+    parameters: ToolParameters

python/assistant-ui-sync-server-api/tests/__init__.py

@@ -0,0 +1 @@
+# Tests for assistant-ui package

python/assistant-ui-sync-server-api/tests/test_client.py

@@ -0,0 +1,510 @@
+"""Tests for the assistant-ui-sync-server-api client."""
+
+import pytest
+import httpx
+import warnings
+from unittest.mock import AsyncMock, Mock, patch
+from assistant_ui import AssistantClient
+from assistant_ui.types import Message, AssistantTransportCommand
+
+
+@pytest.mark.asyncio
+async def test_basic_chat():
+    """Test basic async chat functionality."""
+    client = AssistantClient(
+        base_url="https://api.example.com",
+        headers={"Authorization": "Bearer test"}
+    )
+    
+    messages: list[Message] = [
+        {
+            "role": "user",
+            "content": [{"type": "text", "text": "Hello"}],
+        }
+    ]
+    
+    mock_response = AsyncMock()
+    mock_response.is_success = True
+    mock_response.status_code = 200
+    mock_response.json.return_value = {"response": "Hi there!"}
+    mock_response.text = "Success"
+    
+    with patch.object(client, "_ensure_async_client") as mock_ensure:
+        mock_client = AsyncMock()
+        mock_client.request = AsyncMock(return_value=mock_response)
+        mock_ensure.return_value = mock_client
+        
+        thread = client.threads("thread-123")
+
+        # Suppress deprecation warning for this test
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore", DeprecationWarning)
+            response = await thread.chat(
+                messages=messages,
+                system="You are helpful"
+            )
+
+        assert response.status_code == 200
+
+        # Check the request was made correctly
+        call_args = mock_client.request.call_args
+        payload = call_args[1]["json"]
+
+        assert payload["threadId"] == "thread-123"
+        assert payload["messages"] == messages
+        assert payload["system"] == "You are helpful"
+        # Commands should also be present due to conversion
+        assert "commands" in payload
+    
+    await client.close()
+
+
+@pytest.mark.asyncio
+async def test_chat_with_tools():
+    """Test chat with tools."""
+    async with AssistantClient(base_url="https://api.example.com") as client:
+        messages: list[Message] = []
+        tools = {"search": {"description": "Search the web"}}
+        
+        mock_response = AsyncMock()
+        mock_response.is_success = True
+        mock_response.status_code = 200
+        
+        with patch.object(client, "_ensure_async_client") as mock_ensure:
+            mock_client = AsyncMock()
+            mock_client.request = AsyncMock(return_value=mock_response)
+            mock_ensure.return_value = mock_client
+            
+            thread = client.threads("thread-456")
+            await thread.chat(messages=messages, tools=tools)
+            
+            # Check tools were included
+            call_args = mock_client.request.call_args
+            assert call_args[1]["json"]["tools"] == tools
+
+
+@pytest.mark.asyncio
+async def test_cancel():
+    """Test cancel functionality."""
+    client = AssistantClient(base_url="https://api.example.com")
+    
+    mock_response = AsyncMock()
+    mock_response.is_success = True
+    mock_response.status_code = 200
+    
+    with patch.object(client, "_ensure_async_client") as mock_ensure:
+        mock_client = AsyncMock()
+        mock_client.request = AsyncMock(return_value=mock_response)
+        mock_ensure.return_value = mock_client
+        
+        thread = client.threads("thread-789")
+        await thread.cancel()
+        
+        # Check cancel request was made
+        mock_client.request.assert_called_once_with(
+            "POST",
+            "/api/cancel",
+            headers={"Content-Type": "application/json"},
+            json={"threadId": "thread-789"}
+        )
+    
+    await client.close()
+
+
+def test_sync_chat():
+    """Test synchronous chat."""
+    with AssistantClient(base_url="https://api.example.com") as client:
+        messages: list[Message] = [
+            {
+                "role": "assistant",
+                "content": [
+                    {"type": "text", "text": "I can help."},
+                    {
+                        "type": "tool-call",
+                        "toolCallId": "call-123",
+                        "toolName": "search",
+                        "args": {"query": "weather"}
+                    }
+                ],
+            }
+        ]
+        
+        mock_response = Mock()
+        mock_response.is_success = True
+        mock_response.status_code = 200
+        mock_response.text = "Success"
+        
+        with patch.object(client, "_ensure_sync_client") as mock_ensure:
+            mock_client = Mock()
+            mock_client.request = Mock(return_value=mock_response)
+            mock_ensure.return_value = mock_client
+            
+            thread = client.threads("thread-sync")
+            response = thread.chat_sync(messages=messages)
+            
+            assert response.status_code == 200
+
+
+@pytest.mark.asyncio
+async def test_async_headers():
+    """Test with async header function."""
+    async def get_headers():
+        return {"Authorization": "Bearer dynamic-token"}
+    
+    client = AssistantClient(
+        base_url="https://api.example.com",
+        headers=get_headers
+    )
+    
+    mock_response = AsyncMock()
+    mock_response.is_success = True
+    mock_response.status_code = 200
+    
+    with patch.object(client, "_ensure_async_client") as mock_ensure:
+        mock_client = AsyncMock()
+        mock_client.request = AsyncMock(return_value=mock_response)
+        mock_ensure.return_value = mock_client
+        
+        thread = client.threads("thread-async-headers")
+        await thread.chat(messages=[])
+        
+        # Check dynamic headers were used
+        call_args = mock_client.request.call_args
+        assert call_args[1]["headers"]["Authorization"] == "Bearer dynamic-token"
+    
+    await client.close()
+
+
+def test_sync_with_async_headers_raises():
+    """Test that sync methods raise error with async headers."""
+    async def get_headers():
+        return {"Authorization": "Bearer token"}
+    
+    client = AssistantClient(
+        base_url="https://api.example.com",
+        headers=get_headers
+    )
+    
+    thread = client.threads("thread-123")
+    
+    with pytest.raises(ValueError) as exc_info:
+        thread.chat_sync(messages=[])
+    
+    assert "async header functions" in str(exc_info.value)
+
+
+@pytest.mark.asyncio
+async def test_error_handling():
+    """Test error handling."""
+    client = AssistantClient(base_url="https://api.example.com")
+    
+    mock_response = AsyncMock()
+    mock_response.is_success = False
+    mock_response.status_code = 500
+    mock_response.text = "Internal Server Error"
+    
+    with patch.object(client, "_ensure_async_client") as mock_ensure:
+        mock_client = AsyncMock()
+        mock_client.request = AsyncMock(return_value=mock_response)
+        mock_ensure.return_value = mock_client
+        
+        thread = client.threads("thread-error")
+        
+        with pytest.raises(Exception) as exc_info:
+            await thread.chat(messages=[])
+        
+        assert "Request failed with status 500" in str(exc_info.value)
+    
+    await client.close()
+
+
+@pytest.mark.asyncio
+async def test_context_manager():
+    """Test async context manager."""
+    async with AssistantClient(base_url="https://api.example.com") as client:
+        assert client._async_client is None  # Not created until first use
+        
+        # Force client creation
+        mock_response = AsyncMock()
+        mock_response.is_success = True
+        
+        with patch("httpx.AsyncClient") as mock_async_client_class:
+            mock_instance = AsyncMock()
+            mock_instance.request = AsyncMock(return_value=mock_response)
+            mock_async_client_class.return_value = mock_instance
+            
+            thread = client.threads("test")
+            await thread.chat(messages=[])
+            
+            assert client._async_client is not None
+    
+    # After context exit, client should be closed
+    assert client._async_client is None
+
+
+def test_sync_context_manager():
+    """Test sync context manager."""
+    with AssistantClient(base_url="https://api.example.com") as client:
+        assert client._sync_client is None  # Not created until first use
+        
+        # Force client creation
+        mock_response = Mock()
+        mock_response.is_success = True
+        
+        with patch("httpx.Client") as mock_client_class:
+            mock_instance = Mock()
+            mock_instance.request = Mock(return_value=mock_response)
+            mock_client_class.return_value = mock_instance
+            
+            thread = client.threads("test")
+            thread.chat_sync(messages=[])
+            
+            assert client._sync_client is not None
+    
+    # After context exit, client should be closed
+    assert client._sync_client is None
+
+
+@pytest.mark.asyncio
+async def test_all_parameters():
+    """Test chat with all possible parameters."""
+    client = AssistantClient(
+        base_url="https://api.example.com",
+        headers={"X-API-Key": "test-key"},
+        timeout=30.0
+    )
+    
+    messages: list[Message] = [
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": "Analyze this:"},
+                {"type": "image", "image": "https://example.com/img.jpg"},
+                {"type": "file", "data": "https://example.com/doc.pdf", "mimeType": "application/pdf"}
+            ],
+            "unstable_id": "msg-123",
+        }
+    ]
+    
+    tools = {
+        "analyze": {
+            "description": "Analyze content",
+            "parameters": {
+                "type": "object",
+                "properties": {"content": {"type": "string"}},
+                "required": ["content"]
+            }
+        }
+    }
+    
+    mock_response = AsyncMock()
+    mock_response.is_success = True
+    mock_response.status_code = 200
+    
+    with patch.object(client, "_ensure_async_client") as mock_ensure:
+        mock_client = AsyncMock()
+        mock_client.request = AsyncMock(return_value=mock_response)
+        mock_ensure.return_value = mock_client
+        
+        thread = client.threads("thread-all-params")
+        await thread.chat(
+            messages=messages,
+            tools=tools,
+            system="You are an analyzer",
+            unstable_assistantMessageId="asst-456",
+            runConfig={"model": "gpt-5.4-nano"},
+            state={"session": "abc"},
+            custom_field="custom_value"
+        )
+        
+        # Check all parameters were included
+        call_args = mock_client.request.call_args
+        payload = call_args[1]["json"]
+        
+        assert payload["threadId"] == "thread-all-params"
+        assert payload["system"] == "You are an analyzer"
+        assert payload["messages"] == messages
+        assert payload["tools"] == tools
+        assert payload["unstable_assistantMessageId"] == "asst-456"
+        assert payload["runConfig"] == {"model": "gpt-5.4-nano"}
+        assert payload["state"] == {"session": "abc"}
+        assert payload["custom_field"] == "custom_value"
+
+    await client.close()
+
+
+@pytest.mark.asyncio
+async def test_chat_with_commands():
+    """Test chat with commands instead of messages."""
+    client = AssistantClient(
+        base_url="https://api.example.com",
+        headers={"Authorization": "Bearer test"}
+    )
+
+    commands: list[AssistantTransportCommand] = [
+        {
+            "type": "add-message",
+            "message": {
+                "role": "user",
+                "parts": [{"type": "text", "text": "Hello with commands"}]
+            }
+        },
+        {
+            "type": "add-tool-result",
+            "toolCallId": "call-789",
+            "toolName": "search",
+            "result": {"data": "search results"},
+            "isError": False,
+            "artifact": None
+        }
+    ]
+
+    mock_response = AsyncMock()
+    mock_response.is_success = True
+    mock_response.status_code = 200
+
+    with patch.object(client, "_ensure_async_client") as mock_ensure:
+        mock_client = AsyncMock()
+        mock_client.request = AsyncMock(return_value=mock_response)
+        mock_ensure.return_value = mock_client
+
+        thread = client.threads("thread-cmd")
+        response = await thread.chat(commands=commands)
+
+        assert response.status_code == 200
+
+        # Check the request was made with commands
+        call_args = mock_client.request.call_args
+        payload = call_args[1]["json"]
+
+        assert payload["threadId"] == "thread-cmd"
+        assert payload["commands"] == commands
+        assert "messages" not in payload
+
+    await client.close()
+
+
+@pytest.mark.asyncio
+async def test_chat_with_deprecated_messages():
+    """Test that messages parameter shows deprecation warning."""
+    client = AssistantClient(base_url="https://api.example.com")
+
+    messages: list[Message] = [
+        {
+            "role": "user",
+            "content": [{"type": "text", "text": "Using deprecated API"}],
+        }
+    ]
+
+    mock_response = AsyncMock()
+    mock_response.is_success = True
+    mock_response.status_code = 200
+
+    with patch.object(client, "_ensure_async_client") as mock_ensure:
+        mock_client = AsyncMock()
+        mock_client.request = AsyncMock(return_value=mock_response)
+        mock_ensure.return_value = mock_client
+
+        thread = client.threads("thread-deprecated")
+
+        # Check deprecation warning
+        with warnings.catch_warnings(record=True) as w:
+            warnings.simplefilter("always")
+            await thread.chat(messages=messages)
+
+            assert len(w) == 1
+            assert issubclass(w[0].category, DeprecationWarning)
+            assert "messages' parameter is deprecated" in str(w[0].message)
+
+        # Check both messages and commands were sent
+        call_args = mock_client.request.call_args
+        payload = call_args[1]["json"]
+
+        assert payload["messages"] == messages
+        assert "commands" in payload
+        assert len(payload["commands"]) > 0
+
+    await client.close()
+
+
+@pytest.mark.asyncio
+async def test_chat_messages_and_commands_together():
+    """Test using both messages and commands together."""
+    client = AssistantClient(base_url="https://api.example.com")
+
+    messages: list[Message] = [
+        {
+            "role": "assistant",
+            "content": [{"type": "text", "text": "Previous message"}],
+        }
+    ]
+
+    commands: list[AssistantTransportCommand] = [
+        {
+            "type": "add-message",
+            "message": {
+                "role": "user",
+                "parts": [{"type": "text", "text": "New command"}]
+            }
+        }
+    ]
+
+    mock_response = AsyncMock()
+    mock_response.is_success = True
+    mock_response.status_code = 200
+
+    with patch.object(client, "_ensure_async_client") as mock_ensure:
+        mock_client = AsyncMock()
+        mock_client.request = AsyncMock(return_value=mock_response)
+        mock_ensure.return_value = mock_client
+
+        thread = client.threads("thread-both")
+
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore")
+            await thread.chat(messages=messages, commands=commands)
+
+        # Check both were merged
+        call_args = mock_client.request.call_args
+        payload = call_args[1]["json"]
+
+        assert payload["messages"] == messages
+        assert "commands" in payload
+        # Should have original command plus converted message
+        assert len(payload["commands"]) >= 2
+
+    await client.close()
+
+
+def test_sync_chat_with_commands():
+    """Test synchronous chat with commands."""
+    with AssistantClient(base_url="https://api.example.com") as client:
+        commands: list[AssistantTransportCommand] = [
+            {
+                "type": "add-message",
+                "message": {
+                    "role": "assistant",
+                    "parts": [{"type": "text", "text": "Sync command test"}]
+                }
+            }
+        ]
+
+        mock_response = Mock()
+        mock_response.is_success = True
+        mock_response.status_code = 200
+
+        with patch.object(client, "_ensure_sync_client") as mock_ensure:
+            mock_client = Mock()
+            mock_client.request = Mock(return_value=mock_response)
+            mock_ensure.return_value = mock_client
+
+            thread = client.threads("thread-sync-cmd")
+            response = thread.chat_sync(commands=commands)
+
+            assert response.status_code == 200
+
+            # Check commands were sent
+            call_args = mock_client.request.call_args
+            payload = call_args[1]["json"]
+
+            assert payload["commands"] == commands

python/state-test/README.md

@@ -0,0 +1,41 @@
+# State Management Test
+
+This is a test project for the `assistant-stream` state management functionality. It demonstrates various state operations and updates over time.
+
+## Features
+
+- **Simple State Test**: Basic state updates with primitive values
+- **Complex Test**: Nested state updates with objects and arrays
+- **String Operations**: String concatenation and method testing
+- **List Operations**: List manipulation with append, extend, and other operations
+- **Dictionary Operations**: Dictionary manipulation with various methods
+
+## Setup
+
+1. Install the requirements:
+
+   ```
+   pip install -r requirements.txt
+   ```
+
+2. Run the server:
+
+   ```
+   python server.py
+   ```
+
+3. Open your browser to [http://localhost:8000](http://localhost:8000)
+
+4. Click on the different test buttons to see state updates in action
+
+## Implementation Details
+
+This test server demonstrates the following state management features:
+
+- Primitive values (strings, numbers, booleans)
+- Nested state objects
+- String operations (concatenation, methods)
+- List operations (append, extend, indexing)
+- Dictionary operations (get, setdefault, keys/values)
+
+Each test endpoint updates state over time with various operations to showcase the functionality of the `StateProxy` and `StateManager` classes.

python/state-test/pyproject.toml

@@ -0,0 +1,14 @@
+[project]
+name = "state-test"
+version = "0.1.0"
+description = ""
+authors = [
+    {name = "Simon Farshid",email = "simon.farshid@outlook.com"}
+]
+readme = "README.md"
+requires-python = ">=3.12"
+dependencies = []
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"

python/state-test/requirements.txt

@@ -0,0 +1,3 @@
+fastapi>=0.68.0
+uvicorn>=0.15.0
+../assistant-stream

python/state-test/server.py

@@ -0,0 +1,205 @@
+from fastapi import FastAPI
+from fastapi.middleware.cors import CORSMiddleware
+from assistant_stream import create_run, RunController
+from assistant_stream.serialization import DataStreamResponse
+
+import asyncio
+
+app = FastAPI()
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+
+@app.post("/simple-test")
+async def simple_test():
+    async def run(controller: RunController):
+        # Check that state is initially None
+        assert controller.state is None, "Initial state should be None"
+        
+        # Initialize state with a direct assignment
+        controller.state = {
+            "counter": 0,
+            "message": "Starting simple test"
+        }
+
+        # Update state over time
+        for i in range(1, 6):
+            controller.state["counter"] = i
+            controller.state["message"] = f"Counter updated to {i}"
+
+        # Add a boolean and number
+        controller.state["completed"] = True
+        controller.state["total"] = 5.5
+        
+        # Demonstrate setting root state again
+        controller.state = {
+            "final": True,
+            "summary": "Test completed successfully"
+        }
+
+        controller.append_text("hi")
+
+    return DataStreamResponse(create_run(run))
+
+
+@app.post("/complex-test")
+async def complex_test():
+    async def run(controller: RunController):
+        # Demonstrate that state starts as None
+        print(f"Initial state: {controller.state}")
+        
+        # Initialize nested state directly
+        controller.state = {
+            "user": {
+                "name": "John",
+                "preferences": {
+                    "theme": "dark",
+                    "notifications": True
+                }
+            },
+            "messages": []
+        }
+        
+        # Add messages
+        controller.state["messages"].append("Hello")
+        controller.state["messages"].append("World")
+        
+        # Set state back to None to demonstrate nullifying
+        controller.state = {}
+        controller.state["user"] = {
+            "name": "Test User",
+            "settings": {"theme": "light", "notifications": True},
+        }
+        controller.state["stats"] = {"visits": 0, "actions": []}
+
+        # Update nested state
+        for i in range(1, 6):
+            await asyncio.sleep(0.01)
+            controller.state["stats"]["visits"] = i
+            controller.state["stats"]["actions"].append(f"action_{i}")
+
+            # Toggle theme every other iteration
+            if i % 2 == 0:
+                controller.state["user"]["settings"]["theme"] = "dark"
+            else:
+                controller.state["user"]["settings"]["theme"] = "light"
+
+        # Final update
+        await asyncio.sleep(0.01)
+        controller.state["completed"] = True
+
+    return DataStreamResponse(create_run(run))
+
+
+@app.post("/string-test")
+async def string_test():
+    async def run(controller: RunController):
+        # Initialize string
+        controller.state["message"] = "Hello"
+
+        # Append to string using +=
+        await asyncio.sleep(1)
+        controller.state["message"] += " world"
+
+        # Append more text
+        await asyncio.sleep(1)
+        controller.state["message"] += "!"
+
+        # String methods (these should work through the proxy)
+        await asyncio.sleep(1)
+        controller.state["uppercase"] = controller.state["message"].upper()
+
+        await asyncio.sleep(1)
+        controller.state["contains_world"] = "world" in controller.state["message"]
+
+        await asyncio.sleep(1)
+        controller.state["length"] = len(controller.state["message"])
+
+    return DataStreamResponse(create_run(run))
+
+
+@app.post("/list-test")
+async def list_test():
+    async def run(controller: RunController):
+        # Initialize list
+        controller.state["items"] = []
+
+        # Append items one by one
+        for i in range(5):
+            await asyncio.sleep(1)
+            controller.state["items"].append(f"item_{i}")
+
+        # Extend list with multiple items
+        await asyncio.sleep(1)
+        controller.state["items"] += ["batch_1", "batch_2"]
+
+        # Access by index
+        await asyncio.sleep(1)
+        controller.state["first_item"] = controller.state["items"][0]
+
+        # Get length
+        await asyncio.sleep(1)
+        controller.state["count"] = len(controller.state["items"])
+
+        # Check membership
+        await asyncio.sleep(1)
+        controller.state["contains_batch"] = "batch_1" in controller.state["items"]
+
+        # Clear the list
+        await asyncio.sleep(1)
+        controller.state["items"].clear()
+
+    return DataStreamResponse(create_run(run))
+
+
+@app.post("/dict-test")
+async def dict_test():
+    async def run(controller: RunController):
+        # Initialize dictionary
+        controller.state["config"] = {}
+
+        # Add items one by one
+        await asyncio.sleep(1)
+        controller.state["config"]["theme"] = "light"
+
+        await asyncio.sleep(1)
+        controller.state["config"]["language"] = "en"
+
+        # Set default value
+        await asyncio.sleep(1)
+        controller.state["config"].setdefault("notifications", True)
+
+        # Try to set default for existing key (should not change)
+        await asyncio.sleep(1)
+        controller.state["config"].setdefault("theme", "dark")
+        controller.state["theme_unchanged"] = (
+            controller.state["config"]["theme"] == "light"
+        )
+
+        # Get with default
+        await asyncio.sleep(1)
+        controller.state["missing_with_default"] = controller.state["config"].get(
+            "missing", "default_value"
+        )
+
+        # Keys, values, items
+        await asyncio.sleep(1)
+        controller.state["all_keys"] = list(controller.state["config"].keys())
+        controller.state["all_values"] = list(controller.state["config"].values())
+
+        # Clear the dictionary
+        await asyncio.sleep(1)
+        controller.state["config"].clear()
+
+    return DataStreamResponse(create_run(run))
+
+
+if __name__ == "__main__":
+    import uvicorn
+
+    uvicorn.run(app, host="0.0.0.0", port=8000)

python/state-test/uv.lock

@@ -0,0 +1,8 @@
+version = 1
+revision = 3
+requires-python = ">=3.12"
+
+[[package]]
+name = "state-test"
+version = "0.1.0"
+source = { editable = "." }