From 468f27319e0de8e8ce4c84fcb821e0bdaa66834b Mon Sep 17 00:00:00 2001
From: pradystar <pradeepn@splunk.com>
Date: Thu, 22 Oct 2026 14:38:53 -0700
Subject: [PATCH 1/2] fix(tracing): restore caller context after async callback
 completion (SAO-18234)

- Prevent stale parents and missing traces across repeated invocations
- Capture terminal LangGraph replies without modifying checkpoint history
- Add lifecycle and integration regression coverage
---
 CHANGELOG.md                                  |   7 +
 splunk-ao-adk/tests/test_observer.py          |  38 ++
 .../handlers/langchain/async_handler.py       |  13 +-
 src/splunk_ao/handlers/langchain/handler.py   |  13 +-
 src/splunk_ao/handlers/langchain/utils.py     |  20 +-
 src/splunk_ao/logger/logger.py                |  85 ++++-
 tests/test_handler_context_completion.py      | 354 ++++++++++++++++++
 tests/test_openai_agents.py                   |  34 ++
 8 files changed, 541 insertions(+), 23 deletions(-)
 create mode 100644 tests/test_handler_context_completion.py

diff --git a/CHANGELOG.md b/CHANGELOG.md
index fec4614c..d779612a 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Fixed
+
+- Restore completed SDK operations to their caller context when framework callbacks finish in a copied async task,
+  preventing stale parents and missing telemetry on subsequent LangGraph invocations.
+- LangChain callbacks record the terminal response of a LangGraph root instead of accumulated checkpoint history,
+  without modifying graph state or reducing nested tool-node outputs.
+
 ## [0.6.0]
 
 ### Added
diff --git a/splunk-ao-adk/tests/test_observer.py b/splunk-ao-adk/tests/test_observer.py
index 1f0346a8..37f2fa6e 100644
--- a/splunk-ao-adk/tests/test_observer.py
+++ b/splunk-ao-adk/tests/test_observer.py
@@ -1,8 +1,11 @@
 """Unit tests for SplunkAOObserver extraction methods."""
 
+import asyncio
 from unittest.mock import MagicMock, patch
+from uuid import UUID
 
 import pytest
+from opentelemetry import trace
 
 from splunk_ao.logger import SplunkAOLogger
 from splunk_ao_adk.observer import SplunkAOObserver
@@ -16,6 +19,41 @@ def observer() -> SplunkAOObserver:
     return SplunkAOObserver(ingestion_hook=lambda r: None)
 
 
+@pytest.mark.asyncio
+async def test_observer_completion_in_copied_task_restores_caller_context(monkeypatch: pytest.MonkeyPatch) -> None:
+    # Given: ADK's normal exporter path uses the shared base handler and logger.
+    for name in ("SPLUNK_AO_REALM", "SPLUNK_AO_O11Y_TOKEN", "SPLUNK_AO_O11Y_API_TOKEN"):
+        monkeypatch.delenv(name, raising=False)
+    monkeypatch.setenv("SPLUNK_AO_CONSOLE_URL", "http://fake.test:8088")
+    monkeypatch.setenv("SPLUNK_AO_API_KEY", "api-1234567890")
+    sink = MagicMock()
+    sink.force_flush.return_value = True
+    logger = SplunkAOLogger(project_id="project-id", agent_stream_id="stream-id", _sink=sink)
+    baseline = trace.get_current_span().get_span_context()
+    try:
+        with patch("splunk_ao_adk.observer.splunk_ao_context.get_logger_instance", return_value=logger):
+            adapter = SplunkAOObserver()
+        for _ in range(2):
+            run_id = adapter.on_run_start(MockInvocationContext(), MockContent(text="question"))
+
+            async def finish(run: UUID = run_id) -> None:
+                adapter.on_run_end(run, "answer")
+
+            # When: completion runs in a copied context, as an async framework may do.
+            await asyncio.create_task(finish())
+
+            # Then: the original context is clear and the next invocation can emit.
+            assert trace.get_current_span().get_span_context() == baseline
+            assert logger.current_parent() is None
+        assert sink.emit.call_count == 2
+        first, second = [call.args[0] for call in sink.emit.call_args_list]
+        assert first.parent is None
+        assert second.parent is None
+        assert first.context.trace_id != second.context.trace_id
+    finally:
+        logger.terminate()
+
+
 class TestExtractInvocationMetadata:
     """Tests for _extract_invocation_metadata method."""
 
diff --git a/src/splunk_ao/handlers/langchain/async_handler.py b/src/splunk_ao/handlers/langchain/async_handler.py
index 2e30f728..1e40c926 100644
--- a/src/splunk_ao/handlers/langchain/async_handler.py
+++ b/src/splunk_ao/handlers/langchain/async_handler.py
@@ -7,7 +7,13 @@
 
 from splunk_ao.handlers.base_async_handler import SplunkAOAsyncBaseHandler
 from splunk_ao.handlers.langchain.handler import SplunkAOCallback
-from splunk_ao.handlers.langchain.utils import get_agent_name, is_agent_node, parse_llm_result, update_root_to_agent
+from splunk_ao.handlers.langchain.utils import (
+    get_agent_name,
+    is_agent_node,
+    normalize_langgraph_output,
+    parse_llm_result,
+    update_root_to_agent,
+)
 from splunk_ao.logger import SplunkAOLogger
 from splunk_ao.schema.handlers import NODE_TYPE
 from splunk_ao.schema.trace import TracesIngestRequest
