chore: import upstream snapshot with attribution

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"},
+    ]