fix(tracing): pin eval ids for plain spans and keep them until the upload succeeds

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>

Author: mohammadatallah-scale

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

@@ -256,9 +256,11 @@ async def on_spans_end(self, spans: list[Span]) -> None:
         for span in spans:
             sgp_span = _build_sgp_span(span, self.env_vars)
             sgp_span.end_time = span.end_time.isoformat()  # type: ignore[union-attr]
-            sgp_evals.release_span(span.id)
             sgp_spans.append(sgp_span)
         await client.spans.upsert_batch(items=[s.to_request_params() for s in sgp_spans])
+        # Released only after the upload so a queue retry rebuilds the span with the same ids.
+        for span in spans:
+            sgp_evals.release_span(span.id)
         _metrics.record_export_success(
             event_type="end", span_count=len(spans), processor="sgp"
         )

src/agentex/lib/core/tracing/sgp_evals.py

@@ -81,9 +81,7 @@ def _lookup(span: Span) -> dict[str, Any]:
 
 
 def capture_for_span(span: Span) -> None:
-    """Pin the span's attrs at start. No-op while no eval task is registered in this process."""
-    if not _attrs_by_task:
-        return
+    """Pin the span's attrs at start, empty included, so an eval task registered later cannot claim it."""
     with _lock:
         _attrs_by_span[span.id] = _lookup(span)
         while len(_attrs_by_span) > _MAX_TASKS:

tests/lib/core/tracing/test_sgp_evals.py

@@ -29,7 +29,11 @@
 from agentex.lib.core.clients.temporal.types import ConflictWorkflowPolicy
 from agentex.lib.core.temporal.workers.worker import AgentexWorker
 from agentex.lib.core.temporal.services.temporal_task_service import TemporalTaskService
-from agentex.lib.core.tracing.processors.sgp_tracing_processor import SGPSyncTracingProcessor, _sgp_metadata
+from agentex.lib.core.tracing.processors.sgp_tracing_processor import (
+    SGPSyncTracingProcessor,
+    SGPAsyncTracingProcessor,
+    _sgp_metadata,
+)
 
 EVAL_METADATA = {
     "sgp_evals": "generation-unit",
@@ -149,6 +153,34 @@ def test_plain_span_is_not_stamped_by_a_run_registered_after_it_started(self) ->
 
         assert _sgp_metadata(span) == {"k": 1}
 
+    def test_plain_span_started_with_an_empty_registry_is_not_stamped_later(self) -> None:
+        span = _span(trace_id="task-1", data={"k": 1})
+        sgp_evals.capture_for_span(span)
+
+        sgp_evals.register_task_metadata("task-1", EVAL_METADATA)
+
+        assert _sgp_metadata(span) == {"k": 1}
+
+    async def test_async_processor_keeps_the_capture_when_the_upload_fails(self) -> None:
+        sgp_evals.register_task_metadata("task-1", EVAL_METADATA)
+        span = _span(data={})
+        sgp_evals.capture_for_span(span)
+        span.end_time = datetime.now(UTC)
+        module = "agentex.lib.core.tracing.processors.sgp_tracing_processor"
+        with patch(f"{module}.tracing.init"), patch(f"{module}.EnvironmentVariables"):
+            processor = SGPAsyncTracingProcessor(SGPTracingProcessorConfig(sgp_api_key="", sgp_account_id=""))
+        client = Mock()
+        client.spans.upsert_batch = AsyncMock(side_effect=[RuntimeError("boom"), None])
+        with patch.object(processor, "_get_client", return_value=client):
+            with pytest.raises(RuntimeError):
+                await processor.on_spans_end([span])
+            sgp_evals.unregister_task("task-1")
+            assert sgp_evals.attrs_for_span(span) == EXPECTED_ATTRS
+
+            await processor.on_spans_end([span])
+
+        assert sgp_evals.attrs_for_span(span) == {}
+
     def test_release_drops_the_capture(self) -> None:
         sgp_evals.register_task_metadata("task-1", EVAL_METADATA)
         span = _span(data={})