Trace every chat turn from agent setup to the last event

StreamProcessor starts the trace and mints the request id before the
agent is built, so pre-fetch retrieval and compression are inside it and
side-channel LLM calls share the id. complete_stream activates it in the
SSE pump thread, binds the message and conversation, and writes it once
however the stream ends: paused, failed, abandoned or superseded (dropped).
user_logs rows now carry request_id and message_id.

docsgpt/api/answer/routes/answer.py

@@ -87,7 +87,7 @@ class AnswerResource(Resource, BaseAnswerResource):
         if error := self.validate_request(data):
             return error
         decoded_token = getattr(request, "decoded_token", None)
-        processor = StreamProcessor(data, decoded_token)
+        processor = StreamProcessor(data, decoded_token, trace_source="answer")
         try:
             # ---- Continuation mode ----
             if data.get("tool_actions"):
@@ -115,6 +115,7 @@ class AnswerResource(Resource, BaseAnswerResource):
                     decoded_token=processor.decoded_token,
                     agent_id=processor.agent_id,
                     model_id=processor.model_id,
+                    trace=processor.trace,
                     _continuation={
                         "messages": messages,
                         "tools_dict": tools_dict,
@@ -156,6 +157,8 @@ class AnswerResource(Resource, BaseAnswerResource):
                     is_shared_usage=processor.is_shared_usage,
                     shared_token=processor.shared_token,
                     model_id=processor.model_id,
+                    request_id=processor.request_id,
+                    trace=processor.trace,
                 )
 
             stream_result = self.process_response_stream(stream)

docsgpt/api/answer/routes/base.py

@@ -1,4 +1,5 @@
 import datetime
+import inspect
 import json
 import logging
 import threading
@@ -9,6 +10,7 @@ from typing import Any, Dict, Generator, List, Optional
 from flask import jsonify, make_response, Response
 from flask_restx import Namespace
 
