Serve traces to the Logs UI

GET /api/traces returns a Logs row's stored traces, scoped to the caller or
to an agent they own. get_user_logs rows now carry the ids that find their
trace and a merged summary (duration, LLM and tool calls, tokens, retrieval
time) from one batched lookup per page; searches, MCP searches and graph
builds, which have no log row of their own, are listed from their traces.

docsgpt/api/user/analytics/routes.py

@@ -14,6 +14,10 @@ from docsgpt.api.user.base import (
 )
 from docsgpt.storage.db.redaction import redact_secrets
 from docsgpt.storage.db.repositories.agents import AgentsRepository
+from docsgpt.storage.db.repositories.request_traces import (
+    REF_FIELDS as TRACE_REF_FIELDS,
+    RequestTracesRepository,
+)
 from docsgpt.storage.db.repositories.token_usage import TokenUsageRepository
 from docsgpt.storage.db.session import db_readonly
 
@@ -90,6 +94,154 @@ def _resolve_agent(conn, api_key_id, user_id):
     return agent, api_key, agent_pg_id
 
 
+def _trace_branch(name: str, sources_sql: str, scope: str) -> dict:
+    """A ``get_user_logs`` branch listing stored traces of the given sources."""
+    return {
+        "name": name,
+        "level": "CASE WHEN t.status = 'error' THEN 'error' ELSE 'info' END",
+        # A search's query sits in its retrieval span's preview; graph builds
+        # are named after their source.
+        "summary": "COALESCE(t.spans->0->'preview'->>'query', t.name, t.source)",
+        "where": [f"t.source IN {sources_sql}", scope],
+        "sql": f"""
+        SELECT '{name}' AS event_type,
+               CAST(t.id AS text) AS id,
+               t.user_id AS user_id,
+               t.started_at AS timestamp,
+               {{level}} AS level,
+               t.source AS action,
+               {{summary}} AS summary,
+               jsonb_build_object(
+                   'status', t.status,
+                   'source', t.source,
+                   'duration_ms', t.duration_ms
+               ) AS payload
+        FROM request_traces t
+        WHERE {{where}}
+        """,
+    }
+
+
+def _trace_ref(item: dict):
+    """The ``(field, value)`` that finds a Logs row's trace, if it has one."""
+    event_type = item["event_type"]
+    row_id = item["id"].split("-", 1)[1]
+    if event_type == "chat":
+        return ("request_id", item.get("request_id")) if item.get("request_id") else None
+    if event_type in ("system", "webhook"):
+        return ("activity_id", item.get("activity_id")) if item.get("activity_id") else None
+    if event_type == "workflow":
+        return ("workflow_run_id", row_id)
+    if event_type == "schedule":
+        # The scheduler records the run under its run id.
+        return ("request_id", row_id)
+    if event_type in ("search", "graph"):
+        return ("id", row_id)
+    return None
+
+
+def _merge_trace_summaries(traces: list) -> dict:
+    """One Logs-row summary over every trace for it (a turn plus its resumes)."""
+    totals: dict = {}
+    for trace in traces:
+        for key, value in (trace.get("summary") or {}).items():
+            if isinstance(value, (int, float)) and not isinstance(value, bool):
+                totals[key] = totals.get(key, 0) + value
+    if "retrieval_ms" in totals:
+        totals["retrieval_ms"] = round(totals["retrieval_ms"], 1)
+    return {
+        "count": len(traces),
+        "duration_ms": sum(int(t.get("duration_ms") or 0) for t in traces),
+        "status": traces[-1].get("status"),
+        "started_at": traces[0].get("started_at"),
+        "summary": totals,
+    }
+
+
+def _attach_trace_summaries(conn, items: list, *, user_id, agent_id) -> None:
+    """Add a ``trace`` summary to each Logs row that has a stored trace.
+
+    One batched lookup per link field for the whole page, rather than a join
+    inside the UNION, so the timeline query is unchanged. Rows without a
+    trace (older than the feature, or tracing disabled) get no key.
+    """
+    refs: dict = {}
+    for item in items:
+        ref = _trace_ref(item)
+        if ref:
+            refs.setdefault(ref[0], set()).add(str(ref[1]))
+    if not refs:
+        return
+    found = RequestTracesRepository(conn).summaries_for_refs(
+        refs, user_id=user_id, agent_id=agent_id
+    )
+    for item in items:
+        ref = _trace_ref(item)
+        if not ref:
+            continue
+        traces = found.get(ref[0], {}).get(str(ref[1]))
+        if traces:
+            item["trace"] = {
+                "ref": {"field": ref[0], "value": str(ref[1])},
+                **_merge_trace_summaries(traces),
+            }
+
+
+@analytics_ns.route("/traces")
+class GetTraces(Resource):
+    @api.doc(
+        description=(
+            "Stored execution traces (span timelines) for one Logs row. Pass exactly "
+            "one of message_id, request_id, activity_id, workflow_run_id or id; "
+            "api_key_id scopes to an agent you own."
+        ),
+        params={
+            "message_id": "Assistant message id",
+            "request_id": "Request id (chat turn, scheduled run)",
+            "activity_id": "Agent activity id (webhook and system rows)",
+            "workflow_run_id": "Workflow run id",
+            "id": "Trace id",
+            "api_key_id": "Agent id to scope to",
+        },
+    )
+    def get(self):
+        decoded_token = request.decoded_token
+        if not decoded_token:
+            return make_response(jsonify({"success": False}), 401)
+        user = decoded_token.get("sub")
+        refs = [
+            (field, request.args.get(field))
+            for field in TRACE_REF_FIELDS
+            if request.args.get(field)
+        ]
+        if len(refs) != 1:
+            return make_response(
+                jsonify(
+                    {
+                        "success": False,
+                        "message": "Pass exactly one of: " + ", ".join(TRACE_REF_FIELDS),
+                    }
+                ),
+                400,
+            )
+        field, value = refs[0]
+        api_key_id = request.args.get("api_key_id")
+        try:
+            with db_readonly() as conn:
+                agent, _api_key, agent_pg_id = _resolve_agent(conn, api_key_id, user)
+                if api_key_id and agent is None:
+                    return make_response(jsonify({"success": True, "traces": []}), 200)
+                traces = RequestTracesRepository(conn).list_by_ref(
+                    field, value, user_id=user, agent_id=agent_pg_id
+                )
+        except Exception as err:
+            current_app.logger.error(f"Error getting traces: {err}", exc_info=True)
+            return make_response(jsonify({"success": False}), 400)
+        for trace in traces:
+            trace.pop("_id", None)
+        return make_response(jsonify({"success": True, "traces": traces}), 200)
+
+
 @analytics_ns.route("/get_message_analytics")
 class GetMessageAnalytics(Resource):
     get_message_analytics_model = api.model(
@@ -701,7 +853,9 @@ class GetUserLogs(Resource):
             "event_type": fields.String(
                 required=False,
                 description="Filter by event source",
-                enum=["chat", "schedule", "webhook", "workflow", "system"],
+                enum=[
+                    "chat", "schedule", "webhook", "workflow", "system", "search", "graph",
+                ],
             ),
             "search": fields.String(
                 required=False, description="Substring filter on the summary"
@@ -741,6 +895,8 @@ class GetUserLogs(Resource):
             "webhook",
             "workflow",
             "system",
+            "search",
+            "graph",
         ):
             return make_response(
                 jsonify({"success": False, "message": "Invalid option"}), 400
@@ -831,6 +987,7 @@ class GetUserLogs(Resource):
                         "wr.user_id = :user_id",
                         "wr.workflow_id = CAST(:agent_workflow_id AS uuid)",
                     ]
+                    trace_scope = "t.agent_id = CAST(:agent_pg_id AS uuid)"
                 else:
                     chat_where = ["l.user_id = :user_id"]
                     webhook_where = [
@@ -852,6 +1009,7 @@ class GetUserLogs(Resource):
                         "r.status IN ('success', 'completed', 'failed', 'timeout', 'skipped')",
                     ]
                     workflow_where = ["wr.user_id = :user_id"]
+                    trace_scope = "t.user_id = :user_id"
 
                 # One normalized timeline over five event sources.
                 # ``payload`` carries the per-type detail; the outer query
@@ -892,7 +1050,8 @@ class GetUserLogs(Resource):
                                {summary} AS summary,
                                jsonb_build_object(
                                    'endpoint', s.endpoint,
-                                   'stacks', s.stacks
+                                   'stacks', s.stacks,
+                                   'activity_id', s.activity_id
                                ) AS payload
                         FROM stack_logs s
                         WHERE {where}
@@ -913,7 +1072,8 @@ class GetUserLogs(Resource):
                                {summary} AS summary,
                                jsonb_build_object(
                                    'endpoint', s.endpoint,
-                                   'stacks', s.stacks
+                                   'stacks', s.stacks,
+                                   'activity_id', s.activity_id
                                ) AS payload
                         FROM stack_logs s
                         WHERE {where}
@@ -989,6 +1149,11 @@ class GetUserLogs(Resource):
                         WHERE {where}
                         """,
                     },
+                    # Runs with no log row of their own are listed from their
+                    # stored trace: searches (/api/search and MCP) and graph
+                    # builds.
+                    _trace_branch("search", "('search', 'mcp')", trace_scope),
+                    _trace_branch("graph", "('graph_extraction')", trace_scope),
                 ]
 
                 if level:
@@ -1062,6 +1227,8 @@ class GetUserLogs(Resource):
                             "tool_calls": payload.get("tool_calls"),
                             "agent_id": payload.get("agent_id"),
                             "attachments": payload.get("attachments"),
+                            "request_id": payload.get("request_id"),
+                            "message_id": payload.get("message_id"),
                         }
                     )
                 elif m["event_type"] in ("system", "webhook"):
@@ -1072,6 +1239,15 @@ class GetUserLogs(Resource):
                             # before write-time redaction still carry the
                             # reflected provider/user secrets in ``stacks``.
                             "stacks": redact_secrets(payload.get("stacks")),
+                            "activity_id": payload.get("activity_id"),
+                        }
+                    )
+                elif m["event_type"] in ("search", "graph"):
+                    item.update(
+                        {
+                            "status": payload.get("status"),
+                            "source": payload.get("source"),
+                            "duration_ms": payload.get("duration_ms"),
                         }
                     )
                 elif m["event_type"] == "workflow":
@@ -1104,6 +1280,14 @@ class GetUserLogs(Resource):
                         }
                     )
                 results.append(item)
+            if results:
+                with db_readonly() as conn:
+                    _attach_trace_summaries(
+                        conn,
+                        results,
+                        user_id=user,
+                        agent_id=agent_pg_id if api_key_id else None,
+                    )
         except Exception as err:
             current_app.logger.error(
                 f"Error getting user logs: {err}", exc_info=True

tests/api/user/test_traces_routes.py

@@ -0,0 +1,188 @@
+"""Tests for ``GET /api/traces`` and the trace summaries on ``get_user_logs`` rows."""
+
+from __future__ import annotations
+
+import time
+import uuid
+from contextlib import contextmanager
+from unittest.mock import patch
+
+import pytest
+from flask import Flask
+
+from docsgpt.storage.db.repositories.request_traces import RequestTracesRepository
+
+
+@pytest.fixture
+def app():
+    return Flask(__name__)
+
+
+@contextmanager
+def _patch_db(conn):
+    @contextmanager
+    def _yield_conn():
+        yield conn
+
+    with patch("docsgpt.api.user.analytics.routes.db_readonly", _yield_conn):
+        yield
+
+
+def _trace(pg_conn, **overrides):
+    record = {
+        "id": str(uuid.uuid4()),
+        "request_id": "req-1",
+        "user_id": "owner",
+        "source": "stream",
+        "name": "stream",
+        "status": "ok",
+        "started_at_ns": time.time_ns(),
+        "duration_ms": 900,
+        "span_count": 2,
+        "summary": {"llm_calls": 1, "tool_calls": 1, "input_tokens": 50, "retrieval_ms": 12.5},
+        "spans": [
+            {
+                "id": "s1",
+                "parent_id": None,
+                "kind": "retrieval",
+                "name": "retrieval",
+                "status": "ok",
+                "offset_ms": 0,
+                "duration_ms": 12.5,
+                "attributes": {},
+                "preview": {"query": "how do I deploy"},
+            }
+        ],
+    }
+    record.update(overrides)
+    RequestTracesRepository(pg_conn).insert(record)
+    return record
+
+
+def _get(app, pg_conn, user, query):
+    from docsgpt.api.user.analytics.routes import GetTraces
+
+    with _patch_db(pg_conn), app.test_request_context("/api/traces", query_string=query):
+        from flask import request
+
+        request.decoded_token = {"sub": user} if user else None
+        return GetTraces().get()
+
+
+def _logs(app, pg_conn, user, body):
+    from docsgpt.api.user.analytics.routes import GetUserLogs
+
+    with _patch_db(pg_conn), app.test_request_context(
+        "/api/get_user_logs", method="POST", json=body
+    ):
+        from flask import request
+
+        request.decoded_token = {"sub": user}
+        return GetUserLogs().post()
+
+
+class TestGetTraces:
+    def test_requires_auth(self, app, pg_conn):
+        assert _get(app, pg_conn, None, {"request_id": "x"}).status_code == 401
+
+    def test_requires_exactly_one_ref(self, app, pg_conn):
+        assert _get(app, pg_conn, "owner", {}).status_code == 400
+        assert _get(app, pg_conn, "owner", {"request_id": "a", "id": "b"}).status_code == 400
+
+    def test_returns_own_trace_with_spans(self, app, pg_conn):
+        record = _trace(pg_conn)
+        response = _get(app, pg_conn, "owner", {"request_id": "req-1"})
+        assert response.status_code == 200
+        (trace,) = response.json["traces"]
+        assert trace["id"] == record["id"]
+        assert trace["spans"][0]["preview"]["query"] == "how do I deploy"
+        assert "_id" not in trace
+
+    def test_other_users_get_nothing(self, app, pg_conn):
+        _trace(pg_conn)
+        response = _get(app, pg_conn, "intruder", {"request_id": "req-1"})
+        assert response.json["traces"] == []
+
+    def test_agent_owner_sees_shared_callers_traces(self, app, pg_conn):
+        from docsgpt.storage.db.repositories.agents import AgentsRepository
+
+        agent = AgentsRepository(pg_conn).create("owner", "a", "published", key="k1")
+        _trace(pg_conn, user_id="caller", agent_id=str(agent["id"]))
+        response = _get(
+            app, pg_conn, "owner", {"request_id": "req-1", "api_key_id": str(agent["id"])}
+        )
+        assert len(response.json["traces"]) == 1
+
+    def test_unowned_agent_scope_is_empty(self, app, pg_conn):
+        from docsgpt.storage.db.repositories.agents import AgentsRepository
+
+        agent = AgentsRepository(pg_conn).create("someone", "a", "published", key="k2")
+        _trace(pg_conn, user_id="someone", agent_id=str(agent["id"]))
+        response = _get(
+            app, pg_conn, "owner", {"request_id": "req-1", "api_key_id": str(agent["id"])}
+        )
+        assert response.json["traces"] == []
+
+
+class TestLogsTraceSummaries:
+    def test_chat_row_gets_merged_summary_across_rounds(self, app, pg_conn):
+        from docsgpt.storage.db.repositories.user_logs import UserLogsRepository
+
+        UserLogsRepository(pg_conn).insert(
+            user_id="owner",
+            endpoint="stream_answer",
+            data={"action": "stream_answer", "question": "q", "request_id": "req-1"},
+        )
+        _trace(pg_conn, status="paused")
+        _trace(pg_conn, started_at_ns=time.time_ns() + 1_000_000)
+        response = _logs(app, pg_conn, "owner", {"event_type": "chat"})
+        (row,) = response.json["logs"]
+        assert row["request_id"] == "req-1"
+        trace = row["trace"]
+        assert trace["ref"] == {"field": "request_id", "value": "req-1"}
+        assert trace["count"] == 2
+        assert trace["duration_ms"] == 1800
+        assert trace["status"] == "ok"
+        assert trace["summary"]["llm_calls"] == 2
+        assert trace["summary"]["retrieval_ms"] == 25.0
+
+    def test_row_without_trace_has_no_summary(self, app, pg_conn):
+        from docsgpt.storage.db.repositories.user_logs import UserLogsRepository
+
+        UserLogsRepository(pg_conn).insert(
+            user_id="owner", endpoint="stream_answer", data={"question": "old row"}
+        )
+        (row,) = _logs(app, pg_conn, "owner", {}).json["logs"]
+        assert "trace" not in row
+
+    def test_webhook_row_links_by_activity_id(self, app, pg_conn):
+        from docsgpt.storage.db.repositories.stack_logs import StackLogsRepository
+
+        StackLogsRepository(pg_conn).insert(
+            activity_id="act-1", endpoint="webhook", level="info", user_id="owner", query="{}"
+        )
+        _trace(pg_conn, request_id="task-1", source="webhook", activity_id="act-1")
+        (row,) = _logs(app, pg_conn, "owner", {"event_type": "webhook"}).json["logs"]
+        assert row["activity_id"] == "act-1"
+        assert row["trace"]["ref"] == {"field": "activity_id", "value": "act-1"}
+
+    def test_search_traces_are_listed(self, app, pg_conn):
+        record = _trace(pg_conn, request_id=None, source="mcp", name="mcp", status="error")
+        _trace(pg_conn, request_id=None, source="graph_extraction", name="graph_extraction Docs")
+        search_rows = _logs(app, pg_conn, "owner", {"event_type": "search"}).json["logs"]
+        assert len(search_rows) == 1
+        row = search_rows[0]
+        assert row["id"] == f"search-{record['id']}"
+        assert row["question"] == "how do I deploy"
+        assert row["level"] == "error"
+        assert row["action"] == "mcp"
+        assert row["trace"]["ref"] == {"field": "id", "value": record["id"]}
+        graph_rows = _logs(app, pg_conn, "owner", {"event_type": "graph"}).json["logs"]
+        assert len(graph_rows) == 1
+
+    def test_search_traces_of_other_users_are_hidden(self, app, pg_conn):
+        _trace(pg_conn, request_id=None, source="search", user_id="someone")
+        assert _logs(app, pg_conn, "owner", {"event_type": "search"}).json["logs"] == []
+
+    def test_unknown_event_type_is_rejected(self, app, pg_conn):
+        assert _logs(app, pg_conn, "owner", {"event_type": "bogus"}).status_code == 400