fix(tracing): address greptile review findings

- Extract _build_sgp_span and _add_source_to_span to module-level helpers
  so the sync and async processors share a single implementation, removing
  byte-for-byte duplication that risked drifting between the two classes.
- Strengthen test_sgp_span_input_and_output_propagated_on_end to actually
  assert the updated input/output reach create_span on the end call,
  matching what the test name and docstring promise.

Author: smoreinis

src/agentex/lib/core/tracing/processors/sgp_tracing_processor.py

@@ -27,6 +27,39 @@ def _get_span_type(span: Span) -> str:
     return "STANDALONE"
 
 
+def _add_source_to_span(span: Span, env_vars: EnvironmentVariables) -> None:
+    if span.data is None:
+        span.data = {}
+    if isinstance(span.data, dict):
+        span.data["__source__"] = "agentex"
+        if env_vars.ACP_TYPE is not None:
+            span.data["__acp_type__"] = env_vars.ACP_TYPE
+        if env_vars.AGENT_NAME is not None:
+            span.data["__agent_name__"] = env_vars.AGENT_NAME
+        if env_vars.AGENT_ID is not None:
+            span.data["__agent_id__"] = env_vars.AGENT_ID
+
+
+def _build_sgp_span(span: Span, env_vars: EnvironmentVariables) -> SGPSpan:
+    """Build an SGPSpan from an agentex Span. Idempotent on span_id at the SGP backend."""
+    _add_source_to_span(span, env_vars)
+    sgp_span = cast(
+        SGPSpan,
+        create_span(
+            name=span.name,
+            span_type=_get_span_type(span),
+            span_id=span.id,
+            parent_id=span.parent_id,
+            trace_id=span.trace_id,
+            input=span.input,
+            output=span.output,
+            metadata=span.data,
+        ),
+    )
+    sgp_span.start_time = span.start_time.isoformat()  # type: ignore[union-attr]
+    return sgp_span
+
+
 class SGPSyncTracingProcessor(SyncTracingProcessor):
     def __init__(self, config: SGPTracingProcessorConfig):
         disabled = config.sgp_api_key == "" or config.sgp_account_id == ""
@@ -40,45 +73,14 @@ def __init__(self, config: SGPTracingProcessorConfig):
         )
         self.env_vars = EnvironmentVariables.refresh()
 
-    def _add_source_to_span(self, span: Span) -> None:
-        if span.data is None:
-            span.data = {}
-        if isinstance(span.data, dict):
-            span.data["__source__"] = "agentex"
-            if self.env_vars.ACP_TYPE is not None:
-                span.data["__acp_type__"] = self.env_vars.ACP_TYPE
-            if self.env_vars.AGENT_NAME is not None:
-                span.data["__agent_name__"] = self.env_vars.AGENT_NAME
-            if self.env_vars.AGENT_ID is not None:
-                span.data["__agent_id__"] = self.env_vars.AGENT_ID
-
-    def _build_sgp_span(self, span: Span) -> SGPSpan:
-        """Build an SGPSpan from an agentex Span. Idempotent on span_id at the SGP backend."""
-        self._add_source_to_span(span)
-        sgp_span = cast(
-            SGPSpan,
-            create_span(
-                name=span.name,
-                span_type=_get_span_type(span),
-                span_id=span.id,
-                parent_id=span.parent_id,
-                trace_id=span.trace_id,
-                input=span.input,
-                output=span.output,
-                metadata=span.data,
-            ),
-        )
-        sgp_span.start_time = span.start_time.isoformat()  # type: ignore[union-attr]
-        return sgp_span
-
     @override
     def on_span_start(self, span: Span) -> None:
-        sgp_span = self._build_sgp_span(span)
+        sgp_span = _build_sgp_span(span, self.env_vars)
         sgp_span.flush(blocking=False)
 
     @override
     def on_span_end(self, span: Span) -> None:
-        sgp_span = self._build_sgp_span(span)
+        sgp_span = _build_sgp_span(span, self.env_vars)
         sgp_span.end_time = span.end_time.isoformat()  # type: ignore[union-attr]
         sgp_span.flush(blocking=False)
 
@@ -108,37 +110,6 @@ def __init__(self, config: SGPTracingProcessorConfig):
         )
         self.env_vars = EnvironmentVariables.refresh()
 
-    def _add_source_to_span(self, span: Span) -> None:
-        if span.data is None:
-            span.data = {}
-        if isinstance(span.data, dict):
-            span.data["__source__"] = "agentex"
-            if self.env_vars.ACP_TYPE is not None:
-                span.data["__acp_type__"] = self.env_vars.ACP_TYPE
-            if self.env_vars.AGENT_NAME is not None:
-                span.data["__agent_name__"] = self.env_vars.AGENT_NAME
-            if self.env_vars.AGENT_ID is not None:
-                span.data["__agent_id__"] = self.env_vars.AGENT_ID
-
-    def _build_sgp_span(self, span: Span) -> SGPSpan:
-        """Build an SGPSpan from an agentex Span. Idempotent on span_id at the SGP backend."""
-        self._add_source_to_span(span)
-        sgp_span = cast(
-            SGPSpan,
-            create_span(
-                name=span.name,
-                span_type=_get_span_type(span),
-                span_id=span.id,
-                parent_id=span.parent_id,
-                trace_id=span.trace_id,
-                input=span.input,
-                output=span.output,
-                metadata=span.data,
-            ),
-        )
-        sgp_span.start_time = span.start_time.isoformat()  # type: ignore[union-attr]
-        return sgp_span
-
     @override
     async def on_span_start(self, span: Span) -> None:
         await self.on_spans_start([span])
@@ -152,7 +123,7 @@ async def on_spans_start(self, spans: list[Span]) -> None:
         if not spans:
             return
 
-        sgp_spans = [self._build_sgp_span(span) for span in spans]
+        sgp_spans = [_build_sgp_span(span, self.env_vars) for span in spans]
 
         if self.disabled:
             logger.warning("SGP is disabled, skipping span upsert")
@@ -168,7 +139,7 @@ async def on_spans_end(self, spans: list[Span]) -> None:
 
         sgp_spans: list[SGPSpan] = []
         for span in spans:
-            sgp_span = self._build_sgp_span(span)
+            sgp_span = _build_sgp_span(span, self.env_vars)
             sgp_span.end_time = span.end_time.isoformat()  # type: ignore[union-attr]
             sgp_spans.append(sgp_span)
 

tests/lib/core/tracing/processors/test_sgp_tracing_processor.py

@@ -180,7 +180,8 @@ def capture_create_span(**kwargs):
             captured.append(sgp_span)
             return sgp_span
 
-        with patch(f"{MODULE}.create_span", side_effect=capture_create_span):
+        mock_create_span = MagicMock(side_effect=capture_create_span)
+        with patch(f"{MODULE}.create_span", mock_create_span):
             span = _make_span()
             span.input = {"messages": [{"role": "user", "content": "hello"}]}
             await processor.on_span_start(span)
@@ -199,6 +200,10 @@ def capture_create_span(**kwargs):
         # The end-time SGPSpan should have end_time populated.
         end_span = captured[-1]
         assert end_span.end_time is not None
+        # Verify the updated input/output reached create_span on the end call.
+        end_call_kwargs = mock_create_span.call_args_list[-1].kwargs
+        assert end_call_kwargs["input"]["messages"][-1]["role"] == "assistant"
+        assert end_call_kwargs["output"] == {"response": "hi"}
 
     async def test_on_spans_start_sends_single_upsert_for_batch(self):
         """Given N spans at once, on_spans_start should make ONE upsert_batch HTTP call."""