+from docsgpt import tracing
 from docsgpt.api.answer.services.continuation_service import ContinuationService
 from docsgpt.api.answer.services.conversation_service import (
     ConversationService,
@@ -245,6 +247,59 @@ class BaseAnswerResource:
         return error
 
     def complete_stream(
+        self,
+        *args: Any,
+        request_id: Optional[str] = None,
+        trace: Optional["tracing.Trace"] = None,
+        **kwargs: Any,
+    ) -> Generator[str, None, None]:
+        """Stream the conversation response and record its execution trace.
+
+        Accepts every argument of :meth:`_complete_stream_body`. The trace is
+        activated for the whole stream -- which runs in the SSE pump thread,
+        not the request thread -- and written exactly once when the stream
+        ends, however it ends.
+
+        Args:
+            request_id: The request's id, generated by the route before the
+                agent was built so pre-fetch retrieval shares it. A
+                continuation's saved id takes precedence; absent both, a new
+                one is minted.
+            trace: The trace the route started (it already holds the
+                pre-fetch retrieval). A new one is started when omitted.
+
+        Yields:
+            Server-sent event strings.
+        """
+        bound = inspect.signature(self._complete_stream_body).bind(*args, **kwargs)
+        continuation = bound.arguments.get("_continuation")
+        request_id = (
+            continuation.get("request_id") if continuation else None
+        ) or request_id or str(uuid.uuid4())
+        decoded_token = bound.arguments.get("decoded_token") or {}
+        if not isinstance(trace, tracing.Trace):
+            trace = None
+        if trace is None:
+            agent = bound.arguments.get("agent")
+            trace = tracing.start_trace(
+                source=str(getattr(agent, "endpoint", None) or "stream"),
+                capture_otel_context=False,
+            )
+        if trace is not None:
+            trace.bind(
+                request_id=request_id,
+                user_id=decoded_token.get("sub") if isinstance(decoded_token, dict) else None,
+                agent_id=bound.arguments.get("agent_id"),
+            )
+        with tracing.activate(trace):
+            try:
+                yield from self._complete_stream_body(
+                    *bound.args, request_id=request_id, **bound.kwargs
+                )
+            finally:
+                tracing.flush(trace)
+
+    def _complete_stream_body(
         self,
         question: str,
         agent: Any,
@@ -263,6 +318,7 @@ class BaseAnswerResource:
         model_user_id: Optional[str] = None,
         _continuation: Optional[Dict] = None,
         finalize_tool_pause_as_complete: bool = False,
+        request_id: Optional[str] = None,
     ) -> Generator[str, None, None]:
         """
         Generator function that streams the complete conversation response.
@@ -336,10 +392,10 @@ class BaseAnswerResource:
         pending_pause_event: Optional[dict] = None
 
         # One id shared across the WAL row, primary LLM (token_usage
-        # attribution), the SSE event, and resumed continuations.
+        # attribution), the SSE event, resumed continuations and the trace.
         request_id = (
             _continuation.get("request_id") if _continuation else None
-        ) or str(uuid.uuid4())
+        ) or request_id or str(uuid.uuid4())
 
         # Reserve the placeholder row before the LLM call so a crash
         # mid-stream still leaves the question queryable. Continuations
@@ -405,6 +461,17 @@ class BaseAnswerResource:
         primary_llm = getattr(agent, "llm", None)
         if primary_llm is not None:
             primary_llm._request_id = request_id
+        # Side-channel LLMs built later (guardrail judge, retrievers the agent
+        # creates for its search tools) read the id off the agent.
+        if getattr(agent, "request_id", None) is None:
+            try:
+                agent.request_id = request_id
+            except Exception:
+                logger.debug("Could not stamp request_id on the agent")
+        tracing.bind(
+            message_id=reserved_message_id,
+            conversation_id=str(conversation_id) if conversation_id else None,
+        )
 
         # Flipped to ``streaming`` on the first ``answer``/``sources`` chunk;
         # the reconciler reads ``status`` to tell "never started" from "in
@@ -795,6 +862,9 @@ class BaseAnswerResource:
             # ---- Paused: save continuation state and end stream early ----
             if paused:
                 continuation = getattr(agent, "_pending_continuation", None)
+                trace = tracing.current_trace()
+                if trace is not None:
+                    trace.outcome = tracing.STATUS_PAUSED
 
                 # ---- Stateless-tool-round mode (OpenAI-compatible /v1) ----
                 # OpenAI clients resume by re-POSTing the whole message
@@ -1232,6 +1302,9 @@ class BaseAnswerResource:
                 "tool_calls": tool_calls_for_logging,
                 "attachments": attachment_ids,
                 "timestamp": datetime.datetime.now(datetime.timezone.utc),
+                # Links the Logs row to this turn's execution trace.
+                "request_id": request_id,
+                "message_id": reserved_message_id,
             }
             if is_structured:
                 log_data["structured_output"] = True
@@ -1469,9 +1542,14 @@ class BaseAnswerResource:
             )
             if journal_writer is not None:
                 journal_writer.close()
+            # The user replaced this turn; its trace describes nothing kept.
+            tracing.discard(tracing.current_trace())
             return
         except Exception as e:
             logger.error(f"Error in stream: {str(e)}", exc_info=True)
+            trace = tracing.current_trace()
+            if trace is not None:
+                trace.outcome = tracing.STATUS_ERROR
             # This process took the resume claim, so it owns releasing it. The
             # only other way back is ``revert_stale_resuming``'s 600 s grace,
             # which leaves the user locked out of their own conversation for

docsgpt/api/answer/routes/stream.py

@@ -94,7 +94,7 @@ class StreamResource(Resource, BaseAnswerResource):
         if error := self.validate_request(data, "index" in data):
             return error
         decoded_token = getattr(request, "decoded_token", None)
-        processor = StreamProcessor(data, decoded_token)
+        processor = StreamProcessor(data, decoded_token, trace_source="stream")
 
         try:
             # ---- Continuation mode ----
@@ -130,6 +130,7 @@ class StreamResource(Resource, BaseAnswerResource):
                             agent_id=processor.agent_id,
                             model_id=processor.model_id,
                             model_user_id=processor.model_user_id,
+                            trace=processor.trace,
                             _continuation={
                                 "messages": messages,
                                 "tools_dict": tools_dict,
@@ -181,6 +182,8 @@ class StreamResource(Resource, BaseAnswerResource):
                         shared_token=processor.shared_token,
                         model_id=processor.model_id,
                         model_user_id=processor.model_user_id,
+                        request_id=processor.request_id,
+                        trace=processor.trace,
                     ),
                 ),
                 mimetype="text/event-stream",

docsgpt/api/answer/services/stream_processor.py

@@ -1,9 +1,12 @@
 import datetime
+import functools
 import json
 import logging
+import uuid
 from pathlib import Path
 from typing import Any, Dict, Optional, Set
 
+from docsgpt import tracing
 from docsgpt.agents.agent_creator import AgentCreator
 from docsgpt.agents.default_tools import synthesized_default_tools
 from docsgpt.api.answer.services.compression import CompressionOrchestrator
@@ -104,9 +107,41 @@ def get_prompt(prompt_id: str, prompts_collection=None) -> str:
         raise ValueError(f"Invalid prompt ID: {prompt_id}") from e
 
 
+def _traced_setup(method):
+    """Run a request-setup method inside the request's execution trace.
+
+    Agent setup does real work worth seeing in the trace -- pre-fetch
+    retrieval, history compression -- before ``complete_stream`` runs, so
+    the trace is started here, in the request thread, and handed on.
+    """
+
+    @functools.wraps(method)
+    def wrapper(self, *args, **kwargs):
+        trace = getattr(self, "trace", None)
+        if trace is None:
+            trace = tracing.start_trace(source=getattr(self, "trace_source", "stream"))
+            self.trace = trace
+        with tracing.activate(trace):
+            try:
+                return method(self, *args, **kwargs)
+            finally:
+                if trace is not None:
+                    decoded = getattr(self, "decoded_token", None)
+                    trace.bind(
+                        request_id=getattr(self, "request_id", None),
+                        user_id=decoded.get("sub") if isinstance(decoded, dict) else None,
+                        agent_id=getattr(self, "agent_id", None),
+                    )
+
+    return wrapper
+
+
 class StreamProcessor:
     def __init__(
-        self, request_data: Dict[str, Any], decoded_token: Optional[Dict[str, Any]]
+        self,
+        request_data: Dict[str, Any],
+        decoded_token: Optional[Dict[str, Any]],
+        trace_source: str = "stream",
     ):
         # Legacy attribute retained as None for any external callers that
         # introspect the processor; all DB access uses per-op connections.
@@ -137,6 +172,11 @@ class StreamProcessor:
         self.reserved_message_id: Optional[str] = None
         # Carried through resumes so multi-pause runs keep one request_id.
         self.request_id: Optional[str] = None
+        # The request's execution trace, started by the first traced setup
+        # step and handed to ``complete_stream``; ``trace_source`` names the
+        # entry point it is stored under.
+        self.trace: Optional[tracing.Trace] = None
+        self.trace_source = trace_source
         self.conversation_service = ConversationService()
         self.compression_orchestrator = CompressionOrchestrator(
             self.conversation_service
@@ -161,12 +201,17 @@ class StreamProcessor:
         self._load_conversation_history()
         self._process_attachments()
 
+    @_traced_setup
     def build_agent(self, question: str):
         """One call to go from request data to a ready-to-run agent.
 
         Combines initialize(), pre_fetch_docs(), pre_fetch_tools(), and
-        create_agent() into a single convenience method.
+        create_agent() into a single convenience method. The request id is
+        minted first so pre-fetch retrieval and its side-channel LLM calls
+        share it with the rest of the turn.
         """
+        if not getattr(self, "request_id", None):
+            self.request_id = self.data.get("request_id") or str(uuid.uuid4())
         self.initialize()
 
         agent_type = self.agent_config.get("agent_type", "classic")
@@ -221,6 +266,7 @@ class StreamProcessor:
             tools_data=tools_data,
         )
 
+    @_traced_setup
     def build_continuation_from_messages(self, messages, tool_actions):
         """Rebuild a tool continuation from the request messages (STATELESS).
 
@@ -1118,7 +1164,7 @@ class StreamProcessor:
             user_api_key=self.agent_config["user_api_key"],
             agent_id=self.agent_id,
             decoded_token=self.decoded_token,
-            request_id=self.data.get("request_id"),
+            request_id=self.request_id or self.data.get("request_id"),
         )
 
         def _legacy_classic():
