fixes and enhancements

backend/long_term_memory.py

@@ -2304,6 +2304,13 @@ async def recall_graph_context(
         if relationships:
             parts.append("\nKNOWN RELATIONSHIPS:")
             seen_triples: Set[str] = set()
+            # Cap relationships surfaced in the LLM context.  A 5-entity /
+            # 2-hop traversal on a mature graph can easily return hundreds of
+            # edges; without a cap we blow the prompt budget and the model
+            # loses focus.  Callers who need the full set can call
+            # ``get_relationships`` directly.
+            _rel_cap = max(entity_limit * 6, 20)
+            _rel_count = 0
             for r in relationships:
                 # Filter by minimum confidence when specified
                 if min_confidence > 0.0 and r.get("confidence", 1.0) < min_confidence:
@@ -2316,6 +2323,12 @@ async def recall_graph_context(
                 conf = r.get("confidence", 1.0)
                 conf_str = f" (conf: {conf:.2f})" if conf < 0.8 else ""
                 parts.append(f"  • {triple}{conf_str}")
+                _rel_count += 1
+                if _rel_count >= _rel_cap:
+                    parts.append(
+                        f"  … (+{len(relationships) - _rel_count} more relationships omitted)"
+                    )
+                    break
 
         # Claims
         if include_claims and entity_names:
@@ -2344,8 +2357,16 @@ async def recall_graph_context(
 
         if communities:
             parts.append("\nTHEMATIC CLUSTERS:")
-            for c in communities:
+            # Cap at 5 clusters — the top ones by construction are the most
+            # relevant to the seed entities, and beyond that the summaries
+            # tend to duplicate content already in KNOWN ENTITIES.
+            _community_cap = 5
+            for c in communities[:_community_cap]:
                 parts.append(f"  [{c['topic']}] {c['summary']}")
+            if len(communities) > _community_cap:
+                parts.append(
+                    f"  … (+{len(communities) - _community_cap} more clusters omitted)"
+                )
 
         # Optional: contradictions
         if include_contradictions:

backend/model_backend.py

@@ -965,3 +965,52 @@ def create_model_backend(config: Config) -> ModelBackend:
 
     else:
         raise ValueError(f"Unknown model backend: {backend_type}")
+
+
+# ── Shared model-selection helper ────────────────────────────────────────────
+#
+# Both ``Orchestrator._select_model`` and ``ResearchAgent._select_model_for_step``
+# need to pick a config-driven model name given a backend instance and a
+# heavy/light preference.  Keeping the mapping here avoids the two-implementation
+# drift that ``optimization.py`` inherited before (hardcoded strings that
+# fell out of sync with the ``*_HEAVY_MODEL`` / ``*_LIGHT_MODEL`` env vars).
+def select_model_for_backend(
+    backend: "ModelBackend",
+    config: Any,
+    *,
+    is_heavy: bool,
+) -> Optional[str]:
+    """Return the configured heavy/light model name for ``backend``.
+
+    Parameters
+    ----------
+    backend:
+        Live model backend instance (``OpenAIBackend`` / ``AWSOpenAIBackend`` /
+        etc.).  The concrete type picks which family of ``*_heavy_model`` /
+        ``*_light_model`` config fields to consult.
+    config:
+        A ``Config`` instance carrying the per-backend heavy/light fields.
+    is_heavy:
+        ``True`` → pick the heavy model, ``False`` → pick the light model.
+
+    Returns ``None`` for backends that only support a single model
+    (currently ``HuggingFaceBackend``); callers should fall back to the
+    backend's own default in that case.
+    """
+    if isinstance(backend, OpenAIBackend):
+        return config.openai_heavy_model if is_heavy else config.openai_light_model
+    if isinstance(backend, AzureOpenAIBackend):
+        return config.azure_heavy_model if is_heavy else config.azure_light_model
+    if isinstance(backend, AWSOpenAIBackend):
+        return config.aws_heavy_model if is_heavy else config.aws_light_model
+    if isinstance(backend, GCPVertexAIBackend):
+        return config.gcp_heavy_model if is_heavy else config.gcp_light_model
+    if isinstance(backend, OllamaBackend):
+        return config.ollama_heavy_model if is_heavy else config.ollama_light_model
+    if isinstance(backend, AnthropicBackend):
+        return (
+            config.anthropic_heavy_model if is_heavy else config.anthropic_light_model
+        )
+    if isinstance(backend, HuggingFaceBackend):
+        return None
+    return None

backend/models.py

@@ -123,9 +123,6 @@ class ConfigUpdate(BaseModel):
     analyst_max_tokens: Optional[int] = None
     qa_max_tokens: Optional[int] = None
 
-    # Legacy field kept for backward compat with existing callers
-    model: Optional[str] = None
-
 
 class SynthesisResult(BaseModel):
     summary: str
@@ -177,6 +174,19 @@ def to_dict(self) -> Dict[str, Any]:
             "parallel_group": self.parallel_group,
         }
 
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "ResearchStep":
+        """Reconstruct a ResearchStep from its ``to_dict`` output."""
+        return cls(
+            id=int(data["id"]),
+            name=data.get("name", ""),
+            description=data.get("description", ""),
+            status=StepStatus(data.get("status", StepStatus.PENDING.value)),
+            result=data.get("result") or None,
+            error=data.get("error") or None,
+            parallel_group=data.get("parallel_group"),
+        )
+
 
 class ResearchPlan:
     id: str
@@ -194,3 +204,21 @@ def to_dict(self) -> Dict[str, Any]:
             "goal": self.goal,
             "steps": [step.to_dict() for step in self.steps],
         }
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "ResearchPlan":
+        """Reconstruct a ResearchPlan (preserving ``id``) from its ``to_dict`` output.
+
+        Used by session recovery to rebuild ``Orchestrator._pending_plan`` from
+        a persisted checkpoint so ``approve_plan`` can proceed after a restart.
+        """
+        plan = cls(
+            goal=data.get("goal", ""),
+            steps=[ResearchStep.from_dict(s) for s in data.get("steps", [])],
+        )
+        # Preserve the original plan id so ``approve_plan``'s planId check
+        # against the client's stored id continues to match.
+        pid = data.get("id")
+        if pid:
+            plan.id = pid
+        return plan

backend/optimization.py

@@ -8,6 +8,7 @@
 from model_backend import Message
 from observability import get_logger
 from long_term_memory import AsyncLongTermMemory
+from config import Config
 
 logger = get_logger(__name__)
 
@@ -43,8 +44,9 @@ def __init__(
         model_backend,
         mcp_registry,
         long_term_memory: Optional[AsyncLongTermMemory] = None,
+        config: Optional[Config] = None,
     ):
-        super().__init__(model_backend, mcp_registry)
+        super().__init__(model_backend, mcp_registry, config=config)
         self._long_term_memory = long_term_memory
 
     # ── File I/O helpers ──────────────────────────────────────────────────────