perf: instrument MAP-NAV episode timings

packages/shared-python/shared/services/retrieval/nav/nav_agent.py

@@ -1,7 +1,8 @@
 from __future__ import annotations
 
-import time
+import logging
 import os
+import time
 from typing import Any, List, Optional, Sequence, Tuple
 
 from ._compat import AgentStep, EpisodeResult
@@ -37,6 +38,7 @@
 # Back-compat aliases for tests / callers.
 _evidence_owner_section_id = evidence_owner_section_id
 _unit_score_for_evidence_chunk = unit_score_for_evidence_chunk
+_logger = logging.getLogger(__name__)
 
 
 def _chunks_to_retrieved_nodes(chunks: List[Chunk]) -> List[str]:
@@ -398,6 +400,7 @@ def _run_nav_episode_body(
 
     state = NavState(doc_id=episode_doc, query=query, task_type=task_type)
     steps: List[AgentStep] = []
+    map_started = time.perf_counter()
     if namespace_mode:
         section_ids = list(ts.sections_for_doc(""))
         state.map_scores, state.unit_scores = compute_corpus_map_and_unit_scores(
@@ -408,6 +411,12 @@ def _run_nav_episode_body(
         state.map_scores, state.unit_scores = compute_map_and_unit_scores(
             ts, doc_id=episode_doc, query=query, root_ids=section_ids
         )
+    _logger.info(
+        "retrieval mapnav phase=map_scoring seconds=%.3f documents=%d sections=%d",
+        time.perf_counter() - map_started,
+        len(corpus_ids),
+        len(section_ids),
+    )
     state.highlight_ids = select_map_highlights(
         state.unit_scores, k=int(cfg.collect_top_k)
     )
@@ -434,6 +443,11 @@ def _run_nav_episode_body(
                 }, t0=plan_t0),
             )
         )
+        _logger.info(
+            "retrieval mapnav phase=planner seconds=%.3f subgoals=%d",
+            time.perf_counter() - plan_t0,
+            len(retrieval_plan.subgoals),
+        )
 
     # Checklist: wave orchestration; navigate mode: classic single navigate.
     if cfg.is_checklist and state.retrieval_plan is not None:
@@ -452,6 +466,11 @@ def _run_nav_episode_body(
                 detail=stamp_step_detail(orch_detail, t0=orch_t0),
             )
         )