@@ -1440,6 +1486,7 @@ class StreamProcessor:
             logger.warning(f"Failed to fetch memory tool data: {str(e)}")
             return None
 
+    @_traced_setup
     def resume_from_tool_actions(
         self,
         tool_actions: list,
@@ -1813,7 +1860,7 @@ class StreamProcessor:
                 "llm_name": provider or settings.LLM_PROVIDER,
                 "api_key": system_api_key,
                 "decoded_token": self.decoded_token,
-                "request_id": self.data.get("request_id"),
+                "request_id": self.request_id or self.data.get("request_id"),
             }
 
         elif agent_type == "workflow":

docsgpt/api/v1/routes.py

@@ -257,7 +257,7 @@ def chat_completions():
         internal_data["persist"] = True
 
     try:
-        processor = StreamProcessor(internal_data, decoded_token)
+        processor = StreamProcessor(internal_data, decoded_token, trace_source="v1")
         # Set when this request took the resume claim, so a refusal can release it.
         claimed_conversation_id = None
 
@@ -495,6 +495,8 @@ def _stream_response(
         visibility=visibility,
         _continuation=continuation,
         finalize_tool_pause_as_complete=finalize_stateless_tool_pause,
+        request_id=processor.request_id,
+        trace=getattr(processor, "trace", None),
     )
 
     translation_state = StreamTranslationState()
@@ -572,6 +574,8 @@ def _non_stream_response(
         visibility=visibility,
         _continuation=continuation,
         finalize_tool_pause_as_complete=finalize_stateless_tool_pause,
+        request_id=processor.request_id,
+        trace=getattr(processor, "trace", None),
     )
 
     result = helper.process_response_stream(stream)