Fix Agent365 non-chat inference export and cross-tenant replay

CHANGELOG.md

@@ -27,6 +27,11 @@
   mode `0700` subject to umask and validated before cache initialization.
 
 ### Bugs Fixed
+- Export Agent365 `text_completion` and `generate_content` inference spans instead of filtering them out.
+  ([#316](https://github.com/microsoft/opentelemetry-distro-python/pull/316))
+- Continue Agent365 durable replay for healthy identities when another identity is retrying or blocked,
+  scanning up to 100 records while limiting each pass to 10 delivery attempts.
+  ([#316](https://github.com/microsoft/opentelemetry-distro-python/pull/316))
 - Invalid Agent365 storage configuration or provider initialization failure logs
   a nonfatal error and disables only offline buffering/replay, preserving online
   export. Explicit offline disablement bypasses storage initialization.

src/microsoft/opentelemetry/a365/core/exporters/replay_coordinator.py

@@ -24,6 +24,7 @@
 _logger = logging.getLogger(__name__)
 
 _MAX_RECORDS_PER_PASS = 10
+_MAX_RECORDS_TO_SCAN_PER_PASS = 100
 # A pass leases all records before sending them sequentially. Keep every lease
 # valid for the worst-case full pass so another process sharing the file
 # queue cannot reclaim a record while this coordinator is still sending it.
@@ -137,30 +138,33 @@ def shutdown(self, timeout_seconds: float | None = None) -> bool:
     def run_once(self) -> bool:
         """Claim and process a single bounded replay batch.
 
-        Returns ``True`` only when a maximal batch (``_MAX_RECORDS_PER_PASS``)
-        was claimed *and* every record in it reached a terminal state
-        (delivered or permanently dropped). In that case more records may
-        remain and the loop should run again immediately. It returns ``False``
-        for an empty/partial batch, when the pass was stopped early (retryable
-        failure, shutdown, or unexpected error), or when any record was left
-        behind (gate-blocked or released), so the loop falls back to the
-        periodic cadence and does not busy-spin.
+        Scan up to ``_MAX_RECORDS_TO_SCAN_PER_PASS`` records to find identities
+        whose gates are open, attempting at most ``_MAX_RECORDS_PER_PASS``.
+        Retryable failures update only their identity's gate without stopping the pass.
+        Return ``True`` only when the full attempt budget was deleted, allowing
+        healthy backlogs to drain immediately. Otherwise use the periodic cadence
+        so retained or gate-blocked records do not cause a busy-spin.
         """
-        records = self._storage.claim(_MAX_RECORDS_PER_PASS, _LEASE_SECONDS)
+        records = self._storage.claim(_MAX_RECORDS_TO_SCAN_PER_PASS, _LEASE_SECONDS)
         if not records:
             return False
 
+        attempted_count = 0
         deleted_count = 0
         for index, record in enumerate(records):
             if self._stop_event.is_set():
                 self._release_remaining(records, index)
                 return False
+            if attempted_count >= _MAX_RECORDS_PER_PASS:
+                self._release_remaining(records, index)
+                break
 
             identity = self._identity_for(record)
             if not self._gate.try_acquire(identity):
                 self._release_record(record)
                 continue
 
+            attempted_count += 1
             try:
                 result = self._send(record)
             except ReplayIdentityError as exc:
@@ -216,13 +220,8 @@ def run_once(self) -> bool:
 
             self._gate.record_retryable_failure(identity, result.retry_after)
             self._release_record(record)
-            self._release_remaining(records, index + 1)
-            return False
 
-        # Only request an immediate re-run when we fully drained a maximal
-        # batch; otherwise re-claiming would return records that are still
-        # gated/leased and spin the loop.
-        return deleted_count == len(records) and len(records) >= _MAX_RECORDS_PER_PASS
+        return deleted_count == _MAX_RECORDS_PER_PASS
 
     def _run_loop(self) -> None:
         while not self._stop_event.is_set():

src/microsoft/opentelemetry/a365/core/exporters/utils.py

@@ -70,6 +70,7 @@
         APPLY_GUARDRAIL_OPERATION_NAME,
         InferenceOperationType.CHAT.value,
     }
+    | {operation.name.lower() for operation in InferenceOperationType}
 )
 
 

tests/a365/test_durable_restart.py

@@ -89,6 +89,30 @@ def _queue_size(storage) -> int:
     return len(claimed)
 
 
+def test_replay_releases_unattempted_backlog_for_the_next_pass(tmp_path):
+    exporter = _Agent365Exporter(
+        token_resolver=lambda a, t: "test-token",
+        storage_directory=tmp_path / "queue",
+        enable_durable_delivery=True,
+    )
+    try:
+        exporter._ensure_durable_initialized()
+        identity = IdentityKey("t1", "a1", None, False)
+        for _ in range(15):
+            assert exporter._storage.store(DurableRecord.new(identity, '{"resourceSpans":[]}'))
+        exporter._post_once = MagicMock(  # type: ignore[method-assign]
+            return_value=DeliveryResult(DeliveryDisposition.DELIVERED)
+        )
+
+        assert exporter._replay.run_once() is True
+        assert exporter._post_once.call_count == 10
+        assert exporter._replay.run_once() is False
+        assert exporter._post_once.call_count == 15
+        assert _queue_size(exporter._storage) == 0
+    finally:
+        exporter.shutdown()
+
+
 def test_restart_replays_persisted_record_with_fresh_token(tmp_path):
     storage_dir = tmp_path / "queue"
     storage_dir.mkdir()

tests/a365/test_exporter.py

@@ -605,8 +605,8 @@ def test_export_includes_inference_operation_type_chat_spans(self):
         exporter.shutdown()
 
     @patch.dict(os.environ, {}, clear=True)
-    def test_export_filters_out_unsupported_inference_operation_types(self):
-        """Spans with TextCompletion / GenerateContent are filtered out."""
+    def test_export_filters_out_legacy_non_chat_operation_names(self):
+        """Noncanonical TextCompletion / GenerateContent names remain filtered out."""
         exporter = make_exporter()
         exporter._post_once = MagicMock(return_value=_delivered())
         text_completion_span = _make_span(
@@ -624,17 +624,27 @@ def test_export_filters_out_unsupported_inference_operation_types(self):
 
     @patch.dict(os.environ, {}, clear=True)
     def test_export_does_not_normalize_canonical_operation_names(self):
-        """invoke_agent / execute_tool / output_messages / chat are not rewritten."""
+        """Canonical operation names, including non-chat inference, are exported unchanged."""
         exporter = make_exporter()
         exporter._post_once = MagicMock(return_value=_delivered())
-        for op in ("invoke_agent", "execute_tool", "output_messages", "chat"):
-            with self.subTest(operation_name=op):
-                exporter._post_once.reset_mock()
-                span = _make_span(name=f"{op}_span", trace_id=1, span_id=2, operation_name=op)
-                result = exporter.export([span])
-                self.assertEqual(result, SpanExportResult.SUCCESS)
-                exporter._post_once.assert_called_once()
-        exporter.shutdown()
+        try:
+            for op in (
+                "invoke_agent",
+                "execute_tool",
+                "output_messages",
+                "chat",
+                "text_completion",
+                "generate_content",
+            ):
+                with self.subTest(operation_name=op):
+                    exporter._post_once.reset_mock()
+                    span = _make_span(name=f"{op}_span", trace_id=1, span_id=2, operation_name=op)
+                    result = exporter.export([span])
+                    self.assertEqual(result, SpanExportResult.SUCCESS)
+                    exporter._post_once.assert_called_once()
+                    self.assertEqual(span.attributes["gen_ai.operation.name"], op)
+        finally:
+            exporter.shutdown()
 
 
 # ---------------------------------------------------------------------------

tests/a365/test_replay_coordinator.py

@@ -7,6 +7,7 @@
 
 import threading
 import time
+from dataclasses import replace
 from unittest.mock import MagicMock
 
 from microsoft.opentelemetry.a365.constants import A365_HTTP_TIMEOUT_SECONDS
@@ -86,7 +87,10 @@ def claim(self, limit: int, lease_seconds: float) -> list[DurableRecord]:
             self.block_event.wait()
         with self._lock:
             if self._batches:
-                return self._batches.pop(0)
+                batch = self._batches.pop(0)
+                if len(batch) > limit:
+                    self._batches.insert(0, batch[limit:])
+                return batch[:limit]
             return []
 
     def delete(self, record_id: int) -> bool:
@@ -138,7 +142,7 @@ def test_replay_lease_covers_worst_case_full_pass() -> None:
 
     coordinator.run_once()
 
-    assert storage.claim_limits == [_MAX_RECORDS_PER_PASS]
+    assert storage.claim_limits == [100]
     assert storage.claim_lease_seconds == [_LEASE_SECONDS]
     assert _LEASE_SECONDS > _MAX_RECORDS_PER_PASS * A365_HTTP_TIMEOUT_SECONDS
 
@@ -160,6 +164,68 @@ def test_replay_retains_retryable_record_and_updates_gate() -> None:
     gate.record_retryable_failure.assert_called_once_with(IDENTITY, 45)
 
 
+def test_retryable_identity_allows_other_tenants_and_later_recovers() -> None:
+    sibling = _make_record(2)
+    other_tenant = replace(_make_record(3), tenant_id="tenant-2")
+    storage = FakeStorage([[RECORD, sibling, other_tenant], [RECORD]])
+    now = [0.0]
+    gate = TransmissionGate(clock=lambda: now[0])
+    send = MagicMock(
+        side_effect=[
+            DeliveryResult(DeliveryDisposition.RETRYABLE, retry_after=45),
+            DeliveryResult(DeliveryDisposition.DELIVERED),
+            DeliveryResult(DeliveryDisposition.DELIVERED),
+        ]
+    )
+    coordinator = ReplayCoordinator(storage, gate, send=send)
+
+    assert coordinator.run_once() is False
+    assert [call.args[0].record_id for call in send.call_args_list] == [1, 3]
+    assert storage.released == [1, 2]
+    assert storage.deleted == [3]
+    assert not gate.try_acquire(IDENTITY)
+
+    now[0] = 46.0
+    assert coordinator.run_once() is False
+    assert [call.args[0].record_id for call in send.call_args_list] == [1, 3, 1]
+    assert storage.deleted == [3, 1]
+    assert gate.try_acquire(IDENTITY)
+
+
+def test_replay_reaches_healthy_tenant_beyond_a_blocked_batch() -> None:
+    blocked = [_make_record(i) for i in range(1, 13)]
+    healthy = replace(_make_record(13), tenant_id="tenant-2")
+    storage = FakeStorage([[*blocked, healthy]])
+    gate = TransmissionGate(clock=lambda: 0.0)
+    gate.record_retryable_failure(IDENTITY, retry_after=45)
+    send = MagicMock(return_value=DeliveryResult(DeliveryDisposition.DELIVERED))
+    coordinator = ReplayCoordinator(storage, gate, send=send)
+
+    assert coordinator.run_once() is False
+
+    send.assert_called_once_with(healthy)
+    assert storage.deleted == [13]
+    assert storage.released == list(range(1, 13))
+
+
+def test_full_attempt_batch_with_retryable_record_uses_periodic_backoff() -> None:
+    records = [replace(_make_record(i), tenant_id=f"tenant-{i}") for i in range(1, 11)]
+    storage = FakeStorage([records])
+    gate = TransmissionGate()
+    send = MagicMock(
+        side_effect=[
+            DeliveryResult(DeliveryDisposition.RETRYABLE, retry_after=45),
+            *[DeliveryResult(DeliveryDisposition.DELIVERED) for _ in range(9)],
+        ]
+    )
+    coordinator = ReplayCoordinator(storage, gate, send=send)
+
+    assert coordinator.run_once() is False
+    assert send.call_count == 10
+    assert storage.released == [1]
+    assert storage.deleted == list(range(2, 11))
+
+
 def test_replay_deletes_permanent_record() -> None:
     storage = FakeStorage([[RECORD]])
     gate = MagicMock(spec=TransmissionGate)
@@ -303,15 +369,47 @@ def send(record: DurableRecord) -> DeliveryResult:
     gate.release_probe.assert_called_once_with(IDENTITY)
 
 
-def test_run_once_claims_at_most_ten_records() -> None:
-    storage = FakeStorage([[RECORD]])
+def test_run_once_bounds_scanning_when_all_records_are_blocked() -> None:
+    storage = FakeStorage([[_make_record(i) for i in range(1, 102)]])
     gate = MagicMock(spec=TransmissionGate)
-    gate.try_acquire.return_value = True
-    coordinator = ReplayCoordinator(storage, gate, send=lambda record: DeliveryResult(DeliveryDisposition.DELIVERED))
+    gate.try_acquire.return_value = False
+    send = MagicMock()
+    coordinator = ReplayCoordinator(storage, gate, send=send)
 
-    coordinator.run_once()
+    assert coordinator.run_once() is False
+
+    assert storage.claim_limits == [100]
+    assert gate.try_acquire.call_count == 100
+    assert storage.released == list(range(1, 101))
+    send.assert_not_called()
+
+
+def test_run_once_bounds_attempts_and_releases_unattempted_records() -> None:
+    storage = FakeStorage([[_make_record(i) for i in range(1, 12)]])
+    gate = TransmissionGate()
+    send = MagicMock(return_value=DeliveryResult(DeliveryDisposition.DELIVERED))
+    coordinator = ReplayCoordinator(storage, gate, send=send)
+
+    assert coordinator.run_once() is True
+
+    assert send.call_count == 10
+    assert storage.deleted == list(range(1, 11))
+    assert storage.released == [11]
+
+
+def test_shutdown_during_send_releases_remaining_scanned_records() -> None:
+    storage = FakeStorage([[_make_record(i) for i in range(1, 21)]])
+    gate = TransmissionGate()
+
+    def send(record: DurableRecord) -> DeliveryResult:
+        coordinator._stop_event.set()
+        return DeliveryResult(DeliveryDisposition.DELIVERED)
 
-    assert storage.claim_limits == [10]
+    coordinator = ReplayCoordinator(storage, gate, send=send)
+
+    assert coordinator.run_once() is False
+    assert storage.deleted == [1]
+    assert storage.released == list(range(2, 21))
 
 
 def test_start_and_wake_process_a_later_batch() -> None:

tests/a365/test_utils.py

@@ -130,6 +130,13 @@ def test_groups_by_identity(self):
         self.assertEqual(len(result[("t1", "a1")]), 2)
         self.assertEqual(len(result[("t2", "a2")]), 1)
 
+    def test_includes_canonical_inference_operations(self):
+        for operation_name in ("chat", "text_completion", "generate_content"):
+            with self.subTest(operation_name=operation_name):
+                span = self._make_span("t1", "a1", operation_name)
+
+                self.assertEqual(filter_and_partition_by_identity([span]), {("t1", "a1"): [span]})
+
     def test_skips_missing_tenant(self):
         span = MagicMock()
         span.attributes = {"gen_ai.agent.id": "a1", "gen_ai.operation.name": "invoke_agent"}