Improve retrieval streaming and tracing

app/api.py

@@ -344,6 +344,7 @@ class UIMessageStreamEncoder:
         self.active_reasoning_id = ""
         self.open_tool_call_ids: list[str] = []
         self.announced_tool_call_ids: set[str] = set()
+        self.available_tool_call_ids: set[str] = set()
         self.closed = False
 
     def close_reasoning_block(self) -> list[dict[str, Any]]:
@@ -463,7 +464,7 @@ class UIMessageStreamEncoder:
             )
             args = event.data.get("args")
             args_text = str(event.data.get("args_text", "")).strip()
-            if isinstance(args, dict):
+            if isinstance(args, dict) and args:
                 parts.append(
                     {
                         "type": "tool-input-available",
@@ -472,6 +473,7 @@ class UIMessageStreamEncoder:
                         "input": args,
                     }
                 )
+                self.available_tool_call_ids.add(call_id)
             elif args_text:
                 parts.append(
                     {
@@ -481,6 +483,40 @@ class UIMessageStreamEncoder:
                         "input": {"text": args_text},
                     }
                 )
+                self.available_tool_call_ids.add(call_id)
+            return parts
+
+        if event.type == "tool_call_args_available":
+            call_id = str(event.data.get("call_id", uuid4().hex))
+            tool_name = str(event.data.get("tool_name", "tool"))
+            if call_id not in self.open_tool_call_ids:
+                self.open_tool_call_ids.append(call_id)
+            if call_id not in self.announced_tool_call_ids:
+                parts.append(
+                    {
+                        "type": "tool-input-start",
+                        "toolCallId": call_id,
+                        "toolName": tool_name,
+                    }
+                )
+                self.announced_tool_call_ids.add(call_id)
+            args = event.data.get("args")
+            args_text = str(event.data.get("args_text", "")).strip()
+            if isinstance(args, dict):
+                input_data = args
+            elif args_text:
+                input_data = {"text": args_text}
+            else:
+                return parts
+            parts.append(
+                {
+                    "type": "tool-input-available",
+                    "toolCallId": call_id,
+                    "toolName": tool_name,
+                    "input": input_data,
+                }
+            )
+            self.available_tool_call_ids.add(call_id)
             return parts
 
         if event.type == "source_match":
@@ -507,16 +543,11 @@ class UIMessageStreamEncoder:
             parts.extend(self.close_reasoning_block())
             call_id = str(event.data.get("call_id", uuid4().hex))
             args = event.data.get("args")
-            if (
-                isinstance(args, dict) and args
-            ) or call_id not in self.announced_tool_call_ids:
-                # Providers that stream tool calls incrementally announce the
-                # call before its args are parsed; refresh the input now that
-                # the full args are known. A call id the stream never
-                # announced (e.g. a ToolMessage with a missing id) must also
-                # get an input part first: the AI SDK client throws on an
-                # output for an unknown tool call and drops the rest of the
-                # stream.
+            if call_id not in self.available_tool_call_ids:
+                # The completed-model update normally supplies full arguments
+                # before execution. Keep this completion-time fallback for
+                # providers that omit that update and for orphan ToolMessages:
+                # the AI SDK requires an input part before the output.
                 parts.append(
                     {
                         "type": "tool-input-available",
@@ -526,6 +557,7 @@ class UIMessageStreamEncoder:
                     }
                 )
                 self.announced_tool_call_ids.add(call_id)
