diff --git a/CHANGELOG.md b/CHANGELOG.md index 0237a062..fb759651 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,10 @@ 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, + 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. Turns without a response do not reuse earlier replies. - Nested `@log(span_type="tool")` calls now retain the calling tool as their parent and expose their own active tracing context during execution. Tool, workflow, and agent generators restore the consumer context between yields and keep their own outputs when interleaved. 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..f36755df 100644 --- a/src/splunk_ao/handlers/langchain/utils.py +++ b/src/splunk_ao/handlers/langchain/utils.py @@ -27,6 +27,28 @@ 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): + if isinstance(message, dict): + role = message.get("role", message.get("type")) + else: + role = getattr(message, "role", None) or getattr(message, "type", None) + if isinstance(role, str): + role = role.removesuffix("MessageChunk").lower() + 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 + + 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 +62,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..9853e295 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 @@ -100,7 +101,7 @@ _otel_id_generator = RandomIdGenerator() -@dataclass(frozen=True) +@dataclass class OtelIds: """Stable OTel identity and parentage for one proprietary logger step.""" @@ -108,6 +109,68 @@ class OtelIds: parent_span_context: SpanContext | None exportable: bool session_id: str | None + completed: bool = False + + +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 + 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: + self._span = NonRecordingSpan(ids.span_context) + self._ids = ids + # Copied-task completion can leave stale activations in the caller's context. + while isinstance(previous, _ActiveSpan) and previous._ids.completed: + previous = previous._previous + self._previous = previous + + def _resolve(self) -> otel_trace.Span: + span: otel_trace.Span = self + while isinstance(span, _ActiveSpan): + if not span._ids.completed: + return span._span + span = span._previous + 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) @@ -119,6 +182,7 @@ class ActiveOtelContext: span_context: SpanContext exportable: bool token: Token + ids: OtelIds @dataclass(frozen=True) @@ -151,10 +215,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 +434,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 +563,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 +584,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 +614,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 +675,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 +733,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..cfd4956c --- /dev/null +++ b/tests/test_handler_context_completion.py @@ -0,0 +1,507 @@ +"""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, AIMessageChunk, ChatMessage, HumanMessage, HumanMessageChunk, ToolMessage +from opentelemetry import baggage, context, trace +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 +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.logger.logger import _ActiveSpan +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]) +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]) +@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 = [] + first_turn_depth = None + 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 + activation = trace.get_current_span() + depth = 0 + while isinstance(activation, _ActiveSpan): + depth += 1 + activation = activation._previous + if first_turn_depth is None: + first_turn_depth = depth + assert depth == first_turn_depth + 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("representation", ["objects", "type", "role", "chunks", "chat"]) +@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 == "chunks": + messages = [ + HumanMessageChunk(content=message.content) + if isinstance(message, HumanMessage) + else AIMessageChunk(content=message.content) + for message in history + ] + elif representation == "chat": + messages = [ + ChatMessage(role="user" if isinstance(message, HumanMessage) else "assistant", content=message.content) + for message in history + ] + elif 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]) +@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()