From 5b34b9ed0e33d04458c106d6bc9d6e84bc0c1e9c Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Mon, 5 Jan 2026 10:41:35 +0200 Subject: [PATCH] Track potentially more than one running step in agui emitter --- .../haiku/rag/graph/agui/emitter.py | 17 ++++--- .../haiku/rag/graph/research/graph.py | 10 ++-- tests/graph/agui/test_emitter.py | 46 ++++++++++++++++++- tests/graph/agui/test_stream.py | 2 +- 4 files changed, 60 insertions(+), 15 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/graph/agui/emitter.py b/haiku_rag_slim/haiku/rag/graph/agui/emitter.py index 7e9cc7f9..08040e21 100644 --- a/haiku_rag_slim/haiku/rag/graph/agui/emitter.py +++ b/haiku_rag_slim/haiku/rag/graph/agui/emitter.py @@ -70,7 +70,7 @@ class AGUIEmitter[StateT: BaseModel, ResultT]: self._thread_id = thread_id or str(uuid4()) self._run_id = run_id or str(uuid4()) self._last_state: StateT | None = None - self._current_step: str | None = None + self._active_steps: set[str] = set() self._use_deltas = use_deltas @property @@ -112,14 +112,17 @@ class AGUIEmitter[StateT: BaseModel, ResultT]: Args: step_name: Name of the step being started """ - self._current_step = step_name + self._active_steps.add(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(_serialize_event(StepFinishedEvent(step_name=self._current_step))) - self._current_step = None + def finish_step(self, step_name: str) -> None: + """Emit StepFinished event for the specified step. + + Args: + step_name: Name of the step being finished + """ + self._active_steps.discard(step_name) + self.emit(_serialize_event(StepFinishedEvent(step_name=step_name))) def log(self, message: str, role: str = "assistant") -> None: """Emit a text message event. diff --git a/haiku_rag_slim/haiku/rag/graph/research/graph.py b/haiku_rag_slim/haiku/rag/graph/research/graph.py index c1dab7b7..91e3536b 100644 --- a/haiku_rag_slim/haiku/rag/graph/research/graph.py +++ b/haiku_rag_slim/haiku/rag/graph/research/graph.py @@ -154,7 +154,7 @@ def build_research_graph( ) finally: if deps.agui_emitter: - deps.agui_emitter.finish_step() + deps.agui_emitter.finish_step("plan") @g.step async def search_one( @@ -256,7 +256,7 @@ def build_research_graph( ) finally: if deps.agui_emitter: - deps.agui_emitter.finish_step() + deps.agui_emitter.finish_step(step_name) @g.step async def get_batch( @@ -354,7 +354,7 @@ def build_research_graph( return should_continue finally: if deps.agui_emitter: - deps.agui_emitter.finish_step() + deps.agui_emitter.finish_step("decide") @g.step async def human_decide( @@ -427,7 +427,7 @@ def build_research_graph( return "synthesize" finally: if deps.agui_emitter: - deps.agui_emitter.finish_step() + deps.agui_emitter.finish_step("human_decide") @g.step async def synthesize( @@ -467,7 +467,7 @@ def build_research_graph( return result.output finally: if deps.agui_emitter: - deps.agui_emitter.finish_step() + deps.agui_emitter.finish_step("synthesize") # Build the graph structure collect_answers = g.join( diff --git a/tests/graph/agui/test_emitter.py b/tests/graph/agui/test_emitter.py index 10c7e136..344d813e 100644 --- a/tests/graph/agui/test_emitter.py +++ b/tests/graph/agui/test_emitter.py @@ -50,7 +50,7 @@ async def test_emitter_step_events(): emitter: AGUIEmitter[TestState, TestResult] = AGUIEmitter() emitter.start_step("test_step") - emitter.finish_step() + emitter.finish_step("test_step") await emitter.close() events = [] @@ -132,7 +132,7 @@ async def test_emitter_activity_events(): emitter.update_activity( "done", {"stepName": "analyze", "message": "Completed"}, message_id="msg-1" ) - emitter.finish_step() + emitter.finish_step("analyze") await emitter.close() @@ -273,3 +273,45 @@ async def test_emitter_concurrent_emission(): # Should have all 10 messages text_events = [e for e in events if e["type"] == "TEXT_MESSAGE_CHUNK"] assert len(text_events) == 10 + + +@pytest.mark.asyncio +async def test_emitter_concurrent_steps(): + """Test that multiple steps can run concurrently and finish independently.""" + emitter: AGUIEmitter[TestState, TestResult] = AGUIEmitter() + + async def run_step(step_name: str, delay: float): + emitter.start_step(step_name) + await asyncio.sleep(delay) + emitter.finish_step(step_name) + + # Start collecting events in background + events: list[dict] = [] + + async def collect(): + async for event in emitter: + events.append(event) + + collector = asyncio.create_task(collect()) + + # Run three steps concurrently with different durations + # Step A finishes last, Step B first, Step C middle + await asyncio.gather( + run_step("step_a", 0.03), + run_step("step_b", 0.01), + run_step("step_c", 0.02), + ) + + await emitter.close() + await collector + + started = [e for e in events if e["type"] == "STEP_STARTED"] + finished = [e for e in events if e["type"] == "STEP_FINISHED"] + + # All three steps should have started + started_names = {e["stepName"] for e in started} + assert started_names == {"step_a", "step_b", "step_c"} + + # All three steps should have finished + finished_names = {e["stepName"] for e in finished} + assert finished_names == {"step_a", "step_b", "step_c"} diff --git a/tests/graph/agui/test_stream.py b/tests/graph/agui/test_stream.py index 04d9a0b3..34ce71e7 100644 --- a/tests/graph/agui/test_stream.py +++ b/tests/graph/agui/test_stream.py @@ -37,7 +37,7 @@ class MockGraph: if deps.agui_emitter: deps.agui_emitter.start_step("mock_step") deps.agui_emitter.update_activity("working", {"message": "Doing work"}) - deps.agui_emitter.finish_step() + deps.agui_emitter.finish_step("mock_step") return self.result