Skip gather_context when session context exists

haiku_rag_slim/haiku/rag/agents/research/graph.py

@@ -83,7 +83,7 @@ async def _plan_step_logic(
     has_session_context = bool(state.context.session_context)
     effective_plan_prompt = (
         build_prompt(PLAN_PROMPT_WITH_CONTEXT, config)
-        if has_prior_answers
+        if has_prior_answers or has_session_context
         else plan_prompt
     )
 
@@ -98,17 +98,20 @@ async def _plan_step_logic(
 
     search_filter = state.search_filter
 
-    @plan_agent.tool
+    # Only register gather_context tool when we don't have existing context
-    async def gather_context(
+    if not has_prior_answers and not has_session_context:
-        ctx2: RunContext[ResearchDependencies],
+
-        query: str,
+        @plan_agent.tool
-        limit: int | None = None,
+        async def gather_context(
-    ) -> str:
+            ctx2: RunContext[ResearchDependencies],
-        results = await ctx2.deps.client.search(
+            query: str,
-            query, limit=limit, filter=search_filter
+            limit: int | None = None,
-        )
+        ) -> str:
-        results = await ctx2.deps.client.expand_context(results)
+            results = await ctx2.deps.client.search(
-        return "\n\n".join(r.content for r in results)
+                query, limit=limit, filter=search_filter
+            )
+            results = await ctx2.deps.client.expand_context(results)
+            return "\n\n".join(r.content for r in results)
 
     # Build prompt with existing context if available
     if has_prior_answers:

haiku_rag_slim/haiku/rag/agents/research/prompts.py

@@ -27,23 +27,21 @@ You have access to context that may include:
 - <background>: Domain context for the conversation
 - <prior_answers>: Previous Q&A pairs with confidence scores
 
-Review this first - if prior answers already answer the question completely,
+Review the provided context first. Use <background> to understand the domain.
-you may return an empty sub_questions list. Only create sub-questions to
+If <prior_answers> exist and already answer the question completely,
-fill genuine gaps.
+return an empty sub_questions list. Only create sub-questions to fill gaps.
 
 Responsibilities:
-1. Review prior_answers to understand what's already known
+1. Review provided context to understand what's already known
 2. Identify gaps that need additional research
 3. Propose minimal sub-questions only for missing information
 
 Plan requirements:
-- If prior answers fully answer the question, return an empty sub_questions list.
+- If existing context fully answers the question, return an empty sub_questions list.
 - Only create new sub-questions for genuine gaps in existing knowledge.
 - sub_questions must be a list of plain strings (max 3).
 - Each sub_question must be standalone and self-contained.
-- Prioritize the highest-value gaps first.
+- Prioritize the highest-value gaps first."""
-
-Use the gather_context tool once on the main question before planning."""
 
 SEARCH_PROMPT = """You are a search and question-answering specialist.
 

tests/agents/research/test_plan_prompt_selection.py

@@ -0,0 +1,35 @@
+from haiku.rag.agents.research.prompts import PLAN_PROMPT, PLAN_PROMPT_WITH_CONTEXT
+
+
+def test_plan_prompt_with_context_does_not_instruct_gather_context():
+    """PLAN_PROMPT_WITH_CONTEXT should not instruct to use gather_context.
+
+    When session context already exists, we don't need to gather context again.
+    """
+    assert "gather_context" not in PLAN_PROMPT_WITH_CONTEXT
+
+
+def test_plan_prompt_instructs_gather_context():
+    """PLAN_PROMPT should instruct to use gather_context for initial planning."""
+    assert "gather_context" in PLAN_PROMPT
+
+
+def test_prompt_selection_uses_context_prompt_with_session_context():
+    """When session_context exists, should use PLAN_PROMPT_WITH_CONTEXT.
+
+    This tests the logic pattern used in _plan_step_logic.
+    """
+    # Simulate the selection logic from graph.py
+    has_prior_answers = False
+    has_session_context = True
+
+    # Current buggy behavior would select plan_prompt (the one with gather_context)
+    # Expected behavior: use PLAN_PROMPT_WITH_CONTEXT when session_context exists
+    effective_plan_prompt = (
+        PLAN_PROMPT_WITH_CONTEXT
+        if has_prior_answers or has_session_context
+        else PLAN_PROMPT
+    )
+
+    # Since session_context exists, we should use PLAN_PROMPT_WITH_CONTEXT
+    assert effective_plan_prompt == PLAN_PROMPT_WITH_CONTEXT