fix(utils): retain spans only for active usage scopes

quantmind/utils/usage.py

@@ -31,7 +31,6 @@
 from __future__ import annotations
 
 import threading
-from collections import defaultdict
 from collections.abc import Iterator
 from contextlib import contextmanager
 from dataclasses import dataclass, field
@@ -107,15 +106,21 @@ class UsageScope:
 
 
 class SpanCollector(TracingProcessor):
-    """Process-global sink that buckets finished spans by ``trace_id``.
+    """Process-global sink for spans owned by active usage scopes.
 
-    Bucketing by ``trace_id`` is what makes concurrent scopes and ``batch_run``
-    fan-out safe: each scope only ever reads and evicts its own trace.
+    Each scope opens a bucket with ``start`` and removes it with ``pop``.
+    Unrelated traces and late spans cannot create buckets, so only active
+    scopes retain spans. Concurrent scopes read and evict only their own trace.
     """
 
     def __init__(self) -> None:
         self._lock = threading.Lock()
-        self._by_trace: dict[str, list[Any]] = defaultdict(list)
+        self._by_trace: dict[str, list[Any]] = {}
+
+    def start(self, trace_id: str) -> None:
+        """Collect finished spans for one scope until ``pop`` is called."""
+        with self._lock:
+            self._by_trace[trace_id] = []
 
     def on_trace_start(self, trace: Any) -> None:  # noqa: D102 - no-op sink
         return None
@@ -127,9 +132,11 @@ def on_span_start(self, span: Any) -> None:  # noqa: D102 - no-op sink
         return None
 
     def on_span_end(self, span: Any) -> None:
-        """Store the finished span under its trace bucket (thread-safe)."""
+        """Store a finished span only while its scope is active."""
         with self._lock:
-            self._by_trace[span.trace_id].append(span)
+            spans = self._by_trace.get(span.trace_id)
+            if spans is not None:
+                spans.append(span)
 
     def shutdown(self) -> None:  # noqa: D102 - nothing to flush
         return None
@@ -138,7 +145,7 @@ def force_flush(self) -> None:  # noqa: D102 - synchronous, nothing queued
         return None
 
     def pop(self, trace_id: str) -> list[Any]:
-        """Return and remove every span collected for ``trace_id``."""
+        """End collection and return this scope's spans under the same lock."""
         with self._lock:
             return self._by_trace.pop(trace_id, [])
 
