Trace guardrails, workflow steps and research phases

Guardrail evaluations are traced when they call a remote check or fire; a
firing guardrail drops every content preview from the stored trace. Each
workflow node and each research phase is a step span, and the workflow run
id links the trace to its Logs row.

docsgpt/agents/research_agent.py

@@ -4,6 +4,7 @@ import os
 import time
 from typing import Dict, Generator, List, Optional
 
+from docsgpt import tracing
 from docsgpt.agents.base import BaseAgent
 from docsgpt.agents.tool_executor import ToolExecutor
 from docsgpt.agents.tools.graph_search import add_graph_search_tool
@@ -51,6 +52,15 @@ SYNTHESIS_PROMPT = _load_prompt("synthesis.txt")
 # ---------------------------------------------------------------------------
 
 
+
+def _phase_span(phase: str, **attributes):
+    """Open a trace span for one research phase (clarify, plan, a step, synthesis)."""
+    return tracing.start_span(
+        tracing.KIND_STEP,
+        f"research {phase}",
+        attributes={"docsgpt.research.phase": phase.split(" ")[0], **attributes},
+    )
+
 class CitationManager:
     """Tracks and deduplicates citations across research steps."""
 
@@ -169,7 +179,8 @@ class ResearchAgent(BaseAgent):
 
         # Phase 0: Clarification (skip if user is responding to a prior clarification)
         if not self._is_follow_up():
-            clarification = self._clarification_phase(query)
+            with _phase_span("clarify"):
+                clarification = self._clarification_phase(query)
             if clarification:
                 yield {"metadata": {"is_clarification": True}}
                 yield {"answer": clarification}
@@ -182,7 +193,11 @@ class ResearchAgent(BaseAgent):
 
         # Phase 1: Planning (with adaptive depth)
         yield {"type": "research_progress", "data": {"status": "planning"}}
-        plan, complexity = self._planning_phase(query)
+        with _phase_span("plan") as plan_span:
+            plan, complexity = self._planning_phase(query)
+            plan_span.set(
+                **{"docsgpt.research.steps": len(plan or []), "docsgpt.research.complexity": complexity}
+            )
 
         if not plan:
             logger.warning("ResearchAgent: Planning produced no steps, falling back")
@@ -222,7 +237,9 @@ class ResearchAgent(BaseAgent):
                 },
             }
 
-            report = self._research_step(step_query, tools_dict)
+            with _phase_span(f"step {step_num}", **{"docsgpt.research.step": step_num}) as step_span:
+                step_span.preview("query", step_query)
+                report = self._research_step(step_query, tools_dict)
             intermediate_reports.append({"step": step, "content": report})
 
             yield {
@@ -249,9 +266,18 @@ class ResearchAgent(BaseAgent):
                 "tokens_used": self._tokens_used,
             },
         }
-        yield from self._synthesis_phase(
-            query, plan, intermediate_reports, tools_dict, log_context
-        )
+        synthesis_span = _phase_span("synthesis")
+        try:
+            yield from self._synthesis_phase(
+                query, plan, intermediate_reports, tools_dict, log_context
+            )
+        except GeneratorExit:
+            synthesis_span.end(tracing.STATUS_CANCELLED)
+            raise
+        except Exception as exc:
+            synthesis_span.end(error=exc)
+            raise
+        synthesis_span.end()
 
         # Sources and tool calls
         self.retrieved_docs = self.citations.get_all_docs()

docsgpt/agents/workflows/workflow_engine.py

@@ -6,6 +6,7 @@ import uuid
 from datetime import datetime, timezone
 from typing import Any, Dict, Generator, List, Optional, TYPE_CHECKING
 
+from docsgpt import tracing
 from docsgpt.agents.workflows.cel_evaluator import CelEvaluationError, evaluate_cel
 from docsgpt.agents.workflows.node_agent import WorkflowNodeAgentFactory
 from docsgpt.agents.workflows.schemas import (
@@ -114,6 +115,9 @@ class WorkflowEngine:
         self, initial_inputs: WorkflowState, query: str
     ) -> Generator[Dict[str, str], None, None]:
         self._initialize_state(initial_inputs, query)
+        # A workflow run's Logs row finds its trace by this id; a nested
+        # engine (a workflow inside a workflow) keeps the outermost run's.
+        tracing.bind_if_unset(workflow_run_id=self.workflow_run_id)
 
         # Surface the run id up front so the client can list this run's
         # artifacts (GET /api/artifacts?workflow_run_id=) once it has been
