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)