Update tests

tests/agui/__init__.py

@@ -0,0 +1 @@
+"""Tests for AG-UI implementation."""

tests/agui/test_cli_renderer.py

@@ -0,0 +1,158 @@
+"""Tests for AGUIConsoleRenderer."""
+
+import pytest
+from pydantic import BaseModel
+from rich.console import Console
+
+from haiku.rag.agui.cli_renderer import AGUIConsoleRenderer
+from haiku.rag.agui.events import (
+    emit_activity,
+    emit_run_error,
+    emit_run_finished,
+    emit_run_started,
+    emit_state_snapshot,
+    emit_step_finished,
+    emit_step_started,
+    emit_text_message,
+)
+
+
+class SimpleState(BaseModel):
+    """Simple state for testing."""
+
+    value: int
+
+
+async def async_gen(items):
+    """Helper to create async generator from list."""
+    for item in items:
+        yield item
+
+
+@pytest.mark.asyncio
+async def test_renderer_basic_flow():
+    """Test basic event rendering flow."""
+    console = Console(file=None, force_terminal=False)  # Don't actually print
+    renderer = AGUIConsoleRenderer(console)
+
+    events = [
+        emit_run_started("t1", "r1"),
+        emit_state_snapshot(SimpleState(value=1)),
+        emit_step_started("plan"),
+        emit_activity("m1", "planning", "Planning research"),
+        emit_run_finished("t1", "r1", {"status": "complete"}),
+    ]
+
+    result = await renderer.render(async_gen(events))
+
+    assert result == {"status": "complete"}
+
+
+@pytest.mark.asyncio
+async def test_renderer_state_diff():
+    """Test that renderer computes state diffs correctly."""
+    renderer = AGUIConsoleRenderer()
+
+    # First state
+    old_state = {"value": 1, "text": "old", "nested": {"a": 1, "b": 2}}
+    # Second state with changes
+    new_state = {"value": 2, "text": "old", "nested": {"a": 1, "b": 3, "c": 4}}
+
+    diff = renderer._compute_diff(old_state, new_state)
+
+    assert diff == {
+        "value": 2,
+        "nested": {"b": 3, "c": 4},  # Only changed/new fields in nested
+    }
+
+
+@pytest.mark.asyncio
+async def test_renderer_state_diff_no_changes():
+    """Test that no diff is computed when states are equal."""
+    renderer = AGUIConsoleRenderer()
+
+    state = {"value": 1, "text": "test"}
+    diff = renderer._compute_diff(state, state)
+
+    assert diff == {}
+
+
+@pytest.mark.asyncio
+async def test_renderer_handles_all_event_types():
+    """Test that renderer handles all event types without errors."""
+    renderer = AGUIConsoleRenderer()
+
+    events = [
+        emit_run_started("t1", "r1"),
+        emit_state_snapshot(SimpleState(value=1)),
+        emit_step_started("step1"),
+        emit_step_finished("step1"),
+        emit_text_message("message"),
+        emit_activity("m1", "work", "Working"),
+        emit_run_error("error occurred"),
+        emit_run_finished("t1", "r1", {"result": "done"}),
+    ]
+
+    # Should not raise any exceptions
+    result = await renderer.render(async_gen(events))
+    assert result == {"result": "done"}
+
+
+@pytest.mark.asyncio
+async def test_renderer_initial_state():
+    """Test that initial state is rendered."""
+    renderer = AGUIConsoleRenderer()
+
+    events = [
+        emit_state_snapshot(SimpleState(value=1)),
+        emit_state_snapshot(SimpleState(value=2)),
+    ]
+
+    await renderer.render(async_gen(events))
+
+    # After processing, internal state should be the last state
+    assert renderer._state == {"value": 2}
+
+
+@pytest.mark.asyncio
+async def test_renderer_no_result():
+    """Test renderer when no RUN_FINISHED event."""
+    renderer = AGUIConsoleRenderer()
+
+    events = [
+        emit_run_started("t1", "r1"),
+        emit_step_started("step1"),
+    ]
+
+    result = await renderer.render(async_gen(events))
+    assert result is None
+
+
+@pytest.mark.asyncio
+async def test_renderer_nested_state_diff():
+    """Test nested state diff computation."""
+    renderer = AGUIConsoleRenderer()
+
+    old = {"level1": {"level2": {"value": 1, "text": "old"}, "other": "same"}}
+
+    new = {"level1": {"level2": {"value": 2, "text": "old"}, "other": "same"}}
+
+    diff = renderer._compute_diff(old, new)
+
+    # Should only show the changed nested value
+    assert diff == {"level1": {"level2": {"value": 2}}}
+
+
+@pytest.mark.asyncio
+async def test_renderer_with_empty_state():
+    """Test renderer handles empty state gracefully."""
+    renderer = AGUIConsoleRenderer()
+
+    # Create a state with no changes
+    events = [
+        emit_step_started("step1"),  # No state snapshot
+        emit_run_finished("t1", "r1", {"result": "ok"}),
+    ]
+
+    result = await renderer.render(async_gen(events))
+    assert result == {"result": "ok"}

