chore: import upstream snapshot with attribution

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