learningcircuit--local-deep-research/tests/strategies/test_langgraph_subagent_context.py

"""The LangGraph ``research_subtopic`` subagent pool must propagate the lead
thread's search context (which carries the user's DB password) into its stdlib
``ThreadPoolExecutor`` workers.

Without it, a subagent that re-creates an engine / registers the user's
document collections runs ``get_user_db_session`` with no password and fails to
open the per-user ENCRYPTED database — surfacing as
"Unknown search engine 'collection_…'". Sibling strategies close this gap with
``@preserve_research_context``; this strategy captures the context on the lead
thread and re-sets it per pool task.

The test drives the REAL ``_make_research_subtopic_tool`` pool wiring; only the
leaf subagent internals (``_make_web_search_tool``, ``build_fetch_tool``,
``create_agent``) are stubbed, so the assertion is purely about what context a
worker thread observes.
"""

from types import SimpleNamespace
from unittest.mock import MagicMock, patch

import pytest

import local_deep_research.advanced_search_system.strategies.langgraph_agent_strategy as strat  # noqa: E501
from local_deep_research.utilities.thread_context import (
    clear_search_context,
    get_search_context,
    set_search_context,
)


@pytest.fixture(autouse=True)
def _clean_search_context():
    # Never let a context leak into or out of these tests.
    clear_search_context()
    yield
    clear_search_context()


def _invoke_one_subtopic(seen):
    """Build the real research_subtopic tool with stubbed subagent internals,
    then invoke it once. ``seen`` collects the search context observed inside
    the pool worker (via the stubbed web-search tool factory)."""

    def fake_web_search_tool(*_a, **_k):
        # Runs INSIDE the pool worker — record the context it can see.
        seen.append(get_search_context())
        return SimpleNamespace(name="web_search")

    agent = MagicMock()
    agent.invoke.return_value = {"messages": [SimpleNamespace(content="ok")]}

    with (
        patch.object(strat, "_make_web_search_tool", fake_web_search_tool),
        patch.object(strat, "build_fetch_tool", lambda *_a, **_k: None),
        patch("langchain.agents.create_agent", return_value=agent),
    ):
        tool = strat._make_research_subtopic_tool(
            search_engine_name="collection_abc",
            model=MagicMock(),
            settings_snapshot={"search.tool": "collection_abc"},
            collector=MagicMock(),
            max_sub_iterations=1,
        )
        tool.invoke({"subtopics": ["what is X?"]})


class TestSubagentSearchContextPropagation:
    def test_password_context_reaches_pool_worker(self):
        # Lead thread has the user's DB password in the search context...
        set_search_context({"user_password": "secret", "research_id": "r1"})
        seen = []
        _invoke_one_subtopic(seen)
        # ...and the pool worker sees it (propagated), not None. This fails
        # against the pre-fix code where the ContextVar was never carried in.
        assert seen, "subagent web_search tool was never built"
        assert seen[0] is not None
        assert seen[0].get("user_password") == "secret"

    def test_no_context_does_not_crash(self):
        # No lead-thread context (programmatic/no-encryption): the pool worker
        # simply sees None and the run proceeds without error.
        seen = []
        _invoke_one_subtopic(seen)
        assert seen and seen[0] is None