tests/agui/test_emitter.py

@@ -0,0 +1,231 @@
+"""Tests for AGUIEmitter."""
+
+import asyncio
+
+import pytest
+from pydantic import BaseModel
+
+from haiku.rag.agui.emitter import AGUIEmitter
+
+
+class TestState(BaseModel):
+    """Test state model."""
+
+    value: int
+    text: str
+
+
+class TestResult(BaseModel):
+    """Test result model."""
+
+    status: str
+
+
+@pytest.mark.asyncio
+async def test_emitter_lifecycle():
+    """Test emitter lifecycle events."""
+    emitter: AGUIEmitter[TestState, TestResult] = AGUIEmitter()
+
+    initial_state = TestState(value=1, text="initial")
+    emitter.start_run(initial_state)
+
+    events = []
+    async for event in emitter:
+        events.append(event)
+        if event["type"] == "RUN_STARTED":
+            # Close after getting run started
+            result = TestResult(status="complete")
+            emitter.finish_run(result)
+            await emitter.close()
+
+    assert len(events) >= 3  # RUN_STARTED, STATE_SNAPSHOT, RUN_FINISHED
+    assert events[0]["type"] == "RUN_STARTED"
+    assert events[-1]["type"] == "RUN_FINISHED"
+    assert events[-1]["result"] == {"status": "complete"}
+
+
+@pytest.mark.asyncio
+async def test_emitter_step_events():
+    """Test step lifecycle events."""
+    emitter: AGUIEmitter[TestState, TestResult] = AGUIEmitter()
+
+    emitter.start_step("test_step")
+    emitter.finish_step()
+    await emitter.close()
+
+    events = []
+    async for event in emitter:
+        events.append(event)
+
+    step_events = [e for e in events if e["type"] in ("STEP_STARTED", "STEP_FINISHED")]
+    assert len(step_events) == 2
+    assert step_events[0]["type"] == "STEP_STARTED"
+    assert step_events[0]["stepName"] == "test_step"
+    assert step_events[1]["type"] == "STEP_FINISHED"
+    assert step_events[1]["stepName"] == "test_step"
+
+
+@pytest.mark.asyncio
+async def test_emitter_state_updates():
+    """Test state update events."""
+    emitter: AGUIEmitter[TestState, TestResult] = AGUIEmitter()
+
+    state1 = TestState(value=1, text="first")
+    state2 = TestState(value=2, text="second")
+
+    emitter.update_state(state1)
+    emitter.update_state(state2)
+    await emitter.close()
+
+    events = []
+    async for event in emitter:
+        events.append(event)
+
+    state_events = [e for e in events if e["type"] == "STATE_SNAPSHOT"]
+    assert len(state_events) == 2
+    assert state_events[0]["snapshot"] == {"value": 1, "text": "first"}
+    assert state_events[1]["snapshot"] == {"value": 2, "text": "second"}
+
+
+@pytest.mark.asyncio
+async def test_emitter_activity_events():
+    """Test activity events."""
+    emitter: AGUIEmitter[TestState, TestResult] = AGUIEmitter()
+
+    emitter.update_activity("processing", "Processing data")
+    emitter.update_activity("done", "Completed", message_id="msg-1")
+    await emitter.close()
+
+    events = []
+    async for event in emitter:
+        events.append(event)
+
+    activity_events = [e for e in events if e["type"] == "ACTIVITY_SNAPSHOT"]
+    assert len(activity_events) == 2
+    assert activity_events[0]["activityType"] == "processing"
+    assert activity_events[0]["content"] == "Processing data"
+    assert activity_events[1]["messageId"] == "msg-1"
+    assert activity_events[1]["activityType"] == "done"
+
+
+@pytest.mark.asyncio
+async def test_emitter_text_messages():
+    """Test text message events."""
+    emitter: AGUIEmitter[TestState, TestResult] = AGUIEmitter()
+
+    emitter.log("Test message", role="assistant")
+    emitter.log("Another message", role="user")
+    await emitter.close()
+
+    events = []
+    async for event in emitter:
+        events.append(event)
+
+    text_events = [e for e in events if e["type"] == "TEXT_MESSAGE_CHUNK"]
+    assert len(text_events) == 2
+    assert text_events[0]["delta"] == "Test message"
+    assert text_events[0]["role"] == "assistant"
+    assert text_events[1]["delta"] == "Another message"
+    assert text_events[1]["role"] == "user"
+
+
+@pytest.mark.asyncio
+async def test_emitter_error():
+    """Test error event emission."""
+    emitter: AGUIEmitter[TestState, TestResult] = AGUIEmitter()
+
+    error = ValueError("Test error")
+    emitter.error(error, code="TEST_ERROR")
+    await emitter.close()
+
+    events = []
+    async for event in emitter:
+        events.append(event)
+
+    error_events = [e for e in events if e["type"] == "RUN_ERROR"]
+    assert len(error_events) == 1
+    assert error_events[0]["message"] == "Test error"
+    assert error_events[0]["code"] == "TEST_ERROR"
+
+
+@pytest.mark.asyncio
+async def test_emitter_thread_and_run_ids():
+    """Test thread and run ID management."""
+    emitter: AGUIEmitter[TestState, TestResult] = AGUIEmitter(
+        thread_id="thread-1", run_id="run-1"
+    )
+
+    assert emitter.thread_id == "thread-1"
+    assert emitter.run_id == "run-1"
+
+    initial_state = TestState(value=1, text="test")
+    emitter.start_run(initial_state)
+    await emitter.close()
+
+    events = []
+    async for event in emitter:
+        events.append(event)
+
+    run_started = [e for e in events if e["type"] == "RUN_STARTED"][0]
+    assert run_started["threadId"] == "thread-1"
+    assert run_started["runId"] == "run-1"
+
+
+@pytest.mark.asyncio
+async def test_emitter_generates_thread_id():
+    """Test that thread ID is generated from state hash when not provided."""
+    emitter: AGUIEmitter[TestState, TestResult] = AGUIEmitter()
+
+    initial_state = TestState(value=42, text="test")
+    emitter.start_run(initial_state)
+
+    # Thread ID should be generated deterministically from state
+    assert emitter.thread_id is not None
+    assert len(emitter.thread_id) > 0
+
+    await emitter.close()
+    async for _ in emitter:
+        pass
+
+
+@pytest.mark.asyncio
+async def test_emitter_closes_properly():
+    """Test that emitter closes and stops iteration."""
+    emitter: AGUIEmitter[TestState, TestResult] = AGUIEmitter()
+
+    emitter.log("Message 1")
+    await emitter.close()
+
+    # Attempting to iterate after close should work and stop
+    events = []
+    async for event in emitter:
+        events.append(event)
+
+    # Should have received the message and then stopped
+    assert len(events) == 1
+
+
+@pytest.mark.asyncio
+async def test_emitter_concurrent_emission():
+    """Test that multiple events can be emitted concurrently."""
+    emitter: AGUIEmitter[TestState, TestResult] = AGUIEmitter()
+
+    async def emit_many():
+        for i in range(10):
+            emitter.log(f"Message {i}")
+            await asyncio.sleep(0.001)  # Simulate some work
+        await emitter.close()
+
+    # Start emission in background
+    emit_task = asyncio.create_task(emit_many())
+
+    # Collect events
+    events = []
+    async for event in emitter:
+        events.append(event)
+
+    await emit_task
+
+    # Should have all 10 messages
+    text_events = [e for e in events if e["type"] == "TEXT_MESSAGE_CHUNK"]
+    assert len(text_events) == 10