@@ -140,6 +144,16 @@ class WorkflowEngine:
                 break
             log_entry = self._create_log_entry(node)
             self._last_node_tool_calls = []
+            step_span = tracing.start_span(
+                tracing.KIND_STEP,
+                f"workflow_step {node.title or node.type.value}",
+                attributes={
+                    "docsgpt.workflow.node_id": node.id,
+                    "docsgpt.workflow.node_type": node.type.value,
+                    "docsgpt.workflow.node_title": node.title or None,
+                    "docsgpt.workflow_run_id": self.workflow_run_id,
+                },
+            )
 
             yield {
                 "type": "workflow_step",
@@ -150,7 +164,12 @@ class WorkflowEngine:
             }
 
             try:
-                yield from self._execute_node(node)
+                try:
+                    yield from self._execute_node(node)
+                except GeneratorExit:
+                    step_span.end(tracing.STATUS_CANCELLED)
+                    raise
+                step_span.end()
                 log_entry["status"] = ExecutionStatus.COMPLETED.value
                 self._finalize_log_entry(log_entry, pre_state)
 
@@ -167,6 +186,7 @@ class WorkflowEngine:
                     "output": node_output,
                 }
             except Exception as e:
+                step_span.end(error=e)
                 logger.error(f"Error executing node {node.id}: {e}", exc_info=True)
                 log_entry["status"] = ExecutionStatus.FAILED.value
                 log_entry["error"] = str(e)

docsgpt/guardrails/engine.py

@@ -7,6 +7,7 @@ from threading import Thread
 from time import monotonic
 from typing import List, Optional
 
+from docsgpt import tracing
 from docsgpt.guardrails.base import ScanContext
 from docsgpt.guardrails.config import GuardrailsConfig
 from docsgpt.guardrails.guardrail_creator import GuardrailCreator
@@ -26,6 +27,35 @@ logger = logging.getLogger(__name__)
 _MAX_WORKERS = 8
 
 
+def _start_guardrail_span(stage: Stage, controls) -> "tracing.Span":
+    return tracing.start_span(
+        tracing.KIND_GUARDRAIL,
+        f"guardrail {stage.value}",
+        attributes={
+            "docsgpt.guardrail.stage": stage.value,
+            "docsgpt.guardrail.checks": [c.check for c in controls],
+        },
+    )
+
+
+def _describe_decision(span, decision: StageDecision) -> None:
+    """Record a stage decision on its span; a firing guardrail drops trace previews."""
+    triggered = [v.check for v in decision.triggered]
+    span.set(
+        **{
+            "docsgpt.guardrail.triggered": triggered or None,
+            "docsgpt.guardrail.blocked": decision.blocked,
+            "docsgpt.guardrail.redacted": decision.redacted,
+            "docsgpt.guardrail.categories": decision.categories() or None,
+            "docsgpt.guardrail.unevaluated": [v.check for v in decision.unevaluated] or None,
+        }
+    )
+    if triggered or decision.blocked or decision.redacted:
+        # The scanned text (or text near it) sits in other spans' previews:
+        # the retrieval query, tool results, the answer. Keep none of it.
+        tracing.mark_content_blocked()
+
+
 class GuardrailEngine:
     """Stateless evaluator bound to one agent's resolved config."""
 
@@ -93,14 +123,24 @@ class GuardrailEngine:
             return decision
 
         if any(self._needs_deadline(c) for c in controls):
-            decision.verdicts = self._run_concurrent(controls, text, stage)
+            # Remote checks (e.g. an LLM judge) cost real time, so they are
+            # always traced; the judge's own LLM call nests under this span.
+            with _start_guardrail_span(stage, controls) as span:
+                decision.verdicts = self._run_concurrent(controls, text, stage)
+                self._reduce(decision)
+                _describe_decision(span, decision)
         else:
             # Bounded local checks are pattern matches measured in
             # microseconds. Running them inline keeps the streaming hot loop
             # free of thread churn.
             decision.verdicts = [self._run_control(c, text, stage) for c in controls]
+            self._reduce(decision)
+            # The output guard evaluates every streamed segment; tracing each
+            # clean local scan would bury the trace, so only a firing is kept.
+            if not decision.clean or decision.unevaluated:
+                with _start_guardrail_span(stage, controls) as span:
+                    _describe_decision(span, decision)
 
-        self._reduce(decision)
         self._record(decision)
         return decision
 
