Add OpenRouter planning prepass for tool contexts

Author: jck411

src/backend/chat/orchestrator.py

@@ -25,7 +25,7 @@
 from .image_reflection import reflect_assistant_images
 from .mcp_registry import MCPServerConfig, MCPToolAggregator
 from .streaming import SseEvent, StreamingHandler
-from .tool_context_planner import ToolContextPlanner
+from .tool_context_planner import merge_model_tool_plan, ToolContextPlanner
 
 _TOOL_RATIONALE_INSTRUCTION = (
     "Before each tool call, emit numbered one-sentence rationales in order (e.g.,"
@@ -41,6 +41,42 @@
 _MAX_RANKED_TOOLS = 5
 
 
+def _compact_tool_digest(
+    digest: dict[str, list[dict[str, Any]]] | None,
+) -> dict[str, list[dict[str, Any]]]:
+    if not digest:
+        return {}
+
+    compact: dict[str, list[dict[str, Any]]] = {}
+    for context, entries in digest.items():
+        if not isinstance(entries, list):
+            continue
+        filtered: list[dict[str, Any]] = []
+        for entry in entries:
+            if not isinstance(entry, dict):
+                continue
+            name = entry.get("name")
+            if not isinstance(name, str) or not name.strip():
+                continue
+            compact_entry: dict[str, Any] = {"name": name.strip()}
+            description = entry.get("description")
+            if isinstance(description, str) and description.strip():
+                compact_entry["description"] = description.strip()
+            parameters = entry.get("parameters")
+            if isinstance(parameters, dict) and parameters:
+                compact_entry["parameters"] = parameters
+            server = entry.get("server")
+            if isinstance(server, str) and server.strip():
+                compact_entry["server"] = server.strip()
+            score = entry.get("score")
+            if isinstance(score, (int, float)):
+                compact_entry["score"] = float(score)
+            filtered.append(compact_entry)
+        if filtered:
+            compact[context] = filtered
+    return compact
+
+
 def _iter_attachment_ids(content: Any) -> Iterable[str]:
     if isinstance(content, list):
         for item in content:
@@ -273,38 +309,81 @@ async def process_stream(
             conversation,
             capability_digest=capability_digest,
         )
-        plan_payload = plan.to_dict()
         contexts = plan.contexts_for_attempt(0)
         ranked_tool_names: list[str] = []
         ranked_digest: dict[str, list[dict[str, Any]]] | None = None
-        if contexts:
-            digest_for_contexts = getattr(self._mcp_client, "get_capability_digest", None)
-            if callable(digest_for_contexts):
-                try:
-                    ranked_digest = digest_for_contexts(
-                        contexts, limit=_MAX_RANKED_TOOLS, include_global=False
-                    )
-                except Exception as exc:  # pragma: no cover - defensive fallback
-                    logger.debug(
-                        "Failed to obtain ranked capability digest for contexts %s: %s",
-                        contexts,
-                        exc,
-                    )
+        digest_for_contexts = getattr(self._mcp_client, "get_capability_digest", None)
+        if contexts and callable(digest_for_contexts):
+            try:
+                ranked_digest = digest_for_contexts(
+                    contexts, limit=_MAX_RANKED_TOOLS, include_global=False
+                )
+            except Exception as exc:  # pragma: no cover - defensive fallback
+                logger.debug(
+                    "Failed to obtain ranked capability digest for contexts %s: %s",
+                    contexts,
+                    exc,
+                )
+                ranked_digest = {}
+        plan_request_digest = _compact_tool_digest(ranked_digest)
+        if contexts or plan.broad_search:
+            try:
+                planner_response = await self._client.request_tool_plan(
+                    request=request,
+                    conversation=conversation,
+                    tool_digest=plan_request_digest,
+                )
+            except Exception as exc:  # pragma: no cover - remote planner is best effort
+                logger.debug("Remote tool planning failed: %s", exc)
+            else:
+                merged_plan = merge_model_tool_plan(plan, planner_response)
+                plan = merged_plan
+                contexts = plan.contexts_for_attempt(0)
+                if contexts and callable(digest_for_contexts):
+                    try:
+                        ranked_digest = digest_for_contexts(
+                            contexts, limit=_MAX_RANKED_TOOLS, include_global=False
+                        )
+                    except Exception as exc:  # pragma: no cover - defensive fallback
+                        logger.debug(
+                            "Failed to obtain ranked capability digest for contexts %s: %s",
+                            contexts,
+                            exc,
+                        )
+                        ranked_digest = {}
+                else:
                     ranked_digest = {}
-            if ranked_digest:
-                seen_names: set[str] = set()
-                for context in contexts:
-                    entries = ranked_digest.get(context) or []
-                    for entry in entries:
-                        if not isinstance(entry, dict):
-                            continue
-                        name = entry.get("name")
-                        if not isinstance(name, str) or not name:
-                            continue
-                        if name in seen_names:
-                            continue
-                        seen_names.add(name)
-                        ranked_tool_names.append(name)
+
+        if plan.candidate_tools:
+            seen_names: set[str] = set()
+            ordered_contexts = list(contexts)
+            if "__all__" in plan.candidate_tools and "__all__" not in ordered_contexts:
+                ordered_contexts.append("__all__")
+            for context in ordered_contexts:
+                candidates = plan.candidate_tools.get(context) or []
+                for candidate in candidates:
+                    name = candidate.name.strip()
+                    if not name or name in seen_names:
+                        continue
+                    seen_names.add(name)
+                    ranked_tool_names.append(name)
+
+        if ranked_digest:
+            seen_names = set(ranked_tool_names)
+            for context in contexts:
+                entries = ranked_digest.get(context) or []
+                for entry in entries:
+                    if not isinstance(entry, dict):
+                        continue
+                    name = entry.get("name")
+                    if not isinstance(name, str) or not name:
+                        continue
+                    if name in seen_names:
+                        continue
+                    seen_names.add(name)
+                    ranked_tool_names.append(name)
+
+        plan_payload = plan.to_dict()
         request_event_id: str | None = None
         if isinstance(request.metadata, dict):
             candidate = request.metadata.get("client_request_id")