+        _logger.info(
+            "retrieval mapnav phase=orchestration seconds=%.3f waves=%d",
+            time.perf_counter() - orch_t0,
+            len(orch_detail.get("waves", [])),
+        )
     else:
         navigate(
             ts,
@@ -464,13 +483,19 @@ def _run_nav_episode_body(
             steps_out=steps,
         )
 
+    evidence_started = time.perf_counter()
     fill = pack_nav_evidence(
         _dedupe_scored(list(state.collected)),
         ts,
         state,
         cfg,
         budget_chars=budget_chars,
     )
+    _logger.info(
+        "retrieval mapnav phase=evidence_pack seconds=%.3f chunks=%d",
+        time.perf_counter() - evidence_started,
+        len(fill.kept_chunks),
+    )
     scored_chunks = list(fill.scored_chunks)
     retrieval_seconds = time.perf_counter() - retrieval_t0
     composed = ""

packages/shared-python/shared/services/retrieval/nav/nav_llm.py

@@ -20,7 +20,9 @@
 
 from __future__ import annotations
 
+import logging
 import os
+import time
 from contextlib import contextmanager
 from contextvars import ContextVar
 from typing import Any, Callable, Dict, Iterator, Optional, Sequence
@@ -30,6 +32,7 @@
 NavChatBackend = Callable[..., Dict[str, Any]]
 
 _backend: Optional[NavChatBackend] = None
+_logger = logging.getLogger(__name__)
 
 _DS_DEFAULT_MODEL = "deepseek-v4-flash"
 _DEFAULT_PLANNER_THINK_MAX = 16384
@@ -190,25 +193,34 @@ def nav_chat(
         raise NavTokenLimit(used=nav_tokens_used(), limit=nav_token_limit())
 
     merged_extra = _merge_thinking_extra(extra, role=thinking_role, model=model)
+    call_started = time.perf_counter()
 
     if _backend is not None:
-        result = _backend(
-            purpose=purpose,
-            messages=list(messages),
-            model=model,
-            temperature=temperature,
-            max_tokens=max_tokens,
-            response_format=response_format,
-            extra=merged_extra,
-            thinking_role=thinking_role,
-            context=context,
-            api_key_env=api_key_env,
-            base_url_env=base_url_env,
-            timeout=timeout,
-            usage_tag=usage_tag,
-        )
-        record_episode_tokens((result or {}).get("usage"))
-        return result
+        try:
+            result = _backend(
+                purpose=purpose,
+                messages=list(messages),
+                model=model,
+                temperature=temperature,
+                max_tokens=max_tokens,
+                response_format=response_format,
+                extra=merged_extra,
+                thinking_role=thinking_role,
+                context=context,
+                api_key_env=api_key_env,
+                base_url_env=base_url_env,
+                timeout=timeout,
+                usage_tag=usage_tag,
+            )
+            record_episode_tokens((result or {}).get("usage"))
+            return result
+        finally:
+            _logger.info(
+                "retrieval mapnav llm_call purpose=%s role=%s seconds=%.3f",
+                purpose,
+                thinking_role,
+                time.perf_counter() - call_started,
+            )
 
     from ._compat import cached_chat_completion  # type: ignore
     from ._compat import (  # type: ignore
@@ -230,16 +242,24 @@ def nav_chat(
             f"(model={model!r}; need DS_KEY for deepseek-* or OPENAI_API_KEY)."
         )
     client = make_openai_client(api_key=key, base_url=base_url, timeout=timeout)
-    cached = cached_chat_completion(
-        client,
-        purpose=purpose,
-        model=model,
-        messages=list(messages),
-        temperature=temperature,
-        max_tokens=max_tokens,
-        response_format=response_format,
-        extra=merged_extra,
-    )
+    try:
+        cached = cached_chat_completion(
+            client,
+            purpose=purpose,
+            model=model,
+            messages=list(messages),
+            temperature=temperature,
+            max_tokens=max_tokens,
+            response_format=response_format,
+            extra=merged_extra,
+        )
+    finally:
+        _logger.info(
+            "retrieval mapnav llm_call purpose=%s role=%s seconds=%.3f",
+            purpose,
+            thinking_role,
+            time.perf_counter() - call_started,
+        )
     if usage_tag:
         record_usage(usage_tag, cached.get("usage"))
     record_episode_tokens(cached.get("usage"))

packages/shared-python/shared/services/retrieval/nav/nav_orchestrate.py

@@ -8,14 +8,14 @@
 
 from __future__ import annotations
 
-from .nav_token_budget import stamp_step_detail
-
+import logging
 import re
 import time
 from contextlib import contextmanager
 from dataclasses import asdict
 from typing import Any, Dict, Iterator, List, Optional, Sequence, Set, Tuple
 
+from .nav_token_budget import stamp_step_detail
 from .nav_navigate import navigate
 from .nav_plan import (
     RetrievalPlan,
@@ -29,6 +29,7 @@
 from .nav_verify import apply_bindings_from_result, build_subgoal_result
 
 _SLOT_STRIP_RE = re.compile(r"\{\{\s*[^}]+\s*\}\}")
+_logger = logging.getLogger(__name__)
 
 
 def ready_subgoal_ids(
@@ -483,8 +484,18 @@ def _run_one(sid: str, working_state: NavState, out_steps: Optional[List[Any]])
             )
 
         # Serial wave execution (parallel fan-out retired with ThreadPoolExecutor).
+        wave_started = time.perf_counter()
         for sid in ready:
-            outputs.append(_run_one(sid, state, steps_out))
+            subgoal_started = time.perf_counter()
+            try:
+                outputs.append(_run_one(sid, state, steps_out))
+            finally:
+                _logger.info(
+                    "retrieval mapnav harvest subgoal=%s wave=%d seconds=%.3f",
+                    sid,
+                    wave_idx,
+                    time.perf_counter() - subgoal_started,
+                )
 
         # Bookkeeping shared by both decision paths.
         for item in outputs:
@@ -508,9 +519,17 @@ def _run_one(sid: str, working_state: NavState, out_steps: Optional[List[Any]])
                 state.subgoal_attempt_counts.get(sid, 0)
             ) + 1
 
-        control_detail = _apply_plan_control(
-            ts, state, config, plan=plan, outputs=outputs, by_id=by_id, steps_out=steps_out
-        )
+        control_started = time.perf_counter()
+        try:
+            control_detail = _apply_plan_control(
+                ts, state, config, plan=plan, outputs=outputs, by_id=by_id, steps_out=steps_out
+            )
+        finally:
+            _logger.info(
+                "retrieval mapnav plan_control wave=%d seconds=%.3f",
+                wave_idx,
+                time.perf_counter() - control_started,
+            )
         wave_detail["plan_control"] = control_detail
         replan_requested = bool(control_detail.get("replan"))
         if control_detail.get("done"):
@@ -525,6 +544,12 @@ def _run_one(sid: str, working_state: NavState, out_steps: Optional[List[Any]])
                 )
             )
         summary["waves"].append(wave_detail)
+        _logger.info(
+            "retrieval mapnav wave=%d ready=%d seconds=%.3f",
+            wave_idx,
+            len(ready),
+            time.perf_counter() - wave_started,
+        )
 
         if replan_requested:
             cap = int(getattr(config, "max_replans", 0) or 0)