Merge pull request #216 from ggozad/fix/concurrent-search-emitter

Fix concurrent step tracking in AG-UI emitter

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:
+    def finish_step(self, step_name: str) -> None:
-        """Emit StepFinished event for the current step."""
+        """Emit StepFinished event for the specified step.
-        if self._current_step:
+
-            self.emit(_serialize_event(StepFinishedEvent(step_name=self._current_step)))
+        Args:
-            self._current_step = None
+            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.

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(

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

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