Search every agent source in scheduled and webhook runs

Headless runs searched only the agent's primary source, so an agent whose
knowledge sat in its extra sources answered a schedule or webhook without
it. They now take the primary and every extra source through the same
owner-or-sponsor check a chat uses, shared as one helper, and retrieve
through the per-source dispatcher so each source keeps its own settings.

docsgpt/agents/headless_runner.py

@@ -15,15 +15,16 @@ from docsgpt.api.answer.services.prompt_renderer import (
     resolve_prompt_skeleton,
 )
 from docsgpt.api.answer.services.stream_processor import (
+    authorized_agent_sources,
     authorized_prompt_id,
     get_prompt,
 )
-from docsgpt.api.user.resource_access import ref_principal
 from docsgpt.core.settings import settings
 from docsgpt.quotas.service import QuotaExceededError, QuotaService
+from docsgpt.retriever.dispatcher import build_dispatcher
 from docsgpt.retriever.retriever_creator import RetrieverCreator
-from docsgpt.storage.db.repositories.sources import SourcesRepository
 from docsgpt.storage.db.session import db_readonly
+from docsgpt.storage.db.source_config import SourceConfig
 
 logger = logging.getLogger(__name__)
 
@@ -148,20 +149,22 @@ def _run_agent_headless(
         raise QuotaExceededError(exceeded)
 
     retriever_kind = agent_config.get("retriever", "classic")
-    source_id = agent_config.get("source_id") or agent_config.get("source")
-    source_active: Any = {}
-    if source_id:
+    # Every source a chat with this agent searches: the primary and the
+    # extras, each owned or team-shared to the owner, else attached by an
+    # editor who still qualifies.
+    sources_row = dict(agent_config)
+    sources_row["source_id"] = agent_config.get("source_id") or agent_config.get("source")
+    primary, source_docs = None, []
+    if sources_row["source_id"] or sources_row.get("extra_source_ids"):
         with db_readonly() as conn:
-            # Owned or team-shared to the owner, else attached by an editor
-            # who still qualifies; read unscoped once authorized.
-            src_row = (
-                SourcesRepository(conn).get_by_id(str(source_id))
-                if ref_principal(conn, "agent", agent_config, "source", str(source_id))
-                else None
-            )
-        if src_row:
-            source_active = str(src_row["id"])
-            retriever_kind = src_row.get("retriever", retriever_kind)
+            primary, source_docs = authorized_agent_sources(conn, sources_row)
+    if primary:
+        retriever_kind = primary.get("retriever") or retriever_kind
+    per_source = [
+        {"id": str(doc["id"]), "retrieval": SourceConfig.parse(doc.get("config")).retrieval}
+        for doc in source_docs
+    ]
+    source_active: Any = [entry["id"] for entry in per_source] or {}
     source = {"active_docs": source_active}
     # ``chunks=0`` switches retrieval off; only a missing value takes the default.
     raw_chunks = agent_config.get("chunks")
@@ -196,8 +199,7 @@ def _run_agent_headless(
     system_api_key = get_api_key_for_provider(provider or settings.LLM_PROVIDER)
     doc_token_limit = calculate_doc_token_budget(model_id=model_id, user_id=owner)
 
-    retriever = RetrieverCreator.create_retriever(
-        retriever_kind,
+    retriever_kwargs: Dict[str, Any] = dict(
         source=source,
         chat_history=chat_history or [],
         prompt=prompt,
@@ -208,6 +210,13 @@ def _run_agent_headless(
         agent_id=agent_id,
         decoded_token=decoded_token,
     )
+    # Routed per source like a chat's pre-fetch, so each source keeps its own
+    # retriever and retrieval settings.
+    retriever = build_dispatcher(
+        lambda: RetrieverCreator.create_retriever(retriever_kind, **retriever_kwargs),
+        sources=per_source,
+        **retriever_kwargs,
+    )
     retrieved_docs: List[Dict[str, Any]] = []
     try:
         docs = retriever.search(query)

docsgpt/api/answer/services/stream_processor.py

@@ -4,7 +4,7 @@ import json
 import logging
 import uuid
 from pathlib import Path
-from typing import Any, Callable, Dict, Optional, Set, TypeVar
+from typing import Any, Callable, Dict, List, Optional, Set, Tuple, TypeVar
 
 from flask import after_this_request
 
@@ -183,6 +183,39 @@ def _agent_source_doc(conn: Any, sources_repo: Any, agent: dict, source_id: Any)
     return sources_repo.get_by_id(str(source_id))
 
 
+def authorized_agent_sources(conn: Any, agent: dict) -> Tuple[Optional[dict], List[dict]]:
+    """The source rows an agent run retrieves from: primary first, then extras.
+
+    Each is authorized like :func:`_agent_source_doc` (the owner, else the
+    editor who attached it), and a source listed twice appears once.
+
+    Args:
+        conn: An open database connection.
+        agent: The ``agents`` row.
+
+    Returns:
+        The primary source row (None when unset or not usable) and every
+        usable row in run order.
+    """
+    sources_repo = SourcesRepository(conn)
+    primary: Optional[dict] = None
+    rows: List[dict] = []
+    seen: set = set()
+    refs = [(True, agent.get("source_id"))]
+    refs.extend((False, sid) for sid in agent.get("extra_source_ids") or [])
+    for is_primary, sid_raw in refs:
+        if not sid_raw:
+            continue
+        source_doc = _agent_source_doc(conn, sources_repo, agent, sid_raw)
+        if not source_doc or str(source_doc["id"]) in seen:
+            continue
+        if is_primary:
+            primary = source_doc
+        seen.add(str(source_doc["id"]))
+        rows.append(source_doc)
+    return primary, rows
+
+
 def _wiki_write_owner(conn: Any, source_id: str, caller: str) -> Optional[str]:
     """The owner id to write a wiki source as, when ``caller`` may edit it."""
     from docsgpt.api.user.resource_access import resolve
@@ -751,7 +784,6 @@ class StreamProcessor:
             agent = AgentsRepository(conn).find_by_key(api_key)
             if not agent:
                 raise Exception("Invalid API Key, please generate a new key", 401)
-            sources_repo = SourcesRepository(conn)
             # The repo dict uses "user_id" — the streaming path expects
             # a "user" key (legacy Mongo shape) for identity propagation.
             data: Dict[str, Any] = dict(agent)
@@ -761,68 +793,29 @@ class StreamProcessor:
             # ``_configure_source`` ignores an empty ``data["sources"]``,
             # so the primary must appear in the union too — not only in
             # the legacy ``data["source"]`` slot.
-            sources_list: list = []
-            seen: set = set()
-            primary_id = agent.get("source_id")
-            # ``sources`` row may have NULL ``retriever``/``chunks`` —
-            # fall back to the agent's value (``dict.get`` returns None
-            # even when the key exists with value None).
-            if primary_id:
-                source_doc = _agent_source_doc(conn, sources_repo, agent, primary_id)
-                if source_doc:
-                    sid = str(source_doc["id"])
-                    data["source"] = sid
-                    src_retriever = source_doc.get("retriever")
-                    if src_retriever:
-                        data["retriever"] = src_retriever
-                    src_chunks = source_doc.get("chunks")
-                    if src_chunks is not None:
-                        data["chunks"] = src_chunks
-                    sources_list.append(
-                        {
-                            "id": sid,
-                            "retriever": src_retriever or "classic",
-                            "chunks": (
-                                src_chunks if src_chunks is not None
-                                else data.get("chunks", "6")
-                            ),
-                            # Per-source behaviour contract (lenient read).
-                            "retrieval": SourceConfig.parse(
-                                source_doc.get("config")
-                            ).retrieval,
-                        }
-                    )
-                    seen.add(sid)
-                else:
-                    data["source"] = None
-            else:
-                data["source"] = None
-
-            for sid_raw in agent.get("extra_source_ids") or []:
-                if not sid_raw:
-                    continue
-                source_doc = _agent_source_doc(conn, sources_repo, agent, sid_raw)
-                if not source_doc:
-                    continue
-                sid = str(source_doc["id"])
-                if sid in seen:
-                    continue
-                src_retriever = source_doc.get("retriever")
-                src_chunks = source_doc.get("chunks")
-                sources_list.append(
-                    {
-                        "id": sid,
-                        "retriever": src_retriever or "classic",
-                        "chunks": (
-                            src_chunks if src_chunks is not None
-                            else data.get("chunks", "6")
-                        ),
-                        "retrieval": SourceConfig.parse(
-                            source_doc.get("config")
-                        ).retrieval,
-                    }
-                )
-                seen.add(sid)
+            primary, source_docs = authorized_agent_sources(conn, agent)
+        # ``sources`` row may have NULL ``retriever``/``chunks`` — fall back to
+        # the agent's value (``dict.get`` returns None even when the key
+        # exists with value None). The primary's own values win for the agent.
+        data["source"] = str(primary["id"]) if primary else None
+        if primary:
+            if primary.get("retriever"):
+                data["retriever"] = primary["retriever"]
+            if primary.get("chunks") is not None:
+                data["chunks"] = primary["chunks"]
+        sources_list: list = [
+            {
+                "id": str(source_doc["id"]),
+                "retriever": source_doc.get("retriever") or "classic",
+                "chunks": (
+                    source_doc["chunks"] if source_doc.get("chunks") is not None
+                    else data.get("chunks", "6")
+                ),
+                # Per-source behaviour contract (lenient read).
+                "retrieval": SourceConfig.parse(source_doc.get("config")).retrieval,
+            }
+            for source_doc in source_docs
+        ]
         data["sources"] = sources_list
         data["default_model_id"] = data.get("default_model_id", "")
         return data

tests/api/user/test_resource_sponsors.py

@@ -39,6 +39,7 @@ from docsgpt.storage.db.repositories.workflows import WorkflowsRepository
 OWNER, EDITOR, VIEWER, OTHER = "sp-owner", "sp-editor", "sp-viewer", "sp-other"
 
 _DB_MODULES = (
+    "docsgpt.agents.headless_runner",
     "docsgpt.api.user.agents.routes",
     "docsgpt.api.user.workflows.routes",
     "docsgpt.api.user.base",
@@ -376,6 +377,40 @@ class TestRunTime:
             data = processor._get_data_from_api_key(agent["key"])
         assert [s["id"] for s in data["sources"]] == [source, shared]
 
+    def test_headless_run_retrieves_from_every_agent_source(self, app, pg_conn, monkeypatch):
+        """A scheduled or webhook run searches the same sources a chat does."""
+        from unittest.mock import MagicMock
+
+        from docsgpt.agents import headless_runner as hr
+
+        agent, team_id, _, _, source = self._sponsored_agent(app, pg_conn)
+        shared = str(SourcesRepository(pg_conn).create("shared-src", user_id=OTHER)["id"])
+        TeamMembersRepository(pg_conn).add_member(team_id, OWNER)
+        TeamResourceGrantsRepository(pg_conn).grant(team_id, "source", shared, OTHER, OTHER)
+        foreign = str(SourcesRepository(pg_conn).create("foreign-src", user_id="sp-stranger")["id"])
+        AgentsRepository(pg_conn).update_by_id(str(agent["id"]), {"extra_source_ids": [shared, foreign]})
+
+        searched = []
+        retriever = MagicMock(search=MagicMock(return_value=[]))
+
+        def _create_retriever(cls, *_args, **kwargs):
+            searched.append(kwargs["source"].get("active_docs"))
+            return retriever
+
+        run_agent = MagicMock(gen=MagicMock(return_value=iter([{"answer": "ok"}])))
+        run_agent.llm.token_usage = {}
+        monkeypatch.setattr(hr.RetrieverCreator, "create_retriever", classmethod(_create_retriever))
+        monkeypatch.setattr(hr, "ToolExecutor", lambda *a, **kw: MagicMock(headless_denials=[]))
+        monkeypatch.setattr(hr.AgentCreator, "create_agent", classmethod(lambda cls, *a, **kw: run_agent))
+        monkeypatch.setattr(hr.QuotaService, "check", lambda *a, **kw: None)
+        with _patch_db(pg_conn), \
+                patch("docsgpt.core.model_utils.validate_model_id", return_value=True), \
+                patch("docsgpt.core.model_utils.get_provider_from_model_id", return_value="openai"), \
+                patch("docsgpt.core.model_utils.get_api_key_for_provider", return_value="k"), \
+                patch("docsgpt.utils.calculate_doc_token_budget", return_value=1000):
+            hr.run_agent_headless(_row(pg_conn, str(agent["id"])), "q")
+        assert searched == [[source, shared]]
+
     def test_search_service_authorizes_sponsored_source(self, app, pg_conn):
         from docsgpt.services.search_service import _authorized_source_ids
 

tests/tracing/test_entry_points.py

@@ -82,7 +82,8 @@ class TestHeadless:
         assert trace.request_id == "run-1"
         assert trace.user_id == "owner-1"
         assert trace.status == "ok"
-        assert [s.kind for s in trace.spans] == ["retrieval", "agent"]
+        # The per-source dispatcher's span wraps its retriever's own.
+        assert [s.kind for s in trace.spans] == ["retrieval", "retrieval", "agent"]
 
     def test_trace_can_belong_to_the_scheduling_user(self, monkeypatch, flushed):
         """A schedule on a shared agent is the scheduler's run, not the owner's."""

tests/worker/test_agent_workers.py

@@ -375,4 +375,4 @@ class TestRunAgentHeadlessFromWebhook:
         outcome = headless_runner.run_agent_headless(agent_config, "hello")
 
         assert outcome["answer"] == "done"
-        assert captured_source.get("active_docs") == source_id
+        assert captured_source.get("active_docs") == [source_id]