fix: stream stage events live, thread source_chunks, fix LangChain wrapper

- Stage events now yield immediately to the client instead of being
buffered. Only the chunk is held back for output validation.
- source_chunks threaded through _orchestrator_done metadata and passed
to output_validator.validate(), matching /ask behavior.
- LangChain AgentBenchRetriever updated for RetrievalResult wrapper.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

agent_bench/agents/orchestrator.py

@@ -197,6 +197,7 @@ class Orchestrator:
         messages.append(Message(role=Role.USER, content=question))
         tools = self.registry.get_definitions()
         all_sources: list[str] = []
+        all_source_chunks: list[str] = []
         total_cost = 0.0
         total_input_tokens = 0
         total_output_tokens = 0
@@ -275,6 +276,10 @@ class Orchestrator:
 
                 if "sources" in result.metadata:
                     all_sources.extend(result.metadata["sources"])
+                if "source_chunks" in result.metadata:
+                    all_source_chunks.extend(
+                        result.metadata["source_chunks"]
+                    )
 
         # Max iterations hit — force text answer without tools
         # (same pattern as run(): explicit call after loop)
@@ -314,6 +319,7 @@ class Orchestrator:
                 "tokens_in": total_input_tokens,
                 "tokens_out": total_output_tokens,
                 "iterations": iteration if iteration else 1,
+                "source_chunks": all_source_chunks,
             },
         )
 

agent_bench/langchain_baseline/retriever.py

@@ -17,7 +17,7 @@ from langchain_core.retrievers import BaseRetriever
 class AgentBenchRetriever(BaseRetriever):
     """Wraps agent-bench's async Retriever as a LangChain retriever.
 
-    Delegates to Retriever.search() which returns list[SearchResult].
+    Delegates to Retriever.search() which returns RetrievalResult.
     Each SearchResult has .chunk.content, .chunk.source, .chunk.id, .score.
     """
 
@@ -32,7 +32,7 @@ class AgentBenchRetriever(BaseRetriever):
         *,
         run_manager: AsyncCallbackManagerForRetrieverRun,
     ) -> List[LCDocument]:
-        results = await self.retriever.search(query, top_k=self.top_k)
+        retrieval_result = await self.retriever.search(query, top_k=self.top_k)
         return [
             LCDocument(
                 page_content=r.chunk.content,
@@ -42,7 +42,7 @@ class AgentBenchRetriever(BaseRetriever):
                     "score": r.score,
                 },
             )
-            for r in results
+            for r in retrieval_result.results
         ]
 
     def _get_relevant_documents(

agent_bench/serving/routes.py

@@ -250,8 +250,10 @@ async def ask_stream(body: AskRequest, request: Request) -> StreamingResponse:
             "verdict": injection_verdict_data,
         }).to_sse()
 
-        # Buffer orchestrator events for output validation
-        buffered_events: list = []
+        # Stream orchestrator events live. Stage events are yielded
+        # immediately so the dashboard can animate in real time.
+        # Only the chunk content is accumulated for post-stream
+        # output validation (monitor mode).
         full_answer: list[str] = []
         done_meta: dict = {}
         async for event in orchestrator.run_stream(
@@ -261,21 +263,28 @@ async def ask_stream(body: AskRequest, request: Request) -> StreamingResponse:
             strategy=body.retrieval_strategy,
             history=history,
         ):
-            buffered_events.append(event)
+            if event.type == "_orchestrator_done":
+                # Extract metadata, don't yield to client
+                if event.metadata:
+                    done_meta = event.metadata
+                continue
             if event.type == "chunk" and event.content:
                 full_answer.append(event.content)
-            if event.type == "_orchestrator_done" and event.metadata:
-                done_meta = event.metadata
+                # Don't yield chunk yet — validate first
+                continue
+            # Yield stage and sources events live
+            yield event.to_sse()
 
         # --- Security: output validation (post-generation, monitor mode) ---
         answer_text = "".join(full_answer)
         filtered_answer = answer_text
         output_verdict_data: dict = {"passed": True, "violations": []}
         output_blocked = False
+        source_chunks = done_meta.get("source_chunks", [])
         if output_validator:
             out_verdict = output_validator.validate(
                 output=answer_text,
-                retrieved_chunks=[],
+                retrieved_chunks=source_chunks,
             )
             output_verdict_data = {
                 "passed": out_verdict.passed,
@@ -288,15 +297,11 @@ async def ask_stream(body: AskRequest, request: Request) -> StreamingResponse:
                     "The output was filtered for safety."
                 )
 
-        # Yield buffered orchestrator events (stage events + legacy events)
-        # Filter out _orchestrator_done — route handler emits the real done event
-        for event in buffered_events:
-            if event.type == "_orchestrator_done":
-                continue
-            if output_blocked and event.type == "chunk":
-                yield StreamEvent(type="chunk", content=filtered_answer).to_sse()
-            else:
-                yield event.to_sse()
+        # Yield the (possibly filtered) answer chunk
+        yield StreamEvent(
+            type="chunk",
+            content=filtered_answer if output_blocked else answer_text,
+        ).to_sse()
 
         # --- Output validation stage (monitor mode, after chunk) ---
         yield StreamEvent(type="stage", metadata={
@@ -320,7 +325,8 @@ async def ask_stream(body: AskRequest, request: Request) -> StreamingResponse:
         }).to_sse()
 
         # Record metrics and persist session
-        metrics.record(latency_ms=latency_ms, cost_usd=done_meta.get("estimated_cost_usd", 0.0))
+        cost = done_meta.get("estimated_cost_usd", 0.0)
+        metrics.record(latency_ms=latency_ms, cost_usd=cost)
 
         if body.session_id and conversation_store:
             conversation_store.append(body.session_id, "user", body.question)

tests/test_langchain_baseline/test_retriever.py

@@ -5,6 +5,14 @@ from unittest.mock import AsyncMock, MagicMock
 from agent_bench.langchain_baseline.retriever import AgentBenchRetriever
 
 
+def _make_retrieval_result(results):
+    """Wrap a list of mock SearchResults in a RetrievalResult-like object."""
+    rr = MagicMock()
+    rr.results = results
+    rr.pre_rerank_count = 0
+    return rr
+
+
 def _make_mock_retriever(results=None):
     """Create a mock of agent_bench.rag.retriever.Retriever."""
     retriever = MagicMock()
@@ -17,7 +25,9 @@ def _make_mock_retriever(results=None):
         result.score = 0.85
         result.rank = 1
         results = [result]
-    retriever.search = AsyncMock(return_value=results)
+    retriever.search = AsyncMock(
+        return_value=_make_retrieval_result(results),
+    )
     return retriever