tests/agui/test_events.py

@@ -0,0 +1,155 @@
+"""Tests for AG-UI event creation utilities."""
+
+from pydantic import BaseModel
+
+from haiku.rag.agui.events import (
+    emit_activity,
+    emit_run_error,
+    emit_run_finished,
+    emit_run_started,
+    emit_state_snapshot,
+    emit_step_finished,
+    emit_step_started,
+    emit_text_message,
+)
+
+
+class TestState(BaseModel):
+    """Test state model."""
+
+    value: int
+
+
+class TestResult(BaseModel):
+    """Test result model."""
+
+    status: str
+
+
+def test_emit_run_started():
+    """Test RUN_STARTED event creation."""
+    event = emit_run_started("thread-1", "run-1")
+
+    assert event["type"] == "RUN_STARTED"
+    assert event["threadId"] == "thread-1"
+    assert event["runId"] == "run-1"
+    assert "input" not in event
+
+
+def test_emit_run_started_with_input():
+    """Test RUN_STARTED event with input data."""
+    event = emit_run_started("thread-1", "run-1", input_data="test input")
+
+    assert event["type"] == "RUN_STARTED"
+    assert event["input"] == "test input"
+
+
+def test_emit_run_finished():
+    """Test RUN_FINISHED event creation."""
+    result = TestResult(status="complete")
+    event = emit_run_finished("thread-1", "run-1", result)
+
+    assert event["type"] == "RUN_FINISHED"
+    assert event["threadId"] == "thread-1"
+    assert event["runId"] == "run-1"
+    assert event["result"] == {"status": "complete"}
+
+
+def test_emit_run_finished_with_dict():
+    """Test RUN_FINISHED event with dict result."""
+    result = {"status": "complete", "count": 42}
+    event = emit_run_finished("thread-1", "run-1", result)
+
+    assert event["type"] == "RUN_FINISHED"
+    assert event["result"] == result
+
+
+def test_emit_run_error():
+    """Test RUN_ERROR event creation."""
+    event = emit_run_error("Something went wrong")
+
+    assert event["type"] == "RUN_ERROR"
+    assert event["message"] == "Something went wrong"
+    assert "code" not in event
+
+
+def test_emit_run_error_with_code():
+    """Test RUN_ERROR event with error code."""
+    event = emit_run_error("Something went wrong", code="ERR_001")
+
+    assert event["type"] == "RUN_ERROR"
+    assert event["message"] == "Something went wrong"
+    assert event["code"] == "ERR_001"
+
+
+def test_emit_step_started():
+    """Test STEP_STARTED event creation."""
+    event = emit_step_started("plan")
+
+    assert event["type"] == "STEP_STARTED"
+    assert event["stepName"] == "plan"
+
+
+def test_emit_step_finished():
+    """Test STEP_FINISHED event creation."""
+    event = emit_step_finished("plan")
+
+    assert event["type"] == "STEP_FINISHED"
+    assert event["stepName"] == "plan"
+
+
+def test_emit_text_message():
+    """Test TEXT_MESSAGE_CHUNK event creation."""
+    event = emit_text_message("Hello world")
+
+    assert event["type"] == "TEXT_MESSAGE_CHUNK"
+    assert event["delta"] == "Hello world"
+    assert event["role"] == "assistant"
+    assert "messageId" in event
+
+
+def test_emit_text_message_with_role():
+    """Test TEXT_MESSAGE_CHUNK event with custom role."""
+    event = emit_text_message("Hello", role="user")
+
+    assert event["type"] == "TEXT_MESSAGE_CHUNK"
+    assert event["role"] == "user"
+
+
+def test_emit_state_snapshot():
+    """Test STATE_SNAPSHOT event creation."""
+    state = TestState(value=42)
+    event = emit_state_snapshot(state)
+
+    assert event["type"] == "STATE_SNAPSHOT"
+    assert event["snapshot"] == {"value": 42}
+
+
+def test_emit_activity():
+    """Test ACTIVITY_SNAPSHOT event creation."""
+    event = emit_activity("msg-1", "processing", "Working on task")
+
+    assert event["type"] == "ACTIVITY_SNAPSHOT"
+    assert event["messageId"] == "msg-1"
+    assert event["activityType"] == "processing"
+    assert event["content"] == "Working on task"
+
+
+def test_event_structure_consistency():
+    """Test that all events have consistent structure."""
+    events = [
+        emit_run_started("t1", "r1"),
+        emit_run_finished("t1", "r1", {"result": "done"}),
+        emit_run_error("error"),
+        emit_step_started("step1"),
+        emit_step_finished("step1"),
+        emit_text_message("text"),
+        emit_state_snapshot(TestState(value=1)),
+        emit_activity("m1", "type", "content"),
+    ]
+
+    for event in events:
+        assert isinstance(event, dict)
+        assert "type" in event
+        assert isinstance(event["type"], str)
+        assert event["type"].isupper()  # Event types are uppercase