@@ -81,6 +87,10 @@ async def on_chain_start(
         node_type: NODE_TYPE = "chain"
         node_name = SplunkAOCallback._get_node_name(node_type, serialized, kwargs)
 
+        kwargs["_langgraph_root"] = parent_run_id is None and (
+            node_name.lower() == "langgraph" or (kwargs.get("metadata") or {}).get("ls_integration") == "langgraph"
+        )
+
         # If the `name` is `LangGraph` or `Agent`, set the node type to `agent`.
         if is_agent_node(node_name):
             node_type = "agent"
@@ -105,6 +115,7 @@ async def on_chain_end(
         # The input is sent via kwargs in on_chain_end in async streaming mode
         if "inputs" in kwargs:
             kwargs["input"] = serialize_to_str(kwargs["inputs"])
+        outputs = normalize_langgraph_output(outputs, self._handler.get_node(run_id))
         await self._handler.async_end_node(run_id, output=serialize_to_str(outputs), status_code=200, **kwargs)
 
     async def on_agent_finish(self, finish: AgentFinish, *, run_id: UUID, **kwargs: Any) -> Any:
diff --git a/src/splunk_ao/handlers/langchain/handler.py b/src/splunk_ao/handlers/langchain/handler.py
index 25a8475f..23a6ebd2 100644
--- a/src/splunk_ao/handlers/langchain/handler.py
+++ b/src/splunk_ao/handlers/langchain/handler.py
@@ -6,7 +6,13 @@
 from uuid import UUID
 
 from splunk_ao.handlers.base_handler import SplunkAOBaseHandler
-from splunk_ao.handlers.langchain.utils import get_agent_name, is_agent_node, parse_llm_result, update_root_to_agent
+from splunk_ao.handlers.langchain.utils import (
+    get_agent_name,
+    is_agent_node,
+    normalize_langgraph_output,
+    parse_llm_result,
+    update_root_to_agent,
+)
 from splunk_ao.logger import SplunkAOLogger
 from splunk_ao.schema.handlers import LANGCHAIN_NODE_TYPE, NODE_TYPE
 from splunk_ao.schema.trace import TracesIngestRequest
@@ -107,6 +113,10 @@ def on_chain_start(
         node_type: NODE_TYPE = "chain"
         node_name = self._get_node_name(node_type, serialized, kwargs)
 
+        kwargs["_langgraph_root"] = parent_run_id is None and (
+            node_name.lower() == "langgraph" or (kwargs.get("metadata") or {}).get("ls_integration") == "langgraph"
+        )
+
         # If the `name` is `LangGraph` or `Agent`, set the node type to `agent`.
         if is_agent_node(node_name):
             node_type = "agent"
@@ -124,6 +134,7 @@ def on_chain_end(
         """Langchain callback when a chain ends."""
         # Convert UUID7 to UUID4 if needed
         run_id = convert_uuid_if_uuid7(run_id) or run_id
+        outputs = normalize_langgraph_output(outputs, self._handler.get_node(run_id))
         self._handler.end_node(run_id, output=serialize_to_str(outputs), status_code=200)
 
     def on_agent_finish(self, finish: AgentFinish, *, run_id: UUID, **kwargs: Any) -> Any:
diff --git a/src/splunk_ao/handlers/langchain/utils.py b/src/splunk_ao/handlers/langchain/utils.py
index 9d1d092d..bf9d0362 100644
--- a/src/splunk_ao/handlers/langchain/utils.py
+++ b/src/splunk_ao/handlers/langchain/utils.py
@@ -27,6 +27,20 @@ def is_agent_node(node_name: str) -> bool:
     return node_name.lower() in ["langgraph", "agent"]
 
 
+def normalize_langgraph_output(outputs: Any, node: Node | None) -> Any:
+    """Select the graph root's terminal response without changing checkpoint state."""
+    if node is None or not node.span_params.get("_langgraph_root") or not isinstance(outputs, dict):
+        return outputs
+    messages = outputs.get("messages")
+    if not isinstance(messages, list) or not messages:
+        return outputs
+    for message in reversed(messages):
+        role = message.get("role", message.get("type")) if isinstance(message, dict) else getattr(message, "type", None)
+        if role in ("assistant", "ai", "tool", "function"):
+            return {**outputs, "messages": [message]}
+    return outputs
+
+
 def update_root_to_agent(parent_run_id: UUID | None, metadata: dict[str, Any], parent_node: Node | None) -> None:
     """Update the parent node to be an agent if it is a root-level chain and has LangGraph metadata in the children.
 
@@ -40,8 +54,10 @@ def update_root_to_agent(parent_run_id: UUID | None, metadata: dict[str, Any], p
     # Only mark as agent if parent is a root-level chain (no grandparent)
     if parent_run_id and metadata and any(key.startswith("langgraph_") for key in metadata):
         # Only convert to agent if parent is a root-level chain (no parent of its own)
-        if parent_node and parent_node.node_type == "chain" and parent_node.parent_run_id is None:
-            parent_node.node_type = "agent"
+        if parent_node and parent_node.parent_run_id is None:
+            parent_node.span_params["_langgraph_root"] = True
+            if parent_node.node_type == "chain":
+                parent_node.node_type = "agent"
 
 
 @dataclass
diff --git a/src/splunk_ao/logger/logger.py b/src/splunk_ao/logger/logger.py
index 1c4241eb..829e46e1 100644
--- a/src/splunk_ao/logger/logger.py
+++ b/src/splunk_ao/logger/logger.py
@@ -100,7 +100,7 @@
 _otel_id_generator = RandomIdGenerator()
 
 
-@dataclass(frozen=True)
+@dataclass
 class OtelIds:
     """Stable OTel identity and parentage for one proprietary logger step."""
 
@@ -108,6 +108,29 @@ class OtelIds:
     parent_span_context: SpanContext | None
     exportable: bool
     session_id: str | None
+    completed: bool = False
+
+
+class _ActiveSpan(NonRecordingSpan):
+    """Resolve completed SDK activations to their caller, including in copied contexts.
+
+    Only this internal context view changes on completion; span identities and
+    exported ReadableSpans remain immutable. Copies of an execution context share
+    the operation's completion flag, unlike ContextVar assignments and tokens.
+    """
+
+    def __init__(self, ids: OtelIds, previous: otel_trace.Span) -> None:
+        super().__init__(ids.span_context)
+        self._ids = ids
+        self._previous = previous
+
+    def get_span_context(self) -> SpanContext:
+        span: otel_trace.Span = self
+        while isinstance(span, _ActiveSpan):
+            if not span._ids.completed:
+                return span._ids.span_context
+            span = span._previous
+        return span.get_span_context()
 
 
 @dataclass(frozen=True)
@@ -119,6 +142,7 @@ class ActiveOtelContext:
     span_context: SpanContext
     exportable: bool
     token: Token
+    ids: OtelIds
 
 
 @dataclass(frozen=True)
@@ -151,10 +175,10 @@ def _has_active_exportable_span_context() -> bool:
     if state is None or not state.active_contexts:
         return True
 
-    active = state.active_contexts[-1]
-    if active.span_context != current:
-        return True
-    return active.exportable
+    for active in reversed(state.active_contexts):
+        if active.span_context == current and not active.ids.completed:
+            return active.exportable
+    return True
 
 
 class SplunkAOLogger(TracesLogger):
@@ -370,6 +394,30 @@ def _set_current_parent(self, parent: StepWithChildSpans | None) -> None:
         super()._set_current_parent(parent)
         self._sync_otel_context(parent)
 
+    def current_parent(self) -> StepWithChildSpans | None:
+        """Skip operations completed in another execution context."""
+        parent = super().current_parent()
+        original = parent
+        state = _otel_context_state.get()
+        if parent is not None and state is not None:
+            completed = {
+                active.step_id
+                for active in state.active_contexts
+                if active.logger_id == id(self) and active.ids.completed
+            }
+            while parent is not None and parent.id in completed:
+                parent = parent._parent
+            if parent is not original:
+                # A caller may already have started unrelated native OTel work
+                # above a completed SDK activation. Keep that caller's context.
+                caller_context = (
+                    otel_context.get_current() if not isinstance(otel_trace.get_current_span(), _ActiveSpan) else None
+                )
+                self._set_current_parent(parent)
+                if caller_context is not None:
+                    otel_context.attach(caller_context)
+        return parent
+
     def reset_parent_tracking(self) -> None:
         """Clear proprietary and OTel tracking for the current request context."""
         current_parent = self.current_parent()
@@ -475,17 +523,12 @@ def _sync_otel_context_impl(self, current_parent: StepWithChildSpans | None) ->
                 return
 
         base_context = state.base_context if state is not None else otel_context.get_current()
-        remaining_contexts = tuple(active for active in active_contexts if active.logger_id != logger_id)
-        desired_contexts = tuple(
-            (logger_id, step_id, self._otel_ids[step_id].span_context, self._otel_ids[step_id].exportable)
-            for step_id in desired_step_ids
+        remaining_contexts = tuple(
+            active for active in active_contexts if active.logger_id != logger_id and not active.ids.completed
         )
+        desired_contexts = tuple((logger_id, step_id, self._otel_ids[step_id]) for step_id in desired_step_ids)
         contexts_to_restore = (
-            tuple(
-                (active.logger_id, active.step_id, active.span_context, active.exportable)
-                for active in remaining_contexts
-            )
-            + desired_contexts
+            tuple((active.logger_id, active.step_id, active.ids) for active in remaining_contexts) + desired_contexts
         )
 
         detach_failed = False
@@ -501,15 +544,16 @@ def _sync_otel_context_impl(self, current_parent: StepWithChildSpans | None) ->
             otel_context.attach(base_context)
 
         rebuilt_contexts: list[ActiveOtelContext] = []
-        for owner_id, step_id, span_context, exportable in contexts_to_restore:
-            ctx = otel_trace.set_span_in_context(NonRecordingSpan(span_context))
+        for owner_id, step_id, ids in contexts_to_restore:
+            ctx = otel_trace.set_span_in_context(_ActiveSpan(ids, otel_trace.get_current_span()))
             rebuilt_contexts.append(
                 ActiveOtelContext(
                     logger_id=owner_id,
                     step_id=step_id,
-                    span_context=span_context,
-                    exportable=exportable,
+                    span_context=ids.span_context,
+                    exportable=ids.exportable,
                     token=otel_context.attach(ctx),
+                    ids=ids,
                 )
             )
 
@@ -530,6 +574,8 @@ def _discard_otel_identity_tree(self, step_id: uuid.UUID) -> None:
         if ids is None:
             return
 
+        ids.completed = True
+
         child_ids = tuple(
             child_id
             for child_id, child_ids in self._otel_ids.items()
@@ -589,6 +635,7 @@ def _emit_and_release(self, finished_step: BaseStep) -> None:
         except Exception:
             self._logger.warning("Failed to emit completed step %s.", finished_step.id, exc_info=True)
         finally:
+            ids.completed = True
             self._pending_otel_steps.discard(finished_step.id)
             self._otel_ids.pop(finished_step.id, None)
 
@@ -646,7 +693,7 @@ def _activate_handler_step(self, step: BaseStep) -> HandlerStepContext | None:
         if ids is None:
             return None
         prior_context = otel_context.get_current()
-        active_context = otel_trace.set_span_in_context(NonRecordingSpan(ids.span_context), prior_context)
+        active_context = otel_trace.set_span_in_context(_ActiveSpan(ids, otel_trace.get_current_span()), prior_context)
         return HandlerStepContext(
             prior_context=prior_context, token=otel_context.attach(active_context), span_context=ids.span_context
         )
diff --git a/tests/test_handler_context_completion.py b/tests/test_handler_context_completion.py
new file mode 100644
index 00000000..feb26448
--- /dev/null
+++ b/tests/test_handler_context_completion.py
@@ -0,0 +1,354 @@
+"""Regressions for callback completion in a copied async execution context."""
+
+import asyncio
+import json
+from typing import Any
+from uuid import UUID, uuid4
+
+import pytest
+from langchain_core.callbacks.manager import AsyncCallbackManager
+from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
+from opentelemetry import baggage, context, trace
+from opentelemetry.sdk.trace import ReadableSpan
+from opentelemetry.trace import NonRecordingSpan, SpanContext, TraceFlags
+
+from splunk_ao import get_tracing_headers
+from splunk_ao.handlers.base_handler import SplunkAOBaseHandler
+from splunk_ao.handlers.langchain import SplunkAOAsyncCallback, SplunkAOCallback
+from splunk_ao.handlers.langchain.utils import normalize_langgraph_output
+from splunk_ao.logger import SplunkAOLogger
+from splunk_ao.schema.handlers import Node
+
+
+class RecordingSink:
+    def __init__(self) -> None:
+        self.spans: list[ReadableSpan] = []
+
+    def emit(self, span: ReadableSpan) -> None:
+        self.spans.append(span)
+
+    def force_flush(self) -> bool:
+        return True
+
+    def shutdown(self) -> None:
+        pass
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("callback_type", [SplunkAOCallback, SplunkAOAsyncCallback])
+@pytest.mark.parametrize("flush", [False, True])
+@pytest.mark.parametrize("ownership", ["independent", "external", "attached"])
+async def test_real_async_dispatch_restores_caller_and_allows_callback_reuse(
+    callback_type: Any, flush: bool, ownership: str
+) -> None:
+    # Given: real LangChain dispatch, including shielded end callbacks in copied contexts.
+    sink = RecordingSink()
+    logger = SplunkAOLogger(project_id="project-id", agent_stream_id="stream-id", _sink=sink)
+    caller_context = SpanContext(1234, 5678, is_remote=True, trace_flags=TraceFlags(TraceFlags.SAMPLED))
+    initial = baggage.set_baggage("test-key", "retained", context.Context())
+    if ownership == "external":
+        initial = trace.set_span_in_context(NonRecordingSpan(caller_context), initial)
+    token = context.attach(initial)
+    try:
+        if ownership == "attached":
+            logger.start_trace(input="caller")
+            logger.add_workflow_span(input="caller", name="caller")
+        caller = logger.current_parent()
+        baseline = trace.get_current_span().get_span_context()
+        callback = callback_type(
+            splunk_ao_logger=logger, start_new_trace=ownership != "attached", flush_on_chain_end=flush
+        )
+        manager = AsyncCallbackManager(handlers=[callback], inheritable_handlers=[callback])
+        roots = []
+        for turn in range(3):
+            run = await manager.on_chain_start({"name": "LangGraph"}, {"question": str(turn)})
+            root_context = trace.get_current_span().get_span_context()
+            child = await run.get_child().on_chain_start({"name": "child"}, {"question": str(turn)})
+
+            # When: each end is dispatched by LangChain's real shielded task.
+            await child.on_chain_end({"answer": str(turn)})
+            assert trace.get_current_span().get_span_context() == root_context
+            assert get_tracing_headers()["traceparent"].split("-")[2] == f"{root_context.span_id:016x}"
+            await run.on_chain_end({"answer": str(turn)})
+
+            # Then: the original task has its caller context, and exactly two spans completed.
+            assert trace.get_current_span().get_span_context() == baseline
+            assert logger.current_parent() is caller
+            assert baggage.get_baggage("test-key") == "retained"
+            assert not callback._handler._nodes
+            assert not callback._handler._active_steps
+            child_span, root_span = sink.spans[-2:]
+            assert child_span.parent == root_span.context
+            roots.append(root_span)
+        assert len(sink.spans) == 6
+        assert len({span.context.span_id for span in sink.spans}) == 6
+        if ownership == "independent":
+            assert len({span.context.trace_id for span in roots}) == 3
+            assert all(span.parent is None for span in roots)
+        else:
+            assert all(span.parent == baseline for span in roots)
+    finally:
+        logger.terminate()
+        context.detach(token)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("failed", [False, True])
+async def test_manual_completion_in_another_task_restores_context(failed: bool) -> None:
+    # Given: a manual operation whose conclusion may execute in a copied task.
+    logger = SplunkAOLogger(project_id="project-id", agent_stream_id="stream-id", _sink=RecordingSink())
+    baseline = trace.get_current_span().get_span_context()
+    try:
+        logger.start_trace(input="request")
+        logger.add_agent_span(input="request", name="agent")
+
+        # When: completion occurs in a child task, as framework callbacks may do.
+        async def finish() -> None:
+            logger.conclude(
+                output="failure" if failed else "done", status_code=500 if failed else 200, conclude_all=True
+            )
+
+        await asyncio.create_task(finish())
+
+        # Then: completed operations cannot stay current or parent the next request.
+        assert trace.get_current_span().get_span_context() == baseline
+        assert logger.current_parent() is None
+        logger.start_trace(input="next request")
+        assert logger.current_parent() is not None
+    finally:
+        logger.terminate()
+
+
+@pytest.mark.asyncio
+async def test_reconciling_completed_parent_preserves_new_caller_span() -> None:
+    # Given: an SDK operation completed in another task before the application starts new work.
+    logger = SplunkAOLogger(project_id="project-id", agent_stream_id="stream-id", _sink=RecordingSink())
+    logger.start_trace(input="request")
+    logger.add_agent_span(input="request")
+
+    async def finish() -> None:
+        logger.conclude(output="done", conclude_all=True)
+
+    await asyncio.create_task(finish())
+    caller = SpanContext(4321, 8765, is_remote=False, trace_flags=TraceFlags(TraceFlags.SAMPLED))
+    token = context.attach(trace.set_span_in_context(NonRecordingSpan(caller)))
+    try:
+        # When: a logger read reconciles an obsolete parent beneath a new external span.
+        assert logger.current_parent() is None
+
+        # Then: it leaves that application-owned span active and inherits it on the next invocation.
+        assert trace.get_current_span().get_span_context() == caller
+        operation = logger.start_trace(input="next")
+        assert logger._otel_ids[operation.id].parent_span_context == caller
+    finally:
+        logger.terminate()
+        context.detach(token)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("callback_type", [SplunkAOCallback, SplunkAOAsyncCallback])
+@pytest.mark.parametrize("attached", [False, True])
+async def test_checkpointed_root_exports_current_reply_without_changing_history(
+    callback_type: Any, attached: bool
+) -> None:
+    # Given: a new turn's input differs from the beginning of checkpointed output history.
+    sink = RecordingSink()
+    logger = SplunkAOLogger(project_id="project-id", agent_stream_id="stream-id", _sink=sink)
+    callback = callback_type(splunk_ao_logger=logger, start_new_trace=not attached, flush_on_chain_end=False)
+    manager = AsyncCallbackManager(handlers=[callback], metadata={"ls_integration": "langgraph"})
+    reply = AIMessage(content=[{"type": "text", "text": "current reply"}])
+    history = [
+        HumanMessage(content="old question"),
+        AIMessage(content="old reply"),
+        HumanMessage(content="new question"),
+        reply,
+    ]
+    output = {"messages": history, "other_state": "preserved"}
+    try:
+        if attached:
+            logger.start_trace(input="caller")
+            logger.add_workflow_span(input="caller")
+        caller = logger.current_parent()
+        run = await manager.on_chain_start({"name": "LangGraph"}, {"messages": [history[-2]]})
+
+        # When: a graph root completes with the full persisted state.
+        await run.on_chain_end(output)
+
+        # Then: its exported output is the current response, without changing graph memory.
+        root = sink.spans[-1]
+        messages = json.loads(root.attributes["gen_ai.output.messages"])
+        assert len(messages) == 1
+        assert messages[0]["role"] == "assistant"
+        assert messages[0]["parts"] == [{"type": "text", "content": "current reply"}]
+        assert output["messages"] is history
+        assert len(history) == 4
+        assert output["other_state"] == "preserved"
+        assert logger.current_parent() is caller
+    finally:
+        logger.terminate()
+
+
+@pytest.mark.parametrize("graph", [False, True])
+def test_non_root_tool_messages_are_not_reduced(graph: bool) -> None:
+    # Given: a tool node can return several results, even inside LangGraph.
+    sink = RecordingSink()
+    logger = SplunkAOLogger(project_id="project-id", agent_stream_id="stream-id", _sink=sink)
+    callback = SplunkAOCallback(splunk_ao_logger=logger, flush_on_chain_end=False)
+    root_id, node_id = uuid4(), uuid4()
+    results = [ToolMessage(content="one", tool_call_id="1"), ToolMessage(content="two", tool_call_id="2")]
+    try:
+        callback.on_chain_start({"name": "LangGraph" if graph else "plain chain"}, {}, run_id=root_id)
+        callback.on_chain_start({"name": "tools"}, {}, run_id=node_id, parent_run_id=root_id)
+
+        # When: a nested node returns multiple tool responses.
+        callback.on_chain_end({"messages": results}, run_id=node_id, parent_run_id=root_id)
+
+        # Then: both tool responses remain in its telemetry output.
+        messages = json.loads(sink.spans[0].attributes["gen_ai.output.messages"])
+        assert len(messages) == 2
+        assert [message["role"] for message in messages] == ["tool", "tool"]
+        callback.on_chain_end({"answer": "done"}, run_id=root_id)
+    finally:
+        logger.terminate()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("termination", ["error", "cancellation", "unfinished_child"])
+async def test_shielded_failure_restores_context_and_discards_unfinished_children(termination: str) -> None:
+    # Given: a root operation with an unfinished child in the originating task.
+    sink = RecordingSink()
+    logger = SplunkAOLogger(project_id="project-id", agent_stream_id="stream-id", _sink=sink)
+    callback = SplunkAOAsyncCallback(splunk_ao_logger=logger, flush_on_chain_end=False)
+    manager = AsyncCallbackManager(handlers=[callback], inheritable_handlers=[callback])
+    baseline = trace.get_current_span().get_span_context()
+    try:
+        run = await manager.on_chain_start({"name": "LangGraph"}, {"question": "request"})
+        await run.get_child().on_chain_start({"name": "unfinished child"}, {})
+
+        # When: the framework finishes or reports failure from its shielded task.
+        if termination == "unfinished_child":
+            await run.on_chain_end({"answer": "done"})
+        else:
+            error = asyncio.CancelledError() if termination == "cancellation" else ValueError("failure")
+            await run.on_chain_error(error)
+
+        # Then: the uncompleted child is discarded, and the root emits exactly once.
+        assert trace.get_current_span().get_span_context() == baseline
+        assert logger.current_parent() is None
+        assert len(sink.spans) == 1
+        assert not logger._otel_ids
+        assert not callback._handler._active_steps
+        next_run = await manager.on_chain_start({"name": "LangGraph"}, {})
+        await next_run.on_chain_end({"answer": "next"})
+        assert len(sink.spans) == 2
+        assert sink.spans[0].context.trace_id != sink.spans[1].context.trace_id
+    finally:
+        logger.terminate()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("sampled", [False, True])
+async def test_shared_base_handler_completion_preserves_external_parent_and_sampling(sampled: bool) -> None:
+    # Given: a non-LangChain adapter using the shared base handler under a remote parent.
+    sink = RecordingSink()
+    logger = SplunkAOLogger(project_id="project-id", agent_stream_id="stream-id", _sink=sink)
+    caller = SpanContext(1234, 5678, is_remote=True, trace_flags=TraceFlags(1 if sampled else 0))
+    token = context.attach(trace.set_span_in_context(NonRecordingSpan(caller)))
+    handler = SplunkAOBaseHandler(splunk_ao_logger=logger, flush_on_chain_end=False)
+    try:
+        for _ in range(2):
+            root_id, child_id = uuid4(), uuid4()
+            handler.start_node("agent", None, root_id, input="request", name="agent")
+            handler.start_node("tool", root_id, child_id, input="arguments", name="tool")
+
+            async def finish(child: UUID = child_id, root: UUID = root_id) -> None:
+                handler.end_node(child, output="tool result")
+                handler.end_node(root, output="answer")
+
+            # When: a framework ends its callbacks in another task.
+            await asyncio.create_task(finish())
+
+            # Then: caller ownership and remote sampling decisions are preserved.
+            assert trace.get_current_span().get_span_context() == caller
+            assert logger.current_parent() is None
+        assert len(sink.spans) == (4 if sampled else 0)
+        if sampled:
+            assert all(span.context.trace_id == caller.trace_id for span in sink.spans)
+            assert all(span.parent == caller for span in sink.spans if span.name == "invoke_agent agent")
+    finally:
+        logger.terminate()
+        context.detach(token)
+
+
+@pytest.mark.parametrize(
+    "terminal",
+    [
+        AIMessage(content=""),
+        AIMessage(content="", tool_calls=[{"name": "search", "args": {}, "id": "call"}]),
+        ToolMessage(content="direct result", tool_call_id="call"),
+    ],
+)
+def test_graph_root_normalization_preserves_terminal_message_content(terminal: Any) -> None:
+    node = Node(node_type="agent", run_id=uuid4(), span_params={"_langgraph_root": True})
+    outputs = {"messages": [HumanMessage(content="old"), AIMessage(content="old reply"), terminal]}
+
+    normalized = normalize_langgraph_output(outputs, node)
+
+    assert normalized["messages"] == [terminal]
+    assert normalized["messages"][0] is terminal
+    assert len(outputs["messages"]) == 3
+
+
+@pytest.mark.parametrize("outputs", [{"messages": []}, {"result": "arbitrary state"}, "", None])
+def test_graph_root_normalization_preserves_empty_and_generic_outputs(outputs: Any) -> None:
+    node = Node(node_type="agent", run_id=uuid4(), span_params={"_langgraph_root": True})
+    assert normalize_langgraph_output(outputs, node) is outputs
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("callback_type", [SplunkAOCallback, SplunkAOAsyncCallback])
+@pytest.mark.parametrize("checkpointed", [False, True])
+@pytest.mark.parametrize("flush", [False, True])
+async def test_real_graph_repeated_turns_preserve_memory_and_emit_each_reply(
+    callback_type: Any, checkpointed: bool, flush: bool
+) -> None:
+    # Given: the same real graph can run with or without persisted conversation history.
+    graph_module = pytest.importorskip("langgraph.graph")
+    memory_module = pytest.importorskip("langgraph.checkpoint.memory")
+    sink = RecordingSink()
+    logger = SplunkAOLogger(project_id="project-id", agent_stream_id="stream-id", _sink=sink)
+    callback = callback_type(splunk_ao_logger=logger, start_new_trace=True, flush_on_chain_end=flush)
+    baseline = trace.get_current_span().get_span_context()
+
+    async def assistant(state: dict) -> dict:
+        count = sum(isinstance(message, HumanMessage) for message in state["messages"])
+        return {"messages": [AIMessage(content=f"reply {count}")]}
+
+    builder = graph_module.StateGraph(graph_module.MessagesState)
+    builder.add_node("assistant", assistant)
+    builder.add_edge(graph_module.START, "assistant")
+    builder.add_edge("assistant", graph_module.END)
+    graph = builder.compile(checkpointer=memory_module.InMemorySaver() if checkpointed else None)
+    config = {"configurable": {"thread_id": "conversation"}, "callbacks": [callback]}
+    try:
+        for turn in range(1, 6):
+            before = len(sink.spans)
+
+            # When: one callback processes five independently invoked turns in the same thread.
+            result = await graph.ainvoke({"messages": [HumanMessage(content=f"question {turn}")]}, config=config)
+
+            # Then: state persists only with checkpoints, but every invocation gets its own trace.
+            count = turn if checkpointed else 1
+            assert len(result["messages"]) == count * 2
+            assert result["messages"][-1].content == f"reply {count}"
+            roots = [span for span in sink.spans[before:] if span.parent is None]
+            assert len(roots) == 1
+            output = json.loads(roots[0].attributes["gen_ai.output.messages"])
+            assert len(output) == 1
+            assert output[0]["parts"] == [{"type": "text", "content": f"reply {count}"}]
+            assert trace.get_current_span().get_span_context() == baseline
+            assert logger.current_parent() is None
+        roots = [span for span in sink.spans if span.parent is None]
+        assert len({root.context.trace_id for root in roots}) == 5
+    finally:
+        logger.terminate()
diff --git a/tests/test_openai_agents.py b/tests/test_openai_agents.py
index 725a0964..27d84871 100644
--- a/tests/test_openai_agents.py
+++ b/tests/test_openai_agents.py
@@ -20,6 +20,7 @@
     set_trace_processors,
 )
 from agents.tracing import ResponseSpanData, agent_span, generation_span, response_span
+from opentelemetry import trace as otel_trace
 from opentelemetry.sdk.trace import ReadableSpan
 from pydantic import BaseModel
 from pytest import MonkeyPatch, mark
@@ -464,6 +465,39 @@ def test_openai_agents_public_tracing_lifecycle_exports_and_cleans_state() -> No
         logger.terminate()
 
 
+@pytest.mark.asyncio
+async def test_openai_agents_completion_in_copied_task_releases_original_context() -> None:
+    # Given: the independent OpenAI Agents processor shares logger context helpers.
+    sink = RecordingSink()
+    logger = SplunkAOLogger(project_id="project-id", agent_stream_id="stream-id", _sink=sink)
+    processor = SplunkAOTracingProcessor(splunk_ao_logger=logger, flush_on_trace_end=False)
+    baseline = otel_trace.get_current_span().get_span_context()
+    set_trace_processors([processor])
+    try:
+        for turn in range(2):
+            framework_trace = agents_tracing.trace("Agent trace", trace_id=f"trace_copied_{turn}")
+            framework_trace.start()
+            root = agent_span("agent", span_id=f"span_copied_{turn}", parent=framework_trace)
+            root.start()
+
+            async def finish(agent: Any = root, owner: Any = framework_trace) -> None:
+                agent.finish()
+                owner.finish()
+
+            # When: framework completion happens in a copied execution context.
+            await asyncio.create_task(finish())
+
+            # Then: the original task restores its caller and the processor is reusable.
+            assert otel_trace.get_current_span().get_span_context() == baseline
+            assert logger.current_parent() is None
+            assert processor._trace_states == {}
+        assert len(sink.spans) == 2
+        assert sink.spans[0].context.trace_id != sink.spans[1].context.trace_id
+    finally:
+        set_trace_processors([])
+        logger.terminate()
+
+
 @pytest.mark.asyncio
 async def test_openai_agents_concurrent_public_traces_keep_state_isolated() -> None:
     sink = RecordingSink()

From 8c00acf252a6156c383835255942e3239aebc0fe Mon Sep 17 00:00:00 2001
From: pradystar <pradeepn@splunk.com>
Date: Sat, 24 Oct 2026 08:36:39 -0700
Subject: [PATCH 2/2] fix: preserve caller span recording and prevent stale
 turn outputs

---
 CHANGELOG.md                              |   4 +-
 src/splunk_ao/handlers/langchain/utils.py |   3 +
 src/splunk_ao/logger/logger.py            |  49 +++++++-
 tests/test_handler_context_completion.py  | 135 +++++++++++++++++++++-
 4 files changed, 181 insertions(+), 10 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index d779612a..cca5d34a 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -10,9 +10,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Fixed
 
 - Restore completed SDK operations to their caller context when framework callbacks finish in a copied async task,
-  preventing stale parents and missing telemetry on subsequent LangGraph invocations.
+  preserving caller span recording and preventing stale parents and missing telemetry on subsequent LangGraph invocations.
 - LangChain callbacks record the terminal response of a LangGraph root instead of accumulated checkpoint history,
-  without modifying graph state or reducing nested tool-node outputs.
+  without modifying graph state or reducing nested tool-node outputs. Turns without a response do not reuse earlier replies.
 
 ## [0.6.0]
 
diff --git a/src/splunk_ao/handlers/langchain/utils.py b/src/splunk_ao/handlers/langchain/utils.py
index bf9d0362..e809f7be 100644
--- a/src/splunk_ao/handlers/langchain/utils.py
+++ b/src/splunk_ao/handlers/langchain/utils.py
@@ -38,6 +38,9 @@ def normalize_langgraph_output(outputs: Any, node: Node | None) -> Any:
         role = message.get("role", message.get("type")) if isinstance(message, dict) else getattr(message, "type", None)
         if role in ("assistant", "ai", "tool", "function"):
             return {**outputs, "messages": [message]}
+        if role in ("human", "user"):
+            # A turn with no response must not reuse a reply from checkpoint history.
+            return {**outputs, "messages": []}
     return outputs
 
 
diff --git a/src/splunk_ao/logger/logger.py b/src/splunk_ao/logger/logger.py
index 829e46e1..af37c724 100644
--- a/src/splunk_ao/logger/logger.py
+++ b/src/splunk_ao/logger/logger.py
@@ -6,7 +6,7 @@
 import logging
 import time
 import uuid
-from collections.abc import Callable
+from collections.abc import Callable, Mapping
 from contextvars import ContextVar, Token
 from dataclasses import dataclass
 from datetime import datetime
@@ -16,6 +16,7 @@
 from opentelemetry import trace as otel_trace
 from opentelemetry.sdk.trace.id_generator import RandomIdGenerator
 from opentelemetry.trace import NonRecordingSpan, SpanContext, TraceFlags, TraceState
+from opentelemetry.util.types import Attributes, AttributeValue
 from pydantic import PrivateAttr
 
 from galileo_core.helpers.execution import async_run
@@ -111,7 +112,7 @@ class OtelIds:
     completed: bool = False
 
 
-class _ActiveSpan(NonRecordingSpan):
+class _ActiveSpan(otel_trace.Span):
     """Resolve completed SDK activations to their caller, including in copied contexts.
 
     Only this internal context view changes on completion; span identities and
@@ -120,17 +121,53 @@ class _ActiveSpan(NonRecordingSpan):
     """
 
     def __init__(self, ids: OtelIds, previous: otel_trace.Span) -> None:
-        super().__init__(ids.span_context)
+        self._span = NonRecordingSpan(ids.span_context)
         self._ids = ids
         self._previous = previous
 
-    def get_span_context(self) -> SpanContext:
+    def _resolve(self) -> otel_trace.Span:
         span: otel_trace.Span = self
         while isinstance(span, _ActiveSpan):
             if not span._ids.completed:
-                return span._ids.span_context
+                return span._span
             span = span._previous
-        return span.get_span_context()
+        return span
+
+    def get_span_context(self) -> SpanContext:
+        return self._resolve().get_span_context()
+
+    def is_recording(self) -> bool:
+        return self._resolve().is_recording()
+
+    def set_attribute(self, key: str, value: AttributeValue) -> None:
+        self._resolve().set_attribute(key, value)
+
+    def set_attributes(self, attributes: Mapping[str, AttributeValue]) -> None:
+        self._resolve().set_attributes(attributes)
+
+    def add_event(self, name: str, attributes: Attributes = None, timestamp: int | None = None) -> None:
+        self._resolve().add_event(name, attributes, timestamp)
+
+    def add_link(self, context: SpanContext, attributes: Attributes = None) -> None:
+        self._resolve().add_link(context, attributes)
+
+    def update_name(self, name: str) -> None:
+        self._resolve().update_name(name)
+
+    def set_status(self, status: otel_trace.Status | otel_trace.StatusCode, description: str | None = None) -> None:
+        self._resolve().set_status(status, description)
+
+    def record_exception(
+        self,
+        exception: BaseException,
+        attributes: Attributes = None,
+        timestamp: int | None = None,
+        escaped: bool = False,
+    ) -> None:
+        self._resolve().record_exception(exception, attributes, timestamp, escaped)
+
+    def end(self, end_time: int | None = None) -> None:
+        self._resolve().end(end_time)
 
 
 @dataclass(frozen=True)
diff --git a/tests/test_handler_context_completion.py b/tests/test_handler_context_completion.py
index feb26448..0edd78dd 100644
--- a/tests/test_handler_context_completion.py
+++ b/tests/test_handler_context_completion.py
@@ -9,8 +9,10 @@
 from langchain_core.callbacks.manager import AsyncCallbackManager
 from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
 from opentelemetry import baggage, context, trace
-from opentelemetry.sdk.trace import ReadableSpan
-from opentelemetry.trace import NonRecordingSpan, SpanContext, TraceFlags
+from opentelemetry.sdk.trace import ReadableSpan, TracerProvider
+from opentelemetry.sdk.trace.export import SimpleSpanProcessor
+from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
+from opentelemetry.trace import NonRecordingSpan, SpanContext, StatusCode, TraceFlags
 
 from splunk_ao import get_tracing_headers
 from splunk_ao.handlers.base_handler import SplunkAOBaseHandler
@@ -34,6 +36,63 @@ def shutdown(self) -> None:
         pass
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("callback_type", [SplunkAOCallback, SplunkAOAsyncCallback])
+async def test_copied_completion_preserves_caller_span_recording(callback_type: Any) -> None:
+    # Given: callbacks run beneath a real recording span on an application-owned provider.
+    provider = TracerProvider()
+    exporter = InMemorySpanExporter()
+    provider.add_span_processor(SimpleSpanProcessor(exporter))
+    tracer = provider.get_tracer(__name__)
+    sink = RecordingSink()
+    logger = SplunkAOLogger(project_id="project-id", agent_stream_id="stream-id", _sink=sink)
+    callback = callback_type(splunk_ao_logger=logger, flush_on_chain_end=False)
+    manager = AsyncCallbackManager(handlers=[callback], inheritable_handlers=[callback])
+    try:
+        with tracer.start_as_current_span("server", end_on_exit=False) as server:
+            run = await manager.on_chain_start({"name": "LangGraph"}, {})
+            child = await run.get_child().on_chain_start({"name": "child"}, {})
+            # Open SDK operations must not forward mutations to their caller.
+            assert not trace.get_current_span().is_recording()
+            trace.get_current_span().set_attribute("during_sdk", True)
+
+            # When: both end callbacks execute in LangChain's copied tasks.
+            await child.on_chain_end({"answer": "child"})
+            assert not trace.get_current_span().is_recording()
+            await run.on_chain_end({"answer": "done"})
+
+            # Then: the current span records to the caller before any logger accessor reconciles it.
+            current = trace.get_current_span()
+            assert current.get_span_context() == server.get_span_context()
+            assert current.is_recording()
+            current.set_attribute("after_sdk", True)
+            current.set_attributes({"batch_attribute": "retained"})
+            current.add_event("after_sdk", {"event_attribute": "retained"}, timestamp=123)
+            current.record_exception(ValueError("local failure"), escaped=True)
+            current.update_name("renamed server")
+            current.set_status(StatusCode.ERROR, "local status")
+            link = SpanContext(1234, 5678, is_remote=True)
+            current.add_link(link, {"link_attribute": "retained"})
+            with tracer.start_as_current_span("subsequent child") as subsequent:
+                assert subsequent.parent == server.get_span_context()
+            current.end()
+            assert not current.is_recording()
+            assert logger.current_parent() is None
+        recorded = next(span for span in exporter.get_finished_spans() if span.name == "renamed server")
+        assert recorded.attributes == {"after_sdk": True, "batch_attribute": "retained"}
+        assert recorded.status.status_code == StatusCode.ERROR
+        assert recorded.status.description == "local status"
+        assert [event.name for event in recorded.events] == ["after_sdk", "exception"]
+        assert recorded.events[0].timestamp == 123
+        assert recorded.events[0].attributes == {"event_attribute": "retained"}
+        assert recorded.events[1].attributes["exception.message"] == "local failure"
+        assert recorded.links[0].context == link
+        assert recorded.links[0].attributes == {"link_attribute": "retained"}
+    finally:
+        logger.terminate()
+        provider.shutdown()
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize("callback_type", [SplunkAOCallback, SplunkAOAsyncCallback])
 @pytest.mark.parametrize("flush", [False, True])
@@ -299,12 +358,84 @@ def test_graph_root_normalization_preserves_terminal_message_content(terminal: A
     assert len(outputs["messages"]) == 3
 
 
+@pytest.mark.parametrize("representation", ["objects", "type", "role"])
+@pytest.mark.parametrize("response", [False, True])
+def test_graph_root_normalization_stops_at_current_turn_input(representation: str, response: bool) -> None:
+    # Given: persisted history followed by a new input, optionally with a current response.
+    node = Node(node_type="agent", run_id=uuid4(), span_params={"_langgraph_root": True})
+    history = [HumanMessage(content="old question"), AIMessage(content="old reply"), HumanMessage(content="new input")]
+    if response:
+        history.append(AIMessage(content="new reply"))
+    messages: list[Any] = list(history)
+    if representation == "type":
+        messages = [{"type": message.type, "content": message.content} for message in history]
+    elif representation == "role":
+        messages = [
+            {"role": "user" if message.type == "human" else "assistant", "content": message.content}
+            for message in history
+        ]
+    outputs = {"messages": messages, "other_state": "retained"}
+
+    # When: the root finishes, including turns that produce no response.
+    normalized = normalize_langgraph_output(outputs, node)
+
+    # Then: output never crosses into an earlier turn or presents input as a response.
+    assert normalized["messages"] == ([messages[-1]] if response else [])
+    assert normalized["other_state"] == "retained"
+    assert outputs["messages"] is messages
+    assert len(messages) == (4 if response else 3)
+
+
 @pytest.mark.parametrize("outputs", [{"messages": []}, {"result": "arbitrary state"}, "", None])
 def test_graph_root_normalization_preserves_empty_and_generic_outputs(outputs: Any) -> None:
     node = Node(node_type="agent", run_id=uuid4(), span_params={"_langgraph_root": True})
     assert normalize_langgraph_output(outputs, node) is outputs
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("callback_type", [SplunkAOCallback, SplunkAOAsyncCallback])
+@pytest.mark.parametrize("interrupted", [False, True])
+async def test_real_graph_turn_without_response_does_not_export_previous_reply(
+    callback_type: Any, interrupted: bool
+) -> None:
+    # Given: a checkpointed conversation whose second turn ends without a response.
+    graph_module = pytest.importorskip("langgraph.graph")
+    memory_module = pytest.importorskip("langgraph.checkpoint.memory")
+    types_module = pytest.importorskip("langgraph.types")
+    sink = RecordingSink()
+    logger = SplunkAOLogger(project_id="project-id", agent_stream_id="stream-id", _sink=sink)
+    callback = callback_type(splunk_ao_logger=logger, flush_on_chain_end=False)
+
+    async def assistant(state: dict) -> dict:
+        if sum(isinstance(message, HumanMessage) for message in state["messages"]) == 1:
+            return {"messages": [AIMessage(content="old reply")]}
+        if interrupted:
+            types_module.interrupt("awaiting approval")
+        return {}
+
+    builder = graph_module.StateGraph(graph_module.MessagesState)
+    builder.add_node("assistant", assistant)
+    builder.add_edge(graph_module.START, "assistant")
+    builder.add_edge("assistant", graph_module.END)
+    graph = builder.compile(checkpointer=memory_module.InMemorySaver())
+    config = {"configurable": {"thread_id": "no-response"}, "callbacks": [callback]}
+    try:
+        await graph.ainvoke({"messages": [HumanMessage(content="old question")]}, config=config)
+        before = len(sink.spans)
+
+        # When: a terminal branch or interrupt produces no current-turn message.
+        result = await graph.ainvoke({"messages": [HumanMessage(content="new question")]}, config=config)
+
+        # Then: history remains intact, but root telemetry contains no stale reply or input-as-output.
+        roots = [span for span in sink.spans[before:] if span.parent is None]
+        assert len(roots) == 1
+        assert json.loads(roots[0].attributes["gen_ai.output.messages"]) == []
+        assert [message.content for message in result["messages"]] == ["old question", "old reply", "new question"]
+        assert logger.current_parent() is None
+    finally:
+        logger.terminate()
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize("callback_type", [SplunkAOCallback, SplunkAOAsyncCallback])
 @pytest.mark.parametrize("checkpointed", [False, True])