Trace unexecuted tool calls in one place; tidy span lifecycles

Paused, denied, skipped and client-run tool calls all go through
trace_unexecuted_tool_call. The research synthesis, continuation and
workflow step spans use with-blocks instead of hand-rolled GeneratorExit
handling. Fixes blank lines and adds a missing hint and docstrings.

docsgpt/agents/base.py

@@ -6,11 +6,9 @@ from abc import ABC, abstractmethod
 from datetime import datetime, timezone
 from typing import Any, Dict, Generator, List, Optional
 
-from docsgpt import tracing
 from docsgpt.agents.tool_executor import (
     ToolExecutor,
-    finish_tool_span,
-    record_tool_span_start,
+    trace_unexecuted_tool_call,
     result_status,
     truncate_tool_result,
 )
@@ -504,18 +502,10 @@ class BaseAgent(ABC):
             pending_tool_calls: The pending tool call descriptors from the pause.
             tool_actions: Client-provided actions resolving the pending calls.
         """
-        span = start_agent_span(self, continuation=True)
-        completed = False
-        try:
+        with start_agent_span(self, continuation=True):
             yield from self._gen_continuation_inner(
                 messages, tools_dict, pending_tool_calls, tool_actions, reasoning_content
             )
-            completed = True
-        except Exception as exc:
-            span.end(error=exc)
-            raise
-        finally:
-            span.end(None if completed else tracing.STATUS_CANCELLED)
 
     def _gen_continuation_inner(
         self,
@@ -617,7 +607,7 @@ class BaseAgent(ABC):
                     "arguments": args,
                     "status": "denied",
                 }
-                finish_tool_span(record_tool_span_start(tc), {**denied_data, "error": comment or None})
+                trace_unexecuted_tool_call(tc, {**denied_data, "error": comment or None})
                 yield {"type": "tool_call", "data": denied_data}
 
             elif "result" in action:
@@ -650,10 +640,7 @@ class BaseAgent(ABC):
                     "result": truncate_tool_result(result_str),
                     "status": result_status(result),
                 }
-                finish_tool_span(
-                    record_tool_span_start(tc, **{"docsgpt.client_executed": True}),
-                    client_data,
-                )
+                trace_unexecuted_tool_call(tc, client_data, **{"docsgpt.client_executed": True})
                 yield {"type": "tool_call", "data": client_data}
 
         # Resume the LLM loop with the updated messages

docsgpt/agents/research_agent.py

@@ -47,12 +47,6 @@ STEP_PROMPT = _load_prompt("step.txt")
 SYNTHESIS_PROMPT = _load_prompt("synthesis.txt")
 
 
-# ---------------------------------------------------------------------------
-# CitationManager
-# ---------------------------------------------------------------------------
-
-
-
 def _phase_span(phase: str, **attributes: Any) -> Any:
     """Open a trace span for one research phase (clarify, plan, a step, synthesis)."""
     return tracing.start_span(
@@ -61,6 +55,12 @@ def _phase_span(phase: str, **attributes: Any) -> Any:
         attributes={"docsgpt.research.phase": phase.split(" ")[0], **attributes},
     )
 
+
+# ---------------------------------------------------------------------------
+# CitationManager
+# ---------------------------------------------------------------------------
+
+
 class CitationManager:
     """Tracks and deduplicates citations across research steps."""
 
@@ -266,18 +266,10 @@ class ResearchAgent(BaseAgent):
                 "tokens_used": self._tokens_used,
             },
         }
-        synthesis_span = _phase_span("synthesis")
-        try:
+        with _phase_span("synthesis"):
             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/tool_executor.py

@@ -57,6 +57,22 @@ _TOOL_SPAN_STATUS = {
 }
 
 
+def trace_unexecuted_tool_call(call: Any, data: Dict[str, Any], **attributes: Any) -> None:
+    """Record a tool call that this process did not execute.
+
+    The single entry point for calls that never reach :meth:`ToolExecutor.execute`
+    (which traces executed calls itself): calls paused for approval or for the
+    client, denied ones, ones skipped at the context limit, and results a client
+    sends back on resume. A new pause or refusal path records its call here.
+
+    Args:
+        call: The tool call (anything with ``name`` and ``id``).
+        data: The tool-call record emitted for it (``status``, ``arguments``, ...).
+        **attributes: Extra span attributes.
+    """
+    finish_tool_span(record_tool_span_start(call, **attributes), data)
+
+
 def finish_tool_span(span: Any, data: Dict[str, Any]) -> None:
     """Close a tool span from the ``tool_calls`` entry the call produced.
 