tests/agui/test_server.py

@@ -0,0 +1,250 @@
+"""Tests for AG-UI server."""
+
+import pytest
+from pydantic import BaseModel
+from starlette.testclient import TestClient
+
+from haiku.rag.agui.server import RunAgentInput, create_agui_app, format_sse_event
+from haiku.rag.config.models import AGUIConfig
+
+
+class SimpleState(BaseModel):
+    """Simple state for testing."""
+
+    question: str
+
+
+class SimpleResult(BaseModel):
+    """Simple result for testing."""
+
+    answer: str
+
+
+class MockGraph:
+    """Mock graph that returns immediately."""
+
+    async def run(self, state, deps):  # type: ignore[no-untyped-def]
+        """Return a simple result."""
+        return SimpleResult(answer=f"Answer to: {state.question}")
+
+
+def test_run_agent_input_parsing():
+    """Test RunAgentInput model parsing."""
+    data = {
+        "threadId": "thread-1",
+        "runId": "run-1",
+        "state": {"question": "What is AI?"},
+        "messages": [],
+        "config": {},
+    }
+
+    input_data = RunAgentInput(**data)
+
+    assert input_data.thread_id == "thread-1"
+    assert input_data.run_id == "run-1"
+    assert input_data.state == {"question": "What is AI?"}
+
+
+def test_run_agent_input_defaults():
+    """Test RunAgentInput with defaults."""
+    input_data = RunAgentInput()  # type: ignore[call-arg]
+
+    assert input_data.thread_id is None
+    assert input_data.run_id is None
+    assert input_data.state == {}
+    assert input_data.messages == []
+    assert input_data.config == {}
+
+
+def test_format_sse_event():
+    """Test SSE event formatting."""
+    event = {"type": "TEST_EVENT", "data": "test"}
+
+    sse = format_sse_event(event)
+
+    assert sse.startswith("data: ")
+    assert sse.endswith("\n\n")
+    assert '{"type": "TEST_EVENT"' in sse
+
+
+def test_create_agui_app_basic():
+    """Test basic app creation."""
+    config = AGUIConfig(
+        host="localhost",
+        port=8000,
+        cors_origins=["http://localhost"],
+    )
+
+    def graph_factory():
+        return MockGraph()
+
+    def state_factory(input_state):
+        return SimpleState(question=input_state.get("question", ""))
+
+    def deps_factory(input_config):
+        from dataclasses import dataclass
+
+        @dataclass
+        class SimpleDeps:
+            agui_emitter: None = None
+
+        return SimpleDeps()
+
+    app = create_agui_app(
+        graph_factory=graph_factory,  # type: ignore[arg-type]
+        state_factory=state_factory,
+        deps_factory=deps_factory,
+        config=config,
+    )
+
+    # Should return a Starlette app
+    assert app is not None
+    assert hasattr(app, "routes")
+
+
+def test_server_health_endpoint():
+    """Test health check endpoint."""
+    config = AGUIConfig()
+
+    def graph_factory():
+        return MockGraph()
+
+    def state_factory(input_state):
+        return SimpleState(question="")
+
+    def deps_factory(input_config):
+        from dataclasses import dataclass
+
+        @dataclass
+        class SimpleDeps:
+            agui_emitter: None = None
+
+        return SimpleDeps()
+
+    app = create_agui_app(
+        graph_factory=graph_factory,  # type: ignore[arg-type]
+        state_factory=state_factory,
+        deps_factory=deps_factory,
+        config=config,
+    )
+
+    client = TestClient(app)
+    response = client.get("/health")
+
+    assert response.status_code == 200
+    assert response.json() == {"status": "healthy"}
+
+
+@pytest.mark.asyncio
+async def test_server_stream_endpoint():
+    """Test AG-UI streaming endpoint."""
+    config = AGUIConfig()
+
+    def graph_factory():
+        return MockGraph()
+
+    def state_factory(input_state):
+        question = input_state.get("question", "")
+        return SimpleState(question=question)
+
+    def deps_factory(input_config):
+        from dataclasses import dataclass
+
+        @dataclass
+        class SimpleDeps:
+            agui_emitter: None = None
+
+        return SimpleDeps()
+
+    app = create_agui_app(
+        graph_factory=graph_factory,  # type: ignore[arg-type]
+        state_factory=state_factory,
+        deps_factory=deps_factory,
+        config=config,
+    )
+
+    client = TestClient(app)
+
+    request_data = {
+        "threadId": "test-1",
+        "runId": "run-1",
+        "state": {"question": "What is pydantic-graph?"},
+        "messages": [],
+        "config": {},
+    }
+
+    response = client.post("/v1/agent/stream", json=request_data)
+
+    assert response.status_code == 200
+    assert response.headers["content-type"] == "text/event-stream; charset=utf-8"
+
+    # Read the streamed events
+    events = []
+    for line in response.iter_lines():
+        if line.startswith("data: "):
+            import json
+
+            event_data = line[6:]  # Remove "data: " prefix
+            event = json.loads(event_data)
+            events.append(event)
+
+    # Should have received multiple events
+    assert len(events) > 0
+
+    # Should have RUN_STARTED and RUN_FINISHED
+    event_types = [e["type"] for e in events]
+    assert "RUN_STARTED" in event_types
+    assert "RUN_FINISHED" in event_types
+
+
+def test_server_cors_headers():
+    """Test CORS middleware is configured."""
+    config = AGUIConfig(
+        cors_origins=["http://example.com"],
+        cors_credentials=True,
+    )
+
+    def graph_factory():
+        return MockGraph()
+
+    def state_factory(input_state):
+        return SimpleState(question="")
+
+    def deps_factory(input_config):
+        from dataclasses import dataclass
+
+        @dataclass
+        class SimpleDeps:
+            agui_emitter: None = None
+
+        return SimpleDeps()
+
+    app = create_agui_app(
+        graph_factory=graph_factory,  # type: ignore[arg-type]
+        state_factory=state_factory,
+        deps_factory=deps_factory,
+        config=config,
+    )
+
+    client = TestClient(app)
+
+    # GET request with Origin header should get CORS headers
+    response = client.get("/health", headers={"Origin": "http://example.com"})
+
+    assert response.status_code == 200
+    # CORS middleware should add access-control headers
+    assert (
+        "access-control-allow-origin" in response.headers or response.status_code == 200
+    )
+
+
+def test_agui_config_defaults():
+    """Test AGUIConfig default values."""
+    config = AGUIConfig()
+
+    assert config.host == "0.0.0.0"
+    assert config.port == 8000
+    assert config.cors_origins == ["*"]
+    assert config.cors_credentials is True
+    assert "GET" in config.cors_methods
+    assert "POST" in config.cors_methods

