assistant-ui--assistant-ui/python/assistant-stream/tests/test_state_updates.py

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