From 17cff0f381146f1bc4c968c5c5714c03a82f6202 Mon Sep 17 00:00:00 2001 From: pradystar Date: Wed, 5 Aug 2026 22:58:13 -0700 Subject: [PATCH 1/2] fix(logger): harden OTel context bookkeeping --- src/splunk_ao/logger/logger.py | 152 +++++++++++++++++--- tests/conftest.py | 5 +- tests/test_logger_otel_context.py | 229 +++++++++++++++++++++++++++++- 3 files changed, 358 insertions(+), 28 deletions(-) diff --git a/src/splunk_ao/logger/logger.py b/src/splunk_ao/logger/logger.py index 8ebd8937..7ff64c44 100644 --- a/src/splunk_ao/logger/logger.py +++ b/src/splunk_ao/logger/logger.py @@ -136,6 +136,7 @@ class ActiveOtelContext: logger_id: int step_id: uuid.UUID span_context: SpanContext + previous_context: otel_context.Context token: Token @@ -229,6 +230,8 @@ class SplunkAOLogger(TracesLogger): _task_handler: ThreadPoolTaskHandler _trace_completion_submitted: bool _otel_ids: dict[uuid.UUID, OtelIds] = PrivateAttr(default_factory=dict) + _otel_children_by_parent: dict[uuid.UUID, set[uuid.UUID]] = PrivateAttr(default_factory=dict) + _otel_parent_by_child: dict[uuid.UUID, uuid.UUID] = PrivateAttr(default_factory=dict) _pending_otel_steps: set[uuid.UUID] = PrivateAttr(default_factory=set) def __init__( @@ -368,7 +371,9 @@ def __init__( "User must provide project_name or project_id to SplunkAOLogger, or set it as an environment variable." ) if self.experiment_id is None and self.agent_stream_name is None and self.agent_stream_id is None: - raise SplunkAOLoggerException("agent_stream or agent_stream_id is required to initialize SplunkAOLogger.") + raise SplunkAOLoggerException( + "agent_stream or agent_stream_id is required to initialize SplunkAOLogger." + ) if local_metrics: self.local_metrics = local_metrics @@ -438,11 +443,16 @@ def _record_otel_ids( self, step: BaseStep, parent_step: BaseStep | None = None, parent_span_context: SpanContext | None = None ) -> OtelIds | None: """Assign stable OTel identity without disrupting proprietary logging.""" + existing_ids = self._otel_ids.get(step.id) + if existing_ids is not None: + return existing_ids + + parent_step_id = parent_step.id if parent_step is not None else None try: if parent_step is not None: parent_ids = self._otel_ids.get(parent_step.id) if parent_ids is None: - raise RuntimeError(f"Missing OTel context for parent step {parent_step.id}.") + raise RuntimeError(f"Missing OTel context for parent step {parent_step_id}.") parent_span_context = parent_ids.span_context elif parent_span_context is None: active_context = otel_trace.get_current_span().get_span_context() @@ -459,14 +469,43 @@ def _record_otel_ids( span_context=self._assign_otel_context(otel_trace_id, trace_state), parent_span_context=parent_span_context, ) - self._otel_ids[step.id] = ids + self._insert_otel_identity(step.id, ids, parent_step_id) return ids except Exception: + self._rollback_otel_identity_insert(step.id, parent_step_id) self._logger.warning( "Failed to assign OTel identity for step %s; continuing proprietary logging.", step.id, exc_info=True ) return None + def _insert_otel_identity(self, step_id: uuid.UUID, ids: OtelIds, parent_step_id: uuid.UUID | None) -> None: + """Record one identity and its explicit proprietary parent edge.""" + self._otel_ids[step_id] = ids + if parent_step_id is None: + return + + self._otel_parent_by_child[step_id] = parent_step_id + self._otel_children_by_parent.setdefault(parent_step_id, set()).add(step_id) + + def _rollback_otel_identity_insert(self, step_id: uuid.UUID, parent_step_id: uuid.UUID | None) -> None: + """Best-effort rollback for a partially inserted OTel identity.""" + with contextlib.suppress(Exception): + self._otel_ids.pop(step_id, None) + + recorded_parent_id = parent_step_id + with contextlib.suppress(Exception): + recorded_parent_id = self._otel_parent_by_child.pop(step_id, parent_step_id) + + if recorded_parent_id is None: + return + + with contextlib.suppress(Exception): + children = self._otel_children_by_parent.get(recorded_parent_id) + if children is not None: + children.discard(step_id) + if not children: + self._otel_children_by_parent.pop(recorded_parent_id, None) + def _open_otel_step_ids(self, current_parent: StepWithChildSpans | None) -> tuple[uuid.UUID, ...]: """Return the root-to-current proprietary chain that has OTel identities.""" path: list[uuid.UUID] = [] @@ -525,11 +564,11 @@ def _sync_otel_context_impl(self, current_parent: StepWithChildSpans | None) -> detach_failed = False for active in reversed(active_contexts): - try: - active.token.var.reset(active.token) - except (RuntimeError, ValueError): - # ContextVar tokens cannot be reset from a copied execution - # context. Restore the recorded base before rebuilding below. + otel_context.detach(active.token) + if otel_context.get_current() is not active.previous_context: + # Public detach catches its own errors. Verify the exact + # previously active Context so copied-context failures and + # swallowed no-ops can still trigger deterministic recovery. detach_failed = True if detach_failed: @@ -537,10 +576,15 @@ def _sync_otel_context_impl(self, current_parent: StepWithChildSpans | None) -> rebuilt_contexts: list[ActiveOtelContext] = [] for owner_id, step_id, span_context in contexts_to_restore: - ctx = otel_trace.set_span_in_context(NonRecordingSpan(span_context)) + previous_context = otel_context.get_current() + ctx = otel_trace.set_span_in_context(NonRecordingSpan(span_context), previous_context) rebuilt_contexts.append( ActiveOtelContext( - logger_id=owner_id, step_id=step_id, span_context=span_context, token=otel_context.attach(ctx) + logger_id=owner_id, + step_id=step_id, + span_context=span_context, + previous_context=previous_context, + token=otel_context.attach(ctx), ) ) @@ -556,18 +600,73 @@ def _discard_otel_subtree(self, step: BaseStep) -> None: self._discard_otel_identity_tree(step.id) def _discard_otel_identity_tree(self, step_id: uuid.UUID) -> None: - """Remove one identity and all identities parented to it.""" - ids = self._otel_ids.pop(step_id, None) - if ids is None: - return + """Iteratively remove one identity and its indexed descendants.""" + worklist = [step_id] + visited: set[uuid.UUID] = set() + first_error: Exception | None = None + + while worklist: + current_step_id = worklist.pop() + if current_step_id in visited: + continue + visited.add(current_step_id) + + child_ids: tuple[uuid.UUID, ...] = () + try: + child_ids = tuple(self._otel_children_by_parent.get(current_step_id, ())) + self._otel_children_by_parent.pop(current_step_id, None) + except Exception as exc: + first_error = first_error or exc + worklist.extend(child_ids) - child_ids = tuple( - child_id - for child_id, child_ids in self._otel_ids.items() - if child_ids.parent_span_context == ids.span_context - ) - for child_id in child_ids: - self._discard_otel_identity_tree(child_id) + for child_id in child_ids: + try: + self._otel_parent_by_child.pop(child_id, None) + except Exception as exc: + first_error = first_error or exc + + try: + self._remove_otel_identity(current_step_id) + except Exception as exc: + first_error = first_error or exc + + if first_error is not None: + raise first_error + + def _remove_otel_identity(self, step_id: uuid.UUID) -> None: + """Remove one identity and unlink it from its proprietary parent.""" + first_error: Exception | None = None + parent_step_id: uuid.UUID | None = None + + try: + parent_step_id = self._otel_parent_by_child.get(step_id) + self._otel_parent_by_child.pop(step_id, None) + except Exception as exc: + first_error = exc + + if parent_step_id is not None: + try: + siblings = self._otel_children_by_parent.get(parent_step_id) + if siblings is not None: + siblings.discard(step_id) + if not siblings: + self._otel_children_by_parent.pop(parent_step_id, None) + except Exception as exc: + first_error = first_error or exc + + try: + self._otel_ids.pop(step_id, None) + except Exception as exc: + first_error = first_error or exc + + if first_error is not None: + raise first_error + + def _clear_otel_identities(self) -> None: + """Clear stable identities and both structural indexes.""" + self._otel_ids.clear() + self._otel_children_by_parent.clear() + self._otel_parent_by_child.clear() def _release_otel_context(self, finished_step: BaseStep) -> None: """Release OTel bookkeeping without disrupting proprietary completion.""" @@ -618,7 +717,14 @@ def _emit_and_release(self, finished_step: BaseStep) -> None: self._logger.warning("Failed to emit completed step %s.", finished_step.id, exc_info=True) finally: self._pending_otel_steps.discard(finished_step.id) - self._otel_ids.pop(finished_step.id, None) + try: + self._remove_otel_identity(finished_step.id) + except Exception: + self._logger.warning( + "Failed to release OTel identity for completed step %s; continuing proprietary logging.", + finished_step.id, + exc_info=True, + ) def _emit_pending_descendants(self, finished_step: BaseStep) -> None: """Emit pending descendants in post-order before their enclosing parent.""" @@ -2457,7 +2563,7 @@ def terminate(self) -> None: self._logger.warning("SplunkAOLogger.terminate: sink shutdown failed: %s", exc) finally: self._set_current_parent(None) - self._otel_ids.clear() + self._clear_otel_identities() self._pending_otel_steps.clear() self.traces = [] diff --git a/tests/conftest.py b/tests/conftest.py index cc3a8ba1..018ad4ed 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -146,9 +146,8 @@ def _clear_otel_test_context() -> None: detach_failed = False for active in reversed(state.active_contexts): - try: - active.token.var.reset(active.token) - except (RuntimeError, ValueError): + otel_context.detach(active.token) + if otel_context.get_current() is not active.previous_context: detach_failed = True if detach_failed: diff --git a/tests/test_logger_otel_context.py b/tests/test_logger_otel_context.py index a9481b26..01228f5d 100644 --- a/tests/test_logger_otel_context.py +++ b/tests/test_logger_otel_context.py @@ -1,6 +1,8 @@ import asyncio import atexit -from collections.abc import Callable, Generator +import inspect +import uuid +from collections.abc import Callable, Generator, ItemsView from unittest.mock import Mock import pytest @@ -8,7 +10,32 @@ from opentelemetry.trace import SpanContext, TraceFlags from splunk_ao.logger import SplunkAOLogger -from splunk_ao.logger.logger import _otel_context_state +from splunk_ao.logger.logger import OtelIds, _otel_context_state + + +class ItemsForbiddenIdentityMap(dict[uuid.UUID, OtelIds]): + """Identity map that fails if subtree cleanup attempts a full scan.""" + + def items(self) -> ItemsView[uuid.UUID, OtelIds]: + raise AssertionError("subtree cleanup must not scan unrelated OTel identities") + + +class FailingChildSet(set[uuid.UUID]): + """Child set that simulates a partially failed structural insertion.""" + + def add(self, element: uuid.UUID) -> None: + raise RuntimeError("structural insertion failure") + + +def assert_otel_indexes_consistent(logger: SplunkAOLogger) -> None: + """Assert that both structural indexes describe the same live edges.""" + expected_children: dict[uuid.UUID, set[uuid.UUID]] = {} + for child_id, parent_id in logger._otel_parent_by_child.items(): + assert child_id in logger._otel_ids + assert parent_id in logger._otel_ids + expected_children.setdefault(parent_id, set()).add(child_id) + + assert logger._otel_children_by_parent == expected_children @pytest.fixture(autouse=True) @@ -237,6 +264,204 @@ def test_deep_parent_chain_restores_each_open_context(make_logger: Callable[[], assert not trace.get_current_span().get_span_context().is_valid +def test_context_reconciliation_detaches_in_lifo_order_and_verifies_previous_context( + make_logger: Callable[[], SplunkAOLogger], monkeypatch: pytest.MonkeyPatch +) -> None: + # Given: two managed contexts with their exact pre-attach contexts recorded. + logger = make_logger() + root = logger.start_trace(input="q") + logger.add_workflow_span(input="workflow") + state = _otel_context_state.get() + assert state is not None + expected = list(reversed(state.active_contexts)) + actual_detach = context.detach + detached: list[uuid.UUID] = [] + restored_exact_context: list[bool] = [] + + def recording_detach(token: object) -> None: + active = next(item for item in expected if item.token is token) + actual_detach(token) + detached.append(active.step_id) + restored_exact_context.append(context.get_current() is active.previous_context) + + # When: the child concludes and the logger reconciles back to the root. + with monkeypatch.context() as patch: + patch.setattr(context, "detach", recording_detach) + logger.conclude(output="workflow-output") + + # Then: public detach is LIFO, every reset is verified, and the root remains current. + assert detached == [active.step_id for active in expected] + assert restored_exact_context == [True, True] + assert trace.get_current_span().get_span_context() == logger._otel_ids[root.id].span_context + logger.conclude(output="trace-output") + + +def test_context_reconciliation_recovers_when_public_detach_is_a_silent_noop( + make_logger: Callable[[], SplunkAOLogger], monkeypatch: pytest.MonkeyPatch +) -> None: + # Given: an open root and child, while the public detach API silently does nothing. + logger = make_logger() + root = logger.start_trace(input="q") + logger.add_workflow_span(input="workflow") + detach = Mock(return_value=None) + + # When: the child concludes and both no-op detach calls must be detected. + with monkeypatch.context() as patch: + patch.setattr(context, "detach", detach) + logger.conclude(output="workflow-output") + + # Then: the recorded base is restored before rebuilding the desired root context. + assert detach.call_count == 2 + assert trace.get_current_span().get_span_context() == logger._otel_ids[root.id].span_context + state = _otel_context_state.get() + assert state is not None + assert [active.step_id for active in state.active_contexts if active.logger_id == id(logger)] == [root.id] + logger.conclude(output="trace-output") + assert not trace.get_current_span().get_span_context().is_valid + + +def test_context_reconciliation_uses_no_contextvar_token_internals() -> None: + # Given/When: the context reconciliation implementation is inspected directly. + source = inspect.getsource(SplunkAOLogger._sync_otel_context_impl) + + # Then: only the public OTel detach API is used. + assert "otel_context.detach(" in source + assert "token.var" not in source + + +@pytest.mark.parametrize(("shape", "size"), [("deep", 1_500), ("wide", 1_000)]) +def test_identity_tree_cleanup_handles_deep_and_wide_subtrees_without_recursion( + make_logger: Callable[[], SplunkAOLogger], shape: str, size: int +) -> None: + # Given: a structural identity index much deeper or wider than normal traces. + logger = make_logger() + root = Mock(id=uuid.uuid4()) + assert logger._record_otel_ids(root) is not None + parent = root + for _ in range(size): + child = Mock(id=uuid.uuid4()) + assert logger._record_otel_ids(child, parent_step=root if shape == "wide" else parent) is not None + parent = child + assert_otel_indexes_consistent(logger) + + # When: the root identity subtree is discarded iteratively. + logger._discard_otel_identity_tree(root.id) + + # Then: every node and edge is removed without recursive traversal. + assert logger._otel_ids == {} + assert logger._otel_children_by_parent == {} + assert logger._otel_parent_by_child == {} + + +def test_partial_identity_cleanup_preserves_siblings_and_never_scans_unrelated_ids( + make_logger: Callable[[], SplunkAOLogger], +) -> None: + # Given: a selected subtree, a sibling, and a completely unrelated request tree. + logger = make_logger() + root = Mock(id=uuid.uuid4()) + selected = Mock(id=uuid.uuid4()) + selected_child = Mock(id=uuid.uuid4()) + sibling = Mock(id=uuid.uuid4()) + unrelated_root = Mock(id=uuid.uuid4()) + unrelated_child = Mock(id=uuid.uuid4()) + for step, parent in ( + (root, None), + (selected, root), + (selected_child, selected), + (sibling, root), + (unrelated_root, None), + (unrelated_child, unrelated_root), + ): + assert logger._record_otel_ids(step, parent_step=parent) is not None + logger._otel_ids = ItemsForbiddenIdentityMap(logger._otel_ids) + + # When: only the selected subtree is discarded. + logger._discard_otel_identity_tree(selected.id) + + # Then: cleanup follows indexed UUID edges and leaves unrelated identities untouched. + assert set(logger._otel_ids) == {root.id, sibling.id, unrelated_root.id, unrelated_child.id} + assert logger._otel_children_by_parent == {root.id: {sibling.id}, unrelated_root.id: {unrelated_child.id}} + assert_otel_indexes_consistent(logger) + + +def test_partial_identity_insertion_rolls_back_and_does_not_interrupt_logging( + make_logger: Callable[[], SplunkAOLogger], +) -> None: + # Given: a root whose structural child insertion fails after identity creation. + logger = make_logger() + root = logger.start_trace(input="q") + logger._otel_children_by_parent[root.id] = FailingChildSet() + + # When: a proprietary child is created normally. + workflow = logger.add_workflow_span(input="workflow") + + # Then: proprietary logging continues and all partial OTel index mutations are rolled back. + assert logger.current_parent() is workflow + assert workflow in root.spans + assert workflow.id not in logger._otel_ids + assert workflow.id not in logger._otel_parent_by_child + assert root.id not in logger._otel_children_by_parent + assert_otel_indexes_consistent(logger) + logger.conclude(output="workflow-output") + logger.conclude(output="trace-output") + + +def test_completed_leaf_release_unlinks_identity_edge_before_parent_cleanup() -> None: + # Given: a normal OTLP logger with one active trace envelope. + sink = Mock() + sink.force_flush.return_value = True + logger = SplunkAOLogger(project_id="project-id", agent_stream_id="agent-stream-id", _sink=sink) + root = logger.start_trace(input="q") + + # When: a completed leaf is emitted immediately. + leaf = logger.add_llm_span(input="prompt", output="answer", model="model") + + # Then: its identity and parent edge are removed in O(1), and parent cleanup empties all indexes. + assert leaf.id not in logger._otel_ids + assert leaf.id not in logger._otel_parent_by_child + assert root.id not in logger._otel_children_by_parent + assert_otel_indexes_consistent(logger) + logger.conclude(output="trace-output") + assert logger._otel_ids == {} + assert logger._otel_children_by_parent == {} + assert logger._otel_parent_by_child == {} + logger.terminate() + + +def test_reset_flush_and_termination_clear_visible_context_and_structural_indexes( + make_logger: Callable[[], SplunkAOLogger], +) -> None: + # Given: an open trace with a nested proprietary parent. + logger = make_logger() + logger.start_trace(input="reset") + logger.add_workflow_span(input="workflow") + + # When/Then: reset removes that request subtree and its visible managed context. + logger.reset_parent_tracking() + assert not trace.get_current_span().get_span_context().is_valid + assert logger._otel_ids == {} + assert logger._otel_children_by_parent == {} + assert logger._otel_parent_by_child == {} + + # When/Then: legacy-hook flush concludes and clears a new subtree as well. + logger.start_trace(input="flush") + logger.add_workflow_span(input="workflow") + logger.flush() + assert not trace.get_current_span().get_span_context().is_valid + assert logger._otel_ids == {} + assert logger._otel_children_by_parent == {} + assert logger._otel_parent_by_child == {} + + # When/Then: termination clears unfinished identities without exporting them. + logger.start_trace(input="terminate") + logger.add_workflow_span(input="workflow") + logger.terminate() + assert not trace.get_current_span().get_span_context().is_valid + assert logger._otel_ids == {} + assert logger._otel_children_by_parent == {} + assert logger._otel_parent_by_child == {} + + def test_single_llm_trace_assigns_both_contexts_and_cleans_up( make_logger: Callable[[], SplunkAOLogger], monkeypatch: pytest.MonkeyPatch ) -> None: From 48ff2c9bc734cd9b25eaefb8439876cce76e9706 Mon Sep 17 00:00:00 2001 From: pradystar Date: Wed, 5 Aug 2026 23:45:38 -0700 Subject: [PATCH 2/2] fix(logger): recover from OTel identity removal failures --- src/splunk_ao/logger/logger.py | 16 +++++++++++++--- tests/test_logger_otel_context.py | 29 +++++++++++++++++++++++++++++ 2 files changed, 42 insertions(+), 3 deletions(-) diff --git a/src/splunk_ao/logger/logger.py b/src/splunk_ao/logger/logger.py index 7ff64c44..edf828e3 100644 --- a/src/splunk_ao/logger/logger.py +++ b/src/splunk_ao/logger/logger.py @@ -648,9 +648,19 @@ def _remove_otel_identity(self, step_id: uuid.UUID) -> None: try: siblings = self._otel_children_by_parent.get(parent_step_id) if siblings is not None: - siblings.discard(step_id) - if not siblings: - self._otel_children_by_parent.pop(parent_step_id, None) + try: + siblings.discard(step_id) + except Exception as exc: + first_error = first_error or exc + remaining_siblings = set(siblings) + remaining_siblings.discard(step_id) + if remaining_siblings: + self._otel_children_by_parent[parent_step_id] = remaining_siblings + else: + self._otel_children_by_parent.pop(parent_step_id, None) + else: + if not siblings: + self._otel_children_by_parent.pop(parent_step_id, None) except Exception as exc: first_error = first_error or exc diff --git a/tests/test_logger_otel_context.py b/tests/test_logger_otel_context.py index 01228f5d..1e8eadbf 100644 --- a/tests/test_logger_otel_context.py +++ b/tests/test_logger_otel_context.py @@ -27,6 +27,13 @@ def add(self, element: uuid.UUID) -> None: raise RuntimeError("structural insertion failure") +class FailingDiscardSet(set[uuid.UUID]): + """Child set that simulates a partially failed structural removal.""" + + def discard(self, element: uuid.UUID) -> None: + raise RuntimeError("structural removal failure") + + def assert_otel_indexes_consistent(logger: SplunkAOLogger) -> None: """Assert that both structural indexes describe the same live edges.""" expected_children: dict[uuid.UUID, set[uuid.UUID]] = {} @@ -406,6 +413,28 @@ def test_partial_identity_insertion_rolls_back_and_does_not_interrupt_logging( logger.conclude(output="trace-output") +def test_partial_identity_removal_rebuilds_forward_index_and_does_not_interrupt_logging( + make_logger: Callable[[], SplunkAOLogger], +) -> None: + # Given: a live parent/child edge whose forward child set fails during removal. + logger = make_logger() + root = logger.start_trace(input="q") + workflow = logger.add_workflow_span(input="workflow") + logger._otel_children_by_parent[root.id] = FailingDiscardSet({workflow.id}) + + # When: normal conclusion removes the child identity. + parent = logger.conclude(output="workflow-output") + + # Then: logging continues and the stale forward edge is removed by reconstruction. + assert parent is root + assert logger.current_parent() is root + assert workflow.id not in logger._otel_ids + assert workflow.id not in logger._otel_parent_by_child + assert root.id not in logger._otel_children_by_parent + assert_otel_indexes_consistent(logger) + logger.conclude(output="trace-output") + + def test_completed_leaf_release_unlinks_identity_edge_before_parent_cleanup() -> None: # Given: a normal OTLP logger with one active trace envelope. sink = Mock()