tests/agui/test_stream.py

@@ -0,0 +1,199 @@
+"""Tests for stream_graph function."""
+
+from dataclasses import dataclass
+
+import pytest
+from pydantic import BaseModel
+
+from haiku.rag.agui.emitter import AGUIEmitter
+from haiku.rag.agui.stream import stream_graph
+
+
+class TestState(BaseModel):
+    """Test state model."""
+
+    value: int
+
+
+@dataclass
+class TestDeps:
+    """Test dependencies."""
+
+    agui_emitter: AGUIEmitter | None = None
+
+
+class MockGraph:
+    """Mock graph for testing."""
+
+    def __init__(self, result: str | dict[str, str | int] = "done"):
+        self.result = result
+        self.run_called = False
+
+    async def run(self, state, deps):  # type: ignore[no-untyped-def]
+        """Mock run method."""
+        self.run_called = True
+
+        # Emit some events through the emitter
+        if deps.agui_emitter:
+            deps.agui_emitter.start_step("mock_step")
+            deps.agui_emitter.update_activity("working", "Doing work")
+            deps.agui_emitter.finish_step()
+
+        return self.result
+
+
+@pytest.mark.asyncio
+async def test_stream_graph_basic():
+    """Test basic graph streaming."""
+    graph = MockGraph(result="success")
+    state = TestState(value=1)
+    deps = TestDeps()
+
+    events = []
+    async for event in stream_graph(graph, state, deps):
+        events.append(event)
+
+    # Should have collected events
+    assert len(events) > 0
+
+    # Should have RUN_STARTED and RUN_FINISHED
+    event_types = [e["type"] for e in events]
+    assert "RUN_STARTED" in event_types
+    assert "RUN_FINISHED" in event_types
+
+    # Graph should have been executed
+    assert graph.run_called
+
+
+@pytest.mark.asyncio
+async def test_stream_graph_emits_initial_state():
+    """Test that initial state is emitted."""
+    graph = MockGraph()
+    state = TestState(value=42)
+    deps = TestDeps()
+
+    events = []
+    async for event in stream_graph(graph, state, deps):
+        events.append(event)
+
+    # Should have initial state snapshot
+    state_snapshots = [e for e in events if e["type"] == "STATE_SNAPSHOT"]
+    assert len(state_snapshots) > 0
+    # First snapshot should be initial state
+    assert state_snapshots[0]["snapshot"] == {"value": 42}
+
+
+@pytest.mark.asyncio
+async def test_stream_graph_emits_step_events():
+    """Test that step events from graph are emitted."""
+    graph = MockGraph()
+    state = TestState(value=1)
+    deps = TestDeps()
+
+    events = []
+    async for event in stream_graph(graph, state, deps):
+        events.append(event)
+
+    # Should have step events from MockGraph
+    step_started = [e for e in events if e["type"] == "STEP_STARTED"]
+    assert len(step_started) > 0
+    assert step_started[0]["stepName"] == "mock_step"
+
+
+@pytest.mark.asyncio
+async def test_stream_graph_emits_activity():
+    """Test that activity events from graph are emitted."""
+    graph = MockGraph()
+    state = TestState(value=1)
+    deps = TestDeps()
+
+    events = []
+    async for event in stream_graph(graph, state, deps):
+        events.append(event)
+
+    # Should have activity events from MockGraph
+    activities = [e for e in events if e["type"] == "ACTIVITY_SNAPSHOT"]
+    assert len(activities) > 0
+    assert activities[0]["content"] == "Doing work"
+
+
+@pytest.mark.asyncio
+async def test_stream_graph_handles_error():
+    """Test that graph errors are captured and emitted."""
+
+    class ErrorGraph:
+        async def run(self, state, deps):
+            raise ValueError("Test error")
+
+    graph = ErrorGraph()
+    state = TestState(value=1)
+    deps = TestDeps()
+
+    events = []
+    async for event in stream_graph(graph, state, deps):
+        events.append(event)
+
+    # Should have error event
+    errors = [e for e in events if e["type"] == "RUN_ERROR"]
+    assert len(errors) > 0
+    assert "Test error" in errors[0]["message"]
+
+
+@pytest.mark.asyncio
+async def test_stream_graph_closes_emitter():
+    """Test that emitter is properly closed."""
+
+    class NeverReturnsGraph:
+        async def run(self, state, deps):
+            # Don't return anything
+            pass
+
+    graph = NeverReturnsGraph()
+    state = TestState(value=1)
+    deps = TestDeps()
+
+    events = []
+    try:
+        async for event in stream_graph(graph, state, deps):
+            events.append(event)
+    except RuntimeError:
+        # Expected - graph didn't return a result
+        pass
+
+    # Should have error event about no result
+    errors = [e for e in events if e["type"] == "RUN_ERROR"]
+    assert len(errors) > 0
+
+
+@pytest.mark.asyncio
+async def test_stream_graph_without_emitter_support():
+    """Test error when deps doesn't support agui_emitter."""
+
+    @dataclass
+    class BadDeps:
+        pass
+
+    graph = MockGraph()
+    state = TestState(value=1)
+    deps = BadDeps()
+
+    with pytest.raises(TypeError, match="agui_emitter"):
+        async for _ in stream_graph(graph, state, deps):  # type: ignore[arg-type]
+            pass
+
+
+@pytest.mark.asyncio
+async def test_stream_graph_result_in_finish_event():
+    """Test that graph result is included in RUN_FINISHED event."""
+    graph = MockGraph(result={"status": "complete", "count": 42})
+    state = TestState(value=1)
+    deps = TestDeps()
+
+    events = []
+    async for event in stream_graph(graph, state, deps):
+        events.append(event)
+
+    # Find RUN_FINISHED event
+    finished = [e for e in events if e["type"] == "RUN_FINISHED"]
+    assert len(finished) == 1
+    assert finished[0]["result"] == {"status": "complete", "count": 42}

