Trace scheduled, webhook, search, MCP and graph-extraction runs

run_agent_headless records each unattended run under its endpoint; the
scheduler passes its run id and the webhook worker its task id so Logs rows
can find their trace, while the LLM's own request id stays untouched for
quota counts. /api/search and MCP search_docs record their retrieval, and a
graph build records every extraction call under one step.

docsgpt/agents/headless_runner.py

@@ -5,6 +5,7 @@ from __future__ import annotations
 import logging
 from typing import Any, Dict, Iterable, List, Optional
 
+from docsgpt import tracing
 from docsgpt.agents.agent_creator import AgentCreator
 from docsgpt.agents.tool_executor import ToolExecutor
 from docsgpt.api.answer.services.prompt_renderer import (
@@ -69,12 +70,57 @@ def run_agent_headless(
     endpoint: str = "headless",
     chat_history: Optional[List[Dict[str, Any]]] = None,
     conversation_id: Optional[str] = None,
+    request_id: Optional[str] = None,
 ) -> Dict[str, Any]:
     """Run an agent with no live client; returns a structured outcome dict.
 
+    The run is recorded as one execution trace under ``endpoint`` as its
+    source. ``request_id`` links that trace to the caller's own record (the
+    scheduler passes its run id, the webhook worker its task id); it is kept
+    off the LLM's token-usage rows, whose request ids drive request counts.
+
     Raises:
         QuotaExceededError: If the agent owner's usage quota is exhausted.
     """
+    trace = tracing.start_trace(
+        source=endpoint,
+        request_id=request_id,
+        user_id=_resolve_owner(agent_config),
+        agent_id=_resolve_agent_id(agent_config),
+        conversation_id=conversation_id,
+    )
+    status = None
+    with tracing.activate(trace):
+        try:
+            outcome = _run_agent_headless(
+                agent_config,
+                query,
+                tool_allowlist=tool_allowlist,
+                model_id_override=model_id_override,
+                endpoint=endpoint,
+                chat_history=chat_history,
+                conversation_id=conversation_id,
+            )
+            if outcome.get("error"):
+                status = tracing.STATUS_ERROR
+            return outcome
+        except BaseException:
+            status = tracing.STATUS_ERROR
+            raise
+        finally:
+            tracing.flush(trace, status)
+
+
+def _run_agent_headless(
+    agent_config: Dict[str, Any],
+    query: str,
+    *,
+    tool_allowlist: Optional[Iterable[str]] = None,
+    model_id_override: Optional[str] = None,
+    endpoint: str = "headless",
+    chat_history: Optional[List[Dict[str, Any]]] = None,
+    conversation_id: Optional[str] = None,
+) -> Dict[str, Any]:
     from docsgpt.core.model_utils import (
         get_api_key_for_provider,
         get_default_model_id,

docsgpt/api/user/scheduler_worker.py

@@ -277,6 +277,8 @@ def execute_scheduled_run_body(run_id: str, celery_task_id: Optional[str]) -> Di
             endpoint="schedule",
             conversation_id=schedule.get("origin_conversation_id"),
             chat_history=chat_history,
+            # Links the run's execution trace to its Logs row.
+            request_id=str(run_id),
         )
     except SoftTimeLimitExceeded:
         timed_out = True

docsgpt/graphrag/extraction.py

@@ -23,6 +23,7 @@ import logging
 import re
 from typing import Any, Callable, Dict, List, Optional
 
+from docsgpt import tracing
 from docsgpt.core.model_utils import (
     get_api_key_for_provider,
     get_provider_from_model_id,
@@ -353,7 +354,9 @@ def extract_graph_for_source(
         if pool is not None:
             # ``map`` yields in submission order, so chunks are still applied in
             # the order they were given and a run stays reproducible.
-            prepared = pool.map(_prepare, items)
+            # Pool threads don't inherit context; carry the trace in so each
+            # chunk's extraction LLM call is recorded.
+            prepared = pool.map(tracing.wrap(_prepare), items)
         else:
             prepared = (_prepare(item) for item in items)
         missed = []

docsgpt/mcp_server.py

@@ -53,7 +53,7 @@ async def search_docs(query: str, chunks: int = 5) -> list[dict]:
     if not api_key:
         raise PermissionError("Missing Bearer token")
     try:
-        return await asyncio.to_thread(search, api_key, query, chunks)
+        return await asyncio.to_thread(search, api_key, query, chunks, source="mcp")
     except InvalidAPIKey as exc:
         raise PermissionError("Invalid API key") from exc
     except SearchFailed:

docsgpt/services/search_service.py

@@ -10,10 +10,12 @@ from __future__ import annotations
 import logging
 from typing import Any, Dict, List, Optional
 
+from docsgpt import tracing
 from docsgpt.core.settings import settings
 from docsgpt.retriever.fanout import fetch_per_source
 from docsgpt.storage.db.repositories.agents import AgentsRepository
 from docsgpt.storage.db.session import db_readonly
+from docsgpt.tracing.retrieval import describe_documents, start_retrieval_span
 from docsgpt.vectorstore.vector_creator import VectorCreator
 
 logger = logging.getLogger(__name__)
@@ -219,14 +221,21 @@ def _search_sources(
     return results[:chunks]
 
 
-def search(api_key: str, query: str, chunks: int = 5) -> List[Dict[str, Any]]:
+def search(
+    api_key: str, query: str, chunks: int = 5, *, source: str = "search"
+) -> List[Dict[str, Any]]:
     """Resolve an agent by API key and search its sources.
 
+    Every search that reaches the sources is recorded as an execution trace
+    owned by the agent's owner, under ``source``.
+
     Args:
         api_key: Agent API key (the opaque string stored on
             ``agents.key`` in Postgres).
         query: Free-text search query.
         chunks: Max number of hits to return.
+        source: Trace source name: ``search`` for ``/api/search``, ``mcp``
+            for the MCP ``search_docs`` tool.
 
     Returns:
         List of hit dicts with ``text``, ``title``, ``source`` keys.
@@ -256,4 +265,20 @@ def search(api_key: str, query: str, chunks: int = 5) -> List[Dict[str, Any]]:
     if not source_ids:
         return []
 
-    return _search_sources(query, source_ids, chunks)
+    trace = tracing.start_trace(
+        source=source,
+        user_id=agent.get("user_id"),
+        agent_id=str(agent.get("id")) if agent.get("id") else None,
+    )
+    with tracing.activate(trace):
+        try:
+            with start_retrieval_span(
+                f"retrieval {source}",
+                sources=source_ids,
+                **{"docsgpt.top_k": chunks},
+            ) as span:
+                results = _search_sources(query, source_ids, chunks)
+                describe_documents(span, results, query=query)
+            return results
+        finally:
+            tracing.flush(trace)

docsgpt/worker.py

@@ -16,6 +16,7 @@ from urllib.parse import urljoin, urlsplit
 
 import requests
 
+from docsgpt import tracing
 from docsgpt.core.settings import settings
 from docsgpt.events.publisher import publish_user_event
 from docsgpt.parser.chunking_creator import ChunkerCreator
@@ -2129,6 +2130,7 @@ def agent_webhook_worker(self, agent_id, payload):
             input_data,
             tool_allowlist=_webhook_tool_allowlist(agent_config),
             endpoint="webhook",
+            request_id=getattr(getattr(self, "request", None), "id", None),
         )
         result = {
             "answer": outcome.get("answer", ""),
@@ -2942,16 +2944,36 @@ def extract_graph_worker(self, source_id, user):
         },
     )
 
+    trace = tracing.start_trace(
+        source="graph_extraction",
+        name=f"graph_extraction {source.get('name') or source_id}",
+        request_id=getattr(self.request, "id", None),
+        user_id=user,
+    )
     try:
-        summary = extract_graph_for_source(
-            source_id,
-            user,
-            chunks,
-            config=cfg,
-            request_id=getattr(self.request, "id", None),
-            progress_cb=_progress,
-        )
+        with tracing.activate(trace), tracing.span(
+            tracing.KIND_STEP,
+            "graph_extraction",
+            attributes={"docsgpt.source_id": source_id, "docsgpt.chunk_count": total},
+        ) as span:
+            summary = extract_graph_for_source(
+                source_id,
+                user,
+                chunks,
+                config=cfg,
+                request_id=getattr(self.request, "id", None),
+                progress_cb=_progress,
+            )
+            if isinstance(summary, dict):
+                span.set(
+                    **{
+                        "docsgpt.graph.nodes": summary.get("nodes"),
+                        "docsgpt.graph.edges": summary.get("edges"),
+                        "docsgpt.graph.chunks_processed": summary.get("chunks_processed"),
+                    }
+                )
     except Exception as e:
+        tracing.flush(trace, tracing.STATUS_ERROR)
         _publish_graph_event(
             user,
             source_id,
@@ -2960,6 +2982,7 @@ def extract_graph_worker(self, source_id, user):
         )
         raise
 
+    tracing.flush(trace)
     _publish_graph_event(
         user,
         source_id,