@@ -184,6 +191,7 @@ def usage_scope(
     trace_id = gen_trace_id()
     scope = UsageScope(trace_id=trace_id)
     try:
+        collector.start(trace_id)
         with trace(workflow_name, trace_id=trace_id, group_id=group_id):
             yield scope
     finally:

tests/utils/test_usage.py

@@ -1,13 +1,23 @@
 """Tests for per-run usage collection from Agents SDK spans.
 
-Spans are hand-built so the suite runs fully offline: no model call, no real
-trace export. ``usage_scope`` is driven with a patched ``trace`` and spans
-pushed straight into the collector, simulating what the SDK would record.
+Hand-built spans exercise aggregation. Real SDK traces exercise collection
+with an isolated provider: no model call, network request, or trace export.
 """
 
+import asyncio
 import unittest
+from contextvars import Context
 from typing import Any
-from unittest.mock import patch
+from unittest.mock import Mock, patch
+
+from agents.tracing import (
+    TracingProcessor,
+    generation_span,
+    get_current_trace,
+    trace,
+)
+from agents.tracing import create as tracing_create
+from agents.tracing.provider import DefaultTraceProvider
 
 from quantmind.utils import usage
 from quantmind.utils.usage import RunUsage, SpanCollector, usage_scope
@@ -210,6 +220,8 @@ class SpanCollectorTests(unittest.TestCase):
 
     def test_buckets_by_trace_id_and_pop_clears(self) -> None:
         collector = SpanCollector()
+        collector.start("t1")
+        collector.start("t2")
         collector.on_span_end(_agent("a1", "x", trace_id="t1"))
         collector.on_span_end(_agent("a2", "y", trace_id="t1"))
         collector.on_span_end(_agent("a3", "z", trace_id="t2"))
@@ -302,5 +314,208 @@ def test_usage_populated_even_when_block_raises(self) -> None:
         self.assertEqual(scope_ref["run"].usage.total_tokens, 10)
 
 
+def _sdk_generation(input_tokens: int = 10, output_tokens: int = 5) -> Any:
+    return generation_span(
+        model="synthetic-model",
+        usage={"input_tokens": input_tokens, "output_tokens": output_tokens},
+    )
+
+
+class RealSDKUsageScopeTests(unittest.TestCase):
+    """Real callbacks without changing the SDK's global provider/exporters."""
+
+    def setUp(self) -> None:
+        # DefaultTraceProvider starts empty; only this test's processors are
+        # registered. Its import and the create-module patch are test-only SDK
+        # internals; public trace/span context managers still run unchanged.
+        provider = DefaultTraceProvider()
+        provider.set_disabled(False)
+        self.observer = Mock(spec=TracingProcessor)
+        provider.register_processor(self.observer)
+        for patcher in (
+            patch.object(
+                tracing_create, "get_trace_provider", return_value=provider
+            ),
+            patch.object(
+                usage,
+                "add_trace_processor",
+                side_effect=provider.register_processor,
+            ),
+            patch.object(usage, "_collector", None),
+        ):
+            patcher.start()
+            self.addCleanup(patcher.stop)
+        self.addCleanup(provider.shutdown)
+        self.collector = usage._get_collector()
+        self.context = Context()  # Do not inherit an ambient trace or span.
+
+    def test_unrelated_traces_are_not_retained_and_observers_still_run(
+        self,
+    ) -> None:
+        def scenario() -> None:
+            with usage_scope() as run:
+                with _sdk_generation():
+                    pass
+            for _ in range(3):
+                with trace("unrelated"), _sdk_generation(99, 1):
+                    pass
+
+            self.assertEqual(run.usage.total_tokens, 15)
+            self.assertEqual(run.usage.requests, 1)
+            self.assertEqual(self.observer.on_span_end.call_count, 4)
+            self.assertEqual(self.collector._by_trace, {})
+
+        self.context.run(scenario)
+
+    def test_late_span_does_not_recreate_a_closed_scope_bucket(self) -> None:
+        def scenario() -> None:
+            with usage_scope() as run:
+                with _sdk_generation():
+                    pass
+                late_span = _sdk_generation(99, 1)
+                late_span.start()
+
+            snapshot = run.usage
+            late_span.finish()  # SDK permits span_end after trace_end.
+            self.assertEqual(snapshot.total_tokens, 15)
+            self.assertEqual(snapshot.requests, 1)
+            self.assertIs(run.usage, snapshot)
+            self.assertEqual(self.observer.on_span_end.call_count, 2)
+            self.assertEqual(self.collector._by_trace, {})
+
+        self.context.run(scenario)
+
+    def test_nested_scopes_restore_the_existing_sdk_trace(self) -> None:
+        def scenario() -> None:
+            with trace("external") as external, _sdk_generation(99, 1):
+                with usage_scope() as outer:
+                    with _sdk_generation():
+                        pass
+                    with usage_scope() as inner:
+                        with _sdk_generation(2, 1):
+                            pass
+                    self.assertEqual(
+                        get_current_trace().trace_id, outer.trace_id
+                    )
+                    with _sdk_generation():
+                        pass
+                self.assertIs(get_current_trace(), external)
+
+            self.assertNotEqual(outer.trace_id, inner.trace_id)
+            self.assertEqual(outer.usage.total_tokens, 30)
+            self.assertEqual(outer.usage.requests, 2)
+            self.assertEqual(inner.usage.total_tokens, 3)
+            self.assertEqual(inner.usage.requests, 1)
+            self.assertIsNone(get_current_trace())
+            self.assertEqual(self.collector._by_trace, {})
+
+        self.context.run(scenario)
+
+    def test_concurrent_scopes_keep_their_own_usage(self) -> None:
+        async def scenario() -> None:
+            ready = asyncio.Event()
+            entered = 0
+
+            async def worker(tokens: int) -> Any:
+                nonlocal entered
+                with usage_scope() as run:
+                    with _sdk_generation(tokens, 1):
+                        entered += 1
+                        if entered == 2:
+                            ready.set()
+                        await ready.wait()
+                        self.assertEqual(
+                            get_current_trace().trace_id, run.trace_id
+                        )
+                return run
+
+            first, second = await asyncio.gather(worker(10), worker(20))
+            self.assertNotEqual(first.trace_id, second.trace_id)
+            self.assertEqual(first.usage.total_tokens, 11)
+            self.assertEqual(second.usage.total_tokens, 21)
+            self.assertEqual(first.usage.requests, 1)
+            self.assertEqual(second.usage.requests, 1)
+            self.assertEqual(self.collector._by_trace, {})
+
+        self.context.run(asyncio.run, scenario())
+
+    def test_child_tasks_count_only_when_finished_before_scope_exit(
+        self,
+    ) -> None:
+        async def scenario() -> None:
+            started = asyncio.Event()
+            release = asyncio.Event()
+
+            async def child() -> None:
+                with _sdk_generation(3, 2):
+                    started.set()
+                    await release.wait()
+
+            with usage_scope() as joined:
+                release.set()
+                await asyncio.create_task(child())
+            self.assertEqual(joined.usage.total_tokens, 5)
+            self.assertEqual(joined.usage.requests, 1)
+
+            started.clear()
+            release.clear()
+            with usage_scope() as exited:
+                with _sdk_generation():
+                    pass
+                task = asyncio.create_task(child())
+                await started.wait()
+            snapshot = exited.usage
+            release.set()
+            await task
+            self.assertEqual(snapshot.total_tokens, 15)
+            self.assertEqual(snapshot.requests, 1)
+            self.assertIs(exited.usage, snapshot)
+            self.assertEqual(self.collector._by_trace, {})
+
+        self.context.run(asyncio.run, scenario())
+
+    def test_block_and_trace_entry_exceptions_release_the_scope(self) -> None:
+        def scenario() -> None:
+            with self.assertRaisesRegex(ValueError, "block failed"):
+                with usage_scope() as run, _sdk_generation():
+                    raise ValueError("block failed")
+            self.assertEqual(run.usage.total_tokens, 15)
+            self.assertEqual(run.usage.requests, 1)
+            self.assertEqual(self.collector._by_trace, {})
+
+            with patch.object(usage, "trace") as broken_trace:
+                broken_trace.return_value.__enter__.side_effect = RuntimeError(
+                    "trace entry failed"
+                )
+                with self.assertRaisesRegex(RuntimeError, "trace entry failed"):
+                    with usage_scope():
+                        self.fail("trace entry should have raised")
+            self.assertEqual(self.collector._by_trace, {})
+
+        self.context.run(scenario)
+
+    def test_cancellation_propagates_and_releases_the_scope(self) -> None:
+        async def scenario() -> None:
+            started = asyncio.Event()
+            scopes = []
+
+            async def worker() -> None:
+                with usage_scope() as run, _sdk_generation():
+                    scopes.append(run)
+                    started.set()
+                    await asyncio.Event().wait()
+
+            task = asyncio.create_task(worker())
+            await started.wait()
+            task.cancel()
+            with self.assertRaises(asyncio.CancelledError):
+                await task
+            self.assertEqual(scopes[0].usage.total_tokens, 15)
+            self.assertEqual(scopes[0].usage.requests, 1)
+            self.assertEqual(self.collector._by_trace, {})
+
+        self.context.run(asyncio.run, scenario())
+
+
 if __name__ == "__main__":
     unittest.main()