tests/test_research_graph.py

@@ -1,8 +1,13 @@
 import asyncio
 
+import pytest
+from pydantic_ai.models.test import TestModel
+
+from haiku.rag.agui.stream import stream_graph
+from haiku.rag.client import HaikuRAG
 from haiku.rag.research.dependencies import ResearchContext
 from haiku.rag.research.graph import build_research_graph
-from haiku.rag.research.state import ResearchState
+from haiku.rag.research.state import ResearchDeps, ResearchState
 
 
 def test_build_graph_and_state():
@@ -24,3 +29,57 @@ def test_async_loop_available():
     # Ensure an event loop can be created in test env
     loop = asyncio.new_event_loop()
     loop.close()
+
+
+@pytest.mark.asyncio
+async def test_graph_end_to_end_with_test_model(monkeypatch, temp_db_path):
+    """Test research graph with mocked LLM using AG-UI events."""
+
+    # Mock get_model to return TestModel which generates valid schema-compliant data
+    def test_model_factory(_provider, _model):
+        return TestModel()
+
+    monkeypatch.setattr("haiku.rag.graph_common.utils.get_model", test_model_factory)
+    monkeypatch.setattr("haiku.rag.research.graph.get_model", test_model_factory)
+
+    graph = build_research_graph()
+
+    state = ResearchState(
+        context=ResearchContext(original_question="What is haiku.rag?"),
+        max_iterations=1,
+        confidence_threshold=0.5,
+        max_concurrency=2,
+    )
+
+    # Use real client but with TestModel for LLM calls
+    client = HaikuRAG(temp_db_path)
+    deps = ResearchDeps(client=client)
+
+    events = []
+    result = None
+    async for event in stream_graph(graph, state, deps):
+        events.append(event)
+        if event["type"] == "RUN_FINISHED":
+            result = event["result"]
+        elif event["type"] == "RUN_ERROR":
+            pytest.fail(f"Graph execution failed: {event['message']}")
+
+    # TestModel will generate valid structured output for each node
+    assert result is not None, (
+        f"No result. Events collected: {[e['type'] for e in events]}"
+    )
+    # Result is serialized as dict in AG-UI events
+    assert isinstance(result, dict)
+    assert "title" in result
+    assert isinstance(result["title"], str)
+    assert "executive_summary" in result
+    assert "main_findings" in result
+
+    # Verify AG-UI events were emitted
+    event_types = [e["type"] for e in events]
+    assert "RUN_STARTED" in event_types
+    assert "RUN_FINISHED" in event_types
+    assert "STATE_SNAPSHOT" in event_types
+    assert "STEP_STARTED" in event_types
+
+    client.close()