+                self.available_tool_call_ids.add(call_id)
             output = {
                 "text": str(event.data.get("output_text", "")),
                 "matches": [

app/chat_service.py

@@ -688,6 +688,46 @@ def collect_retrieval_source_matches(payload: str) -> list[SourceMatch]:
     return matches
 
 
+def _source_match_record(match: SourceMatch) -> dict[str, Any]:
+    """Store lightweight source metadata alongside a persistently capped tool."""
+    return {
+        "doc_id": match.doc_id,
+        "title": match.title,
+        "url": match.url,
+        "source_key": match.source_key,
+        "source_label": match.source_label,
+        "score": match.score,
+        "group": match.group,
+        "path": match.path,
+    }
+
+
+def _source_matches_from_records(records: Any) -> list[SourceMatch]:
+    """Restore source metadata retained outside a truncated retrieval payload."""
+    if not isinstance(records, list):
+        return []
+    matches: list[SourceMatch] = []
+    for record in records:
+        if not isinstance(record, dict):
+            continue
+        try:
+            matches.append(
+                SourceMatch(
+                    doc_id=str(record.get("doc_id", "")),
+                    title=str(record.get("title", "")),
+                    url=str(record.get("url", "")),
+                    source_key=str(record.get("source_key", "")),
+                    source_label=str(record.get("source_label", "")),
+                    score=float(record.get("score", 0.0)),
+                    group=str(record.get("group", "")),
+                    path=str(record.get("path", "")),
+                )
+            )
+        except (TypeError, ValueError):
+            continue
+    return matches
+
+
 def _record_evidence(
     target: dict[str, SourceMatch], matches: list[SourceMatch]
 ) -> None:
@@ -1358,6 +1398,15 @@ class StableToolOutputCapMiddleware(AgentMiddleware):
             "sha256": hashlib.sha256(raw).hexdigest(),
             "max_bytes": self.max_bytes,
         }
+        if getattr(result, "name", "") == "retrieve_tutor_context":
+            # Capping inserts a head/tail marker and intentionally makes the
+            # JSON content unparsable. Preserve only lightweight source
+            # metadata so citations and the UI's chunk count stay accurate
+            # without keeping the discarded chunk bodies in the checkpoint.
+            retrieval_matches = collect_retrieval_source_matches(text)
+            metadata["retrieval_matches"] = [
+                _source_match_record(match) for match in retrieval_matches
+            ]
         turn = _turn_id_for(request)
         record_turn_signal(turn, "tool_outputs_capped", 1)
         record_turn_signal(turn, "tool_output_original_bytes", len(raw))
@@ -2894,7 +2943,11 @@ async def stream_chat(request: ChatRequest) -> AsyncIterator[ChatEvent]:
                     tool_name = str(getattr(message, "name", "tool"))
                     call_matches: list[SourceMatch] = []
                     if tool_name == "retrieve_tutor_context":
-                        call_matches = collect_retrieval_source_matches(payload)
+                        call_matches = _source_matches_from_records(
+                            cap_metadata.get("retrieval_matches")
+                        )
+                        if not call_matches:
+                            call_matches = collect_retrieval_source_matches(payload)
                         _record_evidence(retrieval_evidence, call_matches)
 
                     tool_call = tool_calls_by_id.get(
@@ -3009,10 +3062,44 @@ async def stream_chat(request: ChatRequest) -> AsyncIterator[ChatEvent]:
                 if getattr(message, "tool_calls", None):
                     # The completed message has fully parsed args; streamed
                     # fragments may have announced the call with empty args.
+                    # Publish the completed arguments now, before the tools
+                    # node returns, so a long-running search shows its query
+                    # while it is running rather than only with its result.
                     for tool_call in message.tool_calls:
                         call_id = str(tool_call.get("id") or "")
-                        if call_id:
-                            tool_calls_by_id[call_id] = tool_call
+                        if not call_id:
+                            continue
+                        previous = tool_calls_by_id.get(call_id)
+                        tool_calls_by_id[call_id] = tool_call
+                        if previous is None:
+                            if first_token_at is None:
+                                first_token_at = time.monotonic()
+                            yield ChatEvent(
+                                "tool_call_started",
+                                {
+                                    "message_id": message_id,
+                                    "call_id": call_id,
+                                    "tool_name": str(tool_call.get("name", "tool")),
+                                    "args": tool_call.get("args"),
+                                    "args_text": format_tool_args(
+                                        tool_call.get("args")
+                                    ),
+                                },
+                            )
+                            continue
+                        args = tool_call.get("args")
+                        if args == previous.get("args"):
+                            continue
+                        yield ChatEvent(
+                            "tool_call_args_available",
+                            {
+                                "message_id": message_id,
+                                "call_id": call_id,
+                                "tool_name": str(tool_call.get("name", "tool")),
+                                "args": args,
+                                "args_text": format_tool_args(args),
+                            },
+                        )
                     continue
                 completed_answer = message_content_to_text(message.content)
     finally:

app/chroma_rag.py

@@ -11,14 +11,17 @@ import re
 import threading
 import time
 from collections import Counter, deque
+from contextlib import contextmanager
 from dataclasses import asdict, dataclass
 from pathlib import Path
-from typing import Any, Iterable
+from typing import Any, Iterable, Iterator
 from urllib.parse import unquote, urlparse
 
 import chromadb
 import cohere
 import tiktoken
+from langsmith import trace, traceable
+from langsmith.run_helpers import get_current_run_tree
 from tqdm.auto import tqdm
 
 logger = logging.getLogger(__name__)
@@ -164,6 +167,78 @@ class SearchResult:
     retrieval_method: str = ""
 
 
+RETRIEVAL_STAGE_NAMES = (
+    "embed_ms",
+    "chroma_ms",
+    "dense_hydration_ms",
+    "bm25_ms",
+    "fusion_ms",
+    "rerank_ms",
+    "token_budget_ms",
+)
+
+RETRIEVAL_STAGE_TRACE_NAMES = {
+    "embed_ms": "Cohere Embed",
+    "chroma_ms": "Chroma Vector Search",
+    "dense_hydration_ms": "Dense Result Hydration",
+    "bm25_ms": "BM25 Search",
+    "fusion_ms": "RRF Fusion",
+    "rerank_ms": "Cohere Rerank",
+    "token_budget_ms": "Token Budget",
+}
+
+
+@contextmanager
+def _measure_retrieval_stage(
+    timings: dict[str, float],
+    stage: str,
+    *,
+    trace_inputs: dict[str, Any] | None = None,
+) -> Iterator[None]:
+    started_at = time.perf_counter()
+    with trace(
+        RETRIEVAL_STAGE_TRACE_NAMES[stage],
+        run_type="chain",
+        inputs=trace_inputs or {},
+        metadata={"retrieval_stage": stage},
+    ) as stage_run:
+        try:
+            yield
+        except BaseException:
+            elapsed_ms = (time.perf_counter() - started_at) * 1000
+            timings[stage] = elapsed_ms
+            stage_run.add_metadata({"duration_ms": round(elapsed_ms, 2)})
+            raise
+        else:
+            elapsed_ms = (time.perf_counter() - started_at) * 1000
+            timings[stage] = elapsed_ms
+            stage_run.end(metadata={"duration_ms": round(elapsed_ms, 2)})
+
+
+def _trace_retrieval_inputs(inputs: dict[str, Any]) -> dict[str, Any]:
+    allowed_sources = inputs.get("allowed_sources")
+    return {
+        "query": str(inputs.get("query", "")),
+        "requested_source_count": len(set(allowed_sources or [])),
+        "token_budget": inputs.get("token_budget"),
+    }
+
+
+def _trace_retrieval_outputs(results: Any) -> dict[str, Any]:
+    if not isinstance(results, list):
+        return {"result_type": type(results).__name__}
+    return {
+        "result_count": len(results),
+        "sources": sorted(
+            {
+                result.source
+                for result in results
+                if isinstance(result, SearchResult) and result.source
+            }
+        ),
+    }
+
+
 @dataclass(slots=True)
 class BM25Index:
     records: list[ChunkRecord]
@@ -842,7 +917,14 @@ def _wait_for_cohere_retry(
     else:
         delay = min(window_seconds, max(15.0, 2.0**attempt))
 
-    time.sleep(delay + random.uniform(0.5, 2.0))
+    sleep_seconds = delay + random.uniform(0.5, 2.0)
+    logger.warning(
+        "cohere_embed_retry attempt=%d delay_seconds=%.2f status_code=%s",
+        attempt,
+        sleep_seconds,
+        getattr(exc, "status_code", 429),
+    )
+    time.sleep(sleep_seconds)
 
 
 def _cohere_embeddings_list(response: Any) -> list[list[float]]:
@@ -1253,8 +1335,50 @@ class LocalChromaRetriever:
             self._document_dict: dict[str, dict[str, Any]] = pickle.load(handle)
 
         self._bm25_index = load_bm25_index(self._bm25_index_path)
+        document_sources = {
+            str(document.get("source", "")).strip()
+            for document in self._document_dict.values()
+            if isinstance(document, dict) and document.get("source")
+        }
+        bm25_sources = (
+            {
+                str(record.metadata.get("source", "")).strip()
+                for record in self._bm25_index.records
+                if record.metadata.get("source")
+            }
+            if self._bm25_index is not None
+            else set()
+        )
+        # These artifacts are loaded once by the process-cached retriever. Use
+        # their actual source set instead of a configured UI default so a
+        # complete source selection can safely skip Chroma's redundant,
+        # expensive all-record metadata filter.
+        self._indexed_sources = frozenset(document_sources | bm25_sources)
         self._cohere = cohere.ClientV2(api_key=cohere_api_key)
 
+    def _effective_allowed_sources(
+        self, allowed_sources: list[str] | None
+    ) -> tuple[list[str] | None, str]:
+        if allowed_sources is None:
+            return None, "unfiltered"
+
+        requested_sources = frozenset(allowed_sources)
+        if not requested_sources:
+            # The normal chat path disables local tools for an explicit empty
+            # source selection. Preserve the existing helper semantics here.
+            return allowed_sources, "empty_selection"
+        if self._indexed_sources and self._indexed_sources.issubset(requested_sources):
+            return None, "all_sources_omitted"
+        if len(requested_sources) == 1:
+            return allowed_sources, "single_source"
+        return allowed_sources, "source_subset"
+
+    @traceable(
+        name="Hybrid Retrieval Pipeline",
+        run_type="retriever",
+        process_inputs=_trace_retrieval_inputs,
+        process_outputs=_trace_retrieval_outputs,
+    )
     def search(
         self,
         query: str,
@@ -1262,67 +1386,200 @@ class LocalChromaRetriever:
         allowed_sources: list[str] | None = None,
         token_budget: int | None = None,
     ) -> list[SearchResult]:
-        dense_hits = self._dense_search(query, allowed_sources=allowed_sources)
-        bm25_hits = self._bm25_search(query, allowed_sources=allowed_sources)
-        fused_hits = reciprocal_rank_fusion(
-            [hits for hits in (dense_hits, bm25_hits) if hits],
-            rrf_k=self._rrf_k,
-            top_k=self._fusion_top_k,
+        started_at = time.perf_counter()
+        timings = {stage: 0.0 for stage in RETRIEVAL_STAGE_NAMES}
+        effective_sources, filter_mode = self._effective_allowed_sources(
+            allowed_sources
         )
-        if not fused_hits:
-            return []
+        requested_source_count = len(set(allowed_sources or []))
+        counts = {
+            "dense_hit_count": 0,
+            "bm25_hit_count": 0,
+            "fused_hit_count": 0,
+            "reranked_hit_count": 0,
+            "result_count": 0,
+        }
+        status = "success"
 
-        reranked = rerank_results(
-            self._cohere,
-            query,
-            fused_hits,
-            model=self._rerank_model,
-            top_n=self._rerank_top_k,
-        )
-        return self._apply_token_budget(reranked, token_budget)
+        try:
+            dense_hits = self._dense_search(
+                query,
+                allowed_sources=effective_sources,
+                timings=timings,
+            )
+            counts["dense_hit_count"] = len(dense_hits)
+
+            with _measure_retrieval_stage(
+                timings,
+                "bm25_ms",
+                trace_inputs={
+                    "top_k": self._bm25_top_k,
+                    "filter_applied": effective_sources is not None,
+                    "source_count": len(set(effective_sources or [])),
+                },
+            ):
+                bm25_hits = self._bm25_search(query, allowed_sources=effective_sources)
+            counts["bm25_hit_count"] = len(bm25_hits)
+
+            with _measure_retrieval_stage(
+                timings,
+                "fusion_ms",
+                trace_inputs={
+                    "dense_hit_count": len(dense_hits),
+                    "bm25_hit_count": len(bm25_hits),
+                    "top_k": self._fusion_top_k,
+                    "rrf_k": self._rrf_k,
+                },
+            ):
+                fused_hits = reciprocal_rank_fusion(
+                    [hits for hits in (dense_hits, bm25_hits) if hits],
+                    rrf_k=self._rrf_k,
+                    top_k=self._fusion_top_k,
+                )
+            counts["fused_hit_count"] = len(fused_hits)
+            if not fused_hits:
+                return []
+
+            with _measure_retrieval_stage(
+                timings,
+                "rerank_ms",
+                trace_inputs={
+                    "model": self._rerank_model,
+                    "candidate_count": len(fused_hits),
+                    "top_n": self._rerank_top_k,
+                },
+            ):
+                reranked = rerank_results(
+                    self._cohere,
+                    query,
+                    fused_hits,
+                    model=self._rerank_model,
+                    top_n=self._rerank_top_k,
+                )
+            counts["reranked_hit_count"] = len(reranked)
+
+            with _measure_retrieval_stage(
+                timings,
+                "token_budget_ms",
+                trace_inputs={
+                    "candidate_count": len(reranked),
+                    "token_budget": (
+                        self._token_budget if token_budget is None else token_budget
+                    ),
+                },
+            ):
+                results = self._apply_token_budget(reranked, token_budget)
+            counts["result_count"] = len(results)
+            return results
+        except Exception:
+            status = "error"
+            raise
+        finally:
+            total_ms = (time.perf_counter() - started_at) * 1000
+            timing_metadata: dict[str, Any] = {
+                "status": status,
+                "filter_mode": filter_mode,
+                "requested_source_count": requested_source_count,
+                "indexed_source_count": len(self._indexed_sources),
+                "total_ms": round(total_ms, 2),
+                "stage_ms": {
+                    stage: round(timings[stage], 2) for stage in RETRIEVAL_STAGE_NAMES
+                },
+                **counts,
+            }
+            current_run = get_current_run_tree()
+            if current_run is not None:
+                current_run.add_metadata({"retrieval_timing": timing_metadata})
+
+            logger.info(
+                "retrieval_timing status=%s filter_mode=%s "
+                "requested_sources=%d indexed_sources=%d total_ms=%.2f "
+                "embed_ms=%.2f chroma_ms=%.2f dense_hydration_ms=%.2f "
+                "bm25_ms=%.2f fusion_ms=%.2f rerank_ms=%.2f "
+                "token_budget_ms=%.2f dense_hits=%d bm25_hits=%d "
+                "fused_hits=%d reranked_hits=%d results=%d",
+                status,
+                filter_mode,
+                requested_source_count,
+                len(self._indexed_sources),
+                total_ms,
+                timings["embed_ms"],
+                timings["chroma_ms"],
+                timings["dense_hydration_ms"],
+                timings["bm25_ms"],
+                timings["fusion_ms"],
+                timings["rerank_ms"],
+                timings["token_budget_ms"],
+                counts["dense_hit_count"],
+                counts["bm25_hit_count"],
+                counts["fused_hit_count"],
+                counts["reranked_hit_count"],
+                counts["result_count"],
+            )
 
     def _dense_search(
         self,
         query: str,
         *,
         allowed_sources: list[str] | None = None,
+        timings: dict[str, float] | None = None,
     ) -> list[SearchResult]:
-        query_embedding = embed_texts(
-            self._cohere,
-            [query],
-            input_type="search_query",
-            model=self._embed_model,
-        )[0]
+        stage_timings = timings if timings is not None else {}
+        with _measure_retrieval_stage(
+            stage_timings,
+            "embed_ms",
+            trace_inputs={"model": self._embed_model, "input_count": 1},
+        ):
+            query_embedding = embed_texts(
+                self._cohere,
+                [query],
+                input_type="search_query",
+                model=self._embed_model,
+            )[0]
 
         where = build_where_filter(allowed_sources)
-        raw_results = self._collection.query(
-            query_embeddings=[query_embedding],
-            n_results=self._dense_top_k,
-            where=where,
-            include=["documents", "metadatas", "distances"],
-        )
-
-        chunk_ids = _flatten_query_results(raw_results.get("ids"))
-        documents = _flatten_query_results(raw_results.get("documents"))
-        metadatas = _flatten_query_results(raw_results.get("metadatas"))
-        distances = _flatten_query_results(raw_results.get("distances"))
+        with _measure_retrieval_stage(
+            stage_timings,
+            "chroma_ms",
+            trace_inputs={
+                "collection": self._collection_name,
+                "top_k": self._dense_top_k,
+                "where": where,
+            },
+        ):
+            raw_results = self._collection.query(
+                query_embeddings=[query_embedding],
+                n_results=self._dense_top_k,
+                where=where,
+                include=["documents", "metadatas", "distances"],
+            )
 
-        dense_hits: list[SearchResult] = []
-        for chunk_id, chunk_text, metadata, distance in zip(
-            chunk_ids, documents, metadatas, distances, strict=False
+        with _measure_retrieval_stage(
+            stage_timings,
+            "dense_hydration_ms",
+            trace_inputs={"requested_top_k": self._dense_top_k},
         ):
-            if metadata is None:
-                continue
+            chunk_ids = _flatten_query_results(raw_results.get("ids"))
+            documents = _flatten_query_results(raw_results.get("documents"))
+            metadatas = _flatten_query_results(raw_results.get("metadatas"))
+            distances = _flatten_query_results(raw_results.get("distances"))
+
+            dense_hits: list[SearchResult] = []
+            for chunk_id, chunk_text, metadata, distance in zip(
+                chunk_ids, documents, metadatas, distances, strict=False
+            ):
+                if metadata is None:
+                    continue
 
-            dense_hits.append(
-                self._search_result_from_metadata(
-                    chunk_id=str(chunk_id),
-                    score=_distance_to_score(distance),
-                    chunk_text=str(chunk_text),
-                    metadata=dict(metadata),
-                    retrieval_method="dense",
+                dense_hits.append(
+                    self._search_result_from_metadata(
+                        chunk_id=str(chunk_id),
+                        score=_distance_to_score(distance),
+                        chunk_text=str(chunk_text),
+                        metadata=dict(metadata),
+                        retrieval_method="dense",
+                    )
                 )
-            )
         return dense_hits
 
     def _bm25_search(

frontend/components/chat-message.test.tsx

@@ -54,6 +54,7 @@ describe("ChatMessage activity rendering", () => {
     expect(items[0]?.textContent).toContain("Continuation of the first thought");
     expect(items[1]?.textContent).toContain("Hybrid search");
     expect(items[1]?.textContent).toContain("agent memory");
+    expect(items[1]?.textContent).toContain("3 chunks");
     expect(items[2]?.textContent).toContain("Second thought after retrieval");
     expect(screen.getByText("Final answer")).toBeTruthy();
   });

frontend/components/chat-message.tsx

@@ -463,6 +463,7 @@ function ToolRow({ part }: { part: TutorMessagePart }) {
     ? outputObject.matches.length
     : 0;
   const resultSummary = formatToolResultSummary({
+    toolType: part.type,
     outputText,
     matchCount,
     state: part.state,
@@ -600,11 +601,13 @@ function formatToolStateBadge(
 }
 
 function formatToolResultSummary({
+  toolType,
   outputText,
   matchCount,
   state,
   errorText,
 }: {
+  toolType: string;
   outputText: string;
   matchCount: number;
   state?: string;
@@ -614,7 +617,10 @@ function formatToolResultSummary({
     return "";
   }
   if (matchCount > 0) {
-    return `${matchCount} match${matchCount === 1 ? "" : "es"}`;
+    const isRetrieval =
+      toolType.replace(/^tool-/, "") === "retrieve_tutor_context";
+    const noun = isRetrieval ? "chunk" : "match";
+    return `${matchCount} ${noun}${matchCount === 1 ? "" : "s"}`;
   }
   if (outputText) {
     const lineCount = outputText.split("\n").length;

tests/conftest.py

@@ -0,0 +1,13 @@
+from __future__ import annotations
+
+import os
+
+
+# app.config loads the repository's .env during test collection. Without an
+# explicit override, ordinary unit tests upload synthetic LangGraph and
+# retriever runs into the production LangSmith project. Keep normal pytest
+# hermetic; the opt-in live E2E mode deliberately retains tracing.
+if os.getenv("RUN_LIVE_API_E2E") != "1":
+    os.environ["LANGSMITH_TRACING"] = "false"
+    os.environ["LANGSMITH_TRACING_V2"] = "false"
+    os.environ["LANGCHAIN_TRACING_V2"] = "false"

tests/manual_e2e_langsmith.md

@@ -176,6 +176,10 @@ Expected result:
 - The stream includes `tool-input-*`, `tool-output-available`, `text-delta`,
   `source-url`, `source-document`, `data-source`, and `finish` parts.
 - Tool calls include `retrieve_tutor_context` and `run_kb_command`.
+- Every `retrieve_tutor_context` tool run has a nested
+  `Hybrid Retrieval Pipeline` span. Expanding it in the waterfall shows
+  separate Cohere embed, Chroma, dense hydration, BM25, RRF, Cohere rerank,
+  and token-budget runs.
 - The answer has inline citations, not only a final sources list.
 - `data-source` parts include the source cards the frontend will render.
 
@@ -266,6 +270,45 @@ jq -r '
 ' "/tmp/ai_tutor_trace_${TRACE_ID}.json"
 ```
 
+Hybrid retrieval latency breakdown:
+
+```bash
+jq -r '
+  .runs[]
+  | select(.name == "Hybrid Retrieval Pipeline")
+  | (.custom_metadata.retrieval_timing
+      // .extra.metadata.retrieval_timing
+      // {}) as $timing
+  | [
+      (.inputs.query // ""),
+      ($timing.status // ""),
+      ($timing.filter_mode // ""),
+      ($timing.total_ms // ""),
+      ($timing.stage_ms.embed_ms // ""),
+      ($timing.stage_ms.chroma_ms // ""),
+      ($timing.stage_ms.dense_hydration_ms // ""),
+      ($timing.stage_ms.bm25_ms // ""),
+      ($timing.stage_ms.fusion_ms // ""),
+      ($timing.stage_ms.rerank_ms // ""),
+      ($timing.stage_ms.token_budget_ms // "")
+    ]
+  | @tsv
+' "/tmp/ai_tutor_trace_${TRACE_ID}.json"
+```
+
+The columns are query, status, filter mode, total, embed, Chroma, dense-result
+hydration, BM25, RRF, rerank, and token-budget latency, all in milliseconds.
+The backend logs the same values as one `retrieval_timing` line without logging
+the raw student query. A complete source selection should report
+`filter_mode=all_sources_omitted`; real source subsets should report
+`single_source` or `source_subset`.
+
+In the LangSmith waterfall, expand `retrieve_tutor_context`, then
+`Hybrid Retrieval Pipeline`, to see the same stages as individually timed child
+runs. Their names are `Cohere Embed`, `Chroma Vector Search`,
+`Dense Result Hydration`, `BM25 Search`, `RRF Fusion`, `Cohere Rerank`, and
+`Token Budget`.
+
 LLM calls:
 
 ```bash
@@ -295,9 +338,12 @@ For latency debugging, compare these numbers:
 - Root trace duration.
 - Count of `run_kb_command` calls.
 - Total shell/tool duration.
+- The `Hybrid Retrieval Pipeline` stage breakdown for every retrieval call.
 - Count of LLM/chat model calls.
 - Longest LLM call duration.
 - Whether the trace ended in `success`, `error`, or cancellation.
 
-If shell/tool duration is tiny but total trace duration is large, the bottleneck
-is model inference or repeated model turns, not filesystem access.
+If `Hybrid Retrieval Pipeline` is fast but total trace duration is large, the
+bottleneck is model inference or repeated model turns. If retrieval is slow,
+its child runs and metadata identify whether the time was spent in Cohere
+embedding, Chroma, BM25/RRF, Cohere reranking, or token budgeting.

tests/test_api.py

@@ -612,6 +612,60 @@ class ApiTestCase(unittest.TestCase):
         self.assertEqual(matches[0]["score"], 0.9)
         self.assertEqual(matches[0]["path"], "raw/docs/peft/lora.md")
 
+    def test_completed_tool_args_refresh_while_the_tool_is_running(self) -> None:
+        encoder = UIMessageStreamEncoder()
+        encoder.encode(ChatEvent("message_started", {"message_id": "m1"}))
+
+        started = encoder.encode(
+            ChatEvent(
+                "tool_call_started",
+                {
+                    "message_id": "m1",
+                    "call_id": "call_1",
+                    "tool_name": "retrieve_tutor_context",
+                    "args": {},
+                },
+            )
+        )
+        self.assertEqual([part["type"] for part in started], ["tool-input-start"])
+
+        refreshed = encoder.encode(
+            ChatEvent(
+                "tool_call_args_available",
+                {
+                    "message_id": "m1",
+                    "call_id": "call_1",
+                    "tool_name": "retrieve_tutor_context",
+                    "args": {"query": "secure API key storage"},
+                },
+            )
+        )
+        self.assertEqual(
+            [part["type"] for part in refreshed],
+            ["tool-input-available"],
+        )
+        self.assertEqual(
+            refreshed[0]["input"],
+            {"query": "secure API key storage"},
+        )
+
+        completed = encoder.encode(
+            ChatEvent(
+                "tool_call_completed",
+                {
+                    "message_id": "m1",
+                    "call_id": "call_1",
+                    "tool_name": "retrieve_tutor_context",
+                    "args": {"query": "secure API key storage"},
+                    "output_text": "payload",
+                },
+            )
+        )
+        self.assertNotIn(
+            "tool-input-available",
+            [part["type"] for part in completed],
+        )
+
     def test_chat_rejects_oversized_query(self) -> None:
         from app.api import MAX_QUERY_CHARS
 

tests/test_chat_service.py

@@ -117,6 +117,90 @@ class FakeStreamingAgent(FakeAgent):
         }
 
 
+class FakeIncrementalRetrievalAgent(FakeAgent):
+    """DeepSeek-like stream: tool id/name first, complete args at model end."""
+
+    async def astream(self, *args, **kwargs):
+        yield {
+            "type": "messages",
+            "data": (
+                AIMessageChunk(
+                    content="",
+                    tool_calls=[
+                        {
+                            "id": "call_retrieval",
+                            "name": "retrieve_tutor_context",
+                            "args": {},
+                        }
+                    ],
+                ),
+                {"langgraph_node": "model"},
+            ),
+        }
+        yield {
+            "type": "updates",
+            "data": {
+                "model": {
+                    "messages": [
+                        AIMessage(
+                            content="",
+                            tool_calls=[
+                                {
+                                    "id": "call_retrieval",
+                                    "name": "retrieve_tutor_context",
+                                    "args": {"query": "secure API key storage"},
+                                }
+                            ],
+                        )
+                    ]
+                }
+            },
+        }
+        yield {
+            "type": "updates",
+            "data": {
+                "tools": {
+                    "messages": [
+                        ToolMessage(
+                            content=(
+                                '{"query":"secure API key storage","matches":[...\n\n'
+                                "[... tool output truncated ...]\n\n...]}"
+                            ),
+                            name="retrieve_tutor_context",
+                            tool_call_id="call_retrieval",
+                            additional_kwargs={
+                                "stable_tool_cap": {
+                                    "retrieval_matches": [
+                                        {
+                                            "doc_id": "doc-keys",
+                                            "title": "Manage API keys",
+                                            "url": "https://example.com/keys",
+                                            "source_key": "langchain",
+                                            "source_label": "LangChain Docs",
+                                            "score": 0.9,
+                                            "group": "docs",
+                                            "path": "",
+                                        }
+                                    ]
+                                }
+                            },
+                        )
+                    ]
+                }
+            },
+        }
+        yield {
+            "type": "updates",
+            "data": {
+                "model": {
+                    "messages": [
+                        AIMessage(content="Store secrets outside source control.")
+                    ]
+                }
+            },
+        }
+
+
 class FakeAnswerAgent(FakeAgent):
     """Agent that streams nothing and answers with one final AI message."""
 
@@ -871,6 +955,51 @@ class ChatServiceTestCase(unittest.TestCase):
         )
         self.assertTrue(build_agent_mock.call_args.kwargs["include_thoughts"])
 
+    def test_stream_chat_publishes_completed_tool_args_before_result(self) -> None:
+        agent = FakeIncrementalRetrievalAgent([])
+        self.addCleanup(_drop_thread_record, "thread_incremental_tool")
+        request = ChatRequest(
+            query="Where should I store API keys?",
+            source_keys=("langchain",),
+            model_name=DEEPSEEK_DIRECT_MODEL_NAME,
+            include_reasoning=False,
+            enabled_tools=(),
+        )
+
+        async def collect_events():
+            return [event async for event in stream_chat(request)]
+
+        with (
+            patch("app.chat_service.build_agent", return_value=agent),
+            patch(
+                "app.chat_service.new_thread_id",
+                return_value="thread_incremental_tool",
+            ),
+        ):
+            events = asyncio.run(collect_events())
+
+        event_types = [event.type for event in events]
+        self.assertLess(
+            event_types.index("tool_call_started"),
+            event_types.index("tool_call_args_available"),
+        )
+        self.assertLess(
+            event_types.index("tool_call_args_available"),
+            event_types.index("tool_call_completed"),
+        )
+        args_event = next(
+            event for event in events if event.type == "tool_call_args_available"
+        )
+        self.assertEqual(
+            args_event.data["args"],
+            {"query": "secure API key storage"},
+        )
+        completed = next(
+            event for event in events if event.type == "tool_call_completed"
+        )
+        self.assertEqual(len(completed.data["matches"]), 1)
+        self.assertEqual(completed.data["matches"][0]["doc_id"], "doc-keys")
+
     def test_stream_chat_resolves_shell_citation_after_final_answer(self) -> None:
         agent = FakeStreamingAgent([])
         self.addCleanup(_drop_thread_record, "thread_rg")

tests/test_chroma_rag.py

@@ -5,7 +5,7 @@ import pickle
 import tempfile
 import unittest
 from pathlib import Path
-from unittest.mock import patch
+from unittest.mock import ANY, Mock, patch
 
 import chromadb
 import tiktoken
@@ -351,11 +351,123 @@ class TokenBudgetTestCase(unittest.TestCase):
         self.assertEqual(len(kept_all), 3)
 
 
+class SourceNormalizationAndTimingTestCase(unittest.TestCase):
+    def _retriever(self) -> LocalChromaRetriever:
+        retriever = LocalChromaRetriever.__new__(LocalChromaRetriever)
+        retriever._indexed_sources = frozenset({"langchain", "peft"})
+        retriever._rrf_k = 60
+        retriever._bm25_top_k = 30
+        retriever._fusion_top_k = 30
+        retriever._rerank_model = "rerank-test"
+        retriever._rerank_top_k = 5
+        retriever._token_budget = 100_000
+        retriever._cohere = object()
+        return retriever
+
+    def _result(self) -> SearchResult:
+        return SearchResult(
+            chunk_id="chunk-1",
+            doc_id="doc-1",
+            title="Doc",
+            url="https://example.com",
+            source="peft",
+            retrieve_doc=False,
+            tokens=10,
+            score=0.9,
+            content="content",
+            chunk_content="content",
+            retrieval_method="dense",
+        )
+
+    def test_complete_indexed_source_selection_omits_filter(self) -> None:
+        retriever = self._retriever()
+
+        effective, mode = retriever._effective_allowed_sources(["langchain", "peft"])
+        self.assertIsNone(effective)
+        self.assertEqual(mode, "all_sources_omitted")
+
+        # Unknown requested keys do not change the result set, so a superset
+        # still safely covers every source present in the retrieval artifacts.
+        effective, mode = retriever._effective_allowed_sources(
+            ["langchain", "peft", "future-source"]
+        )
+        self.assertIsNone(effective)
+        self.assertEqual(mode, "all_sources_omitted")
+
+    def test_partial_source_selection_keeps_filter(self) -> None:
+        retriever = self._retriever()
+
+        sources = ["peft"]
+        effective, mode = retriever._effective_allowed_sources(sources)
+
+        self.assertEqual(effective, sources)
+        self.assertEqual(mode, "single_source")
+
+    def test_search_logs_and_traces_stage_timings_after_normalization(self) -> None:
+        retriever = self._retriever()
+        result = self._result()
+        retriever._dense_search = Mock(return_value=[result])
+        retriever._bm25_search = Mock(return_value=[])
+        retriever._apply_token_budget = Mock(return_value=[result])
+
+        class _FakeRun:
+            def __init__(self) -> None:
+                self.metadata: dict[str, object] = {}
+
+            def add_metadata(self, metadata: dict[str, object]) -> None:
+                self.metadata.update(metadata)
+
+        fake_run = _FakeRun()
+        with (
+            patch(
+                "app.chroma_rag.rerank_results",
+                return_value=[result],
+            ),
+            patch(
+                "app.chroma_rag.get_current_run_tree",
+                return_value=fake_run,
+            ),
+            self.assertLogs("app.chroma_rag", level="INFO") as captured_logs,
+        ):
+            results = retriever.search(
+                "adapter configuration",
+                allowed_sources=["langchain", "peft"],
+            )
+
+        self.assertEqual(results, [result])
+        retriever._dense_search.assert_called_once_with(
+            "adapter configuration",
+            allowed_sources=None,
+            timings=ANY,
+        )
+        retriever._bm25_search.assert_called_once_with(
+            "adapter configuration",
+            allowed_sources=None,
+        )
+        timing = fake_run.metadata["retrieval_timing"]
+        assert isinstance(timing, dict)
+        self.assertEqual(timing["filter_mode"], "all_sources_omitted")
+        self.assertEqual(timing["requested_source_count"], 2)
+        self.assertEqual(timing["indexed_source_count"], 2)
+        self.assertEqual(timing["result_count"], 1)
+        self.assertIn("stage_ms", timing)
+        self.assertTrue(
+            any(
+                "filter_mode=all_sources_omitted" in message and "chroma_ms=" in message
+                for message in captured_logs.output
+            )
+        )
+
+
 class CollectionOpenTestCase(unittest.TestCase):
-    def _write_document_dict(self, directory: str) -> str:
+    def _write_document_dict(
+        self,
+        directory: str,
+        documents: dict[str, dict[str, object]] | None = None,
+    ) -> str:
         path = Path(directory) / "document_dict_test.pkl"
         with open(path, "wb") as handle:
-            pickle.dump({}, handle)
+            pickle.dump(documents or {}, handle)
         return str(path)
 
     def test_init_fails_loudly_when_collection_missing(self) -> None:
@@ -383,7 +495,13 @@ class CollectionOpenTestCase(unittest.TestCase):
             chromadb.PersistentClient(path=temp_dir).create_collection(
                 name="test-collection"
             )
-            document_dict_path = self._write_document_dict(temp_dir)
+            document_dict_path = self._write_document_dict(
+                temp_dir,
+                {
+                    "doc-1": {"source": "peft"},
+                    "doc-2": {"source": "langchain"},
+                },
+            )
 
             retriever = LocalChromaRetriever(
                 db_path=temp_dir,
@@ -393,6 +511,10 @@ class CollectionOpenTestCase(unittest.TestCase):
             )
 
             self.assertEqual(retriever._collection.name, "test-collection")
+            self.assertEqual(
+                retriever._indexed_sources,
+                frozenset({"peft", "langchain"}),
+            )
 
 
 if __name__ == "__main__":

tests/test_memory_variants.py

@@ -9,6 +9,7 @@ no model client, no API keys, no vector DB.
 from __future__ import annotations
 
 import hashlib
+import json
 import unittest
 from types import SimpleNamespace
 from unittest import mock
@@ -300,6 +301,43 @@ class StableToolOutputCapTests(unittest.TestCase):
             signals["tool_output_retained_bytes"],
         )
 
+    def test_cap_preserves_retrieval_chunk_metadata(self) -> None:
+        raw = json.dumps(
+            {
+                "query": "secure API key storage",
+                "matches": [
+                    {
+                        "chunk_id": "chunk-1",
+                        "doc_id": "doc-1",
+                        "title": "Manage API keys",
+                        "url": "https://example.com/keys",
+                        "source": "langchain",
+                        "retrieve_doc": True,
+                        "tokens": 1200,
+                        "score": 0.9,
+                        "content": "x" * 5_000,
+                        "chunk_content": "API key guidance",
+                        "heading_path": "Testing > Manage API keys",
+                        "retrieval_method": "hybrid",
+                    }
+                ],
+            }
+        )
+        result = StableToolOutputCapMiddleware(2_048)._cap(
+            make_request([], "retrieval-cap"),
+            ToolMessage(
+                content=raw,
+                name="retrieve_tutor_context",
+                tool_call_id="call-retrieval",
+            ),
+        )
+
+        self.assertIn("tool output truncated", result.content)
+        retained = result.additional_kwargs["stable_tool_cap"]["retrieval_matches"]
+        self.assertEqual(len(retained), 1)
+        self.assertEqual(retained[0]["doc_id"], "doc-1")
+        self.assertEqual(retained[0]["source_key"], "langchain")
+
 
 class ExperimentCompactionMiddlewareTests(unittest.TestCase):
     class FakeModel: