Do not duplicate ResearchContext in ResearchState

src/haiku/rag/app.py

@@ -222,9 +222,9 @@ class HaikuRAGApp:
                     self.console.print()
 
                 graph = build_research_graph()
+                context = ResearchContext(original_question=question)
                 state = ResearchState(
-                    question=question,
+                    context=context,
-                    context=ResearchContext(original_question=question),
                     max_iterations=max_iterations,
                     confidence_threshold=confidence_threshold,
                     max_concurrency=max_concurrency,

src/haiku/rag/research/nodes/evaluate.py

@@ -55,8 +55,8 @@ class EvaluateNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]):
         for insight in output.key_insights:
             state.context.add_insight(insight)
         for new_q in output.new_questions:
-            if new_q not in state.sub_questions:
+            if new_q not in state.context.sub_questions:
-                state.sub_questions.append(new_q)
+                state.context.sub_questions.append(new_q)
 
         state.last_eval = output
         state.iterations += 1

src/haiku/rag/research/nodes/plan.py

@@ -45,7 +45,7 @@ class PlanNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]):
 
         prompt = (
             "Plan a focused research approach for the main question.\n\n"
-            f"Main question: {state.question}"
+            f"Main question: {state.context.original_question}"
         )
 
         agent_deps = ResearchDependencies(
@@ -55,12 +55,16 @@ class PlanNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]):
             stream=deps.stream,
         )
         plan_result = await plan_agent.run(prompt, deps=agent_deps)
-        state.sub_questions = list(plan_result.output.sub_questions)
+        state.context.sub_questions = list(plan_result.output.sub_questions)
 
         log(deps, state, "\n[bold green]✅ Research Plan Created:[/bold green]")
-        log(deps, state, f"   [bold]Main Question:[/bold] {state.question}")
+        log(
+            deps,
+            state,
+            f"   [bold]Main Question:[/bold] {state.context.original_question}",
+        )
         log(deps, state, "   [bold]Sub-questions:[/bold]")
-        for i, sq in enumerate(state.sub_questions, 1):
+        for i, sq in enumerate(state.context.sub_questions, 1):
             log(deps, state, f"      {i}. {sq}")
 
         return SearchDispatchNode(self.provider, self.model)

src/haiku/rag/research/nodes/search.py

@@ -24,7 +24,7 @@ class SearchDispatchNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]):
     ) -> BaseNode[ResearchState, ResearchDeps, ResearchReport]:
         state = ctx.state
         deps = ctx.deps
-        if not state.sub_questions:
+        if not state.context.sub_questions:
             from haiku.rag.research.nodes.evaluate import EvaluateNode
 
             return EvaluateNode(self.provider, self.model)
@@ -32,8 +32,8 @@ class SearchDispatchNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]):
         # Take up to max_concurrency questions and answer them concurrently
         take = max(1, state.max_concurrency)
         batch: list[str] = []
-        while state.sub_questions and len(batch) < take:
+        while state.context.sub_questions and len(batch) < take:
-            batch.append(state.sub_questions.pop(0))
+            batch.append(state.context.sub_questions.pop(0))
 
         async def answer_one(sub_q: str) -> SearchAnswer | None:
             log(

src/haiku/rag/research/state.py

@@ -1,4 +1,4 @@
-from dataclasses import dataclass, field
+from dataclasses import dataclass
 
 from rich.console import Console
 
@@ -23,9 +23,7 @@ class ResearchDeps:
 
 @dataclass
 class ResearchState:
-    question: str
     context: ResearchContext
-    sub_questions: list[str] = field(default_factory=list)
     iterations: int = 0
     max_iterations: int = 3
     max_concurrency: int = 1

src/haiku/rag/research/stream.py

@@ -34,13 +34,13 @@ class ResearchStateSnapshot:
             last_sufficient = state.last_eval.is_sufficient
 
         return cls(
-            question=state.question,
+            question=context.original_question,
-            sub_questions=list(state.sub_questions),
+            sub_questions=list(context.sub_questions),
             iterations=state.iterations,
             max_iterations=state.max_iterations,
             max_concurrency=state.max_concurrency,
             confidence_threshold=state.confidence_threshold,
-            pending_sub_questions=len(state.sub_questions),
+            pending_sub_questions=len(context.sub_questions),
             answered_questions=len(context.qa_responses),
             insights=list(context.insights),
             gaps=list(context.gaps),

tests/test_research_graph.py

@@ -9,7 +9,6 @@ def test_build_graph_and_state():
     assert graph is not None
 
     state = ResearchState(
-        question="What are the key features of haiku.rag?",
         context=ResearchContext(
             original_question="What are the key features of haiku.rag?"
         ),
@@ -17,7 +16,7 @@ def test_build_graph_and_state():
         confidence_threshold=0.8,
     )
     assert state.iterations == 0
-    assert state.sub_questions == []
+    assert state.context.sub_questions == []
 
 
 def test_async_loop_available():

tests/test_research_graph_integration.py

@@ -21,7 +21,6 @@ async def test_graph_end_to_end_with_patched_nodes(monkeypatch):
     graph = build_research_graph()
 
     state = ResearchState(
-        question="What is haiku.rag?",
         context=ResearchContext(original_question="What is haiku.rag?"),
         max_iterations=1,
         confidence_threshold=0.5,
@@ -32,7 +31,7 @@ async def test_graph_end_to_end_with_patched_nodes(monkeypatch):
     )  # client unused in patched nodes
 
     async def fake_plan_run(self, ctx) -> Any:
-        ctx.state.sub_questions = [
+        ctx.state.context.sub_questions = [
             "Describe haiku.rag in one sentence",
             "List core components of haiku.rag",
         ]
@@ -41,8 +40,8 @@ async def test_graph_end_to_end_with_patched_nodes(monkeypatch):
 
     async def fake_search_dispatch_run(self, ctx) -> Any:
         # Answer all pending questions deterministically, then move to evaluation
-        while ctx.state.sub_questions:
+        while ctx.state.context.sub_questions:
-            q = ctx.state.sub_questions.pop(0)
+            q = ctx.state.context.sub_questions.pop(0)
             # pydantic BaseModel kwargs not fully typed for pyright
             ctx.state.context.add_qa_response(
                 SearchAnswer(query=q, answer="A", context=["x"], sources=["s"])  # pyright: ignore[reportCallIssue]