@@ -117,7 +157,7 @@ class GuardrailEngine:
         for control in controls[:_MAX_WORKERS]:
             slot: dict = {"control": control, "verdict": None}
             thread = Thread(
-                target=self._fill_slot,
+                target=tracing.wrap(self._fill_slot),
                 args=(slot, control, text, stage),
                 daemon=True,
                 name=f"guardrail-{control.check}",

docsgpt/tracing/core.py

@@ -5,7 +5,7 @@ scheduled run, a search). Spans are recorded while the request runs and the
 whole trace is written once, when its owner calls :func:`flush`.
 
 Nesting is tracked with one span stack per thread. Only *container* spans
-(agent, tool, retrieval, step) are pushed, so a leaf such as an LLM call can
+(agent, tool, retrieval, rerank, guardrail, step) are pushed, so a leaf such as an LLM call can
 never become the parent of a sibling that starts while it is still open --
 which matters because DocsGPT's agent loop is a chain of suspended
 generators. Ending a span pops anything left above it (marked
@@ -40,7 +40,9 @@ KIND_GUARDRAIL = "guardrail"
 KIND_STEP = "step"
 
 #: Kinds that become the implicit parent of spans started while they are open.
-CONTAINER_KINDS = frozenset({KIND_AGENT, KIND_TOOL, KIND_RETRIEVAL, KIND_RERANK, KIND_STEP})
+CONTAINER_KINDS = frozenset(
+    {KIND_AGENT, KIND_TOOL, KIND_RETRIEVAL, KIND_RERANK, KIND_GUARDRAIL, KIND_STEP}
+)
 
 STATUS_OK = "ok"
 STATUS_ERROR = "error"

tests/agents/test_research_agent.py

@@ -1626,3 +1626,53 @@ class TestResearchAgentClarificationCoverage:
         text = 'Here is my response: {"needs_clarification": false} end.'
         result = agent._parse_clarification_json(text)
         assert result == {"needs_clarification": False}
+
+
+@pytest.mark.unit
+class TestGenInnerTraceSpans:
+    """Each research phase is a step span in the execution trace."""
+
+    def test_phases_become_step_spans(
+        self,
+        agent_base_params,
+        mock_llm,
+        mock_llm_creator,
+        mock_llm_handler_creator,
+        log_context,
+        monkeypatch,
+    ):
+        from docsgpt import tracing
+        from docsgpt.core.settings import settings
+
+        monkeypatch.setattr(settings, "TRACES_ENABLED", True)
+        agent = ResearchAgent(**agent_base_params)
+        plan_steps = [{"query": "s1", "rationale": "r"}, {"query": "s2", "rationale": "r"}]
+
+        def fake_step(query, tools_dict):
+            with tracing.span(tracing.KIND_LLM, "chat m"):
+                return f"report {query}"
+
+        trace = tracing.start_trace(source="stream", capture_otel_context=False)
+        with tracing.activate(trace), \
+             patch.object(agent, "_setup_tools", return_value={}), \
+             patch.object(agent, "_is_follow_up", return_value=False), \
+             patch.object(agent, "_clarification_phase", return_value=None), \
+             patch.object(agent, "_planning_phase", return_value=(plan_steps, "moderate")), \
+             patch.object(agent, "_research_step", side_effect=fake_step), \
+             patch.object(agent, "_synthesis_phase", return_value=iter([{"answer": "final"}])), \
+             patch.object(agent, "_get_truncated_tool_calls", return_value=[]):
+            list(agent._gen_inner("Compare A and B", log_context))
+
+        names = [s.name for s in trace.spans if s.kind == tracing.KIND_STEP]
+        assert names == [
+            "research clarify",
+            "research plan",
+            "research step 1",
+            "research step 2",
+            "research synthesis",
+        ]
+        by_name = {s.name: s for s in trace.spans}
+        assert by_name["research plan"].attributes["docsgpt.research.steps"] == 2
+        llm_spans = [s for s in trace.spans if s.kind == tracing.KIND_LLM]
+        assert llm_spans[0].parent_id == by_name["research step 1"].id
+        assert by_name["research step 1"].previews["query"] == "s1"

tests/agents/test_workflow_engine_coverage.py

@@ -1081,3 +1081,51 @@ class TestNodeDocumentManifest:
         monkeypatch.setattr("docsgpt.storage.db.session.db_readonly", broken_readonly)
         config = AgentNodeConfig(agent_type="classic", input_documents=["*"])
         assert engine._node_document_manifest(config) == ""
+
+
+class TestTraceSpans:
+    """Every executed node is a ``workflow_step`` span; the run id links the trace."""
+
+    @pytest.fixture(autouse=True)
+    def _trace(self, monkeypatch):
+        from docsgpt import tracing
+        from docsgpt.core.settings import settings
+
+        monkeypatch.setattr(settings, "TRACES_ENABLED", True)
+        self.trace = tracing.start_trace(source="stream", capture_otel_context=False)
+        with tracing.activate(self.trace):
+            yield
+
+    @pytest.mark.unit
+    def test_step_spans_and_run_binding(self):
+        nodes = [
+            _make_node("n1", NodeType.START, "Start"),
+            _make_node("n2", NodeType.END, "End", config={"config": {}}),
+        ]
+        engine = WorkflowEngine(
+            _make_graph(nodes, [_make_edge("e1", "n1", "n2")]), _make_agent()
+        )
+        list(engine.execute({}, "hello"))
+        steps = [s for s in self.trace.spans if s.kind == "step"]
+        assert [s.name for s in steps] == ["workflow_step Start", "workflow_step End"]
+        assert all(s.status == "ok" for s in steps)
+        assert steps[0].attributes["docsgpt.workflow.node_type"] == "start"
+        assert self.trace.workflow_run_id == str(engine.workflow_run_id)
+
+    @pytest.mark.unit
+    def test_failed_node_is_an_error_span(self):
+        nodes = [
+            _make_node("n1", NodeType.START),
+            _make_node(
+                "n2",
+                NodeType.STATE,
+                "State",
+                config={"config": {"operations": [{"expression": "bad!!!", "target_variable": "x"}]}},
+            ),
+        ]
+        engine = WorkflowEngine(
+            _make_graph(nodes, [_make_edge("e1", "n1", "n2")]), _make_agent()
+        )
+        list(engine.execute({}, "q"))
+        state_span = [s for s in self.trace.spans if s.name == "workflow_step State"][0]
+        assert state_span.status == "error"

tests/guardrails/test_engine.py

@@ -340,3 +340,61 @@ def test_action_enum_serialises_as_value():
     dumped = config.model_dump(mode="json")
     assert dumped["controls"][0]["action"] == Action.REDACT.value
     assert dumped["controls"][0]["stage"] == Stage.INPUT.value
+
+
+class QuickRemoteCheck(GuardrailCheck):
+    name = "_test_quick_remote"
+    label = "Quick remote"
+    supported_stages = {Stage.INPUT, Stage.OUTPUT}
+    remote = True
+
+    def scan(self, text, stage, context):
+        from docsgpt import tracing
+
+        with tracing.span(tracing.KIND_LLM, "chat judge"):
+            pass
+        return CheckOutcome.clean()
+
+
+class TestTraceSpans:
+    """Guardrail evaluations appear in the execution trace only when they matter."""
+
+    @pytest.fixture(autouse=True)
+    def _trace(self, monkeypatch):
+        from docsgpt import tracing
+        from docsgpt.core.settings import settings
+
+        monkeypatch.setattr(settings, "TRACES_ENABLED", True)
+        GuardrailCreator.register(QuickRemoteCheck.name, QuickRemoteCheck)
+        self.trace = tracing.start_trace(source="stream", capture_otel_context=False)
+        with tracing.activate(self.trace):
+            yield
+        GuardrailCreator.checks.pop(QuickRemoteCheck.name, None)
+
+    def test_clean_local_scan_is_not_traced(self):
+        engine = GuardrailEngine(
+            _config(controls=[{"check": "_test_never", "stage": "output", "action": "block"}])
+        )
+        engine.evaluate("text", Stage.OUTPUT)
+        assert self.trace.spans == []
+        assert self.trace.content_blocked is False
+
+    def test_firing_local_scan_is_traced_and_blocks_previews(self):
+        engine = GuardrailEngine(_config())
+        engine.evaluate("some text", Stage.INPUT)
+        (span,) = self.trace.spans
+        assert span.kind == "guardrail"
+        assert span.attributes["docsgpt.guardrail.blocked"] is True
+        assert span.attributes["docsgpt.guardrail.triggered"] == ["_test_always"]
+        assert self.trace.content_blocked is True
+
+    def test_remote_scan_is_traced_with_judge_nested(self):
+        engine = GuardrailEngine(
+            _config(controls=[{"check": "_test_quick_remote", "stage": "input", "action": "block"}])
+        )
+        decision = engine.evaluate("text", Stage.INPUT)
+        assert decision.clean
+        guard, judge = self.trace.spans
+        assert guard.kind == "guardrail"
+        assert judge.parent_id == guard.id
+        assert self.trace.content_blocked is False