tests/test_research_graph_integration.py

@@ -1,55 +0,0 @@
-import pytest
-from pydantic_ai.models.test import TestModel
-
-from haiku.rag.client import HaikuRAG
-from haiku.rag.research.dependencies import ResearchContext
-from haiku.rag.research.graph import build_research_graph
-from haiku.rag.research.models import ResearchReport
-from haiku.rag.research.state import ResearchDeps, ResearchState
-from haiku.rag.research.stream import stream_research_graph
-
-
-@pytest.mark.asyncio
-async def test_graph_end_to_end_with_test_model(monkeypatch, temp_db_path):
-    """Test research graph with mocked LLM using TestModel."""
-
-    # Mock get_model to return TestModel which generates valid schema-compliant data
-    def test_model_factory(provider, model):
-        return TestModel()
-
-    monkeypatch.setattr("haiku.rag.graph_common.utils.get_model", test_model_factory)
-    monkeypatch.setattr("haiku.rag.research.graph.get_model", test_model_factory)
-
-    graph = build_research_graph()
-
-    state = ResearchState(
-        context=ResearchContext(original_question="What is haiku.rag?"),
-        max_iterations=1,
-        confidence_threshold=0.5,
-        max_concurrency=2,
-    )
-
-    # Use real client but with TestModel for LLM calls
-    client = HaikuRAG(temp_db_path)
-    deps = ResearchDeps(client=client, console=None)
-
-    collected = []
-    report = None
-    async for event in stream_research_graph(graph, state, deps):
-        collected.append(event)
-        if event.type == "report":
-            report = event.report
-            break
-        elif event.type == "error":
-            pytest.fail(f"Graph execution failed: {event.error}")
-
-    # TestModel will generate valid structured output for each node
-    assert report is not None, (
-        f"No report generated. Events collected: {[e.type for e in collected]}"
-    )
-    assert isinstance(report, ResearchReport)
-    assert report.title is not None
-    assert isinstance(report.title, str)
-    assert any(evt.type == "log" for evt in collected)
-
-    client.close()