docsgpt/agents/workflows/workflow_engine.py

@@ -164,12 +164,8 @@ class WorkflowEngine:
             }
 
             try:
-                try:
+                with step_span:
                     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)
 
@@ -186,7 +182,6 @@ 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

@@ -5,7 +5,7 @@ from __future__ import annotations
 import logging
 from threading import Thread
 from time import monotonic
-from typing import List, Optional
+from typing import Any, List, Optional, Sequence
 
 from docsgpt import tracing
 from docsgpt.guardrails.base import ScanContext
@@ -27,7 +27,8 @@ logger = logging.getLogger(__name__)
 _MAX_WORKERS = 8
 
 
-def _start_guardrail_span(stage: Stage, controls) -> "tracing.Span":
+def _start_guardrail_span(stage: Stage, controls: Sequence[Any]) -> "tracing.Span":
+    """Open a ``guardrail`` span for evaluating ``controls`` at ``stage``."""
     return tracing.start_span(
         tracing.KIND_GUARDRAIL,
         f"guardrail {stage.value}",
@@ -39,6 +40,7 @@ def _start_guardrail_span(stage: Stage, controls) -> "tracing.Span":
 
 
 def _content_fired(decision: StageDecision) -> bool:
+    """True when a check matched, or the text was blocked or redacted."""
     return bool(decision.triggered or decision.blocked or decision.redacted)
 
 

docsgpt/llm/handlers/base.py

@@ -5,6 +5,7 @@ from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from typing import Any, Dict, Generator, List, Optional, Union
 
+from docsgpt.agents.tool_executor import trace_unexecuted_tool_call
 from docsgpt.logging import build_stack_data
 
 logger = logging.getLogger(__name__)
@@ -97,17 +98,6 @@ class LLMResponse:
         return bool(self.tool_calls) and self.finish_reason == "tool_calls"
 
 
-def _trace_unexecuted_tool(call: Any, data: Dict[str, Any]) -> None:
-    """Record a tool call that was paused, denied or skipped instead of run.
-
-    Those calls never reach ``ToolExecutor.execute`` (which traces executed
-    calls), so without this the trace would not show them at all.
-    """
-    from docsgpt.agents.tool_executor import finish_tool_span, record_tool_span_start
-
-    finish_tool_span(record_tool_span_start(call), data)
-
-
 def _is_restated_payload(existing: str, incoming: str) -> bool:
     """True when ``incoming`` restates an already-complete ``existing`` payload.
 
@@ -1138,7 +1128,7 @@ class LLMHandler(ABC):
 
                     # Mark remaining tools as skipped
                     for remaining_call in tool_calls[i:]:
-                        _trace_unexecuted_tool(
+                        trace_unexecuted_tool_call(
                             remaining_call,
                             {"tool_name": "system", "action_name": remaining_call.name, "status": "skipped"},
                         )
@@ -1235,7 +1225,7 @@ class LLMHandler(ABC):
                             "error_type", "tool_not_allowed"
                         ),
                     }
-                    _trace_unexecuted_tool(call, denied_data)
+                    trace_unexecuted_tool_call(call, denied_data)
                     yield {"type": "tool_call", "data": denied_data}
                     continue
                 # Yield pause event so the client knows this tool is waiting
@@ -1250,7 +1240,7 @@ class LLMHandler(ABC):
                 # can wire the sticky "don't ask again" button.
                 if pause_info.get("device_id"):
                     pause_data["device_id"] = pause_info["device_id"]
-                _trace_unexecuted_tool(call, pause_data)
+                trace_unexecuted_tool_call(call, pause_data)
                 yield {"type": "tool_call", "data": pause_data}
                 pending_actions.append(pause_info)
                 # Do NOT add messages for pending tools here.

docsgpt/storage/db/repositories/request_traces.py

@@ -33,10 +33,12 @@ _SUMMARY_COLUMNS = (
 
 
 def _dump_jsonb(value: Any) -> str:
+    """Serialize ``value`` for a JSONB parameter, without NUL bytes Postgres rejects."""
     return json.dumps(strip_null_bytes(value), cls=PGNativeJSONEncoder)
 
 
 def _uuid_or_none(value: Optional[str]) -> Optional[str]:
+    """``value`` as a string when it is a UUID, else ``None`` (the column is typed)."""
     return str(value) if value and looks_like_uuid(str(value)) else None