Merge pull request #199 from ggozad/chore/ag_ui-core

Replace custom event classes with `ag_ui.core` types

CHANGELOG.md

@@ -27,6 +27,9 @@
 
 ### Changed
 
+- **AG-UI Events**: Replaced custom event classes with `ag_ui.core` types
+  - Removed `haiku.rag.graph.agui.events` module
+  - Event factory functions (`emit_*`) now wrap official `ag_ui.core` event classes
 - **Chunker Sets Order**: Chunkers now set `chunk.order` directly
 - **Unified Research Graph**: Simplified and unified research and deep QA into a single configurable graph
   - Removed `analyze_insights` node - graph now flows directly from `collect_answers` to `decide`

haiku_rag_slim/haiku/rag/graph/agui/__init__.py

@@ -1,8 +1,8 @@
 """Generic AG-UI protocol support for haiku.rag graphs."""
 
 from haiku.rag.graph.agui.cli_renderer import AGUIConsoleRenderer
-from haiku.rag.graph.agui.emitter import AGUIEmitter
+from haiku.rag.graph.agui.emitter import (
-from haiku.rag.graph.agui.events import (
+    AGUIEmitter,
     AGUIEvent,
     emit_activity,
     emit_activity_delta,
@@ -17,6 +17,9 @@ from haiku.rag.graph.agui.events import (
     emit_text_message_content,
     emit_text_message_end,
     emit_text_message_start,
+    emit_tool_call_args,
+    emit_tool_call_end,
+    emit_tool_call_start,
 )
 from haiku.rag.graph.agui.server import (
     RunAgentInput,
@@ -48,6 +51,9 @@ __all__ = [
     "emit_text_message_content",
     "emit_text_message_end",
     "emit_text_message_start",
+    "emit_tool_call_args",
+    "emit_tool_call_end",
+    "emit_tool_call_start",
     "format_sse_event",
     "stream_graph",
 ]

haiku_rag_slim/haiku/rag/graph/agui/cli_renderer.py

@@ -5,7 +5,7 @@ from typing import Any
 
 from rich.console import Console
 
-from haiku.rag.graph.agui.events import AGUIEvent
+from haiku.rag.graph.agui.emitter import AGUIEvent
 
 
 class AGUIConsoleRenderer:

haiku_rag_slim/haiku/rag/graph/agui/emitter.py

@@ -2,24 +2,39 @@
 
 import asyncio
 import hashlib
+import json
 from collections.abc import AsyncIterator
 from typing import Any
 from uuid import uuid4
 
+from ag_ui.core import (
+    ActivitySnapshotEvent,
+    BaseEvent,
+    RunErrorEvent,
+    RunFinishedEvent,
+    RunStartedEvent,
+    StateDeltaEvent,
+    StateSnapshotEvent,
+    StepFinishedEvent,
+    StepStartedEvent,
+    TextMessageChunkEvent,
+    TextMessageContentEvent,
+    TextMessageEndEvent,
+    TextMessageStartEvent,
+    ToolCallArgsEvent,
+    ToolCallEndEvent,
+    ToolCallStartEvent,
+)
 from pydantic import BaseModel
 
-from haiku.rag.graph.agui.events import (
+from haiku.rag.graph.agui.state import compute_state_delta
-    AGUIEvent,
+
-    emit_activity,
+AGUIEvent = dict[str, Any]
-    emit_run_error,
+
-    emit_run_finished,
+
-    emit_run_started,
+def _serialize_event(event: BaseEvent) -> AGUIEvent:
-    emit_state_delta,
+    """Serialize an ag_ui event to a dict with camelCase keys."""
-    emit_state_snapshot,
+    return event.model_dump(mode="json", by_alias=True, exclude_none=True)
-    emit_step_finished,
-    emit_step_started,
-    emit_text_message,
-)
 
 
 class AGUIEmitter[StateT: BaseModel, ResultT]:
@@ -80,8 +95,14 @@ class AGUIEmitter[StateT: BaseModel, ResultT]:
             self._thread_id = self._generate_thread_id(state_json)
 
         # RunStarted (state snapshot follows immediately with full state)
-        self.emit(emit_run_started(self._thread_id, self._run_id))
+        self.emit(
-        self.emit(emit_state_snapshot(initial_state))
+            _serialize_event(
+                RunStartedEvent(thread_id=self._thread_id, run_id=self._run_id)
+            )
+        )
+        self.emit(
+            _serialize_event(StateSnapshotEvent(snapshot=initial_state.model_dump()))
+        )
         # Store a deep copy to detect future changes
         self._last_state = initial_state.model_copy(deep=True)
 
@@ -92,12 +113,12 @@ class AGUIEmitter[StateT: BaseModel, ResultT]:
             step_name: Name of the step being started
         """
         self._current_step = step_name
-        self.emit(emit_step_started(step_name))
+        self.emit(_serialize_event(StepStartedEvent(step_name=step_name)))
 
     def finish_step(self) -> None:
         """Emit StepFinished event for the current step."""
         if self._current_step:
-            self.emit(emit_step_finished(self._current_step))
+            self.emit(_serialize_event(StepFinishedEvent(step_name=self._current_step)))
             self._current_step = None
 
     def log(self, message: str, role: str = "assistant") -> None:
@@ -107,7 +128,16 @@ class AGUIEmitter[StateT: BaseModel, ResultT]:
             message: The message content
             role: The role of the sender (default: assistant)
         """
-        self.emit(emit_text_message(message, role))
+        message_id = str(uuid4())
+        self.emit(
+            _serialize_event(
+                TextMessageChunkEvent(
+                    message_id=message_id,
+                    role=role,  # type: ignore[arg-type]
+                    delta=message,
+                )
+            )
+        )
 
     def update_state(self, new_state: StateT) -> None:
         """Emit StateDelta or StateSnapshot for state change.
@@ -117,10 +147,13 @@ class AGUIEmitter[StateT: BaseModel, ResultT]:
         """
         if self._use_deltas and self._last_state is not None:
             # Emit delta for incremental updates
-            self.emit(emit_state_delta(self._last_state, new_state))
+            delta = compute_state_delta(self._last_state, new_state)
+            self.emit(_serialize_event(StateDeltaEvent(delta=delta)))
         else:
             # Emit full snapshot for initial state or when deltas disabled
-            self.emit(emit_state_snapshot(new_state))
+            self.emit(
+                _serialize_event(StateSnapshotEvent(snapshot=new_state.model_dump()))
+            )
         # Store a deep copy to detect future changes
         self._last_state = new_state.model_copy(deep=True)
 
@@ -139,7 +172,15 @@ class AGUIEmitter[StateT: BaseModel, ResultT]:
         """
         if message_id is None:
             message_id = str(uuid4())
-        self.emit(emit_activity(message_id, activity_type, content))
+        self.emit(
+            _serialize_event(
+                ActivitySnapshotEvent(
+                    message_id=message_id,
+                    activity_type=activity_type,
+                    content=content,
+                )
+            )
+        )
 
     def finish_run(self, result: ResultT) -> None:
         """Emit RunFinished event.
@@ -147,7 +188,18 @@ class AGUIEmitter[StateT: BaseModel, ResultT]:
         Args:
             result: The final result from the graph
         """
-        self.emit(emit_run_finished(self._thread_id, self._run_id, result))
+        # Convert result to dict if it's a Pydantic model
+        result_data: Any = result
+        if hasattr(result, "model_dump"):
+            result_data = result.model_dump()  # type: ignore[union-attr]
+
+        self.emit(
+            _serialize_event(
+                RunFinishedEvent(
+                    thread_id=self._thread_id, run_id=self._run_id, result=result_data
+                )
+            )
+        )
 
     def error(self, error: Exception, code: str | None = None) -> None:
         """Emit RunError event.
@@ -156,7 +208,7 @@ class AGUIEmitter[StateT: BaseModel, ResultT]:
             error: The exception that occurred
             code: Optional error code
         """
-        self.emit(emit_run_error(str(error), code))
+        self.emit(_serialize_event(RunErrorEvent(message=str(error), code=code)))
 
     def emit(self, event: AGUIEvent) -> None:
         """Put event in queue.
@@ -199,3 +251,132 @@ class AGUIEmitter[StateT: BaseModel, ResultT]:
         # Use hash of input for deterministic thread ID
         hash_obj = hashlib.sha256(input_data.encode("utf-8"))
         return hash_obj.hexdigest()[:16]
+
+
+def emit_text_message_start(message_id: str, role: str = "assistant") -> AGUIEvent:
+    """Create a TextMessageStart event."""
+    return _serialize_event(
+        TextMessageStartEvent(message_id=message_id, role=role)  # type: ignore[arg-type]
+    )
+
+
+def emit_text_message_content(message_id: str, delta: str) -> AGUIEvent:
+    """Create a TextMessageContent event."""
+    return _serialize_event(TextMessageContentEvent(message_id=message_id, delta=delta))
+
+
+def emit_text_message_end(message_id: str) -> AGUIEvent:
+    """Create a TextMessageEnd event."""
+    return _serialize_event(TextMessageEndEvent(message_id=message_id))
+
+
+def emit_tool_call_start(
+    tool_call_id: str,
+    tool_name: str,
+    parent_message_id: str | None = None,
+) -> AGUIEvent:
+    """Create a ToolCallStart event."""
+    return _serialize_event(
+        ToolCallStartEvent(
+            tool_call_id=tool_call_id,
+            tool_call_name=tool_name,
+            parent_message_id=parent_message_id,
+        )
+    )
+
+
+def emit_tool_call_args(tool_call_id: str, args: dict[str, Any]) -> AGUIEvent:
+    """Create a ToolCallArgs event."""
+    return _serialize_event(
+        ToolCallArgsEvent(tool_call_id=tool_call_id, delta=json.dumps(args))
+    )
+
+
+def emit_tool_call_end(tool_call_id: str) -> AGUIEvent:
+    """Create a ToolCallEnd event."""
+    return _serialize_event(ToolCallEndEvent(tool_call_id=tool_call_id))
+
+
+def emit_run_started(thread_id: str, run_id: str) -> AGUIEvent:
+    """Create a RunStarted event."""
+    return _serialize_event(RunStartedEvent(thread_id=thread_id, run_id=run_id))
+
+
+def emit_run_finished(thread_id: str, run_id: str, result: Any) -> AGUIEvent:
+    """Create a RunFinished event."""
+    # Convert result to dict if it's a Pydantic model
+    if hasattr(result, "model_dump"):
+        result = result.model_dump()
+    return _serialize_event(
+        RunFinishedEvent(thread_id=thread_id, run_id=run_id, result=result)
+    )
+
+
+def emit_run_error(message: str, code: str | None = None) -> AGUIEvent:
+    """Create a RunError event."""
+    return _serialize_event(RunErrorEvent(message=message, code=code))
+
+
+def emit_step_started(step_name: str) -> AGUIEvent:
+    """Create a StepStarted event."""
+    return _serialize_event(StepStartedEvent(step_name=step_name))
+
+
+def emit_step_finished(step_name: str) -> AGUIEvent:
+    """Create a StepFinished event."""
+    return _serialize_event(StepFinishedEvent(step_name=step_name))
+
+
+def emit_text_message(content: str, role: str = "assistant") -> AGUIEvent:
+    """Create a TextMessageChunk event (convenience wrapper)."""
+    message_id = str(uuid4())
+    return _serialize_event(
+        TextMessageChunkEvent(
+            message_id=message_id,
+            role=role,  # type: ignore[arg-type]
+            delta=content,
+        )
+    )
+
+
+def emit_state_snapshot(state: BaseModel) -> AGUIEvent:
+    """Create a StateSnapshot event."""
+    return _serialize_event(StateSnapshotEvent(snapshot=state.model_dump()))
+
+
+def emit_state_delta(old_state: BaseModel, new_state: BaseModel) -> AGUIEvent:
+    """Create a StateDelta event with JSON Patch operations."""
+    delta = compute_state_delta(old_state, new_state)
+    return _serialize_event(StateDeltaEvent(delta=delta))
+
+
+def emit_activity(
+    message_id: str,
+    activity_type: str,
+    content: dict[str, Any],
+) -> AGUIEvent:
+    """Create an ActivitySnapshot event."""
+    return _serialize_event(
+        ActivitySnapshotEvent(
+            message_id=message_id,
+            activity_type=activity_type,
+            content=content,
+        )
+    )
+
+
+def emit_activity_delta(
+    message_id: str,
+    activity_type: str,
+    patch: list[dict[str, Any]],
+) -> AGUIEvent:
+    """Create an ActivityDelta event with JSON Patch operations."""
+    from ag_ui.core import ActivityDeltaEvent
+
+    return _serialize_event(
+        ActivityDeltaEvent(
+            message_id=message_id,
+            activity_type=activity_type,
+            patch=patch,
+        )
+    )

haiku_rag_slim/haiku/rag/graph/agui/events.py

@@ -1,313 +0,0 @@
-"""Generic AG-UI event creation utilities for any graph."""
-
-from typing import Any
-from uuid import uuid4
-
-from pydantic import BaseModel
-
-from haiku.rag.graph.agui.state import compute_state_delta
-
-# Type aliases for AG-UI events (actual types from ag_ui.core will be used at runtime)
-AGUIEvent = dict[str, Any]
-
-
-def emit_run_started(
-    thread_id: str, run_id: str, input_data: str | None = None
-) -> dict[str, Any]:
-    """Create a RunStarted event.
-
-    Args:
-        thread_id: Unique identifier for the conversation thread
-        run_id: Unique identifier for this run
-        input_data: Optional input that started the run
-
-    Returns:
-        RunStarted event dict
-    """
-    event: dict[str, Any] = {
-        "type": "RUN_STARTED",
-        "threadId": thread_id,
-        "runId": run_id,
-    }
-    if input_data:
-        event["input"] = input_data
-    return event
-
-
-def emit_run_finished(thread_id: str, run_id: str, result: Any) -> dict[str, Any]:
-    """Create a RunFinished event.
-
-    Args:
-        thread_id: Unique identifier for the conversation thread
-        run_id: Unique identifier for this run
-        result: The final result of the run
-
-    Returns:
-        RunFinished event dict
-    """
-    # Convert result to dict if it's a Pydantic model
-    if hasattr(result, "model_dump"):
-        result = result.model_dump()
-
-    return {
-        "type": "RUN_FINISHED",
-        "threadId": thread_id,
-        "runId": run_id,
-        "result": result,
-    }
-
-
-def emit_run_error(message: str, code: str | None = None) -> dict[str, Any]:
-    """Create a RunError event.
-
-    Args:
-        message: Error message
-        code: Optional error code
-
-    Returns:
-        RunError event dict
-    """
-    event: dict[str, Any] = {
-        "type": "RUN_ERROR",
-        "message": message,
-    }
-    if code:
-        event["code"] = code
-    return event
-
-
-def emit_step_started(step_name: str) -> dict[str, Any]:
-    """Create a StepStarted event.
-
-    Args:
-        step_name: Name of the step being started
-
-    Returns:
-        StepStarted event dict
-    """
-    return {
-        "type": "STEP_STARTED",
-        "stepName": step_name,
-    }
-
-
-def emit_step_finished(step_name: str) -> dict[str, Any]:
-    """Create a StepFinished event.
-
-    Args:
-        step_name: Name of the step that finished
-
-    Returns:
-        StepFinished event dict
-    """
-    return {
-        "type": "STEP_FINISHED",
-        "stepName": step_name,
-    }
-
-
-def emit_text_message(content: str, role: str = "assistant") -> dict[str, Any]:
-    """Create a TextMessageChunk event (convenience wrapper).
-
-    This creates a complete text message in one event.
-
-    Args:
-        content: The message content
-        role: The role of the sender (default: assistant)
-
-    Returns:
-        TextMessageChunk event dict
-    """
-    message_id = str(uuid4())
-    return {
-        "type": "TEXT_MESSAGE_CHUNK",
-        "messageId": message_id,
-        "role": role,
-        "delta": content,
-    }
-
-
-def emit_text_message_start(message_id: str, role: str = "assistant") -> dict[str, Any]:
-    """Create a TextMessageStart event.
-
-    Args:
-        message_id: Unique identifier for this message
-        role: The role of the sender
-
-    Returns:
-        TextMessageStart event dict
-    """
-    return {
-        "type": "TEXT_MESSAGE_START",
-        "messageId": message_id,
-        "role": role,
-    }
-
-
-def emit_text_message_content(message_id: str, delta: str) -> dict[str, Any]:
-    """Create a TextMessageContent event.
-
-    Args:
-        message_id: Identifier for the message being streamed
-        delta: Content chunk to append
-
-    Returns:
-        TextMessageContent event dict
-    """
-    return {
-        "type": "TEXT_MESSAGE_CONTENT",
-        "messageId": message_id,
-        "delta": delta,
-    }
-
-
-def emit_text_message_end(message_id: str) -> dict[str, Any]:
-    """Create a TextMessageEnd event.
-
-    Args:
-        message_id: Identifier for the message being completed
-
-    Returns:
-        TextMessageEnd event dict
-    """
-    return {
-        "type": "TEXT_MESSAGE_END",
-        "messageId": message_id,
-    }
-
-
-def emit_state_snapshot(state: BaseModel) -> dict[str, Any]:
-    """Create a StateSnapshot event.
-
-    Args:
-        state: The complete state to snapshot (any Pydantic BaseModel)
-
-    Returns:
-        StateSnapshot event dict
-    """
-    return {
-        "type": "STATE_SNAPSHOT",
-        "snapshot": state.model_dump(),
-    }
-
-
-def emit_state_delta(old_state: BaseModel, new_state: BaseModel) -> dict[str, Any]:
-    """Create a StateDelta event with JSON Patch operations.
-
-    Args:
-        old_state: Previous state (any Pydantic BaseModel)
-        new_state: Current state (same type as old_state)
-
-    Returns:
-        StateDelta event dict
-    """
-    delta = compute_state_delta(old_state, new_state)
-    return {
-        "type": "STATE_DELTA",
-        "delta": delta,
-    }
-
-
-def emit_activity(
-    message_id: str,
-    activity_type: str,
-    content: dict[str, Any],
-) -> dict[str, Any]:
-    """Create an ActivitySnapshot event.
-
-    Args:
-        message_id: Message ID to associate activity with (required)
-        activity_type: Type of activity (e.g., "planning", "searching")
-        content: Structured payload representing the activity state
-
-    Returns:
-        ActivitySnapshot event dict
-    """
-    return {
-        "type": "ACTIVITY_SNAPSHOT",
-        "messageId": message_id,
-        "activityType": activity_type,
-        "content": content,
-    }
-
-
-def emit_activity_delta(
-    message_id: str,
-    activity_type: str,
-    patch: list[dict[str, Any]],
-) -> dict[str, Any]:
-    """Create an ActivityDelta event with JSON Patch operations.
-
-    Args:
-        message_id: Message ID of the activity being updated
-        activity_type: Type of activity being updated
-        patch: JSON Patch operations to apply
-
-    Returns:
-        ActivityDelta event dict
-    """
-    return {
-        "type": "ACTIVITY_DELTA",
-        "messageId": message_id,
-        "activityType": activity_type,
-        "patch": patch,
-    }
-
-
-def emit_tool_call_start(
-    tool_call_id: str,
-    tool_name: str,
-    parent_message_id: str | None = None,
-) -> dict[str, Any]:
-    """Create a ToolCallStart event.
-
-    Args:
-        tool_call_id: Unique identifier for this tool call
-        tool_name: Name of the tool being called
-        parent_message_id: Optional parent message ID
-
-    Returns:
-        ToolCallStart event dict
-    """
-    event: dict[str, Any] = {
-        "type": "TOOL_CALL_START",
-        "toolCallId": tool_call_id,
-        "toolCallName": tool_name,
-    }
-    if parent_message_id:
-        event["parentMessageId"] = parent_message_id
-    return event
-
-
-def emit_tool_call_args(tool_call_id: str, args: dict[str, Any]) -> dict[str, Any]:
-    """Create a ToolCallArgs event.
-
-    Args:
-        tool_call_id: Identifier for the tool call
-        args: Tool arguments
-
-    Returns:
-        ToolCallArgs event dict
-    """
-    import json
-
-    return {
-        "type": "TOOL_CALL_ARGS",
-        "toolCallId": tool_call_id,
-        "delta": json.dumps(args),
-    }
-
-
-def emit_tool_call_end(tool_call_id: str) -> dict[str, Any]:
-    """Create a ToolCallEnd event.
-
-    Args:
-        tool_call_id: Identifier for the tool call being completed
-
-    Returns:
-        ToolCallEnd event dict
-    """
-    return {
-        "type": "TOOL_CALL_END",
-        "toolCallId": tool_call_id,
-    }

haiku_rag_slim/haiku/rag/graph/agui/server.py

@@ -18,8 +18,7 @@ from starlette.responses import JSONResponse, StreamingResponse
 from starlette.routing import Route
 
 from haiku.rag.config.models import AGUIConfig
-from haiku.rag.graph.agui.emitter import AGUIEmitter
+from haiku.rag.graph.agui.emitter import AGUIEmitter, AGUIEvent
-from haiku.rag.graph.agui.events import AGUIEvent
 from haiku.rag.graph.agui.stream import stream_graph
 
 

haiku_rag_slim/haiku/rag/graph/agui/stream.py

@@ -8,8 +8,7 @@ from typing import Protocol, TypeVar
 from pydantic import BaseModel
 from pydantic_graph.beta import Graph
 
-from haiku.rag.graph.agui.emitter import AGUIEmitter
+from haiku.rag.graph.agui.emitter import AGUIEmitter, AGUIEvent
-from haiku.rag.graph.agui.events import AGUIEvent
 
 StateT = TypeVar("StateT", bound=BaseModel)
 ResultT = TypeVar("ResultT")

haiku_rag_slim/haiku/rag/graph/research/graph.py

@@ -9,7 +9,7 @@ from pydantic_graph.beta.join import reduce_list_append
 
 from haiku.rag.config import Config
 from haiku.rag.config.models import AppConfig
-from haiku.rag.graph.agui.events import (
+from haiku.rag.graph.agui.emitter import (
     emit_text_message_end,
     emit_text_message_start,
     emit_tool_call_args,

tests/graph/agui/test_cli_renderer.py

@@ -5,7 +5,7 @@ from pydantic import BaseModel
 from rich.console import Console
 
 from haiku.rag.graph.agui.cli_renderer import AGUIConsoleRenderer
-from haiku.rag.graph.agui.events import (
+from haiku.rag.graph.agui.emitter import (
     emit_activity,
     emit_run_error,
     emit_run_finished,

tests/graph/agui/test_events.py

@@ -1,201 +0,0 @@
-"""Tests for AG-UI event creation utilities."""
-
-from pydantic import BaseModel
-
-from haiku.rag.graph.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,
-    emit_tool_call_args,
-    emit_tool_call_end,
-    emit_tool_call_start,
-)
-
-
-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", {"message": "Working on task"})
-
-    assert event["type"] == "ACTIVITY_SNAPSHOT"
-    assert event["messageId"] == "msg-1"
-    assert event["activityType"] == "processing"
-    assert event["content"] == {"message": "Working on task"}
-
-
-def test_emit_tool_call_start():
-    """Test TOOL_CALL_START event creation."""
-    event = emit_tool_call_start("call-1", "search_documents")
-
-    assert event["type"] == "TOOL_CALL_START"
-    assert event["toolCallId"] == "call-1"
-    assert event["toolCallName"] == "search_documents"
-    assert "parentMessageId" not in event
-
-
-def test_emit_tool_call_start_with_parent():
-    """Test TOOL_CALL_START event with parent message ID."""
-    event = emit_tool_call_start("call-1", "search", parent_message_id="msg-1")
-
-    assert event["type"] == "TOOL_CALL_START"
-    assert event["toolCallId"] == "call-1"
-    assert event["toolCallName"] == "search"
-    assert event["parentMessageId"] == "msg-1"
-
-
-def test_emit_tool_call_args():
-    """Test TOOL_CALL_ARGS event creation."""
-    import json
-
-    args = {"query": "test query", "limit": 10}
-    event = emit_tool_call_args("call-1", args)
-
-    assert event["type"] == "TOOL_CALL_ARGS"
-    assert event["toolCallId"] == "call-1"
-    assert event["delta"] == json.dumps(args)
-
-
-def test_emit_tool_call_end():
-    """Test TOOL_CALL_END event creation."""
-    event = emit_tool_call_end("call-1")
-
-    assert event["type"] == "TOOL_CALL_END"
-    assert event["toolCallId"] == "call-1"
-
-
-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": "value"}),
-        emit_tool_call_start("c1", "tool"),
-        emit_tool_call_args("c1", {"arg": "value"}),
-        emit_tool_call_end("c1"),
-    ]
-
-    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