From 3bbe45846ceeef4bec8418773d8f9b2e8b99cbb1 Mon Sep 17 00:00:00 2001 From: pradystar Date: Thu, 13 Aug 2026 09:15:34 -0700 Subject: [PATCH 01/13] feat(tracing): migrate SDK propagation to W3C context --- CHANGELOG.md | 26 + src/splunk_ao/__init__.py | 3 +- src/splunk_ao/constants/__init__.py | 6 +- src/splunk_ao/constants/tracing.py | 8 - src/splunk_ao/decorator.py | 43 -- src/splunk_ao/handlers/base_async_handler.py | 51 +- src/splunk_ao/handlers/base_handler.py | 212 ++++++- .../handlers/openai_agents/handler.py | 129 +++- src/splunk_ao/logger/logger.py | 563 ++---------------- src/splunk_ao/logger/task_handler.py | 200 ------- src/splunk_ao/middleware/tracing.py | 200 +------ src/splunk_ao/tracing.py | 65 +- src/splunk_ao/utils/env_helpers.py | 5 +- src/splunk_ao/utils/singleton.py | 20 +- tests/conftest.py | 24 - tests/test_async_base_handler.py | 55 +- tests/test_base_handler.py | 88 ++- tests/test_decorator_distributed.py | 36 +- tests/test_logger_distributed.py | 120 +++- tests/test_logger_otel_context.py | 23 +- tests/test_middleware_tracing.py | 365 ++++-------- tests/test_openai_agents.py | 53 +- tests/test_otel_native_paths.py | 101 +++- tests/test_thread_pool_task_handler.py | 352 ----------- tests/test_tracing.py | 121 ++++ tests/testutils/setup.py | 172 ------ 26 files changed, 1099 insertions(+), 1942 deletions(-) delete mode 100644 src/splunk_ao/constants/tracing.py delete mode 100644 src/splunk_ao/logger/task_handler.py delete mode 100644 tests/test_thread_pool_task_handler.py create mode 100644 tests/test_tracing.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 1b750ee4..a86d3e8f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,32 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Added + +- Added module-level `get_tracing_headers()` and + `extract_tracing_context()` helpers for standard W3C `traceparent` and + `tracestate` propagation. `TracingMiddleware` now extracts and scopes that + OpenTelemetry context for Starlette/FastAPI requests. + +### Changed + +- The temporarily retained `mode="batch"` and `mode="distributed"` values now + use identical scheduled OTLP batch export and W3C propagation. Concluding an + operation ends and queues it; a per-operation `flush()` is not required. +- Callback handlers keep a live real root during framework execution so + outbound W3C context uses the same identity and visible hierarchy that is + exported at commit. + +### Removed + +- **Breaking:** Removed `trace_id=` and `span_id=` from `SplunkAOLogger`, the + logger-level proprietary `get_tracing_headers()` method, and custom + `Splunk-AO-Trace-ID` / `Splunk-AO-Parent-ID` continuation. Use the new + module-level W3C helpers instead. +- Removed the obsolete distributed-only REST streaming worker and task queue; + normal telemetry in both retained modes uses the existing + `BatchSpanProcessor` path. + ## [0.2.1] - 2026-08-07 ### Fixed diff --git a/src/splunk_ao/__init__.py b/src/splunk_ao/__init__.py index 0582d1fc..d850956d 100644 --- a/src/splunk_ao/__init__.py +++ b/src/splunk_ao/__init__.py @@ -80,7 +80,7 @@ SplunkAOFutureError, ValidationError, ) -from splunk_ao.tracing import get_tracing_headers +from splunk_ao.tracing import extract_tracing_context, get_tracing_headers from splunk_ao.types import MetricSpec from splunk_ao.utils.log_config import enable_console_logging @@ -167,6 +167,7 @@ "delete_annotation_queue_field", "delete_api_key", "enable_console_logging", + "extract_tracing_context", "get_agent_control_target", "get_annotation_queue", "get_annotation_queue_records", diff --git a/src/splunk_ao/constants/__init__.py b/src/splunk_ao/constants/__init__.py index c8c96f7b..4040a835 100644 --- a/src/splunk_ao/constants/__init__.py +++ b/src/splunk_ao/constants/__init__.py @@ -7,18 +7,14 @@ DEFAULT_API_URL = "https://api.galileo.ai/" DEFAULT_CONSOLE_URL = "https://app.galileo.ai/" -# HTTP header prefix for all Splunk AO headers -SPLUNK_AO_HEADER_PREFIX = "Splunk-AO" - # Type definitions LoggerModeType = Literal["batch", "distributed"] __all__ = ( + "DEFAULT_AGENT_STREAM_NAME", "DEFAULT_API_URL", "DEFAULT_CONSOLE_URL", - "DEFAULT_AGENT_STREAM_NAME", "DEFAULT_MODE", "DEFAULT_PROJECT_NAME", - "SPLUNK_AO_HEADER_PREFIX", "LoggerModeType", ) diff --git a/src/splunk_ao/constants/tracing.py b/src/splunk_ao/constants/tracing.py deleted file mode 100644 index c62539bb..00000000 --- a/src/splunk_ao/constants/tracing.py +++ /dev/null @@ -1,8 +0,0 @@ -"""Constants for distributed tracing.""" - -from splunk_ao.constants import SPLUNK_AO_HEADER_PREFIX - -# HTTP header names for propagating distributed tracing context -# These headers follow the pattern of namespaced custom headers (Splunk-AO-*) -TRACE_ID_HEADER = f"{SPLUNK_AO_HEADER_PREFIX}-Trace-ID" -PARENT_ID_HEADER = f"{SPLUNK_AO_HEADER_PREFIX}-Parent-ID" diff --git a/src/splunk_ao/decorator.py b/src/splunk_ao/decorator.py index fb02472b..ab56f31e 100644 --- a/src/splunk_ao/decorator.py +++ b/src/splunk_ao/decorator.py @@ -61,7 +61,6 @@ def call_llm(prompt, temperature=0.7): from galileo_core.schemas.logging.trace import Trace from splunk_ao.constants import LoggerModeType from splunk_ao.logger import SplunkAOLogger -from splunk_ao.logger.logger import STUB_TRACE_NAME from splunk_ao.schema.content_blocks import is_content_block_list from splunk_ao.schema.datasets import DatasetRecord from splunk_ao.schema.metrics import LocalMetricConfig @@ -94,10 +93,6 @@ def call_llm(prompt, temperature=0.7): _mode_context: ContextVar[LoggerModeType | None] = ContextVar("mode_context", default=None) _session_id_context: ContextVar[str | None] = ContextVar("session_id_context", default=None) -# Distributed tracing context variables (for middleware) -_trace_id_context: ContextVar[str | None] = ContextVar("trace_id_context", default=None) -_parent_id_context: ContextVar[str | None] = ContextVar("parent_id_context", default=None) - # Context variables for dataset fields (ground truth/reference output) # These allow setting ground truth data that will be attached to all spans # created within the context, enabling scorers that require reference output. @@ -1024,34 +1019,6 @@ def _handle_call_result( status_code = span_params.get("status_code") logger.conclude(output=output, duration_ns=span_params["duration_ns"], status_code=status_code) - # In distributed mode, update parent trace output after concluding a top-level workflow - # This ensures the trace shows the latest workflow's output (last workflow wins) - # Skip stub traces (created from distributed tracing headers - they're managed by the client) - if logger.mode == "distributed" and not stack: - current_parent = logger.current_parent() - if current_parent is not None and isinstance(current_parent, Trace): - is_stub_trace = current_parent.name == STUB_TRACE_NAME - - if not is_stub_trace: - # _coerce_output preserves str and List[ContentBlock], - # serializes everything else (Message, List[Document], etc.) to JSON string. - if output is not None: - current_parent.output = SplunkAOLogger._coerce_output(output) - if redacted_output is not None: - current_parent.redacted_output = SplunkAOLogger._coerce_output(redacted_output) - - # Update trace duration - # Note: In distributed mode, trace.created_at may be set by the server - # Using max() to ensure parent is never shorter than its children. - if current_parent.created_at: - elapsed_ns = convert_time_delta_to_ns(_get_timestamp() - current_parent.created_at) - workflow_ns = span_params.get("duration_ns", 0) - prev_ns = current_parent.metrics.duration_ns or 0 - current_parent.metrics.duration_ns = max(elapsed_ns, workflow_ns, prev_ns) - - if status_code is not None: - current_parent.status_code = status_code - else: # Non-concludable spans (llm, tool, retriever) are added to the parent span_methods = {"llm": "add_llm_span", "tool": "add_tool_span", "retriever": "add_retriever_span"} @@ -1234,12 +1201,6 @@ def get_logger_instance( "experiment_id": experiment_id or _experiment_id_context.get(), "mode": _get_mode_or_default(mode) if mode is not None else _mode_context.get(), } - trace_id_from_context = _trace_id_context.get() - span_id_from_context = _parent_id_context.get() - if trace_id_from_context: - kwargs["trace_id"] = trace_id_from_context - if span_id_from_context: - kwargs["span_id"] = span_id_from_context if ingestion_hook is not None: kwargs["ingestion_hook"] = ingestion_hook @@ -1374,10 +1335,6 @@ def reset(self) -> None: _span_stack_context.set([]) _trace_context.set(None) _session_id_context.set(None) - # Reset distributed tracing context - _trace_id_context.set(None) - _parent_id_context.set(None) - # Clear all stacks _get_or_init_list(_project_stack).clear() _get_or_init_list(_agent_stream_stack).clear() diff --git a/src/splunk_ao/handlers/base_async_handler.py b/src/splunk_ao/handlers/base_async_handler.py index 19d0f91f..68053a96 100644 --- a/src/splunk_ao/handlers/base_async_handler.py +++ b/src/splunk_ao/handlers/base_async_handler.py @@ -21,7 +21,9 @@ class SplunkAOAsyncBaseHandler(SplunkAOBaseHandler): _nodes : dict[UUID, Node] A dictionary of nodes, where the key is the run_id and the value is the node. _start_new_trace : bool - Whether to start a new trace when a chain starts. Set this to `False` to continue using the current trace. + Whether the handler owns and concludes a local trace lifecycle. An + active W3C parent is still inherited when this is ``True``. Set it to + ``False`` only when attaching to caller-owned active logger state. _flush_on_chain_end : bool Whether to flush the trace when a chain ends. """ @@ -42,34 +44,53 @@ async def async_commit(self) -> None: _logger.warning("Unable to add nodes to trace: Root node does not exist") return - owned_trace = None try: - if self._start_new_trace: - owned_trace = self._splunk_ao_logger.start_trace( - input=serialize_to_str(root_node.span_params.get("input", "")), - name=root_node.span_params.get("name"), - metadata=root_node.span_params.get("metadata"), + if self._owned_root is not None: + self._update_owned_root(root_node) + self._log_node_children(root_node) + root_output = self._root_output(root_node) + self._splunk_ao_logger.conclude( + output=serialize_to_str(root_output), + duration_ns=root_node.span_params.get("duration_ns"), + status_code=root_node.span_params.get("status_code"), ) - - self.log_node_tree(root_node) - root_output = root_node.span_params.get("output", "") - - if self._start_new_trace: + if self._owned_trace is not None: + self._conclude_owned_trace( + self._owned_trace, + output=serialize_to_str(root_output), + status_code=root_node.span_params.get("status_code"), + ) + elif self._start_new_trace: + if self._owned_trace is None: + self._start_owned_root(root_node) + if self._owned_trace is None: + return + + self.log_node_tree(root_node) + root_output = self._root_output(root_node) self._conclude_owned_trace( - owned_trace, + self._owned_trace, output=serialize_to_str(root_output), status_code=root_node.span_params.get("status_code"), ) + elif self._has_reusable_caller_root(): + # Preserve the pre-DTA visible tree while leaving the live + # caller operation current and caller-owned. + self.log_node_tree(root_node) + else: + _logger.warning("Unable to commit async handler telemetry: no caller-owned active operation") if self._flush_on_chain_end: await self._splunk_ao_logger.async_flush() except Exception: - if owned_trace is not None: - self._conclude_owned_trace(owned_trace, output="", status_code=500) + self._conclude_owned_state_on_failure() _logger.warning("Failed to commit async handler telemetry", exc_info=True) finally: self._nodes.clear() self._root_node = None + self._owned_trace = None + self._owned_root = None + self._owned_parent = None async def async_end_node(self, run_id: UUID, **kwargs: Any) -> None: """ diff --git a/src/splunk_ao/handlers/base_handler.py b/src/splunk_ao/handlers/base_handler.py index c095e717..63c2d62c 100644 --- a/src/splunk_ao/handlers/base_handler.py +++ b/src/splunk_ao/handlers/base_handler.py @@ -8,6 +8,7 @@ from splunk_ao import splunk_ao_context from splunk_ao.logger import SplunkAOLogger from splunk_ao.schema.handlers import INTEGRATION, NODE_TYPE, Node +from splunk_ao.schema.logged import LoggedAgentSpan, LoggedWorkflowSpan from splunk_ao.schema.trace import TracesIngestRequest from splunk_ao.utils.serialization import convert_to_string_dict, serialize_to_str @@ -25,7 +26,9 @@ class SplunkAOBaseHandler: _nodes : dict[UUID, Node] A dictionary of nodes, where the key is the run_id and the value is the node. _start_new_trace : bool - Whether to start a new trace when a chain starts. Set this to `False` to continue using the current trace. + Whether the handler owns and concludes a local trace lifecycle. An + active W3C parent is still inherited when this is ``True``. Set it to + ``False`` only when attaching to caller-owned active logger state. _flush_on_chain_end : bool Whether to flush the trace when a chain ends. _root_node : Optional[Node] @@ -56,6 +59,9 @@ def __init__( self._nodes: dict[str, Node] = {} self._root_node: Node | None = None self._integration: INTEGRATION = integration + self._owned_trace: Any = None + self._owned_root: Any = None + self._owned_parent: Any = None def commit(self) -> None: """Commit the nodes to the trace using the Splunk AO Logger. Optionally flush the trace.""" @@ -73,36 +79,193 @@ def commit(self) -> None: _logger.warning("Unable to add nodes to trace: Root node does not exist") return - owned_trace = None try: - if self._start_new_trace: - owned_trace = self._splunk_ao_logger.start_trace( - input=SplunkAOLogger._coerce_output(root_node.span_params.get("input", "")), - name=root_node.span_params.get("name"), - metadata=root_node.span_params.get("metadata"), + if self._owned_root is not None: + self._update_owned_root(root_node) + self._log_node_children(root_node) + root_output = self._root_output(root_node) + self._splunk_ao_logger.conclude( + output=serialize_to_str(root_output), + duration_ns=root_node.span_params.get("duration_ns"), + status_code=root_node.span_params.get("status_code"), ) - - self.log_node_tree(root_node) - - # Conclude the trace with the root node's output - root_output = root_node.span_params.get("output", "") - - if self._start_new_trace: + if self._owned_trace is not None: + self._conclude_owned_trace( + self._owned_trace, + output=SplunkAOLogger._coerce_output(root_output), + status_code=root_node.span_params.get("status_code"), + ) + elif self._start_new_trace: + if self._owned_trace is None: + self._start_owned_root(root_node) + if self._owned_trace is None: + return + + # Leaf-only roots cannot remain active in the Path 1 logger: + # leaf creation is an immediate completion operation. Preserve + # the pre-DTA direct leaf topology at commit while the owned + # envelope retains the incoming context. + self.log_node_tree(root_node) + root_output = self._root_output(root_node) self._conclude_owned_trace( - owned_trace, + self._owned_trace, output=SplunkAOLogger._coerce_output(root_output), status_code=root_node.span_params.get("status_code"), ) + elif self._has_reusable_caller_root(): + # Preserve the pre-DTA visible tree: the buffered handler root + # remains a child of the caller-owned live operation. The + # handler concludes only that buffered subtree and returns to + # the same caller operation. + self.log_node_tree(root_node) + else: + _logger.warning("Unable to commit handler telemetry: no caller-owned active operation") if self._flush_on_chain_end: self._splunk_ao_logger.flush() except Exception: - if owned_trace is not None: - self._conclude_owned_trace(owned_trace, output="", status_code=500) + self._conclude_owned_state_on_failure() _logger.warning("Failed to commit handler telemetry", exc_info=True) finally: self._nodes.clear() self._root_node = None + self._owned_trace = None + self._owned_root = None + self._owned_parent = None + + def _start_owned_root(self, root_node: Node) -> None: + """Open the handler-owned envelope and real root at callback start.""" + if self._owned_root is not None: + return + + if not self._start_new_trace: + # Caller-owned mode deliberately leaves the caller's real + # operation current for outbound W3C propagation. The buffered + # handler root is reconstructed beneath it at commit. + self._owned_parent = self._splunk_ao_logger.current_parent() + return + + if self._owned_trace is None: + self._owned_trace = self._splunk_ao_logger.start_trace( + input=SplunkAOLogger._coerce_output(root_node.span_params.get("input", "")), + name=root_node.span_params.get("name"), + metadata=root_node.span_params.get("metadata"), + ) + parent = self._owned_trace + + if parent is None: + return + self._owned_parent = parent + + # A leaf root is still logged at commit. There is no open Path 1 leaf + # identity to keep active without changing its completion semantics. + if root_node.node_type not in ("agent", "chain", "workflow"): + return + + metadata = root_node.span_params.get("metadata") + if metadata is not None: + metadata = convert_to_string_dict(metadata) + step_number = self._step_number(metadata) + root_kwargs = { + "input": serialize_to_str(root_node.span_params.get("input", "")), + "name": root_node.span_params.get("name"), + "metadata": metadata, + "tags": root_node.span_params.get("tags"), + "created_at": root_node.span_params.get("created_at"), + "status_code": root_node.span_params.get("status_code"), + "step_number": step_number, + } + if root_node.node_type == "agent": + self._owned_root = self._splunk_ao_logger.add_agent_span(**root_kwargs) + else: + self._owned_root = self._splunk_ao_logger.add_workflow_span(**root_kwargs) + + if self._owned_root is None: + self._conclude_owned_state_on_failure() + self._owned_trace = None + self._owned_parent = None + + def _update_owned_root(self, root_node: Node) -> None: + """Apply final callback data to the live root before it is completed.""" + if self._owned_root is None: + return + self._sync_owned_root_kind(root_node) + self._owned_root.input = serialize_to_str(root_node.span_params.get("input", "")) + if root_node.span_params.get("name") is not None: + self._owned_root.name = root_node.span_params["name"] + if root_node.span_params.get("tags") is not None: + self._owned_root.tags = root_node.span_params["tags"] + metadata = root_node.span_params.get("metadata") + if metadata is not None: + converted_metadata = convert_to_string_dict(metadata) + self._owned_root.user_metadata = converted_metadata + self._owned_root.step_number = self._step_number(converted_metadata) + + @staticmethod + def _step_number(metadata: dict[str, str] | None) -> int | None: + """Read the LangGraph step number using the pre-DTA conversion rule.""" + if not metadata or not (value := metadata.get("langgraph_step")): + return None + try: + return int(value) + except Exception as exc: + _logger.warning(f"Invalid step number: {value}, exception raised {exc}") + return None + + def _sync_owned_root_kind(self, root_node: Node) -> None: + """Preserve a late LangGraph chain-to-agent classification. + + LangGraph identifies some roots from child metadata after the root + callback has started. Replace only the proprietary model object while + retaining its UUID and preassigned OTel identity, so outbound W3C + context and the eventually exported agent span remain identical. + """ + if root_node.node_type != "agent" or not isinstance(self._owned_root, LoggedWorkflowSpan): + return + if self._splunk_ao_logger.current_parent() is not self._owned_root: + return + + previous = self._owned_root + parent = previous._parent + if parent is None: + return + replacement = LoggedAgentSpan.model_validate(previous.model_dump(exclude={"type"})) + replacement._parent = parent + replacement.spans = previous.spans + parent.spans = [replacement if child is previous else child for child in parent.spans] + self._owned_root = replacement + self._splunk_ao_logger._set_current_parent(replacement) + + def _log_node_children(self, root_node: Node) -> None: + """Log buffered descendants beneath the already-active real root.""" + for child_id in root_node.children: + child_node = self._nodes.get(child_id) + if child_node is not None: + self.log_node_tree(child_node) + else: + _logger.warning(f"Child node {child_id} not found") + + def _root_output(self, root_node: Node) -> Any: + output = root_node.span_params.get("output", "") + if output or not root_node.children: + return output + last_child = self._nodes.get(root_node.children[-1]) + return last_child.span_params.get("output", "") if last_child is not None else "" + + def _has_reusable_caller_root(self) -> bool: + """Return whether start_new_trace=False has caller-owned active state.""" + return self._splunk_ao_logger.current_parent() is not None + + def _conclude_owned_state_on_failure(self) -> None: + """Close only the handler root/envelope and return to caller state.""" + while ( + self._owned_parent is not None + and self._splunk_ao_logger.current_parent() is not None + and self._splunk_ao_logger.current_parent() is not self._owned_parent + ): + self._splunk_ao_logger.conclude(output="", status_code=500) + if self._owned_trace is not None: + self._conclude_owned_trace(self._owned_trace, output="", status_code=500) def _conclude_owned_trace(self, trace: Any, output: Any, status_code: int | None) -> None: current_parent = self._splunk_ao_logger.current_parent() @@ -142,7 +305,7 @@ def log_node_tree(self, node: Node) -> None: _logger.warning(f"Invalid step number: {metadata_step_number}, exception raised {e}") # Log the current node based on its type - if node.node_type == "chain": + if node.node_type in ("chain", "workflow"): self._splunk_ao_logger.add_workflow_span( input=input_, output=output, @@ -286,6 +449,19 @@ def start_node(self, node_type: NODE_TYPE, parent_run_id: UUID | None, run_id: U if not self._root_node: _logger.debug(f"Setting root node to {node_id}") self._root_node = node + try: + self._start_owned_root(node) + except Exception: + self._conclude_owned_state_on_failure() + self._owned_trace = None + self._owned_root = None + self._owned_parent = None + _logger.warning("Failed to start handler root telemetry", exc_info=True) + + # A LangGraph child may have reclassified the existing root from a + # chain to an agent immediately before this callback reached us. + if self._root_node is not None: + self._sync_owned_root_kind(self._root_node) # Add to parent's children if parent exists if parent_run_id: diff --git a/src/splunk_ao/handlers/openai_agents/handler.py b/src/splunk_ao/handlers/openai_agents/handler.py index 6826a880..9c46590a 100644 --- a/src/splunk_ao/handlers/openai_agents/handler.py +++ b/src/splunk_ao/handlers/openai_agents/handler.py @@ -59,6 +59,8 @@ def __init__(self, splunk_ao_logger: SplunkAOLogger | None = None, flush_on_trac self._last_status_code: int | None = None self._first_input: Any = None self._owned_trace: Any = None + self._owned_root: Any = None + self._owned_root_node_id: str | None = None def on_trace_start(self, trace: Trace) -> None: """Called when an OpenAI Agent trace starts.""" @@ -73,6 +75,16 @@ def on_trace_start(self, trace: Trace) -> None: }, ) self._nodes[trace.trace_id] = node + if self._splunk_ao_logger.current_parent() is None: + try: + self._owned_trace = self._splunk_ao_logger.start_trace( + input="Agent Workflow", + name="Trace", + created_at=datetime.fromisoformat(node.span_params["start_time_iso"]), + ) + except Exception: + self._owned_trace = None + _logger.warning("Failed to start OpenAI Agents trace telemetry", exc_info=True) def on_trace_end(self, trace: Trace) -> None: """Called when an OpenAI Agent trace ends.""" @@ -96,6 +108,8 @@ def on_trace_end(self, trace: Trace) -> None: self._last_status_code = None self._first_input = None self._owned_trace = None + self._owned_root = None + self._owned_root_node_id = None def _commit_trace(self, trace: Trace) -> None: if not self._nodes: @@ -103,11 +117,34 @@ def _commit_trace(self, trace: Trace) -> None: return root_node = self._nodes.get(trace.trace_id) - if root_node: - self._log_node_tree(root_node, first_node=True) - else: + if root_node is None: _logger.warning(f"Root node {trace.trace_id} not found") - self._splunk_ao_logger.conclude(output=self._last_output, status_code=self._last_status_code) + return + + live_root_node = self._nodes.get(self._owned_root_node_id) if self._owned_root_node_id else None + if self._owned_trace is not None and self._owned_root is not None and live_root_node is not None: + self._owned_trace.input = self._first_input or "Agent Workflow" + self._owned_trace.metrics.duration_ns = root_node.span_params.get("duration_ns") + self._update_owned_root(live_root_node) + self._log_node_tree(live_root_node, reuse_current=True) + for child_id in root_node.children: + if child_id == self._owned_root_node_id: + continue + child = self._nodes.get(child_id) + if child is not None: + self._log_node_tree(child) + self._splunk_ao_logger.conclude(output=self._last_output, status_code=self._last_status_code) + return + + # A caller-owned operation remains current; log framework operations + # beneath it without ever concluding the caller's parent. + for child_id in root_node.children: + child = self._nodes.get(child_id) + if child is not None: + self._log_node_tree(child) + + if self._owned_trace is not None: + self._splunk_ao_logger.conclude(output=self._last_output, status_code=self._last_status_code) def _conclude_current_trace_on_failure(self) -> None: if self._owned_trace is None: @@ -123,7 +160,7 @@ def _conclude_current_trace_on_failure(self) -> None: if root is self._owned_trace: self._splunk_ao_logger.conclude(output="", status_code=500, conclude_all=True) - def _log_node_tree(self, node: Node, first_node: bool = False) -> None: + def _log_node_tree(self, node: Node, reuse_current: bool = False) -> None: """ Log a node and its children recursively. @@ -131,8 +168,8 @@ def _log_node_tree(self, node: Node, first_node: bool = False) -> None: ---------- node : Node The node to log. - first_node : bool - Whether this is the root trace node. + reuse_current : bool + Whether this node is the live root already created at span start. """ is_workflow_span = False input = node.span_params.get("input", "") @@ -145,26 +182,18 @@ def _log_node_tree(self, node: Node, first_node: bool = False) -> None: # Convert metadata to a dict[str, str] if metadata is not None: metadata = convert_to_string_dict(metadata) - if first_node: - self._owned_trace = self._splunk_ao_logger.add_trace( - input=self._first_input or "Agent Workflow", - output=self._last_output, - duration_ns=node.span_params.get("duration_ns"), - created_at=start_time_iso, - name="Trace", - tags=tags, - ) # Log the current node based on its type - elif node.node_type in ("agent", "chain", "workflow"): - self._splunk_ao_logger.add_workflow_span( - input=input or node.node_type.capitalize() + " Step", - output=output, - name=name, - metadata=metadata, - tags=tags, - created_at=start_time_iso, - duration_ns=node.span_params.get("duration_ns"), - ) + if node.node_type in ("agent", "chain", "workflow"): + if not reuse_current: + self._splunk_ao_logger.add_workflow_span( + input=input or node.node_type.capitalize() + " Step", + output=output, + name=name, + metadata=metadata, + tags=tags, + created_at=start_time_iso, + duration_ns=node.span_params.get("duration_ns"), + ) is_workflow_span = True elif node.node_type in ("llm", "chat"): tools = node.span_params.get("tools") @@ -239,10 +268,47 @@ def _log_node_tree(self, node: Node, first_node: bool = False) -> None: if error: output = error status_code = 500 - self._splunk_ao_logger.conclude(output=serialize_to_str(output), status_code=status_code) + self._splunk_ao_logger.conclude( + output=serialize_to_str(output), + duration_ns=node.span_params.get("duration_ns"), + status_code=status_code, + ) self._last_status_code = status_code self._last_output = output + def _start_owned_root(self, node: Node) -> None: + """Create the first top-level real operation while framework work is active.""" + if self._owned_trace is None or self._owned_root is not None: + return + if node.node_type not in ("agent", "chain", "workflow"): + return + + metadata = node.span_params.get("metadata") + if metadata is not None: + metadata = convert_to_string_dict(metadata) + self._owned_root = self._splunk_ao_logger.add_workflow_span( + input=node.span_params.get("input") or node.node_type.capitalize() + " Step", + name=node.span_params.get("name"), + metadata=metadata, + tags=node.span_params.get("tags"), + created_at=datetime.fromisoformat(node.span_params["start_time_iso"]), + ) + if self._owned_root is not None: + self._owned_root_node_id = str(node.run_id) + + def _update_owned_root(self, node: Node) -> None: + """Apply final OpenAI callback data to the already-live root.""" + if self._owned_root is None: + return + self._owned_root.input = node.span_params.get("input") or node.node_type.capitalize() + " Step" + if node.span_params.get("name") is not None: + self._owned_root.name = node.span_params["name"] + if node.span_params.get("tags") is not None: + self._owned_root.tags = node.span_params["tags"] + metadata = node.span_params.get("metadata") + if metadata is not None: + self._owned_root.user_metadata = convert_to_string_dict(metadata) + def on_span_start(self, span: Span[Any]) -> None: """Called when an OpenAI Agent span starts.""" span_id = span.span_id @@ -320,6 +386,15 @@ def on_span_start(self, span: Span[Any]) -> None: _logger.warning(f"Parent node {parent_id} not found for span {span_id} in trace {trace_id}") return parent_node.children.append(span_id) + if parent_id == trace_id: + try: + self._start_owned_root(node) + except Exception: + self._conclude_current_trace_on_failure() + self._owned_trace = None + self._owned_root = None + self._owned_root_node_id = None + _logger.warning("Failed to start OpenAI Agents root telemetry", exc_info=True) def on_span_end(self, span: Span[Any]) -> None: """Called when an OpenAI Agent span ends.""" diff --git a/src/splunk_ao/logger/logger.py b/src/splunk_ao/logger/logger.py index 8ebd8937..a079b04c 100644 --- a/src/splunk_ao/logger/logger.py +++ b/src/splunk_ao/logger/logger.py @@ -1,7 +1,6 @@ import asyncio import atexit import contextlib -import copy import inspect import json import logging @@ -13,10 +12,6 @@ from datetime import datetime from typing import TYPE_CHECKING, Any, TypeVar, Union -if TYPE_CHECKING: - from splunk_ao.handlers.agent_control import SplunkAOAgentControlBridge - -import backoff from opentelemetry import context as otel_context from opentelemetry import trace as otel_trace from opentelemetry.sdk.trace.id_generator import RandomIdGenerator @@ -32,16 +27,14 @@ LlmSpanAllowedInputType, LlmSpanAllowedOutputType, RetrieverSpan, - Span, StepWithChildSpans, ToolSpan, ) -from galileo_core.schemas.logging.step import BaseStep, Metrics, StepType +from galileo_core.schemas.logging.step import BaseStep, Metrics from galileo_core.schemas.logging.trace import Trace from galileo_core.schemas.shared.traces_logger import TracesLogger from splunk_ao.agent_streams import AgentStreams from splunk_ao.constants import LoggerModeType -from splunk_ao.constants.tracing import PARENT_ID_HEADER, TRACE_ID_HEADER from splunk_ao.converter import SpanConverter from splunk_ao.deployment import DeploymentMode, O11yConfig, StandaloneConfig, resolve_deployment from splunk_ao.exceptions import SplunkAOLoggerException @@ -56,7 +49,6 @@ ) from splunk_ao.exporter.diagnostics import get_export_health from splunk_ao.logger.control import ControlAppliesTo, ControlCheckStage, ControlResult -from splunk_ao.logger.task_handler import ThreadPoolTaskHandler from splunk_ao.projects import Projects from splunk_ao.schema.content_blocks import ( DataContentBlock, @@ -81,42 +73,28 @@ LogRecordsSearchRequest, RetrieverSpanAllowedOutputType, SessionCreateRequest, - SpansIngestRequest, - SpanUpdateRequest, TracesIngestRequest, - TraceUpdateRequest, ) from splunk_ao.traces import Traces -from splunk_ao.utils.decorators import ( - async_warn_catch_exception, - nop_async, - nop_sync, - retry_on_transient_http_error, - warn_catch_exception, -) +from splunk_ao.utils.decorators import async_warn_catch_exception, nop_async, nop_sync, warn_catch_exception from splunk_ao.utils.env_helpers import _get_mode_or_default from splunk_ao.utils.metrics import populate_local_metrics from splunk_ao.utils.retrievers import convert_to_documents from splunk_ao.utils.serialization import serialize_to_str +if TYPE_CHECKING: + from splunk_ao.handlers.agent_control import SplunkAOAgentControlBridge + # Type alias for metadata values that can be auto-converted to strings MetadataValue = str | bool | int | float | None StepT = TypeVar("StepT", bound=BaseStep) -STREAMING_MAX_RETRIES = 5 -STREAMING_MAX_TIME_SECONDS = 70 # Maximum time to spend retrying a single request -DISTRIBUTED_FLUSH_TIMEOUT_SECONDS = 90 # Timeout for waiting on background trace/span update tasks during flush() -# Default upper bound for shutdown wait in terminate(). Independent from -# DISTRIBUTED_FLUSH_TIMEOUT_SECONDS (which guards live flushes); kept aligned -# at 90s so explicit terminate() calls behave like flush() by default. DEFAULT_TERMINATE_TIMEOUT_SECONDS = 90 # Absolute threshold above which a slow SplunkAOLogger.terminate() shutdown is # logged as a warning. Fast-path shutdowns are sub-millisecond; >1s always # indicates a real anomaly (busy-poll, stuck task, in-flight HTTP retry) and # should be visible in CI logs regardless of the configured timeout. _SLOW_SHUTDOWN_WARN_THRESHOLD_SECONDS = 1.0 -STUB_TRACE_NAME = "stub_trace" # Name for stub traces created from distributed tracing headers - _logger = logging.getLogger("splunk_ao.logger") _otel_id_generator = RandomIdGenerator() @@ -127,6 +105,7 @@ class OtelIds: span_context: SpanContext parent_span_context: SpanContext | None + exportable: bool @dataclass(frozen=True) @@ -136,6 +115,7 @@ class ActiveOtelContext: logger_id: int step_id: uuid.UUID span_context: SpanContext + exportable: bool token: Token @@ -150,6 +130,22 @@ class OtelContextState: _otel_context_state: ContextVar[OtelContextState | None] = ContextVar("_otel_context_state", default=None) +def _has_active_exportable_span_context() -> bool: + """Return whether the current OTel context can be propagated as a wire parent.""" + current = otel_trace.get_current_span().get_span_context() + if not current.is_valid: + return False + + state = _otel_context_state.get() + 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 + + class SplunkAOLogger(TracesLogger): """ This class can be used to upload traces to Splunk AO. @@ -218,16 +214,12 @@ class SplunkAOLogger(TracesLogger): agent_stream_id: str | None = None experiment_id: str | None = None session_id: str | None = None - trace_id: str | None = None - span_id: str | None = None local_metrics: list[LocalMetricConfig] | None = None mode: LoggerModeType | None = None _session_external_id: str | None = None _logger = logging.getLogger("splunk_ao.logger") _traces_client: Union["Traces", None] = None - _task_handler: ThreadPoolTaskHandler - _trace_completion_submitted: bool _otel_ids: dict[uuid.UUID, OtelIds] = PrivateAttr(default_factory=dict) _pending_otel_steps: set[uuid.UUID] = PrivateAttr(default_factory=set) @@ -238,8 +230,6 @@ def __init__( agent_stream: str | None = None, agent_stream_id: str | None = None, experiment_id: str | None = None, - trace_id: str | None = None, - span_id: str | None = None, local_metrics: list[LocalMetricConfig] | None = None, mode: str | None = None, ingestion_hook: Callable[[TracesIngestRequest], None] | None = None, @@ -261,25 +251,12 @@ def __init__( Agent stream ID. experiment_id: Optional[str] Experiment ID. Used by the experiment runner. - trace_id: Optional[str] - Trace ID for distributed tracing. This can only be used in "distributed" mode. - - When provided, creates a local stub trace without fetching from the backend. - This allows downstream services to continue a distributed trace without waiting - for backend ingestion. - span_id: Optional[str] - Parent span ID for distributed tracing. This can only be used in "distributed" mode. - - When provided, creates a local stub span without fetching from the backend. - This allows downstream services to continue a distributed trace without waiting - for backend ingestion. local_metrics: Optional[list[LocalMetricConfig]] Local metrics mode: Optional[str] Logger mode: "batch" or "distributed". Defaults to SPLUNK_AO_MODE env var, or "batch" if not set. - Both modes enqueue completed spans for scheduled OTLP batch export, and - flush() only drains the export queue. "distributed" additionally enables - the legacy trace_id/span_id continuation parameters. + Both accepted values enqueue completed spans for the same scheduled OTLP + batch export. The value is retained temporarily for ingestion-hook compatibility. ingestion_hook: Optional[Callable[[TracesIngestRequest], None]] A callable that intercepts trace data before ingestion. This hook is called when the logger is flushed and can be a @@ -290,7 +267,6 @@ def __init__( super().__init__() mode = _get_mode_or_default(mode) self.mode: LoggerModeType = mode - self._task_counter = 0 self._terminated = False self._traces_client = None self._pending_otel_steps = set() @@ -313,31 +289,6 @@ def __init__( self._auto_enable_agent_control_if_available() return - if trace_id or span_id: - if self.mode != "distributed": - raise SplunkAOLoggerException("trace_id or span_id can only be used in distributed mode") - if span_id and not trace_id: - raise SplunkAOLoggerException( - "trace_id is required when span_id is provided. " - "In distributed tracing, both trace_id and span_id must be propagated together." - ) - - # Validate UUIDs to prevent crashes from malformed input - if trace_id: - try: - uuid.UUID(trace_id) - except (ValueError, AttributeError, TypeError) as e: - raise SplunkAOLoggerException(f"Invalid trace_id: '{trace_id}' is not a valid UUID. Error: {e}") - - if span_id: - try: - uuid.UUID(span_id) - except (ValueError, AttributeError, TypeError) as e: - raise SplunkAOLoggerException(f"Invalid span_id: '{span_id}' is not a valid UUID. Error: {e}") - - self.trace_id = trace_id - self.span_id = span_id - if (agent_stream or agent_stream_id) and experiment_id: raise SplunkAOLoggerException("User cannot specify both an agent stream and an experiment.") @@ -368,7 +319,9 @@ def __init__( "User must provide project_name or project_id to SplunkAOLogger, or set it as an environment variable." ) if self.experiment_id is None and self.agent_stream_name is None and self.agent_stream_id is None: - raise SplunkAOLoggerException("agent_stream or agent_stream_id is required to initialize SplunkAOLogger.") + raise SplunkAOLoggerException( + "agent_stream or agent_stream_id is required to initialize SplunkAOLogger." + ) if local_metrics: self.local_metrics = local_metrics @@ -393,12 +346,6 @@ def __init__( else: self._sink = build_span_sink(build_standalone_exporter(StandaloneConfig.from_env(), routing)) - # If continuing an existing distributed trace, create local stubs instead of - # fetching from the backend to avoid race conditions with eventual consistency. - # Note: trace_id/span_id can ONLY be provided in distributed mode for distributed tracing - if self.trace_id: - self._init_distributed_trace_stubs() - # cleans up when the python interpreter closes atexit.register(self.terminate) self._auto_enable_agent_control_if_available() @@ -424,13 +371,13 @@ def reset_parent_tracking(self) -> None: if root is not None: self._discard_otel_subtree(root) - def _assign_otel_context(self, otel_trace_id: int, trace_state: TraceState) -> SpanContext: - """Create a sampled local SpanContext without changing the active context.""" + def _assign_otel_context(self, otel_trace_id: int, trace_flags: TraceFlags, trace_state: TraceState) -> SpanContext: + """Create a local SpanContext without changing the active context.""" return SpanContext( trace_id=otel_trace_id, span_id=_otel_id_generator.generate_span_id(), is_remote=False, - trace_flags=TraceFlags(TraceFlags.SAMPLED), + trace_flags=trace_flags, trace_state=trace_state, ) @@ -450,14 +397,17 @@ def _record_otel_ids( if parent_span_context is None: otel_trace_id = _otel_id_generator.generate_trace_id() + trace_flags = TraceFlags(TraceFlags.SAMPLED) trace_state = TraceState() else: otel_trace_id = parent_span_context.trace_id + trace_flags = parent_span_context.trace_flags trace_state = parent_span_context.trace_state ids = OtelIds( - span_context=self._assign_otel_context(otel_trace_id, trace_state), + span_context=self._assign_otel_context(otel_trace_id, trace_flags, trace_state), parent_span_context=parent_span_context, + exportable=not isinstance(step, Trace), ) self._otel_ids[step.id] = ids return ids @@ -516,10 +466,14 @@ def _sync_otel_context_impl(self, current_parent: StepWithChildSpans | None) -> 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) for step_id in desired_step_ids + (logger_id, step_id, self._otel_ids[step_id].span_context, self._otel_ids[step_id].exportable) + for step_id in desired_step_ids ) contexts_to_restore = ( - tuple((active.logger_id, active.step_id, active.span_context) for active in remaining_contexts) + tuple( + (active.logger_id, active.step_id, active.span_context, active.exportable) + for active in remaining_contexts + ) + desired_contexts ) @@ -536,11 +490,15 @@ 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 in contexts_to_restore: + for owner_id, step_id, span_context, exportable in contexts_to_restore: ctx = otel_trace.set_span_in_context(NonRecordingSpan(span_context)) rebuilt_contexts.append( ActiveOtelContext( - logger_id=owner_id, step_id=step_id, span_context=span_context, token=otel_context.attach(ctx) + logger_id=owner_id, + step_id=step_id, + span_context=span_context, + exportable=exportable, + token=otel_context.attach(ctx), ) ) @@ -605,6 +563,9 @@ def _emit_and_release(self, finished_step: BaseStep) -> None: if isinstance(finished_step, Trace): return + if not bool(ids.span_context.trace_flags & TraceFlags.SAMPLED): + return + self._sink.emit( self._converter.convert_span( span=finished_step, @@ -746,49 +707,6 @@ def _ensure_session_crud_client(self) -> Traces: self._traces_client = self._create_traces_client() return self._traces_client - def _init_distributed_trace_stubs(self) -> None: - """ - Initialize local stub objects for distributed tracing. To only be used in distributed mode. - - When a downstream service receives trace_id/span_id via headers, we create - local stub objects instead of fetching from the backend. This avoids race - conditions as the parent trace/span may not have been ingested yet when the downstream service starts. - - The stubs are placeholders that allow: - 1. Adding new spans to the distributed trace - 2. Proper parent-child relationships via _parent pointers - 3. Correct trace_id in ingestion requests - - Note: trace_id and span_id are already validated as UUIDs in __init__ - """ - stub_trace = LoggedTrace( - input="", - name=STUB_TRACE_NAME, - created_at=datetime.now(), - id=uuid.UUID(self.trace_id), - metrics=Metrics(duration_ns=0), - ) - self.traces.append(stub_trace) - - # Set trace as current parent using parent pointers - stub_trace._parent = None # Root trace has no parent - self._record_otel_ids(stub_trace) - self._set_current_parent(stub_trace) - - if self.span_id: - # If span_id is provided, also add the span (it's the immediate parent) - stub_span = LoggedWorkflowSpan( - input="", - name="stub_parent_span", - created_at=datetime.now(), - id=uuid.UUID(self.span_id), - metrics=Metrics(duration_ns=0), - ) - # Set parent pointer and update current parent - stub_span._parent = stub_trace - self._record_otel_ids(stub_span, parent_step=stub_trace) - self._set_current_parent(stub_span) - def add_trace( self, input: str, @@ -968,228 +886,6 @@ def _get_last_output(node: BaseStep | None) -> tuple[IngestOutputType | None, In return None, None - @nop_sync - @warn_catch_exception(exceptions=(Exception,)) - def _ingest_trace_streaming(self, trace: Trace, is_complete: bool = False) -> None: - traces_ingest_request = TracesIngestRequest( - traces=[copy.deepcopy(trace)], session_id=self.session_id, is_complete=is_complete, reliable=True - ) - - task_id = f"trace-ingest-{trace.id}" - - @backoff.on_exception( - backoff.expo, - Exception, - max_tries=self._max_retries, - max_time=self._max_time, - base=2, - logger=None, - on_backoff=lambda details: ( - self._task_handler.increment_retry(task_id), - self._logger.info(f"Retry #{self._task_handler.get_retry_count(task_id)} for task {task_id}"), - ), - on_giveup=lambda details: self._logger.error( - f"Task {task_id} failed after {details['tries']} attempts: {details.get('exception')}", exc_info=False - ), - ) - @retry_on_transient_http_error - async def ingest_traces_with_backoff(request: Any) -> None: - return await self._traces_client.ingest_traces(request) - - self._task_handler.submit_task( - task_id, lambda: ingest_traces_with_backoff(traces_ingest_request), dependent_on_prev=False - ) - self._logger.info("ingested trace %s.", trace.id) - - @nop_sync - @warn_catch_exception(exceptions=(Exception,)) - def _ingest_span_streaming(self, span: Span) -> None: - parent_step: StepWithChildSpans | None = ( - self.current_parent() - if span.type - not in [ - StepType.trace, - StepType.workflow, - StepType.agent, - ] # TODO: change this to StepWithChildSpans once we fix tool and retriever spans in `core - else self.previous_parent() - ) - if parent_step is None: - raise ValueError("A trace needs to be created in order to add a span.") - - # Use IDs from the current trace and parent step - trace_id = self.traces[0].id - parent_id = parent_step.id - spans_ingest_request = SpansIngestRequest( - spans=[copy.deepcopy(span)], trace_id=trace_id, parent_id=parent_id, reliable=True - ) - - task_id = f"span-ingest-{span.id}" - - @backoff.on_exception( - backoff.expo, - Exception, - max_tries=self._max_retries, - max_time=self._max_time, - base=2, - logger=None, - on_backoff=lambda details: ( - self._task_handler.increment_retry(task_id), - self._logger.info(f"Retry #{self._task_handler.get_retry_count(task_id)} for task {task_id}"), - ), - on_giveup=lambda details: self._logger.error( - f"Task {task_id} failed after {details['tries']} attempts: {details.get('exception')}", exc_info=False - ), - ) - @retry_on_transient_http_error - async def ingest_spans_with_backoff(request: Any) -> None: - return await self._traces_client.ingest_spans(request) - - self._task_handler.submit_task( - task_id, lambda: ingest_spans_with_backoff(spans_ingest_request), dependent_on_prev=False - ) - self._logger.info("ingested span %s.", span.id) - - @nop_sync - @warn_catch_exception(exceptions=(Exception,)) - def _update_trace_streaming(self, trace: Trace, is_complete: bool = False) -> None: - output: str | None = None - if trace.output is not None: - output = trace.output if isinstance(trace.output, str) else serialize_to_str(trace.output) - trace_update_request = TraceUpdateRequest( - trace_id=trace.id, - log_stream_id=self.agent_stream_id, - experiment_id=self.experiment_id, - output=output, - status_code=trace.status_code, - tags=trace.tags, - is_complete=is_complete, - duration_ns=trace.metrics.duration_ns, - reliable=True, - ) - - # Use counter to make each update task unique (same trace can be updated multiple times) - self._task_counter += 1 - task_id = f"trace-update-{trace.id}-{self._task_counter}" - - # Find the most recent trace update task for this specific trace (if any) - # This ensures trace updates for the same trace happen in order - prev_trace_update_task = None - for existing_task_id in reversed(list(self._task_handler._tasks.keys())): - if existing_task_id.startswith(f"trace-update-{trace.id}-"): - prev_trace_update_task = existing_task_id - break - - @backoff.on_exception( - backoff.expo, - Exception, - max_tries=self._max_retries, - max_time=self._max_time, - base=2, - logger=None, - on_backoff=lambda details: ( - self._task_handler.increment_retry(task_id), - self._logger.info(f"Retry #{self._task_handler.get_retry_count(task_id)} for trace update {task_id}"), - ), - on_giveup=lambda details: ( - self._logger.error( - f"Task {task_id} failed after {details['tries']} attempts: {details.get('exception')}", - exc_info=False, - ), - ), - ) - @retry_on_transient_http_error - async def update_trace_with_backoff(request: Any) -> None: - return await self._traces_client.update_trace(request) - - # Submit with dependency on the previous trace update for this trace - if prev_trace_update_task: - self._task_handler.submit_task_with_parent( - task_id, lambda: update_trace_with_backoff(trace_update_request), parent_task_id=prev_trace_update_task - ) - else: - self._task_handler.submit_task( - task_id, lambda: update_trace_with_backoff(trace_update_request), dependent_on_prev=True - ) - - # Mark that we've submitted the trace completion update to prevent duplicates - if is_complete: - self._trace_completion_submitted = True - - self._logger.info("updated trace %s.", trace.id) - - @nop_sync - @warn_catch_exception(exceptions=(Exception,)) - def _update_span_streaming(self, span: Span) -> None: - span_update_request = SpanUpdateRequest( - span_id=span.id, - log_stream_id=self.agent_stream_id, - experiment_id=self.experiment_id, - output=span.output, - status_code=span.status_code, - tags=span.tags, - duration_ns=span.metrics.duration_ns, - reliable=True, - ) - - # Use counter to make each update task unique (same span can be updated multiple times) - self._task_counter += 1 - task_id = f"span-update-{span.id}-{self._task_counter}" - - # Find the most recent update/ingest task for this specific span - # This ensures span updates happen in order - parent_task_id = None - for existing_task_id in reversed(list(self._task_handler._tasks.keys())): - if existing_task_id.startswith(f"span-update-{span.id}-") or existing_task_id == f"span-ingest-{span.id}": - parent_task_id = existing_task_id - break - - # If no previous task found, depend on the span ingest - if not parent_task_id: - parent_task_id = f"span-ingest-{span.id}" - - @backoff.on_exception( - backoff.expo, - Exception, - max_tries=self._max_retries, - max_time=self._max_time, - base=2, - logger=None, - on_backoff=lambda details: ( - self._task_handler.increment_retry(task_id), - self._logger.info( - f"Retry #{self._task_handler.get_retry_count(task_id)} for task {task_id}, waiting {details['wait']:.1f}s" - ), - ), - on_giveup=lambda details: self._logger.error( - f"Task {task_id} failed after {details['tries']} attempts: {details.get('exception')}", exc_info=False - ), - ) - @retry_on_transient_http_error - async def update_span_with_backoff(request: Any) -> None: - return await self._traces_client.update_span(request) - - self._task_handler.submit_task_with_parent( - task_id, lambda: update_span_with_backoff(span_update_request), parent_task_id=parent_task_id - ) - self._logger.info("updated span %s.", span.id) - - @nop_sync - @warn_catch_exception(exceptions=(Exception,)) - def _ingest_step_streaming(self, step: StepWithChildSpans, is_complete: bool = False) -> None: - if isinstance(step, LoggedTrace): - self._ingest_trace_streaming(step, is_complete=is_complete) - else: - self._ingest_span_streaming(step) - - @nop_sync - @warn_catch_exception(exceptions=(Exception,)) - def _update_step_streaming(self, step: StepWithChildSpans, is_complete: bool = False) -> None: - if isinstance(step, LoggedTrace): - self._update_trace_streaming(step, is_complete=is_complete) - else: - self._update_span_streaming(step) - @nop_sync @warn_catch_exception(exceptions=(Exception,)) def previous_parent(self) -> StepWithChildSpans | None: @@ -1198,8 +894,6 @@ def previous_parent(self) -> StepWithChildSpans | None: @nop_sync @warn_catch_exception(exceptions=(Exception,)) def has_active_trace(self) -> bool: - if self.mode == "distributed" and (self.trace_id or self.span_id): - return True current_parent = self.current_parent() # Each logger has its own per-instance ContextVar for parent tracking. # The traces check is a sanity check to ensure consistency. @@ -1222,71 +916,6 @@ def disable_agent_control(self) -> None: if bridge is not None: bridge.unregister() - def get_tracing_headers(self) -> dict[str, str]: - """ - Get tracing headers for distributed tracing. - Returns headers that can be passed to downstream services to continue the distributed trace. - - Returns - ------- - dict[str, str] - Dictionary with the following headers: - - Splunk-AO-Trace-ID: The root trace ID - - Splunk-AO-Parent-ID: The ID of the current parent (trace or span) that downstream - spans should attach to - - Raises - ------ - SplunkAOLoggerException - If not in distributed mode or if no trace has been started. - - Examples - -------- - ```python - logger = SplunkAOLogger(mode="distributed") - logger.start_trace(input="question") - headers = logger.get_tracing_headers() - # headers = { - # "Splunk-AO-Trace-ID": "...", - # "Splunk-AO-Parent-ID": "...", # trace ID as parent - # } - - logger.add_workflow_span(input="workflow", name="orchestrator") - headers = logger.get_tracing_headers() - # headers = { - # "Splunk-AO-Trace-ID": "...", - # "Splunk-AO-Parent-ID": "...", # workflow span ID as parent - # } - - # Pass headers to HTTP request - response = httpx.post(url, headers=headers) - ``` - - Note: Project and log_stream are configured per service (via env vars or logger initialization), - not propagated via headers, following standard distributed tracing patterns. - """ - if self.mode != "distributed": - raise SplunkAOLoggerException( - "get_tracing_headers is only supported in distributed mode for distributed tracing." - ) - - if len(self.traces) == 0: - raise SplunkAOLoggerException("Start trace before getting tracing headers.") - - headers: dict[str, str] = {} - - root_trace = self.traces[-1] - headers[TRACE_ID_HEADER] = str(root_trace.id) - - current_parent = self.current_parent() - - if not current_parent: - raise SplunkAOLoggerException("No parent trace or span found.") - - headers[PARENT_ID_HEADER] = str(current_parent.id) - - return headers - @nop_sync @warn_catch_exception() def start_trace( @@ -2263,72 +1892,6 @@ async def async_flush(self) -> None: if not await asyncio.to_thread(self._sink.force_flush): raise RuntimeError("force_flush timed out; some spans may not have been exported") - @async_warn_catch_exception(exceptions=(Exception,)) - async def _wait_for_all_tasks_async(self, timeout_seconds: int) -> None: - """Wait for all background tasks to complete (async polling). - - Parameters - ---------- - timeout_seconds: int - Maximum time to wait for tasks to complete - """ - start_wait = time.time() - while not self._task_handler.all_tasks_completed(): - if time.time() - start_wait > timeout_seconds: - raise TimeoutError( - f"Flush timeout reached after {timeout_seconds}s. " - "Some trace/span update requests may still be in progress." - ) - await asyncio.sleep(0.1) - - @warn_catch_exception(exceptions=(Exception,)) - def _wait_for_all_tasks_sync(self, timeout_seconds: int) -> None: - """Wait for all background tasks to complete (synchronous polling). - - Parameters - ---------- - timeout_seconds: int - Maximum time to wait for tasks to complete - """ - start_wait = time.time() - while not self._task_handler.all_tasks_completed(): - if time.time() - start_wait > timeout_seconds: - self._logger.warning( - f"Terminate timeout reached after {timeout_seconds}s. " - "Some trace/span update requests may still be in progress." - ) - break - time.sleep(0.1) - - @warn_catch_exception(exceptions=(Exception,)) - def _wait_for_pending_span_ingests(self, timeout_seconds: int) -> None: - """Wait for all pending span ingest tasks to complete. - - Note: Uses time.sleep() for polling even though callers may have @nop_sync. - This briefly blocks but is acceptable since we're just polling task status. - - Parameters - ---------- - timeout_seconds: int - Maximum time to wait for span ingests to complete - """ - pending_span_tasks = [ - task_id - for task_id in self._task_handler._tasks - if task_id.startswith("span-ingest-") and self._task_handler.get_status(task_id) in ["pending", "running"] - ] - - if pending_span_tasks: - start_wait = time.time() - while pending_span_tasks and (time.time() - start_wait) < timeout_seconds: - pending_span_tasks = [ - task_id - for task_id in pending_span_tasks - if self._task_handler.get_status(task_id) in ["pending", "running"] - ] - if pending_span_tasks: - time.sleep(0.1) - @nop_sync @warn_catch_exception(exceptions=(Exception,)) def _auto_conclude_trace(self) -> None: @@ -2343,39 +1906,13 @@ def _auto_conclude_trace(self) -> None: # Use the last trace in self.traces (should be the only active trace) trace = self.traces[-1] - # Don't auto-conclude stub traces - they're owned by the upstream service - # Downstream services that receive distributed tracing headers create stubs - # but should not mark them as complete - if trace.name == STUB_TRACE_NAME: - return - # If there are unconcluded items in the stack, conclude them if self._parent_stack: self._logger.info("Concluding unconcluded spans before flush...") # Get output from last child span if trace has no explicit output output, redacted_output = SplunkAOLogger._get_last_output(trace) # conclude() with conclude_all=True will conclude all unconcluded items in _parent_stack - # This will mark the trace as complete in distributed mode self.conclude(output=output, redacted_output=redacted_output, conclude_all=True) - elif self.mode == "distributed": - if not self._trace_completion_submitted: - # Wait for all span ingests to complete before marking trace complete - self._wait_for_pending_span_ingests(timeout_seconds=DISTRIBUTED_FLUSH_TIMEOUT_SECONDS) - self._update_trace_streaming(trace, is_complete=True) - - async def _flush_distributed(self) -> list[LoggedTrace]: - """Legacy proprietary distributed flush path, unused by OTLP egress.""" - self._auto_conclude_trace() - - # Wait for all pending trace/span update requests to complete - self._logger.info("Waiting for all distributed tracing tasks to complete...") - await self._wait_for_all_tasks_async(timeout_seconds=DISTRIBUTED_FLUSH_TIMEOUT_SECONDS) - self._logger.info("All distributed tracing requests are complete.") - - self.traces = [] - self._set_current_parent(None) - - return [] async def _flush_batch(self) -> list[LoggedTrace]: """Flush in batch mode: conclude unconcluded traces and send all traces to backend.""" diff --git a/src/splunk_ao/logger/task_handler.py b/src/splunk_ao/logger/task_handler.py deleted file mode 100644 index 3864c158..00000000 --- a/src/splunk_ao/logger/task_handler.py +++ /dev/null @@ -1,200 +0,0 @@ -import time -from collections.abc import Awaitable, Callable, Coroutine -from concurrent.futures import Future -from typing import Any, Literal - -from galileo_core.helpers.event_loop_thread_pool import EventLoopThreadPool - -NUM_THREADS = 4 - -TaskStatus = Literal["not_found", "pending", "running", "completed", "failed"] - - -class ThreadPoolTaskHandler: - """A task handler that manages dependencies and executes tasks in a thread pool.""" - - _pool: EventLoopThreadPool - _tasks: dict[str, dict] - _retry_counts: dict[str, int] - - def __init__(self, num_threads: int = NUM_THREADS): - self._tasks = {} - self._retry_counts = {} - self._pool = EventLoopThreadPool(num_threads=num_threads) - - def _handle_task_completion(self, task_id: str) -> None: - """Handle the completion of a task, triggering any children.""" - # Find all child tasks that depend on this task - for _child_task_id, task in list(self._tasks.items()): - if task.get("parent_task_id") == task_id and task.get("callback"): - # Execute the callback which will submit the child task - task["callback"]() - - def _add_or_update_task( - self, - task_id: str, - future: Future | None = None, - start_time: float | None = None, - parent_task_id: str | None = None, - callback: Callable | None = None, - ) -> None: - """ - Track a submitted future. - - Parameters - ---------- - task_id: str - The ID of the task. - future: Optional[Future] - The future to track. - start_time: Optional[float] - The start time of the task. - parent_task_id: Optional[str] - The ID of the parent task. - callback: Optional[Callable] - The callback to run when the task is completed. - """ - self._tasks[task_id] = { - "future": future, - "start_time": start_time, - "parent_task_id": parent_task_id, - "callback": callback, - } - self._retry_counts[task_id] = 0 - - def submit_task( - self, task_id: str, async_fn: Callable[[], Awaitable[Any]] | Coroutine, dependent_on_prev: bool = False - ) -> None: - """ - Submit a task to the thread pool. - - Parameters - ---------- - task_id: str - The ID of the task. - async_fn: Union[Callable[[], Awaitable[Any]], Coroutine] - The async function to submit to the thread pool. - dependent_on_prev: bool - Whether the task depends on the previous task. - """ - - def _submit(*args) -> None: - future = self._pool.submit(async_fn, wait_for_result=False) - future.add_done_callback(lambda f: self._handle_task_completion(task_id)) - self._add_or_update_task(task_id=task_id, future=future, start_time=time.time(), parent_task_id=None) - - if dependent_on_prev: - if not self._tasks: - _submit() - else: - last_task_id = list(self._tasks.keys())[-1] - if self.get_status(last_task_id) == "completed": - _submit() - else: - self._add_or_update_task( - task_id=task_id, future=None, start_time=None, parent_task_id=last_task_id, callback=_submit - ) - else: - _submit() - - def submit_task_with_parent( - self, task_id: str, async_fn: Callable[[], Awaitable[Any]] | Coroutine, parent_task_id: str - ) -> None: - """ - Submit a task that depends on a specific parent task. - - Parameters - ---------- - task_id: str - The ID of the task. - async_fn: Union[Callable[[], Awaitable[Any]], Coroutine] - The async function to submit to the thread pool. - parent_task_id: str - The ID of the parent task this depends on. - """ - - def _submit(*args) -> None: - future = self._pool.submit(async_fn, wait_for_result=False) - future.add_done_callback(lambda f: self._handle_task_completion(task_id)) - self._add_or_update_task(task_id=task_id, future=future, start_time=time.time(), parent_task_id=None) - - if parent_task_id not in self._tasks or self.get_status(parent_task_id) == "completed": - _submit() - else: - self._add_or_update_task( - task_id=task_id, future=None, start_time=None, parent_task_id=parent_task_id, callback=_submit - ) - - def get_children(self, parent_task_id: str) -> list[dict]: - """Get the children of a task.""" - return [task for _, task in self._tasks.items() if task.get("parent_task_id") == parent_task_id] - - def increment_retry(self, task_id: str) -> None: - """ - Increment the retry count for a task. - - Parameters - ---------- - task_id: str - The ID of the task. - """ - self._retry_counts[task_id] = self._retry_counts.get(task_id, 0) + 1 - - def get_status(self, task_id: str) -> TaskStatus: - """ - Returns the status of a task. - - Parameters - ---------- - task_id: str - The ID of the task. - - Returns - ------- - TaskStatus - The status of the task. - """ - if task_id not in self._tasks: - return "not_found" - - task = self._tasks[task_id] - - if task.get("parent_task_id"): - return "pending" - - future = task["future"] - if not future: - return "not_found" - - if not future.done(): - return "running" - - try: - future.result() # This will raise exception if task failed - return "completed" - except Exception: - return "failed" - - def get_result(self, task_id: str) -> Any: - """Get result if task completed, otherwise raises exception.""" - if task_id not in self._tasks: - raise ValueError(f"Task {task_id} not found") - return self._tasks[task_id]["future"].result() - - def get_retry_count(self, task_id: str) -> int: - """Get the retry count for a task.""" - return self._retry_counts.get(task_id, 0) - - def all_tasks_completed(self) -> bool: - """ - Check if all tasks are completed. - - Returns - ------- - bool - True if all tasks are completed, False otherwise. - """ - return all(self.get_status(task_id) not in ["running", "pending"] for task_id in self._tasks) - - def terminate(self) -> None: - self._pool.stop() diff --git a/src/splunk_ao/middleware/tracing.py b/src/splunk_ao/middleware/tracing.py index 57b2725f..ceb677db 100644 --- a/src/splunk_ao/middleware/tracing.py +++ b/src/splunk_ao/middleware/tracing.py @@ -1,62 +1,37 @@ -""" -Distributed tracing middleware for Starlette-based applications. - -This middleware automatically extracts distributed tracing headers from incoming HTTP requests -and makes them available to the Splunk AO logger within request handlers. - -Works with any ASGI framework built on Starlette: -- FastAPI -- Starlette -- Any other Starlette-based framework +"""W3C trace-context middleware for Starlette-based applications. -Example usage with FastAPI: - ```python - from fastapi import FastAPI - from splunk_ao.middleware import TracingMiddleware, get_request_logger +``TracingMiddleware`` extracts ``traceparent`` and ``tracestate`` through the +configured OpenTelemetry propagator. Applications explicitly own the local +Splunk AO logger and trace lifecycle:: - app = FastAPI() app.add_middleware(TracingMiddleware) @app.post("/process") async def process_request(data: dict): - # Logger automatically continues the distributed trace logger = get_request_logger() - logger.add_workflow_span(input=str(data), name="process_workflow") - # ... process request ... - logger.conclude(output="done") - return {"status": "success"} - ``` - -Example usage with Starlette: - ```python - from starlette.applications import Starlette - from starlette.routing import Route - from splunk_ao.middleware import TracingMiddleware, get_request_logger - - async def homepage(request): - logger = get_request_logger() - logger.add_workflow_span(input="homepage", name="homepage_handler") - logger.conclude(output="success") - return {"status": "ok"} - - app = Starlette( - routes=[Route("/", homepage)], - middleware=[TracingMiddleware] - ) - ``` + try: + logger.start_trace(input=str(data), name="request") + logger.add_workflow_span(input=str(data), name="process") + result = await process(data) + logger.conclude(output=str(result)) + logger.conclude(output=str(result)) + return {"result": result} + finally: + logger.terminate() """ import logging from typing import Any, NoReturn -from splunk_ao.constants.tracing import PARENT_ID_HEADER, TRACE_ID_HEADER -from splunk_ao.decorator import _parent_id_context, _trace_id_context +from opentelemetry import context as otel_context + from splunk_ao.logger import SplunkAOLogger +from splunk_ao.tracing import extract_tracing_context _logger = logging.getLogger(__name__) INSTALL_ERR_MSG = ( - "Starlette is not installed. Install optional middleware dependencies with: pip install galileo[middleware]" + "Starlette is not installed. Install optional middleware dependencies with: pip install splunk-ao[middleware]" ) try: @@ -65,7 +40,7 @@ async def homepage(request): from starlette.responses import Response from starlette.types import ASGIApp except ImportError: - # Create stub classes if Starlette is not available + # Keep imports available without the optional Starlette dependency. class BaseHTTPMiddleware: # type: ignore[no-redef] def __init__(self, *args: Any, **kwargs: Any) -> NoReturn: raise ImportError(INSTALL_ERR_MSG) @@ -84,144 +59,27 @@ class ASGIApp: # type: ignore[no-redef] class TracingMiddleware(BaseHTTPMiddleware): - """ - Middleware that extracts distributed tracing headers from incoming requests. - - This middleware looks for the following headers in incoming HTTP requests: - - Splunk-AO-Trace-ID: The root trace ID - - Splunk-AO-Parent-ID: The parent span/trace ID to attach to - - These values are stored in context variables, making them available to request - handlers via the `get_request_logger()` function. - - The middleware is compatible with FastAPI and any Starlette-based framework. - - Note: Project and log_stream are configured per service via environment variables - (SPLUNK_AO_PROJECT and SPLUNK_AO_AGENT_STREAM). They are not propagated via headers, - following standard distributed tracing patterns. - """ + """Attach incoming W3C trace context for the duration of each request.""" def __init__(self, app: ASGIApp) -> None: - """ - Initialize the tracing middleware. - - Parameters - ---------- - app : ASGIApp - The ASGI application - """ super().__init__(app) async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response: - """ - Process the request and extract tracing headers. - - Parameters - ---------- - request : Request - The incoming HTTP request - call_next : RequestResponseEndpoint - The next middleware or route handler - - Returns - ------- - Response - The HTTP response - """ - # Extract tracing headers from request - trace_id = request.headers.get(TRACE_ID_HEADER) - parent_id = request.headers.get(PARENT_ID_HEADER) - - # Store in context variables (thread-safe for async) - trace_id_token = _trace_id_context.set(trace_id) - parent_id_token = _parent_id_context.set(parent_id) - + """Extract, attach, and reliably detach the request's trace context.""" + extracted = extract_tracing_context(request.headers) + token = otel_context.attach(extracted) try: - # Process the request return await call_next(request) finally: - # Clean up context variables - _trace_id_context.reset(trace_id_token) - _parent_id_context.reset(parent_id_token) + otel_context.detach(token) def get_request_logger() -> SplunkAOLogger: - """ - Get a request-scoped SplunkAOLogger configured for distributed mode. - - Distributed mode enables the legacy trace_id/span_id continuation parameters. - Completed spans use the same scheduled OTLP batch export as batch mode. - - This function should be called within a request handler after the TracingMiddleware has - been registered. It creates a new SplunkAOLogger instance per request that automatically - continues the distributed trace from the upstream service. - - The logger is configured using trace context extracted by the middleware: - - Splunk-AO-Trace-ID: Root trace ID - - Splunk-AO-Parent-ID: Parent span/trace ID to attach to - - Project and log_stream are configured per service via environment variables - (SPLUNK_AO_PROJECT and SPLUNK_AO_AGENT_STREAM), not propagated via headers, following - standard distributed tracing patterns. - - If no tracing headers were present in the request, a regular logger is returned - (using SPLUNK_AO_PROJECT and SPLUNK_AO_AGENT_STREAM env vars). - - Note: This creates a new logger per request, unlike the decorator's get_logger_instance() - which uses a singleton pattern. - - Returns - ------- - SplunkAOLogger - A logger instance configured for the current request's trace context - - Examples - -------- - ```python - @app.post("/process") - async def process_request(data: dict): - logger = get_request_logger() - - # This span will be attached to the distributed trace - logger.add_workflow_span(input=str(data), name="process_workflow") - result = await process(data) - logger.conclude(output=str(result)) - - logger.flush() - logger.terminate() - - return {"result": result} - ``` - - ```python - @app.post("/retrieve") - async def retrieve_endpoint(query: str): - # Get logger with trace context from upstream service - logger = get_request_logger() - - # If trace context exists, this creates a workflow span - # Otherwise, it starts a new trace - if logger.trace_id: - logger.add_workflow_span(input=query, name="retrieval_service") - else: - logger.start_trace(input=query, name="retrieval_service") - - results = retrieve(query, logger) - - logger.conclude(output=str(results)) - logger.flush() - logger.terminate() + """Create a request-scoped logger under the middleware's active context. - return {"results": results} - ``` + The caller must explicitly call ``start_trace()`` before adding spans and + should call ``terminate()`` in ``finally``. With an extracted W3C parent, + the first real local operation becomes its descendant; without one, it + starts a new trace. """ - # Get trace context from middleware - trace_id = _trace_id_context.get() - parent_id = _parent_id_context.get() - - # Create logger with trace context - # Project and log_stream come from env vars (SPLUNK_AO_PROJECT, SPLUNK_AO_AGENT_STREAM) - # If parent_id equals trace_id, it means the parent is the root trace itself, - # not a span. In this case, we should pass None as span_id to avoid - # SplunkAOLoggerException when it tries to look up a span with the trace_id. - return SplunkAOLogger(mode="distributed", trace_id=trace_id, span_id=parent_id if parent_id != trace_id else None) + return SplunkAOLogger() diff --git a/src/splunk_ao/tracing.py b/src/splunk_ao/tracing.py index a1092ec8..96f59c4c 100644 --- a/src/splunk_ao/tracing.py +++ b/src/splunk_ao/tracing.py @@ -1,45 +1,44 @@ -"""Utilities for distributed tracing with Splunk AO.""" +"""W3C distributed-tracing utilities for Splunk AO.""" -from splunk_ao.decorator import splunk_ao_context +from collections.abc import Mapping, MutableMapping +from opentelemetry import propagate +from opentelemetry.context import Context -def get_tracing_headers() -> dict[str, str]: - """ - Get tracing headers for distributed tracing (decorator usage). +from splunk_ao.exceptions import SplunkAOLoggerException +from splunk_ao.logger.logger import _has_active_exportable_span_context + + +def get_tracing_headers(carrier: MutableMapping[str, str] | None = None) -> MutableMapping[str, str]: + """Inject the active operation's W3C trace context into a carrier. - Returns headers from the singleton logger context that can be passed to downstream services. - For direct logger instances, use logger.get_tracing_headers() instead. + Parameters + ---------- + carrier : MutableMapping[str, str] | None + Existing carrier to populate. A new dictionary is created when omitted. Returns ------- - dict[str, str] - Dictionary with Splunk-AO-Trace-ID and Splunk-AO-Parent-ID headers + MutableMapping[str, str] + The supplied carrier, or a new dictionary, containing ``traceparent`` + and any other fields owned by the configured OTel propagator. Raises ------ SplunkAOLoggerException - If not in distributed mode or if no trace has been started - - Examples - -------- - Using with decorators to propagate trace context to downstream services: - - ```python - from splunk_ao import log, get_tracing_headers - import httpx - - @log() - async def orchestrator(): - # Get headers to pass to downstream service - headers = get_tracing_headers() - - # Call downstream service with trace context - async with httpx.AsyncClient() as client: - response = await client.post( - "http://service:8000/process", - headers=headers, - json={"data": "..."} - ) - ``` + If there is no active exportable operation span. The internal Splunk AO + trace envelope is not exportable and cannot be used as a wire parent. """ - return splunk_ao_context.get_logger_instance().get_tracing_headers() + if not _has_active_exportable_span_context(): + raise SplunkAOLoggerException("Distributed tracing requires an active exportable operation span") + + if carrier is None: + carrier = {} + propagate.inject(carrier) + return carrier + + +def extract_tracing_context(carrier: Mapping[str, str]) -> Context: + """Extract W3C trace context from an incoming text-map carrier.""" + normalized = {str(key).lower(): value for key, value in carrier.items()} + return propagate.extract(normalized) diff --git a/src/splunk_ao/utils/env_helpers.py b/src/splunk_ao/utils/env_helpers.py index 19132002..4b291a6d 100644 --- a/src/splunk_ao/utils/env_helpers.py +++ b/src/splunk_ao/utils/env_helpers.py @@ -20,8 +20,9 @@ def _get_mode_or_default(mode: str | None) -> LoggerModeType: ------- LoggerModeType The mode value to use: - - "batch": Uses scheduled OTLP batch export (default) - - "distributed": Also enables legacy trace_id/span_id continuation + Both "batch" and the temporarily retained "distributed" value use + scheduled OTLP batch export. The distinction remains only for the + deprecated ingestion-hook compatibility guard. """ if mode is None: mode = getenv("SPLUNK_AO_MODE", DEFAULT_MODE) diff --git a/src/splunk_ao/utils/singleton.py b/src/splunk_ao/utils/singleton.py index 99c00017..eb985dec 100644 --- a/src/splunk_ao/utils/singleton.py +++ b/src/splunk_ao/utils/singleton.py @@ -53,8 +53,6 @@ def _get_key( agent_stream_id: str | None, mode: str, experiment_id: str | None = None, - trace_id: str | None = None, - span_id: str | None = None, ingestion_hook_id: int | None = None, ) -> tuple[str, ...]: """ @@ -74,18 +72,13 @@ def _get_key( The experiment ID. mode: The logger mode. - trace_id: (Optional[str]) - The distributed trace ID. - span_id: (Optional[str]) - The distributed parent span ID. ingestion_hook_id: (Optional[int]) Identity of the temporary ingestion hook compatibility path. Returns ------- Tuple[str, ...] - A tuple key used for caching. Includes trace_id and span_id for proper - isolation of concurrent requests in async web servers. + A tuple key used for caching. """ _logger.debug("current thread is %s", threading.current_thread().name) @@ -121,11 +114,6 @@ def _get_key( destination_key = f"id:{routing.agent_stream_id or ''}" base_key = (*key, deployment.value, project_key, destination_key) - # Add trace_id and span_id to key if present (for distributed tracing) - if trace_id is not None: - base_key = (*base_key, trace_id) - if span_id is not None: - base_key = (*base_key, span_id) if ingestion_hook_id is not None: base_key = (*base_key, str(ingestion_hook_id)) @@ -158,8 +146,6 @@ def get( experiment_id: str | None = None, mode: str | None = None, local_metrics: list[LocalMetricConfig] | None = None, - trace_id: str | None = None, - span_id: str | None = None, ingestion_hook: Callable | None = None, ) -> SplunkAOLogger: """ @@ -197,8 +183,6 @@ def get( agent_stream_id, mode, experiment_id, - trace_id, - span_id, ingestion_hook_id=id(ingestion_hook) if ingestion_hook else None, ) @@ -221,8 +205,6 @@ def get( "experiment_id": experiment_id, "local_metrics": local_metrics, "mode": mode, - "trace_id": trace_id, - "span_id": span_id, "ingestion_hook": ingestion_hook, } # Create the logger with filtered kwargs. diff --git a/tests/conftest.py b/tests/conftest.py index cc3a8ba1..ba59c4bb 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -66,7 +66,6 @@ from splunk_ao.resources.models import DatasetContent, DatasetRow, DatasetRowValuesDict from splunk_ao.resources.models.messages_list_item import MessagesListItem from splunk_ao.utils.singleton import SplunkAOLoggerSingleton -from tests.testutils.setup import setup_thread_pool_request_capture class _TestSpanSink: @@ -380,29 +379,6 @@ def dataset_content_150_rows(): return DatasetContent(rows=rows) -@pytest.fixture -def thread_pool_capture(): - """ - Pytest fixture that provides a function to capture thread pool requests from distributed tracing methods. - - Usage: - def test_distributed_method(thread_pool_capture): - logger = SplunkAOLogger(project="test", agent_stream="test", mode="distributed") - capture = thread_pool_capture(logger) - - logger._ingest_trace_streaming(trace) - - capture.mock_pool.assert_called_once() - request = capture.get_latest_request() - assert isinstance(request, TracesIngestRequest) - """ - - def _capture_factory(logger): - return setup_thread_pool_request_capture(logger) - - return _capture_factory - - @pytest.fixture def enable_galileo_logging(): """Temporarily enable SDK logging for tests that need to capture log output.""" diff --git a/tests/test_async_base_handler.py b/tests/test_async_base_handler.py index 10e4d2f9..f2934e1d 100644 --- a/tests/test_async_base_handler.py +++ b/tests/test_async_base_handler.py @@ -4,6 +4,7 @@ import pytest +from splunk_ao import get_tracing_headers from splunk_ao.handlers.base_async_handler import SplunkAOAsyncBaseHandler from splunk_ao.logger.logger import SplunkAOLogger from tests.testutils.setup import setup_mock_logstreams_client, setup_mock_projects_client, setup_mock_traces_client @@ -30,6 +31,7 @@ def handler(self, splunk_ao_logger: SplunkAOLogger) -> Generator[SplunkAOAsyncBa yield handler # Clean up after each test handler._root_node = None + splunk_ao_logger.terminate() @pytest.mark.asyncio async def test_initialization(self, splunk_ao_logger: SplunkAOLogger) -> None: @@ -82,6 +84,24 @@ async def test_start_node(self, handler: SplunkAOAsyncBaseHandler) -> None: assert handler._root_node assert handler._root_node.run_id == parent_id + @pytest.mark.asyncio + async def test_live_root_is_reused_by_async_commit( + self, handler: SplunkAOAsyncBaseHandler, splunk_ao_logger: SplunkAOLogger + ) -> None: + run_id = uuid.uuid4() + await handler.async_start_node( + node_type="chain", parent_run_id=None, run_id=run_id, name="Root", input="request" + ) + live_root = handler._owned_root + live_context = splunk_ao_logger._otel_ids[live_root.id].span_context + + headers = get_tracing_headers() + await handler.async_end_node(run_id, output="done") + + assert headers["traceparent"].split("-")[2] == format(live_context.span_id, "016x") + assert splunk_ao_logger.traces[0].spans[0] is live_root + assert splunk_ao_logger.current_parent() is None + @pytest.mark.asyncio async def test_end_node(self, handler: SplunkAOAsyncBaseHandler, splunk_ao_logger: SplunkAOLogger) -> None: """Test ending a node and updating its parameters""" @@ -109,9 +129,42 @@ async def test_commit_failure_concludes_handler_owned_trace( run_id = uuid.uuid4() await handler.async_start_node(node_type="chain", parent_run_id=None, run_id=run_id, name="Test", input="test") - with patch.object(handler, "log_node_tree", side_effect=RuntimeError("conversion failed")): + with patch.object(handler, "_log_node_children", side_effect=RuntimeError("conversion failed")): await handler.async_end_node(run_id, output="result") assert splunk_ao_logger.current_parent() is None assert handler._nodes == {} assert handler._root_node is None + + @pytest.mark.asyncio + async def test_start_new_trace_false_preserves_handler_root_and_caller( + self, splunk_ao_logger: SplunkAOLogger + ) -> None: + splunk_ao_logger.start_trace(input="request", name="caller") + caller_operation = splunk_ao_logger.add_workflow_span(input="outer", name="outer") + handler = SplunkAOAsyncBaseHandler( + splunk_ao_logger=splunk_ao_logger, start_new_trace=False, flush_on_chain_end=False + ) + root_id = uuid.uuid4() + child_id = uuid.uuid4() + await handler.async_start_node( + node_type="chain", parent_run_id=None, run_id=root_id, name="handler-root", input="work" + ) + await handler.async_start_node( + node_type="llm", + parent_run_id=root_id, + run_id=child_id, + name="handler-child", + input="prompt", + output="answer", + model="model", + ) + + await handler.async_end_node(root_id, output="done") + + assert splunk_ao_logger.current_parent() is caller_operation + [handler_root] = caller_operation.spans + assert handler_root.name == "handler-root" + assert [child.name for child in handler_root.spans] == ["handler-child"] + splunk_ao_logger.conclude(output="outer done") + splunk_ao_logger.conclude(output="done") diff --git a/tests/test_base_handler.py b/tests/test_base_handler.py index a730d652..826a407d 100644 --- a/tests/test_base_handler.py +++ b/tests/test_base_handler.py @@ -4,8 +4,10 @@ import pytest +from splunk_ao import get_tracing_headers from splunk_ao.handlers.base_handler import SplunkAOBaseHandler from splunk_ao.logger.logger import SplunkAOLogger +from splunk_ao.schema.logged import LoggedAgentSpan from tests.testutils.setup import setup_mock_logstreams_client, setup_mock_projects_client, setup_mock_traces_client @@ -24,9 +26,9 @@ def splunk_ao_logger(self, mock_traces_client: Mock, mock_projects_client: Mock, @pytest.fixture def handler(self, splunk_ao_logger: SplunkAOLogger) -> Generator[SplunkAOBaseHandler, None, None]: """Creates a SplunkAOBaseHandler with a mock logger""" - return SplunkAOBaseHandler(splunk_ao_logger=splunk_ao_logger, flush_on_chain_end=False) - # Reset the root node before each test - # Clean up after each test + handler = SplunkAOBaseHandler(splunk_ao_logger=splunk_ao_logger, flush_on_chain_end=False) + yield handler + splunk_ao_logger.terminate() def test_initialization(self, splunk_ao_logger: SplunkAOLogger) -> None: """Test callback initialization with various parameters""" @@ -77,6 +79,45 @@ def test_start_node(self, handler: SplunkAOBaseHandler) -> None: assert handler._root_node assert handler._root_node.run_id == parent_id + def test_root_is_live_for_w3c_injection_and_reused_at_commit( + self, handler: SplunkAOBaseHandler, splunk_ao_logger: SplunkAOLogger + ) -> None: + run_id = uuid.uuid4() + handler.start_node(node_type="chain", parent_run_id=None, run_id=run_id, name="Root", input="request") + live_root = handler._owned_root + live_context = splunk_ao_logger._otel_ids[live_root.id].span_context + + headers = get_tracing_headers() + handler.end_node(run_id, output="done") + + assert headers["traceparent"].split("-")[2] == format(live_context.span_id, "016x") + assert len(splunk_ao_logger.traces[0].spans) == 1 + assert splunk_ao_logger.traces[0].spans[0] is live_root + assert splunk_ao_logger.current_parent() is None + + def test_late_langgraph_agent_classification_preserves_live_otel_identity( + self, handler: SplunkAOBaseHandler, splunk_ao_logger: SplunkAOLogger + ) -> None: + root_id = uuid.uuid4() + child_id = uuid.uuid4() + root_node = handler.start_node( + node_type="chain", parent_run_id=None, run_id=root_id, name="Root", input="request" + ) + original_root = handler._owned_root + original_ids = splunk_ao_logger._otel_ids[original_root.id] + + root_node.node_type = "agent" + handler.start_node(node_type="llm", parent_run_id=root_id, run_id=child_id, name="Child", input="prompt") + + assert isinstance(handler._owned_root, LoggedAgentSpan) + assert handler._owned_root.id == original_root.id + assert splunk_ao_logger._otel_ids[handler._owned_root.id] is original_ids + assert splunk_ao_logger.current_parent() is handler._owned_root + + handler.end_node(child_id, output="answer", model="model") + handler.end_node(root_id, output="done") + assert isinstance(splunk_ao_logger.traces[0].spans[0], LoggedAgentSpan) + def test_end_node(self, handler: SplunkAOBaseHandler, splunk_ao_logger: SplunkAOLogger) -> None: """Test ending a node and updating its parameters""" # Create a node @@ -146,7 +187,7 @@ def test_commit_failure_concludes_only_handler_owned_trace( run_id = uuid.uuid4() handler.start_node(node_type="chain", parent_run_id=None, run_id=run_id, name="Test", input="test") - with patch.object(handler, "log_node_tree", side_effect=RuntimeError("conversion failed")): + with patch.object(handler, "_log_node_children", side_effect=RuntimeError("conversion failed")): handler.end_node(run_id, output="result") assert splunk_ao_logger.current_parent() is None @@ -154,7 +195,8 @@ def test_commit_failure_concludes_only_handler_owned_trace( assert handler._root_node is None def test_commit_failure_preserves_caller_owned_trace(self, splunk_ao_logger: SplunkAOLogger) -> None: - caller_trace = splunk_ao_logger.start_trace(input="request", name="caller") + splunk_ao_logger.start_trace(input="request", name="caller") + caller_operation = splunk_ao_logger.add_workflow_span(input="outer", name="outer") handler = SplunkAOBaseHandler( splunk_ao_logger=splunk_ao_logger, start_new_trace=False, flush_on_chain_end=False ) @@ -164,5 +206,39 @@ def test_commit_failure_preserves_caller_owned_trace(self, splunk_ao_logger: Spl with patch.object(handler, "log_node_tree", side_effect=RuntimeError("conversion failed")): handler.end_node(run_id, output="result") - assert splunk_ao_logger.current_parent() is caller_trace + assert splunk_ao_logger.current_parent() is caller_operation + splunk_ao_logger.conclude(output="outer done") + splunk_ao_logger.conclude(output="done") + + def test_start_new_trace_false_preserves_handler_root_beneath_caller( + self, splunk_ao_logger: SplunkAOLogger + ) -> None: + splunk_ao_logger.start_trace(input="request", name="caller") + caller_operation = splunk_ao_logger.add_workflow_span(input="outer", name="outer") + handler = SplunkAOBaseHandler( + splunk_ao_logger=splunk_ao_logger, start_new_trace=False, flush_on_chain_end=False + ) + caller_context = splunk_ao_logger._otel_ids[caller_operation.id].span_context + root_id = uuid.uuid4() + child_id = uuid.uuid4() + handler.start_node(node_type="chain", parent_run_id=None, run_id=root_id, name="handler-root", input="work") + headers = get_tracing_headers() + handler.start_node( + node_type="llm", + parent_run_id=root_id, + run_id=child_id, + name="handler-child", + input="prompt", + output="answer", + model="model", + ) + + handler.end_node(root_id, output="done") + + assert splunk_ao_logger.current_parent() is caller_operation + assert headers["traceparent"].split("-")[2] == format(caller_context.span_id, "016x") + [handler_root] = caller_operation.spans + assert handler_root.name == "handler-root" + assert [child.name for child in handler_root.spans] == ["handler-child"] + splunk_ao_logger.conclude(output="outer done") splunk_ao_logger.conclude(output="done") diff --git a/tests/test_decorator_distributed.py b/tests/test_decorator_distributed.py index d01f74ef..3b6a2e77 100644 --- a/tests/test_decorator_distributed.py +++ b/tests/test_decorator_distributed.py @@ -3,12 +3,11 @@ from unittest.mock import Mock, patch import pytest +from opentelemetry import context as otel_context from galileo_core.schemas.shared.document import Document from galileo_core.schemas.shared.multimodal import ContentModality -from splunk_ao import Message, MessageRole, log, splunk_ao_context -from splunk_ao.constants.tracing import PARENT_ID_HEADER, TRACE_ID_HEADER -from splunk_ao.decorator import _parent_id_context, _trace_id_context +from splunk_ao import Message, MessageRole, extract_tracing_context, log, splunk_ao_context from splunk_ao.schema.content_blocks import DataContentBlock, TextContentBlock from splunk_ao.tracing import get_tracing_headers from tests.testutils.setup import setup_mock_logstreams_client, setup_mock_projects_client, setup_mock_traces_client @@ -49,33 +48,40 @@ def test_decorator_get_tracing_headers(reset_context: None, distributed_clients: @log(span_type="workflow") def orchestrator(query: str) -> dict: - trace = splunk_ao_context.get_current_trace() - return {"result": query, "headers": get_tracing_headers(), "trace_id": str(trace.id)} + return {"result": query, "headers": get_tracing_headers()} result = orchestrator("test input") headers = result["headers"] - assert headers[TRACE_ID_HEADER] == result["trace_id"] - assert PARENT_ID_HEADER in headers + assert "traceparent" in headers + assert not any(header.lower().startswith("splunk-ao-") for header in headers) + traceparent = headers["traceparent"].split("-") + [span] = logger._sink.spans + assert int(traceparent[1], 16) == span.context.trace_id + assert int(traceparent[2], 16) == span.context.span_id assert logger.current_parent() is None def test_decorator_respects_incoming_distributed_context(reset_context: None, distributed_clients: Mock) -> None: - trace_id = "12345678-1234-4678-9abc-123456789abc" - parent_id = "87654321-4321-4876-9cba-987654321cba" - _trace_id_context.set(trace_id) - _parent_id_context.set(parent_id) + trace_id = "4bf92f3577b34da6a3ce929d0e0e4736" + parent_id = "00f067aa0ba902b7" logger = init_logger() @log(span_type="workflow") def downstream_service(query: str) -> str: return f"processed: {query}" - assert downstream_service("test input") == "processed: test input" + token = otel_context.attach(extract_tracing_context({"traceparent": f"00-{trace_id}-{parent_id}-01"})) + try: + assert downstream_service("test input") == "processed: test input" + finally: + otel_context.detach(token) + assert logger.mode == "distributed" - assert str(logger.traces[0].id) == trace_id - assert logger.traces[0].name == "stub_trace" - assert (logger._sink.spans[-1].attributes or {})["gen_ai.operation.name"] == "invoke_workflow" + exported = logger._sink.spans[-1] + assert exported.context.trace_id == int(trace_id, 16) + assert exported.parent.span_id == int(parent_id, 16) + assert (exported.attributes or {})["gen_ai.operation.name"] == "invoke_workflow" def test_completed_workflow_is_enqueued_and_flush_does_not_change_ownership( diff --git a/tests/test_logger_distributed.py b/tests/test_logger_distributed.py index 647c1e5b..c1026c5c 100644 --- a/tests/test_logger_distributed.py +++ b/tests/test_logger_distributed.py @@ -1,10 +1,11 @@ from collections.abc import Generator -from unittest.mock import Mock -from uuid import uuid4 +from unittest.mock import Mock, patch import pytest +from opentelemetry import context from opentelemetry.sdk.trace import ReadableSpan +from splunk_ao import extract_tracing_context from splunk_ao.logger import SplunkAOLogger @@ -37,6 +38,23 @@ def operation_names(spans: list[ReadableSpan]) -> list[str | None]: return [(span.attributes or {}).get("gen_ai.operation.name") for span in spans] +@pytest.mark.parametrize("mode", ["batch", "distributed"]) +def test_both_mode_values_construct_the_shared_batch_span_sink(mode: str) -> None: + sink = RecordingSink() + exporter = Mock() + with ( + patch("splunk_ao.logger.logger.build_standalone_exporter", return_value=exporter), + patch("splunk_ao.logger.logger.build_span_sink", return_value=sink) as build_sink, + ): + logger = SplunkAOLogger(project_id="project-id", agent_stream_id="stream-id", mode=mode) + + try: + build_sink.assert_called_once_with(exporter) + assert not hasattr(logger, "_task_handler") + finally: + logger.terminate() + + def test_distributed_mode_uses_otlp_completion_queue(distributed_logger: tuple[SplunkAOLogger, RecordingSink]) -> None: logger, sink = distributed_logger root = logger.start_trace(input="question") @@ -125,46 +143,82 @@ def test_distributed_conclude_all_emits_inner_to_outer( assert len({span.context.trace_id for span in sink.spans if span.context is not None}) == 1 -def test_distributed_continuation_keeps_stub_parent_local() -> None: +@pytest.mark.parametrize("mode", ["batch", "distributed"]) +def test_retained_mode_values_use_identical_w3c_otlp_behavior(mode: str) -> None: sink = RecordingSink() - trace_id = str(uuid4()) - parent_id = str(uuid4()) - logger = SplunkAOLogger( - project_id="project-id", - agent_stream_id="log-stream-id", - mode="distributed", - trace_id=trace_id, - span_id=parent_id, - _sink=sink, - ) + trace_id = "4bf92f3577b34da6a3ce929d0e0e4736" + parent_id = "00f067aa0ba902b7" + token = context.attach(extract_tracing_context({"traceparent": f"00-{trace_id}-{parent_id}-01"})) + logger = SplunkAOLogger(project_id="project-id", agent_stream_id="log-stream-id", mode=mode, _sink=sink) try: - assert str(logger.traces[0].id) == trace_id - assert str(logger.current_parent().id) == parent_id + logger.start_trace(input="request") + operation = logger.add_workflow_span(input="work", name="operation") + operation_context = logger._otel_ids[operation.id] + logger.conclude(output="done") + logger.conclude(output="complete") + + assert operation_context.span_context.trace_id == int(trace_id, 16) + assert operation_context.parent_span_context is not None + assert operation_names(sink.spans) == ["invoke_workflow"] + assert sink.spans[0].parent is not None + assert sink.spans[0].parent.span_id == int(parent_id, 16) + assert sink.force_flush_calls == 0 + assert logger.current_parent() is None + finally: + logger.terminate() + context.detach(token) - logger.add_llm_span(input="prompt", output="answer", model="model") - assert operation_names(sink.spans) == ["chat"] - assert sink.spans[0].parent == logger._otel_ids[logger.current_parent().id].span_context +@pytest.mark.parametrize("mode", ["batch", "distributed"]) +def test_mode_does_not_change_parent_child_topology_or_require_flush(mode: str) -> None: + sink = RecordingSink() + logger = SplunkAOLogger(project_id="project-id", agent_stream_id="log-stream-id", mode=mode, _sink=sink) + try: + logger.start_trace(input="request") + logger.add_workflow_span(input="workflow") + logger.add_agent_span(input="agent") + logger.add_llm_span(input="prompt", output="answer", model="model") + logger.conclude(output="complete", conclude_all=True) + + llm, agent, workflow = sink.spans + assert operation_names(sink.spans) == ["chat", "invoke_agent", "invoke_workflow"] + assert llm.parent == agent.context + assert agent.parent == workflow.context + assert workflow.parent is None + assert sink.force_flush_calls == 0 + assert logger.current_parent() is None finally: logger.terminate() -def test_distributed_terminate_discards_unfinished_stubs() -> None: +def test_unsampled_remote_parent_is_preserved_and_not_exported() -> None: sink = RecordingSink() - logger = SplunkAOLogger( - project_id="project-id", - agent_stream_id="log-stream-id", - mode="distributed", - trace_id=str(uuid4()), - span_id=str(uuid4()), - _sink=sink, + token = context.attach( + extract_tracing_context({"traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-00"}) ) + logger = SplunkAOLogger(project_id="project-id", agent_stream_id="log-stream-id", _sink=sink) + try: + logger.start_trace(input="request") + operation = logger.add_workflow_span(input="work") + operation_context = logger._otel_ids[operation.id].span_context + logger.add_llm_span(input="prompt", output="answer", model="model") + logger.conclude(output="done") + logger.conclude(output="complete") - logger.terminate() - logger.terminate() + assert operation_context.is_valid + assert not bool(operation_context.trace_flags) + assert sink.spans == [] + finally: + logger.terminate() + context.detach(token) - assert sink.spans == [] - assert sink.force_flush_calls == 1 - assert sink.shutdown_calls == 1 - assert logger.current_parent() is None - assert logger._otel_ids == {} + +def test_custom_trace_and_span_id_constructor_arguments_are_removed() -> None: + with pytest.raises(TypeError, match="trace_id"): + SplunkAOLogger(trace_id="4bf92f35-77b3-4da6-a3ce-929d0e0e4736") + with pytest.raises(TypeError, match="span_id"): + SplunkAOLogger(span_id="00f067aa-0ba9-42b7-8000-000000000000") + + +def test_old_logger_header_method_is_removed() -> None: + assert not hasattr(SplunkAOLogger, "get_tracing_headers") diff --git a/tests/test_logger_otel_context.py b/tests/test_logger_otel_context.py index a9481b26..62a84d35 100644 --- a/tests/test_logger_otel_context.py +++ b/tests/test_logger_otel_context.py @@ -106,29 +106,24 @@ def test_completed_leaf_siblings_share_parent_and_do_not_become_active( logger.conclude(output="done") -def test_otel_identity_failure_does_not_interrupt_span_creation_or_legacy_streaming( +def test_otel_identity_failure_does_not_interrupt_span_creation( make_logger: Callable[[], SplunkAOLogger], monkeypatch: pytest.MonkeyPatch ) -> None: logger = make_logger() root = logger.start_trace(input="q") - ingest_step = Mock() warning = Mock() def fail_assignment(*args, **kwargs) -> SpanContext: raise RuntimeError("identity failure") - logger.mode = "distributed" with monkeypatch.context() as patch: patch.setattr(SplunkAOLogger, "_assign_otel_context", fail_assignment) - patch.setattr(SplunkAOLogger, "_ingest_step_streaming", ingest_step) patch.setattr(logger._logger, "warning", warning) span = logger.add_llm_span(input="prompt", output="answer", model="model") - logger.mode = "batch" assert span is not None assert span in root.spans assert span.id not in logger._otel_ids - ingest_step.assert_not_called() assert "Failed to assign OTel identity" in warning.call_args.args[0] logger.conclude(output="done") @@ -139,22 +134,17 @@ def test_otel_sync_failure_does_not_interrupt_parentable_span_creation( ) -> None: logger = make_logger() root = logger.start_trace(input="q") - ingest_step = Mock() warning = Mock() - logger.mode = "distributed" with monkeypatch.context() as patch: patch.setattr(SplunkAOLogger, "_sync_otel_context_impl", Mock(side_effect=RuntimeError("sync failure"))) - patch.setattr(SplunkAOLogger, "_ingest_step_streaming", ingest_step) patch.setattr(logger._logger, "warning", warning) workflow = logger.add_workflow_span(input="workflow") - logger.mode = "batch" assert workflow is not None assert workflow in root.spans assert logger.current_parent() is workflow assert workflow.id in logger._otel_ids - ingest_step.assert_not_called() assert "Failed to synchronize OTel context" in warning.call_args.args[0] logger._sync_otel_context(workflow) @@ -168,22 +158,17 @@ def test_otel_sync_and_release_failures_do_not_interrupt_hook_conclusion( logger = make_logger() root = logger.start_trace(input="q") workflow = logger.add_workflow_span(input="workflow") - update_step = Mock() warning = Mock() - logger.mode = "distributed" with monkeypatch.context() as patch: patch.setattr(SplunkAOLogger, "_sync_otel_context_impl", Mock(side_effect=RuntimeError("sync failure"))) patch.setattr(SplunkAOLogger, "_discard_otel_subtree", Mock(side_effect=RuntimeError("release failure"))) - patch.setattr(SplunkAOLogger, "_update_step_streaming", update_step) patch.setattr(logger._logger, "warning", warning) parent = logger.conclude(output="workflow-output") - logger.mode = "batch" assert parent is root assert logger.current_parent() is root assert workflow.output == "workflow-output" - update_step.assert_not_called() warning_messages = [call.args[0] for call in warning.call_args_list] assert any("Failed to synchronize OTel context" in message for message in warning_messages) assert any("Failed to release OTel context" in message for message in warning_messages) @@ -287,7 +272,7 @@ def test_start_trace_inherits_remote_parent_and_preserves_tracestate(make_logger context.detach(token) -def test_trace_flags_are_always_sampled_with_unsampled_remote_parent(make_logger: Callable[[], SplunkAOLogger]) -> None: +def test_trace_flags_preserve_unsampled_remote_parent(make_logger: Callable[[], SplunkAOLogger]) -> None: remote_context = propagate.extract({"traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-00"}) token = context.attach(remote_context) try: @@ -295,8 +280,8 @@ def test_trace_flags_are_always_sampled_with_unsampled_remote_parent(make_logger root = logger.start_trace(input="q") child = logger.add_llm_span(input="prompt", output="answer", model="model") - assert logger._otel_ids[root.id].span_context.trace_flags == TraceFlags(TraceFlags.SAMPLED) - assert logger._otel_ids[child.id].span_context.trace_flags == TraceFlags(TraceFlags.SAMPLED) + assert logger._otel_ids[root.id].span_context.trace_flags == TraceFlags(TraceFlags.DEFAULT) + assert logger._otel_ids[child.id].span_context.trace_flags == TraceFlags(TraceFlags.DEFAULT) logger.conclude(output="a") finally: context.detach(token) diff --git a/tests/test_middleware_tracing.py b/tests/test_middleware_tracing.py index 81ee7751..e6f5680a 100644 --- a/tests/test_middleware_tracing.py +++ b/tests/test_middleware_tracing.py @@ -1,288 +1,141 @@ -"""Tests for distributed tracing middleware.""" +"""Tests for W3C tracing middleware.""" -from unittest.mock import Mock, patch +from typing import Any +from unittest.mock import Mock import pytest from fastapi import FastAPI from fastapi.testclient import TestClient +from opentelemetry import trace +from opentelemetry.sdk.trace import ReadableSpan -from splunk_ao.constants.tracing import PARENT_ID_HEADER, TRACE_ID_HEADER -from splunk_ao.decorator import _parent_id_context, _trace_id_context +from splunk_ao import get_tracing_headers from splunk_ao.logger import SplunkAOLogger from splunk_ao.middleware import TracingMiddleware, get_request_logger -from tests.testutils.setup import setup_mock_logstreams_client, setup_mock_projects_client, setup_mock_traces_client -@pytest.fixture -def app(): - """Create a test FastAPI app with tracing middleware.""" - app = FastAPI() - app.add_middleware(TracingMiddleware) - - @app.get("/test") - async def test_endpoint(): - logger = get_request_logger() - return { - "trace_id": logger.trace_id, - "span_id": logger.span_id, - "is_logger": isinstance(logger, SplunkAOLogger), # Verify logger was created - } - - @app.get("/test-context") - async def test_context_endpoint(): - # Direct access to context variables - return {"trace_id": _trace_id_context.get(), "parent_id": _parent_id_context.get()} +class RecordingSink: + def __init__(self) -> None: + self.spans: list[ReadableSpan] = [] - @app.get("/test-add-span") - async def test_add_span_endpoint(): - """Endpoint that tries to add a span to test distributed tracing stub creation. + def emit(self, span: ReadableSpan) -> None: + self.spans.append(span) - Note: In distributed tracing (distributed mode), when span_id is provided, - the logger creates local stub objects in __init__ without fetching from backend. - """ - try: - logger = get_request_logger() - # The validation should have happened during get_request_logger() call above - # If we get here without exception, validation passed (or didn't run) - # Try to add a span to confirm logger works - logger.add_workflow_span(input="test", name="test_span") - return { - "error": None, - "success": True, - "trace_id": str(logger.trace_id), - "span_id": str(logger.span_id) if logger.span_id else None, - } - except Exception as e: - return {"error": str(e), "error_type": type(e).__name__, "success": False} + def force_flush(self) -> bool: + return True - return app + def shutdown(self) -> None: + return None @pytest.fixture -def client(app): - """Create a test client.""" - return TestClient(app) - - -@patch("splunk_ao.logger.logger.AgentStreams") -@patch("splunk_ao.logger.logger.Projects") -@patch("splunk_ao.logger.logger.Traces") -def test_middleware_extracts_headers( - mock_traces_client: Mock, mock_projects_client: Mock, mock_logstreams_client: Mock, client: TestClient -): - """Test that middleware correctly extracts tracing headers.""" - setup_mock_traces_client(mock_traces_client) - setup_mock_projects_client(mock_projects_client) - setup_mock_logstreams_client(mock_logstreams_client) - - # Use valid UUID4 values - trace_id = "6c4e3f7e-4a9a-4e7e-8c1f-3a9a3a9a3a9d" - parent_id = "6c4e3f7e-4a9a-4e7e-8c1f-3a9a3a9a3a9e" - - # Override get_span to return a span that belongs to the correct trace - # Also override get_trace to return a trace with the correct trace_id - import datetime - from unittest.mock import AsyncMock - from uuid import UUID - - now = datetime.datetime.now() - mock_span = { - "id": UUID(parent_id), - "trace_id": UUID(trace_id), # Same trace as header - "type": "workflow", - "name": "test-workflow-span", - "input": "test-input", - "output": None, - "created_at": now, - "updated_at": now, - "user_metadata": {}, - "metrics": {}, - "parent_id": None, - } - mock_trace = { - "id": UUID(trace_id), - "name": "test-trace", - "type": "trace", - "input": "test-input", - "output": None, - "created_at": now, - "updated_at": now, - "user_metadata": {}, - "spans": [], - } - mock_instance = mock_traces_client.return_value - mock_instance.get_span = AsyncMock(return_value=mock_span) - mock_instance.get_trace = AsyncMock(return_value=mock_trace) - - response = client.get("/test", headers={TRACE_ID_HEADER: trace_id, PARENT_ID_HEADER: parent_id}) - - assert response.status_code == 200 - data = response.json() - assert data["is_logger"] is True - assert data["trace_id"] == trace_id - assert data["span_id"] == parent_id - - -@patch("splunk_ao.logger.logger.AgentStreams") -@patch("splunk_ao.logger.logger.Projects") -@patch("splunk_ao.logger.logger.Traces") -def test_middleware_handles_missing_headers( - mock_traces_client: Mock, mock_projects_client: Mock, mock_logstreams_client: Mock, client: TestClient -): - """Test that middleware handles requests without tracing headers.""" - setup_mock_traces_client(mock_traces_client) - setup_mock_projects_client(mock_projects_client) - setup_mock_logstreams_client(mock_logstreams_client) - - response = client.get("/test") - - assert response.status_code == 200 - data = response.json() - assert data["is_logger"] is True - assert data["trace_id"] is None - assert data["span_id"] is None - - -@patch("splunk_ao.logger.logger.AgentStreams") -@patch("splunk_ao.logger.logger.Projects") -@patch("splunk_ao.logger.logger.Traces") -def test_middleware_handles_partial_headers( - mock_traces_client: Mock, mock_projects_client: Mock, mock_logstreams_client: Mock, client: TestClient -): - """Test that middleware handles requests with only trace ID. - - Note: In normal operation, both TRACE_ID_HEADER and PARENT_ID_HEADER should be present. +def app_factory(monkeypatch: pytest.MonkeyPatch): + loggers: list[SplunkAOLogger] = [] + sinks: list[RecordingSink] = [] - However, when only TRACE_ID_HEADER is provided (no PARENT_ID_HEADER), the middleware - handles this by creating a logger with the trace_id but no span_id, - allowing the service to start a new root span within the existing trace. - """ - setup_mock_traces_client(mock_traces_client) - setup_mock_projects_client(mock_projects_client) - setup_mock_logstreams_client(mock_logstreams_client) + def logger_factory() -> SplunkAOLogger: + sink = RecordingSink() + logger = SplunkAOLogger(project_id="project-id", agent_stream_id="stream-id", _sink=sink) + loggers.append(logger) + sinks.append(sink) + return logger - trace_id = "12345678-1234-4678-9abc-123456789abc" + monkeypatch.setattr("splunk_ao.middleware.tracing.SplunkAOLogger", logger_factory) - response = client.get("/test", headers={TRACE_ID_HEADER: trace_id}) + def make_app(*, fail: bool = False) -> FastAPI: + app = FastAPI() + app.add_middleware(TracingMiddleware) - assert response.status_code == 200 - data = response.json() - assert data["is_logger"] is True - assert data["trace_id"] == trace_id - assert data["span_id"] is None - - -@patch("splunk_ao.logger.logger.AgentStreams") -@patch("splunk_ao.logger.logger.Projects") -@patch("splunk_ao.logger.logger.Traces") -def test_context_cleanup_after_request( - mock_traces_client: Mock, mock_projects_client: Mock, mock_logstreams_client: Mock, client: TestClient -): - """Test that context variables are cleaned up after request.""" - setup_mock_traces_client(mock_traces_client) - setup_mock_projects_client(mock_projects_client) - setup_mock_logstreams_client(mock_logstreams_client) - - trace_id = "12345678-1234-4678-9abc-123456789abc" - parent_id = "87654321-4321-4876-9cba-987654321cba" - - # First request with headers - response1 = client.get("/test-context", headers={TRACE_ID_HEADER: trace_id, PARENT_ID_HEADER: parent_id}) - assert response1.status_code == 200 - data1 = response1.json() - assert data1["trace_id"] == trace_id - assert data1["parent_id"] == parent_id - - # Second request without headers - should not see previous request's context - response2 = client.get("/test-context") - assert response2.status_code == 200 - data2 = response2.json() - assert data2["trace_id"] is None - assert data2["parent_id"] is None - - -@patch("splunk_ao.logger.logger.AgentStreams") -@patch("splunk_ao.logger.logger.Projects") -@patch("splunk_ao.logger.logger.Traces") -def test_get_request_logger_when_parent_id_equals_trace_id( - mock_traces_client: Mock, mock_projects_client: Mock, mock_logstreams_client: Mock, client: TestClient -): - """Test that get_request_logger handles the case when parent_id equals trace_id. - - When upstream services forward headers immediately after start_trace(), both - Splunk-AO-Trace-ID and Splunk-AO-Parent-ID are identical (the root trace id). - In this case, we should pass None as span_id to avoid SplunkAOLoggerException. - """ - setup_mock_traces_client(mock_traces_client) - setup_mock_projects_client(mock_projects_client) - setup_mock_logstreams_client(mock_logstreams_client) - - trace_id = "12345678-1234-4678-9abc-123456789abc" - # parent_id equals trace_id (first hop after start_trace) - parent_id = trace_id - - response = client.get("/test", headers={TRACE_ID_HEADER: trace_id, PARENT_ID_HEADER: parent_id}) + @app.get("/test") + async def endpoint() -> dict[str, Any]: + logger = get_request_logger() + try: + logger.start_trace(input="request", name="request") + operation = logger.add_workflow_span(input="work", name="operation") + ids = logger._otel_ids[operation.id] + outbound = get_tracing_headers() + if fail: + raise RuntimeError("request failed") + logger.conclude(output="done") + logger.conclude(output="complete") + return { + "trace_id": format(ids.span_context.trace_id, "032x"), + "span_id": format(ids.span_context.span_id, "016x"), + "parent_span_id": ( + format(ids.parent_span_context.span_id, "016x") if ids.parent_span_context is not None else None + ), + "outbound": outbound, + } + finally: + logger.terminate() + + return app + + return make_app, loggers, sinks + + +def test_middleware_and_request_logger_continue_w3c_parent(app_factory) -> None: + make_app, _, sinks = app_factory + trace_id = "4bf92f3577b34da6a3ce929d0e0e4736" + parent_id = "00f067aa0ba902b7" + response = TestClient(make_app()).get( + "/test", headers={"traceparent": f"00-{trace_id}-{parent_id}-01", "tracestate": "vendor=value"} + ) assert response.status_code == 200 data = response.json() - - # When parent_id equals trace_id, span_id should be None (not the trace_id) - # This prevents SplunkAOLogger from trying to look up a span with the trace_id - assert data["is_logger"] is True assert data["trace_id"] == trace_id - assert data["span_id"] is None - - -@patch("splunk_ao.logger.logger.AgentStreams") -@patch("splunk_ao.logger.logger.Projects") -@patch("splunk_ao.logger.logger.Traces") -def test_mismatched_trace_and_span_ids( - mock_traces_client: Mock, mock_projects_client: Mock, mock_logstreams_client: Mock, client: TestClient -): - """Test that get_request_logger creates stubs for distributed tracing. - - In distributed tracing (distributed mode), when trace_id and span_id are provided, - the logger creates local stub objects instead of fetching from the backend. - This allows distributed tracing to work without backend validation, with eventual - consistency handling any mismatches during ingestion retries. - """ - setup_mock_traces_client(mock_traces_client) - setup_mock_projects_client(mock_projects_client) - setup_mock_logstreams_client(mock_logstreams_client) - - trace_id = "12345678-1234-4678-9abc-123456789abc" - parent_id = "99999999-9999-4999-9999-999999999999" - - response = client.get("/test-add-span", headers={TRACE_ID_HEADER: trace_id, PARENT_ID_HEADER: parent_id}) + assert data["outbound"]["tracestate"] == "vendor=value" + [exported] = sinks[0].spans + assert exported.context.trace_id == int(trace_id, 16) + assert exported.parent.span_id == int(parent_id, 16) + + +def test_middleware_without_header_starts_new_trace_and_does_not_leak(app_factory) -> None: + make_app, _, sinks = app_factory + client = TestClient(make_app()) + first = client.get( + "/test", headers={"traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"} + ).json() + second = client.get("/test").json() + + assert first["trace_id"] == "4bf92f3577b34da6a3ce929d0e0e4736" + assert second["trace_id"] != first["trace_id"] + [first_exported] = sinks[0].spans + [second_exported] = sinks[1].spans + assert first_exported.parent.span_id == int("00f067aa0ba902b7", 16) + assert second_exported.parent is None + + +def test_middleware_ignores_proprietary_headers(app_factory) -> None: + make_app, _, sinks = app_factory + response = TestClient(make_app()).get( + "/test", + headers={ + "Splunk-AO-Trace-ID": "4bf92f35-77b3-4da6-a3ce-929d0e0e4736", + "Splunk-AO-Parent-ID": "00f067aa-0ba9-42b7-8000-000000000000", + }, + ) assert response.status_code == 200 - data = response.json() + assert response.json()["trace_id"] != "4bf92f3577b34da6a3ce929d0e0e4736" + [exported] = sinks[0].spans + assert exported.parent is None - # With stubs, logger initialization succeeds without backend validation - # The stub objects are created locally, allowing distributed tracing to work - # even when the parent span hasn't been ingested yet (eventual consistency) - assert data["success"] is True, "Expected logger initialization to succeed with stubs" - assert data["error"] is None - assert data["trace_id"] == trace_id - assert data["span_id"] == parent_id +def test_middleware_detaches_context_on_exception(app_factory) -> None: + make_app, _, _ = app_factory + with pytest.raises(RuntimeError, match="request failed"): + TestClient(make_app(fail=True)).get( + "/test", headers={"traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"} + ) -@patch("splunk_ao.logger.logger.Projects") -@patch("splunk_ao.logger.logger.AgentStreams") -def test_invalid_uuid_headers_raise_exception( - mock_logstreams_client: Mock, mock_projects_client: Mock, app: FastAPI, client: TestClient -): - """Test that invalid UUID headers raise SplunkAOLoggerException.""" - from splunk_ao.exceptions import SplunkAOLoggerException + assert not trace.get_current_span().get_span_context().is_valid - setup_mock_projects_client(mock_projects_client) - setup_mock_logstreams_client(mock_logstreams_client) - # Test with completely invalid trace_id - should raise exception - with pytest.raises(SplunkAOLoggerException, match="Invalid trace_id"): - client.get("/test", headers={TRACE_ID_HEADER: "not-a-valid-uuid"}) +def test_get_request_logger_constructs_plain_request_scoped_logger(monkeypatch: pytest.MonkeyPatch) -> None: + constructor = Mock(return_value=Mock(spec=SplunkAOLogger)) + monkeypatch.setattr("splunk_ao.middleware.tracing.SplunkAOLogger", constructor) - # Test with valid trace_id but invalid parent_id - should raise exception - valid_trace_id = "12345678-1234-4678-9abc-123456789abc" - with pytest.raises(SplunkAOLoggerException, match="Invalid span_id"): - client.get("/test", headers={TRACE_ID_HEADER: valid_trace_id, PARENT_ID_HEADER: "invalid-parent"}) + assert get_request_logger() is constructor.return_value + constructor.assert_called_once_with() diff --git a/tests/test_openai_agents.py b/tests/test_openai_agents.py index 16ed813e..205bc6f9 100644 --- a/tests/test_openai_agents.py +++ b/tests/test_openai_agents.py @@ -20,6 +20,7 @@ from pytest import MonkeyPatch, mark from galileo_core.schemas.logging.span import LlmSpan, ToolSpan +from splunk_ao import get_tracing_headers from splunk_ao.handlers.openai_agents import SplunkAOTracingProcessor from splunk_ao.logger.logger import SplunkAOLogger from splunk_ao.schema.handlers import Node @@ -182,18 +183,13 @@ def test_commit_failure_concludes_handler_owned_trace( processor = SplunkAOTracingProcessor(splunk_ao_logger=logger) trace = MagicMock(trace_id="trace-id", name="Agent trace", metadata={}) processor.on_trace_start(trace) - owned_traces = [] + owned_trace = processor._owned_trace - def fail_after_starting_trace(*args, **kwargs): - processor._owned_trace = logger.add_trace(input="input", name="Trace") - owned_traces.append(processor._owned_trace) - raise RuntimeError("conversion failed") - - with patch.object(processor, "_log_node_tree", side_effect=fail_after_starting_trace): + with patch.object(processor, "_commit_trace", side_effect=RuntimeError("conversion failed")): processor.on_trace_end(trace) assert logger.current_parent() is None - assert owned_traces[0].status_code == 500 + assert owned_trace.status_code == 500 assert processor._nodes == {} assert processor._owned_trace is None @@ -222,6 +218,47 @@ def test_commit_failure_preserves_caller_owned_trace( logger.conclude(output="done") +@patch("splunk_ao.logger.logger.AgentStreams") +@patch("splunk_ao.logger.logger.Projects") +@patch("splunk_ao.logger.logger.Traces") +def test_openai_agents_live_root_is_used_for_outbound_context_and_commit( + mock_traces_client: Mock, mock_projects_client: Mock, mock_logstreams_client: Mock +) -> None: + setup_mock_traces_client(mock_traces_client) + setup_mock_projects_client(mock_projects_client) + setup_mock_logstreams_client(mock_logstreams_client) + logger = SplunkAOLogger(project="test", agent_stream="test", ingestion_hook=lambda _: None) + processor = SplunkAOTracingProcessor(splunk_ao_logger=logger) + trace = MagicMock(trace_id="trace-id", name="Agent trace", metadata={}) + processor.on_trace_start(trace) + node = Node( + node_type="agent", + run_id="root-span-id", + parent_run_id="trace-id", + span_params={ + "input": "question", + "output": "answer", + "name": "Agent root", + "start_time_iso": "2025-01-01T00:00:00+00:00", + "duration_ns": 10, + "status_code": 200, + }, + ) + processor._nodes[str(node.run_id)] = node + processor._nodes[trace.trace_id].children.append(str(node.run_id)) + processor._start_owned_root(node) + live_root = processor._owned_root + live_context = logger._otel_ids[live_root.id].span_context + + headers = get_tracing_headers() + processor.on_trace_end(trace) + + assert headers["traceparent"].split("-")[2] == format(live_context.span_id, "016x") + assert logger.traces[0].spans == [live_root] + assert logger.current_parent() is None + logger.terminate() + + def _create_mock_response_with_tools(tool_calls: list[dict]) -> dict: """Create a mock OpenAI API response with embedded tool calls.""" return { diff --git a/tests/test_otel_native_paths.py b/tests/test_otel_native_paths.py index 4896ee69..d2850c74 100644 --- a/tests/test_otel_native_paths.py +++ b/tests/test_otel_native_paths.py @@ -6,11 +6,14 @@ from opentelemetry import trace from opentelemetry.sdk.resources import Resource from opentelemetry.sdk.trace import Event, ReadableSpan -from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult +from opentelemetry.sdk.trace import TracerProvider as SDKTracerProvider +from opentelemetry.sdk.trace.export import BatchSpanProcessor, SpanExporter, SpanExportResult from opentelemetry.sdk.util.instrumentation import InstrumentationScope from opentelemetry.trace import Link, SpanContext, SpanKind, TraceFlags from opentelemetry.trace.status import Status, StatusCode +from galileo_core.schemas.logging.span import WorkflowSpan +from splunk_ao import get_tracing_headers from splunk_ao.decorator import ( _agent_stream_context, _dataset_input_context, @@ -21,11 +24,13 @@ _session_id_context, ) from splunk_ao.deployment import DeploymentMode, O11yConfig, StandaloneConfig +from splunk_ao.logger import SplunkAOLogger from splunk_ao.otel import ( _TRACE_PROVIDER_CONTEXT_VAR, SplunkAOOTLPExporter, SplunkAOSpanProcessor, add_splunk_ao_span_processor, + start_splunk_ao_span, ) from splunk_ao.shared.exceptions import MissingConfigurationError @@ -451,3 +456,97 @@ def test_processor_construction_does_not_replace_global_provider() -> None: assert trace.get_tracer_provider() is global_before processor.shutdown() + + +def test_sdk_native_otel_span_injects_w3c_context() -> None: + provider = SDKTracerProvider() + _TRACE_PROVIDER_CONTEXT_VAR.set(provider) + try: + with start_splunk_ao_span(WorkflowSpan(input="request", name="operation")) as span: + headers = get_tracing_headers() + + parts = headers["traceparent"].split("-") + assert int(parts[1], 16) == span.get_span_context().trace_id + assert int(parts[2], 16) == span.get_span_context().span_id + finally: + provider.shutdown() + + +def test_caller_owned_otel_span_injects_w3c_context() -> None: + provider = SDKTracerProvider() + tracer = provider.get_tracer("caller-owned") + try: + with tracer.start_as_current_span("operation") as span: + headers = get_tracing_headers() + + parts = headers["traceparent"].split("-") + assert int(parts[1], 16) == span.get_span_context().trace_id + assert int(parts[2], 16) == span.get_span_context().span_id + finally: + provider.shutdown() + + +def test_caller_owned_otel_and_path1_logger_interoperate_bidirectionally() -> None: + exporter = RecordingExporter() + processor = SplunkAOSpanProcessor(_exporter=exporter) + provider = SDKTracerProvider() + provider.add_span_processor(processor) + tracer = provider.get_tracer("caller-owned") + sink = MagicMock() + sink.force_flush.return_value = True + logger = SplunkAOLogger(project_id="project-id", agent_stream_id="stream-id", _sink=sink) + try: + with tracer.start_as_current_span("otel-upstream") as upstream: + logger.start_trace(input="request") + operation = logger.add_workflow_span(input="agent work", name="path1-operation") + operation_context = logger._otel_ids[operation.id].span_context + + with tracer.start_as_current_span("otel-downstream") as downstream: + headers = get_tracing_headers() + + logger.conclude(output="done") + logger.conclude(output="complete") + + [path1_span] = [call.args[0] for call in sink.emit.call_args_list] + assert path1_span.parent == upstream.get_span_context() + assert path1_span.context == operation_context + assert downstream.parent == operation_context + assert headers["traceparent"].split("-")[2] == format(downstream.get_span_context().span_id, "016x") + finally: + logger.terminate() + provider.shutdown() + + by_name = {span.name: span for span in exporter.spans} + assert by_name["otel-downstream"].parent == operation_context + + +def test_sdk_native_and_caller_owned_spans_keep_topology_and_queue_without_flush() -> None: + exporter = RecordingExporter() + processor = SplunkAOSpanProcessor(_exporter=exporter) + provider = SDKTracerProvider() + provider.add_span_processor(processor) + provider_token = _TRACE_PROVIDER_CONTEXT_VAR.set(provider) + tracer = provider.get_tracer("caller-owned") + try: + assert isinstance(processor.processor, BatchSpanProcessor) + with ( + patch.object(processor, "on_end", wraps=processor.on_end) as on_end, + patch.object(processor.processor, "force_flush", wraps=processor.processor.force_flush) as force_flush, + ): + operation = WorkflowSpan(input="request", name="sdk-operation") + with start_splunk_ao_span(operation) as sdk_span: + with tracer.start_as_current_span("caller-child") as caller_span: + pass + + assert on_end.call_count == 2 + force_flush.assert_not_called() + assert caller_span.parent is not None + assert caller_span.parent.span_id == sdk_span.get_span_context().span_id + assert caller_span.get_span_context().trace_id == sdk_span.get_span_context().trace_id + finally: + _TRACE_PROVIDER_CONTEXT_VAR.reset(provider_token) + provider.shutdown() + + by_name = {span.name: span for span in exporter.spans} + assert set(by_name) == {"sdk-operation", "caller-child"} + assert by_name["caller-child"].parent == by_name["sdk-operation"].context diff --git a/tests/test_thread_pool_task_handler.py b/tests/test_thread_pool_task_handler.py deleted file mode 100644 index e5080bf4..00000000 --- a/tests/test_thread_pool_task_handler.py +++ /dev/null @@ -1,352 +0,0 @@ -from concurrent.futures import Future -from unittest.mock import Mock, patch - -import pytest -from src.splunk_ao.logger.task_handler import ThreadPoolTaskHandler - - -class TestThreadPoolTaskHandler: - """Test suite for ThreadPoolTaskHandler class.""" - - @pytest.fixture - def mock_pool(self): - """Mock EventLoopThreadPool.""" - with patch("src.splunk_ao.logger.task_handler.EventLoopThreadPool") as mock_pool_class: - mock_pool = Mock() - mock_pool_class.return_value = mock_pool - yield mock_pool - - @pytest.fixture - def handler(self, mock_pool): - """Create ThreadPoolTaskHandler instance with mocked pool.""" - return ThreadPoolTaskHandler(num_threads=2) - - @pytest.fixture - def mock_future(self): - """Create a mock Future object.""" - future = Mock(spec=Future) - future.done.return_value = False - future.result.return_value = "test_result" - future.add_done_callback = Mock() - return future - - async def dummy_async_func(self) -> str: - """Dummy async function for testing.""" - return "async_result" - - def test_init(self, mock_pool) -> None: - """Test ThreadPoolTaskHandler initialization.""" - handler = ThreadPoolTaskHandler(num_threads=4) - assert handler._tasks == {} - assert handler._retry_counts == {} - assert handler._pool is not None - - def test_submit_task_independent(self, handler, mock_pool, mock_future) -> None: - """Test submitting an independent task (dependent_on_prev=False).""" - mock_pool.submit.return_value = mock_future - - async_fn = self.dummy_async_func - task_id = "test_task_1" - - handler.submit_task(task_id, async_fn, dependent_on_prev=False) - - # Verify task was submitted to pool - mock_pool.submit.assert_called_once_with(async_fn, wait_for_result=False) - - # Verify done callback was added - mock_future.add_done_callback.assert_called_once() - - # Verify task was tracked - assert task_id in handler._tasks - assert handler._tasks[task_id]["future"] == mock_future - assert handler._tasks[task_id]["parent_task_id"] is None - assert handler._retry_counts[task_id] == 0 - - def test_submit_task_dependent_no_previous(self, handler, mock_pool, mock_future) -> None: - """Test submitting a dependent task when no previous task exists.""" - mock_pool.submit.return_value = mock_future - - async_fn = self.dummy_async_func - task_id = "test_task_1" - - handler.submit_task(task_id, async_fn, dependent_on_prev=True) - - # Should submit immediately since no previous task exists - mock_pool.submit.assert_called_once_with(async_fn, wait_for_result=False) - assert task_id in handler._tasks - assert handler._tasks[task_id]["future"] == mock_future - assert handler._tasks[task_id]["parent_task_id"] is None # No parent since no previous tasks - - def test_submit_task_dependent_previous_completed(self, handler, mock_pool, mock_future) -> None: - """Test submitting a dependent task when previous task is completed.""" - # Set up first task - mock_pool.submit.return_value = mock_future - first_task_id = "task_1" - handler.submit_task(first_task_id, self.dummy_async_func, dependent_on_prev=False) - - # Mock the first task as completed - with patch.object(handler, "get_status", return_value="completed"): - second_mock_future = Mock(spec=Future) - mock_pool.submit.return_value = second_mock_future - - second_task_id = "task_2" - handler.submit_task(second_task_id, self.dummy_async_func, dependent_on_prev=True) - - # Should submit immediately since previous task is completed - assert mock_pool.submit.call_count == 2 - assert second_task_id in handler._tasks - assert handler._tasks[second_task_id]["future"] == second_mock_future - - def test_submit_task_dependent_previous_running(self, handler, mock_pool, mock_future) -> None: - """Test submitting a dependent task when previous task is still running.""" - # Set up first task - mock_pool.submit.return_value = mock_future - first_task_id = "task_1" - handler.submit_task(first_task_id, self.dummy_async_func, dependent_on_prev=False) - - # Mock the first task as still running - with patch.object(handler, "get_status", return_value="running"): - second_task_id = "task_2" - handler.submit_task(second_task_id, self.dummy_async_func, dependent_on_prev=True) - - # Should not submit yet, but should be tracked as pending - assert mock_pool.submit.call_count == 1 # Only first task submitted - assert second_task_id in handler._tasks - assert handler._tasks[second_task_id]["future"] is None - assert handler._tasks[second_task_id]["parent_task_id"] == first_task_id - assert handler._tasks[second_task_id]["callback"] is not None - - def test_handle_task_completion_triggers_children(self, handler, mock_pool) -> None: - """Test that task completion triggers child tasks.""" - # Set up parent task - parent_future = Mock(spec=Future) - mock_pool.submit.return_value = parent_future - parent_task_id = "parent_task" - handler.submit_task(parent_task_id, self.dummy_async_func, dependent_on_prev=False) - - # Set up child task (dependent) - with patch.object(handler, "get_status", return_value="running"): - child_task_id = "child_task" - handler.submit_task(child_task_id, self.dummy_async_func, dependent_on_prev=True) - - # Mock child task submission - child_future = Mock(spec=Future) - mock_pool.submit.return_value = child_future - - # Trigger parent task completion - handler._handle_task_completion(parent_task_id) - - # Verify child task was submitted - assert mock_pool.submit.call_count == 2 # Parent + child - assert handler._tasks[child_task_id]["future"] == child_future - - def test_increment_retry(self, handler) -> None: - """Test retry count increment.""" - task_id = "test_task" - - # Initially 0 - assert handler.get_retry_count(task_id) == 0 - - # Increment retry - handler.increment_retry(task_id) - assert handler.get_retry_count(task_id) == 1 - - # Increment again - handler.increment_retry(task_id) - assert handler.get_retry_count(task_id) == 2 - - def test_get_status_not_found(self, handler) -> None: - """Test get_status for non-existent task.""" - assert handler.get_status("non_existent") == "not_found" - - def test_get_status_pending(self, handler, mock_pool, mock_future) -> None: - """Test get_status for pending task.""" - # Set up parent task - mock_pool.submit.return_value = mock_future - parent_task_id = "parent" - handler.submit_task(parent_task_id, self.dummy_async_func, dependent_on_prev=False) - - # Set up child task as pending - with patch.object( - handler, "get_status", side_effect=lambda tid: "running" if tid == parent_task_id else "pending" - ): - child_task_id = "child" - handler.submit_task(child_task_id, self.dummy_async_func, dependent_on_prev=True) - - # Reset side_effect and test actual status - handler.get_status = handler.__class__.get_status.__get__(handler) - assert handler.get_status(child_task_id) == "pending" - - def test_get_status_running(self, handler, mock_pool, mock_future) -> None: - """Test get_status for running task.""" - mock_pool.submit.return_value = mock_future - mock_future.done.return_value = False - - task_id = "running_task" - handler.submit_task(task_id, self.dummy_async_func, dependent_on_prev=False) - - assert handler.get_status(task_id) == "running" - - def test_get_status_completed(self, handler, mock_pool, mock_future) -> None: - """Test get_status for completed task.""" - mock_pool.submit.return_value = mock_future - mock_future.done.return_value = True - mock_future.result.return_value = "success" - - task_id = "completed_task" - handler.submit_task(task_id, self.dummy_async_func, dependent_on_prev=False) - - assert handler.get_status(task_id) == "completed" - - def test_get_status_failed(self, handler, mock_pool, mock_future) -> None: - """Test get_status for failed task.""" - mock_pool.submit.return_value = mock_future - mock_future.done.return_value = True - mock_future.result.side_effect = Exception("Task failed") - - task_id = "failed_task" - handler.submit_task(task_id, self.dummy_async_func, dependent_on_prev=False) - - assert handler.get_status(task_id) == "failed" - - def test_get_result(self, handler, mock_pool, mock_future) -> None: - """Test getting task result.""" - mock_pool.submit.return_value = mock_future - expected_result = "test_result" - mock_future.result.return_value = expected_result - - task_id = "test_task" - handler.submit_task(task_id, self.dummy_async_func, dependent_on_prev=False) - - result = handler.get_result(task_id) - assert result == expected_result - - def test_get_result_not_found(self, handler) -> None: - """Test getting result for non-existent task.""" - with pytest.raises(ValueError, match="Task non_existent not found"): - handler.get_result("non_existent") - - def test_get_children(self, handler, mock_pool, mock_future) -> None: - """Test getting child tasks.""" - # Set up parent task - mock_pool.submit.return_value = mock_future - parent_task_id = "parent" - handler.submit_task(parent_task_id, self.dummy_async_func, dependent_on_prev=False) - - # Manually add two child tasks that depend on the parent - # (This simulates what would happen if we had a more complex dependency system) - def child1_callback(): - return handler._pool.submit(self.dummy_async_func, wait_for_result=False) - - def child2_callback(): - return handler._pool.submit(self.dummy_async_func, wait_for_result=False) - - handler._add_or_update_task( - task_id="child1", future=None, start_time=None, parent_task_id=parent_task_id, callback=child1_callback - ) - - handler._add_or_update_task( - task_id="child2", future=None, start_time=None, parent_task_id=parent_task_id, callback=child2_callback - ) - - children = handler.get_children(parent_task_id) - assert len(children) == 2 - - # Verify both children have the correct parent_task_id - child_task_ids = [ - task_id for task_id, task in handler._tasks.items() if task.get("parent_task_id") == parent_task_id - ] - assert len(child_task_ids) == 2 - assert "child1" in child_task_ids - assert "child2" in child_task_ids - - def test_all_tasks_completed_empty(self, handler) -> None: - """Test all_tasks_completed with no tasks.""" - assert handler.all_tasks_completed() is True - - def test_all_tasks_completed_true(self, handler, mock_pool, mock_future) -> None: - """Test all_tasks_completed when all tasks are done.""" - mock_pool.submit.return_value = mock_future - - with patch.object(handler, "get_status", return_value="completed"): - handler.submit_task("task1", self.dummy_async_func, dependent_on_prev=False) - handler.submit_task("task2", self.dummy_async_func, dependent_on_prev=False) - - assert handler.all_tasks_completed() is True - - def test_all_tasks_completed_false(self, handler, mock_pool, mock_future) -> None: - """Test all_tasks_completed when some tasks are still running.""" - mock_pool.submit.return_value = mock_future - - def mock_status(task_id) -> str: - return "completed" if task_id == "task1" else "running" - - with patch.object(handler, "get_status", side_effect=mock_status): - handler.submit_task("task1", self.dummy_async_func, dependent_on_prev=False) - handler.submit_task("task2", self.dummy_async_func, dependent_on_prev=False) - - assert handler.all_tasks_completed() is False - - def test_terminate(self, handler, mock_pool) -> None: - """Test handler termination.""" - handler.terminate() - mock_pool.stop.assert_called_once() - - def test_complex_dependency_chain(self, handler, mock_pool) -> None: - """Test a complex chain of dependent tasks.""" - # Create multiple futures for different tasks - futures = [Mock(spec=Future) for _ in range(4)] - future_index = 0 - - def get_next_future(*args, **kwargs): - nonlocal future_index - future = futures[future_index] - future_index += 1 - return future - - mock_pool.submit.side_effect = get_next_future - - # Submit task chain: A -> B -> C -> D - task_ids = ["task_A", "task_B", "task_C", "task_D"] - - # Submit first task (independent) - handler.submit_task(task_ids[0], self.dummy_async_func, dependent_on_prev=False) - - # Submit dependent tasks - for i in range(1, len(task_ids)): - with patch.object(handler, "get_status", return_value="running"): - handler.submit_task(task_ids[i], self.dummy_async_func, dependent_on_prev=True) - - # Verify only first task was submitted initially - assert mock_pool.submit.call_count == 1 - - # Simulate completion of each task in sequence - for i in range(len(task_ids) - 1): - # Mock the current task as completed for next submission - mock_pool.submit.side_effect = get_next_future - handler._handle_task_completion(task_ids[i]) - - # Verify next task was submitted - expected_calls = i + 2 # +1 for initial task, +1 for current completion - assert mock_pool.submit.call_count == expected_calls - - def test_done_callback_integration(self, handler, mock_pool) -> None: - """Test that the done callback properly triggers task completion handling.""" - mock_future = Mock(spec=Future) - mock_pool.submit.return_value = mock_future - - task_id = "callback_test" - handler.submit_task(task_id, self.dummy_async_func, dependent_on_prev=False) - - # Get the callback that was registered - callback_calls = mock_future.add_done_callback.call_args_list - assert len(callback_calls) == 1 - registered_callback = callback_calls[0][0][0] - - # Mock _handle_task_completion to verify it's called - with patch.object(handler, "_handle_task_completion") as mock_handle: - # Trigger the callback - registered_callback(mock_future) - - # Verify _handle_task_completion was called with correct task_id - mock_handle.assert_called_once_with(task_id) diff --git a/tests/test_tracing.py b/tests/test_tracing.py new file mode 100644 index 00000000..71ff3bd3 --- /dev/null +++ b/tests/test_tracing.py @@ -0,0 +1,121 @@ +from collections.abc import Generator + +import pytest +from opentelemetry import context, trace +from opentelemetry.sdk.trace import ReadableSpan +from opentelemetry.trace import NonRecordingSpan, SpanContext, TraceFlags, TraceState + +from splunk_ao import extract_tracing_context, get_tracing_headers +from splunk_ao.exceptions import SplunkAOLoggerException +from splunk_ao.logger import SplunkAOLogger +from splunk_ao.logger.logger import _otel_context_state + + +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: + return None + + +@pytest.fixture(autouse=True) +def isolated_context() -> Generator[None, None, None]: + state_token = _otel_context_state.set(None) + token = context.attach(context.Context()) + try: + yield + finally: + context.detach(token) + _otel_context_state.reset(state_token) + + +def make_logger() -> tuple[SplunkAOLogger, RecordingSink]: + sink = RecordingSink() + logger = SplunkAOLogger(project_id="project-id", agent_stream_id="stream-id", _sink=sink) + return logger, sink + + +def attach_span_context(span_context: SpanContext): + return context.attach(trace.set_span_in_context(NonRecordingSpan(span_context))) + + +def test_get_tracing_headers_injects_native_otel_context_into_same_carrier() -> None: + span_context = SpanContext( + trace_id=0x4BF92F3577B34DA6A3CE929D0E0E4736, + span_id=0x00F067AA0BA902B7, + is_remote=False, + trace_flags=TraceFlags.SAMPLED, + trace_state=TraceState([("vendor", "value")]), + ) + token = attach_span_context(span_context) + try: + carrier = {"existing": "header"} + result = get_tracing_headers(carrier) + finally: + context.detach(token) + + assert result is carrier + assert carrier["existing"] == "header" + assert carrier["traceparent"] == "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01" + assert carrier["tracestate"] == "vendor=value" + + +def test_get_tracing_headers_rejects_no_active_operation() -> None: + with pytest.raises(SplunkAOLoggerException, match="active exportable operation"): + get_tracing_headers() + + +def test_get_tracing_headers_rejects_internal_trace_envelope() -> None: + logger, _ = make_logger() + try: + logger.start_trace(input="request") + with pytest.raises(SplunkAOLoggerException, match="active exportable operation"): + get_tracing_headers() + logger.conclude(output="done") + finally: + logger.terminate() + + +def test_get_tracing_headers_uses_real_path1_operation_and_no_routing_headers() -> None: + logger, _ = make_logger() + try: + logger.start_trace(input="request") + operation = logger.add_workflow_span(input="work", name="operation") + operation_context = logger._otel_ids[operation.id].span_context + + headers = get_tracing_headers() + + assert headers["traceparent"].split("-")[2] == format(operation_context.span_id, "016x") + assert set(headers).isdisjoint( + {"project", "projectid", "logstream", "logstreamid", "experimentid", "X-SF-Token", "Splunk-AO-API-Key"} + ) + logger.conclude(output="work") + logger.conclude(output="done") + finally: + logger.terminate() + + +def test_extract_tracing_context_is_case_insensitive_and_preserves_tracestate() -> None: + extracted = extract_tracing_context( + {"TRACEPARENT": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01", "TRACESTATE": "vendor=value"} + ) + span_context = trace.get_current_span(extracted).get_span_context() + + assert span_context.is_valid + assert span_context.trace_id == 0x4BF92F3577B34DA6A3CE929D0E0E4736 + assert span_context.span_id == 0x00F067AA0BA902B7 + assert span_context.is_remote + assert span_context.trace_state.get("vendor") == "value" + + +@pytest.mark.parametrize("carrier", [{}, {"traceparent": "not-valid"}]) +def test_extract_tracing_context_rejects_missing_or_malformed_parent(carrier: dict[str, str]) -> None: + extracted = extract_tracing_context(carrier) + assert not trace.get_current_span(extracted).get_span_context().is_valid diff --git a/tests/testutils/setup.py b/tests/testutils/setup.py index 2866591e..4b32098f 100644 --- a/tests/testutils/setup.py +++ b/tests/testutils/setup.py @@ -1,15 +1,8 @@ -import copy import datetime -from collections.abc import Callable -from contextlib import contextmanager -from typing import Any from unittest.mock import AsyncMock, Mock from uuid import UUID -from pydantic import BaseModel - from splunk_ao.agent_streams import AgentStream -from splunk_ao.logger.logger import SplunkAOLogger from splunk_ao.projects import Project from splunk_ao.resources.models import ExperimentResponse, ProjectType from splunk_ao.resources.models.log_stream_response import LogStreamResponse @@ -17,171 +10,6 @@ from splunk_ao.resources.models.task_type import TaskType -class ThreadPoolTaskInfo(BaseModel): - task_id: str - function_name: str - request: Any - task_func: Callable - kwargs: dict - - -class ThreadPoolRequestCapture: - """Helper class to capture requests submitted to the thread pool in streaming methods.""" - - def __init__(self): - self.captured_tasks: list[ThreadPoolTaskInfo] = [] - self.mock_pool = None - - def capture_request(self, task_id: str, async_fn: Callable, **kwargs) -> None: - """Capture function that extracts requests and function names from ThreadPoolTaskHandler.submit_task calls.""" - - # Extract both the request and function from the lambda - captured_request = None - captured_function_name = None - captured_function = None - - # Get the first function name referenced in the lambda - # Since lambdas follow pattern: lambda: function_name(request) - code = async_fn.__code__ - - if code.co_names: - # The first (and typically only) name in co_names is our function - captured_function_name = code.co_names[0] - - # Extract request from closure and also look for callable functions - if async_fn.__closure__: - for cell in async_fn.__closure__: - cell_content = cell.cell_contents - - # Check for request object - if isinstance(cell_content, BaseModel): - captured_request = copy.deepcopy(cell_content) - - # Check for callable function (if we didn't get function name from co_names) - elif captured_function_name is None and callable(cell_content) and hasattr(cell_content, "__name__"): - captured_function_name = cell_content.__name__ - captured_function = cell_content - - # If we have function name but not function reference, look for it - elif ( - captured_function_name - and callable(cell_content) - and hasattr(cell_content, "__name__") - and cell_content.__name__ == captured_function_name - ): - captured_function = cell_content - - # Create a new lambda that uses the deep-copied request to avoid mutation issues - if captured_function and captured_request: - - def isolated_task_func(): - return captured_function(captured_request) - else: - isolated_task_func = async_fn - - task_info = ThreadPoolTaskInfo( - task_id=task_id, - function_name=captured_function_name, - request=captured_request, - task_func=isolated_task_func, - kwargs=kwargs, - ) - self.captured_tasks.append(task_info) - - def get_request_by_function_name(self, function_name: str) -> Any | None: - return next((task.request for task in self.captured_tasks if task.function_name == function_name), None) - - def get_latest_task(self) -> ThreadPoolTaskInfo | None: - """Get the most recent task info (request + function name).""" - return self.captured_tasks[-1] if self.captured_tasks else None - - def get_all_tasks(self) -> list[ThreadPoolTaskInfo]: - """ - Get all captured tasks. - - Returns: - list[ThreadPoolTaskInfo]: List of captured tasks. - """ - return self.captured_tasks - - def get_all_requests(self) -> list[Any]: - """Get all captured requests.""" - return [task.request for task in self.captured_tasks] - - def get_all_function_names(self) -> list[str]: - """Get all captured function names.""" - return [task.function_name for task in self.captured_tasks] - - def get_task_by_function_name(self, function_name: str) -> ThreadPoolTaskInfo | None: - """Get the first task info that matches the given function name. - - Args: - function_name (str): The name of the function to search for. - - Returns: - Optional[ThreadPoolTaskInfo]: The first task info that matches the given function name, or None if no match is found. - """ - return next((task for task in self.captured_tasks if task.function_name == function_name), None) - - def count_function_calls(self, function_name: str) -> int: - """Count how many times a particular function was called.""" - functions = self.get_all_function_names() - return sum(1 for func_name in functions if func_name == function_name) - - def clear(self) -> None: - """Clear all captured tasks.""" - self.captured_tasks.clear() - - def assert_function_called(self, expected_function_name: str) -> None: - """Assert that the latest captured function matches the expected name.""" - assert expected_function_name in self.get_all_function_names(), ( - f"Expected function '{expected_function_name}' not in '{self.get_all_function_names()}'" - ) - - def assert_functions_called(self, expected_function_names: list[str]) -> None: - """Assert that the captured functions match the expected sequence.""" - actual_functions = self.get_all_function_names() - assert actual_functions == expected_function_names, ( - f"Expected functions {expected_function_names}, but got {actual_functions}" - ) - - -def setup_thread_pool_request_capture(logger: SplunkAOLogger) -> ThreadPoolRequestCapture: - """ - Set up request capture for a logger's thread pool. - - Args: - logger: SplunkAOLogger instance to mock the thread pool for - - Returns: - ThreadPoolRequestCapture instance that can be used to inspect captured requests - """ - capture = ThreadPoolRequestCapture() - mock_pool = Mock(side_effect=capture.capture_request) - logger._task_handler.submit_task = mock_pool - logger._task_handler.submit_task_with_parent = mock_pool - capture.mock_pool = mock_pool - return capture - - -@contextmanager -def capture_streaming_requests(logger): - """ - Context manager for capturing streaming requests. - - Usage: - with capture_streaming_requests(logger) as capture: - logger._ingest_trace_streaming(trace) - request = capture.get_latest_request() - assert isinstance(request, TracesIngestRequest) - """ - capture = setup_thread_pool_request_capture(logger) - try: - yield capture - finally: - capture.clear() - - def setup_mock_projects_client(mock_projects_client: Mock): now = datetime.datetime.now() mock_instance = mock_projects_client.return_value From a06e36292e5bb9be8c0a10a91ddb588ce880cb92 Mon Sep 17 00:00:00 2001 From: pradystar Date: Thu, 13 Aug 2026 11:06:07 -0700 Subject: [PATCH 02/13] feat(tracing): automate context propagation and handler batching --- ARCHITECTURE.md | 22 ++ CHANGELOG.md | 13 + README.md | 53 ++- poetry.lock | 193 ++++++++++- pyproject.toml | 8 +- src/splunk_ao/__init__.py | 2 + src/splunk_ao/decorator.py | 15 +- src/splunk_ao/exporter/sink.py | 14 +- src/splunk_ao/handlers/base_async_handler.py | 19 +- src/splunk_ao/handlers/base_handler.py | 116 ++++++- .../handlers/openai_agents/handler.py | 83 ++++- src/splunk_ao/handlers/span_lifecycle.py | 143 ++++++++ src/splunk_ao/http_instrumentation.py | 119 +++++++ src/splunk_ao/logger/logger.py | 95 ++++- src/splunk_ao/otel.py | 9 +- src/splunk_ao/session_context.py | 107 ++++++ src/splunk_ao/tracing.py | 6 +- tests/conftest.py | 4 +- tests/test_async_base_handler.py | 50 +++ tests/test_base_handler.py | 84 +++++ tests/test_exporter_sink.py | 10 + tests/test_http_instrumentation.py | 325 ++++++++++++++++++ tests/test_logger_otel_egress.py | 17 + tests/test_openai_agents.py | 76 ++++ tests/test_otel_native_paths.py | 16 +- tests/test_tracing.py | 99 +++++- 26 files changed, 1651 insertions(+), 47 deletions(-) create mode 100644 src/splunk_ao/handlers/span_lifecycle.py create mode 100644 src/splunk_ao/http_instrumentation.py create mode 100644 src/splunk_ao/session_context.py create mode 100644 tests/test_http_instrumentation.py diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index faac9a7a..714ddbe1 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -86,10 +86,29 @@ client/server wrapping and message-context propagation. The SDK may add processors, but it must never silently replace the global tracer provider. The application owns the provider and calls `shutdown()`; SDK-owned logger/export resources use their own termination path. +### Distributed context and HTTP transports + +Supported automatic propagation reuses upstream OpenTelemetry instrumentors rather than custom SDK HTTP wrappers. +`instrument_distributed_tracing()` configures FastAPI/Starlette inbound instrumentation and Requests, HTTPX sync/async, +and aiohttp-client outbound instrumentation with a caller-owned provider. It never sets or replaces the global tracer +provider. Manual `get_tracing_headers()` and `TracingMiddleware` remain fallbacks for unsupported transports; supported +automatic instrumentation does not require either one on individual requests. + +An explicit SDK session propagates as the standard W3C baggage member `gen_ai.conversation.id`. Export normalization +may derive the local compatibility attribute `splunk_ao.session.id`, but that attribute is not propagated as a second +baggage member. The SDK must not propagate authentication, deployment, Project, Agent Stream/log stream, experiment, +endpoint, model, workflow, or agent identity in baggage. + ## Span Lifecycle and Export Ownership `SpanSink` owns the SDK's private provider and `BatchSpanProcessor` for internal telemetry. A completed operation is enqueued immediately, allowing scheduled export without an explicit flush. +Without an explicit internal `BatchConfig`, both SDK-owned and caller-owned processors defer to standard OTel +`OTEL_BSP_*` environment configuration. + +Framework handlers register a stable operation identity at each start callback and enqueue that operation at its +matching end callback, so a child can enter `BatchSpanProcessor` while its parent remains active. The deprecated +`ingestion_hook` is the compatibility exception and retains whole-tree payload construction at trace completion. - `flush()` / `async_flush()` drain completed work and do not end an active trace. - `terminate()` drains completed work, shuts down SDK-owned telemetry resources, and releases unfinished state. @@ -156,6 +175,9 @@ receiver failures, and expose bounded acknowledgement health. Do not log secrets These handlers adapt framework events into the internal Path 1 model. Keep optional dependencies lazy and imports free of surprising side effects. CrewAI is excluded on Python 3.14; code and tests must support the unavailable path. +With `start_new_trace=True`, a handler owns and concludes its local envelope and root while still inheriting an active +W3C parent. With `start_new_trace=False`, it attaches below caller-owned active state, emits only its subtree, and leaves +the caller open. ### OpenAI wrapper diff --git a/CHANGELOG.md b/CHANGELOG.md index a86d3e8f..9bfae46a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added +- Added the `distributed-tracing` extra and + `instrument_distributed_tracing()` startup helper for supported upstream + FastAPI/Starlette, Requests, HTTPX, and aiohttp-client instrumentation. +- Explicit SDK sessions now propagate across supported services as standard + W3C `gen_ai.conversation.id` baggage. SDK routing, authentication, and + application identity are not propagated. - Added module-level `get_tracing_headers()` and `extract_tracing_context()` helpers for standard W3C `traceparent` and `tracestate` propagation. `TracingMiddleware` now extracts and scopes that @@ -22,6 +28,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Callback handlers keep a live real root during framework execution so outbound W3C context uses the same identity and visible hierarchy that is exported at commit. +- Normal LangChain, CrewAI, Google ADK, and OpenAI Agents callbacks now enqueue + each completed operation into the existing `BatchSpanProcessor` at that + operation's end callback. The deprecated `ingestion_hook` retains its + whole-tree compatibility behavior. +- Default logger-owned batching now honors standard OpenTelemetry + `OTEL_BSP_*` configuration, matching caller-owned OTel paths. Explicit + internal `BatchConfig` values remain authoritative when supplied. ### Removed diff --git a/README.md b/README.md index 7c4e7d77..17210603 100644 --- a/README.md +++ b/README.md @@ -32,6 +32,7 @@ Install the optional integration dependencies used by your application: ```shell pip install "splunk-ao[openai]" pip install "splunk-ao[langchain]" langchain-openai +pip install "splunk-ao[distributed-tracing]" ``` Other available extras include `crewai`, `middleware`, and `all`. @@ -347,7 +348,8 @@ tracer_provider.shutdown() `add_splunk_ao_span_processor()` reads the same deployment and routing configuration described in [Setup](#setup). It registers a `SplunkAOSpanProcessor`, which uses a standard OpenTelemetry -`BatchSpanProcessor`. +`BatchSpanProcessor`. Default logger-owned and caller-owned processors honor +the standard `OTEL_BSP_*` batch processor environment variables. If your application needs another processor or exporter, register it through the standard OpenTelemetry extension points. For example, to send the same @@ -364,6 +366,42 @@ The application owns a `TracerProvider` that it constructs, so it must call `tracer_provider.shutdown()` during teardown. A `SplunkAOLogger` created elsewhere owns a separate provider and should be terminated separately. +##### Automatic HTTP distributed tracing + +The `distributed-tracing` extra provides a supported one-time setup for +upstream OpenTelemetry FastAPI/Starlette, Requests, HTTPX sync/async, and +aiohttp-client instrumentation: + +```python +from fastapi import FastAPI +from opentelemetry.sdk.trace import TracerProvider + +from splunk_ao import instrument_distributed_tracing +from splunk_ao.otel import add_splunk_ao_span_processor + +provider = TracerProvider() +add_splunk_ao_span_processor(provider) + +app = FastAPI() +instrument_distributed_tracing(tracer_provider=provider, app=app) +``` + +The application owns `provider` and calls `provider.shutdown()` during process +teardown. The helper never replaces the process-global provider. Configure +other OpenTelemetry/OpenInference agent or model instrumentations with the +same provider when they should join the trace. + +After this startup call, supported inbound requests extract W3C context and +supported outbound clients inject it automatically. Application code does not +call `get_tracing_headers()` for each request and does not install +`TracingMiddleware`. Those APIs remain supported for transports and frameworks +outside the automatic support matrix. + +An explicit SDK session propagates as the standard +`gen_ai.conversation.id` W3C baggage member. Project, Agent Stream, +experiment, deployment, authentication, model, workflow, and agent identity +remain local and are never added to baggage by the SDK. + #### Export diagnostics The SDK does not change your application's global logging configuration. @@ -416,6 +454,19 @@ not close a trace that was started by the caller. In short-lived jobs, ensure the owning logger or provider is terminated or shut down after framework work finishes. +For normal OTLP export, LangChain, CrewAI, Google ADK, and OpenAI Agents +handlers enqueue each operation into the existing `BatchSpanProcessor` when +that operation's end callback runs. A child can therefore enter a scheduled +or size-based export batch while its parent remains active. This is immediate +enqueue, not one network request per span, and no per-trace `flush()` is +required. The deprecated `ingestion_hook` compatibility path still constructs +one mutable whole-tree payload at trace completion. + +`start_new_trace` controls ownership, not whether distributed tracing works. +Keep its default `True` for a standalone handler. Set it to `False` only when +the handler is intentionally placed beneath caller-owned active logger state; +the handler emits its own subtree and does not conclude the caller. + #### Datasets Create a dataset: diff --git a/poetry.lock b/poetry.lock index 5edd9823..4cfb87f4 100644 --- a/poetry.lock +++ b/poetry.lock @@ -246,6 +246,23 @@ files = [ [package.extras] test = ["coverage", "mypy", "pexpect", "ruff", "wheel"] +[[package]] +name = "asgiref" +version = "3.12.1" +description = "ASGI specs, helper code, and adapters" +optional = true +python-versions = ">=3.10" +groups = ["main"] +markers = "extra == \"distributed-tracing\" or extra == \"all\"" +files = [ + {file = "asgiref-3.12.1-py3-none-any.whl", hash = "sha256:fe386d1c2bff7259ea95929266d12a8cf9a8b5a1c2598402967d8792e7a7c094"}, + {file = "asgiref-3.12.1.tar.gz", hash = "sha256:59dcb51c272ad209d59bed5708a64a333083e86017d7fcdd67498eeab7784340"}, +] + +[package.extras] +mypy = ["mypy (>=1.14.0)"] +tests = ["pytest", "pytest-asyncio"] + [[package]] name = "attrs" version = "25.3.0" @@ -3123,6 +3140,162 @@ opentelemetry-sdk = ">=1.38.0,<1.39.0" requests = ">=2.7,<3.0" typing-extensions = ">=4.5.0" +[[package]] +name = "opentelemetry-instrumentation" +version = "0.59b0" +description = "Instrumentation Tools & Auto Instrumentation for OpenTelemetry Python" +optional = true +python-versions = ">=3.9" +groups = ["main"] +markers = "extra == \"distributed-tracing\" or extra == \"all\"" +files = [ + {file = "opentelemetry_instrumentation-0.59b0-py3-none-any.whl", hash = "sha256:44082cc8fe56b0186e87ee8f7c17c327c4c2ce93bdbe86496e600985d74368ee"}, + {file = "opentelemetry_instrumentation-0.59b0.tar.gz", hash = "sha256:6010f0faaacdaf7c4dff8aac84e226d23437b331dcda7e70367f6d73a7db1adc"}, +] + +[package.dependencies] +opentelemetry-api = ">=1.4,<2.0" +opentelemetry-semantic-conventions = "0.59b0" +packaging = ">=18.0" +wrapt = ">=1.0.0,<2.0.0" + +[[package]] +name = "opentelemetry-instrumentation-aiohttp-client" +version = "0.59b0" +description = "OpenTelemetry aiohttp client instrumentation" +optional = true +python-versions = ">=3.9" +groups = ["main"] +markers = "extra == \"distributed-tracing\" or extra == \"all\"" +files = [ + {file = "opentelemetry_instrumentation_aiohttp_client-0.59b0-py3-none-any.whl", hash = "sha256:8e7f234d2b6b385d5b1757ba7f13baa71c45ec020e1a08ba7f7649da7958470a"}, + {file = "opentelemetry_instrumentation_aiohttp_client-0.59b0.tar.gz", hash = "sha256:665ee520f2fb5f44d6d600918eb7bbd2c29b355e0d0deda49991db857adee51f"}, +] + +[package.dependencies] +opentelemetry-api = ">=1.12,<2.0" +opentelemetry-instrumentation = "0.59b0" +opentelemetry-semantic-conventions = "0.59b0" +opentelemetry-util-http = "0.59b0" +wrapt = ">=1.0.0,<2.0.0" + +[package.extras] +instruments = ["aiohttp (>=3.0,<4.0)"] + +[[package]] +name = "opentelemetry-instrumentation-asgi" +version = "0.59b0" +description = "ASGI instrumentation for OpenTelemetry" +optional = true +python-versions = ">=3.9" +groups = ["main"] +markers = "extra == \"distributed-tracing\" or extra == \"all\"" +files = [ + {file = "opentelemetry_instrumentation_asgi-0.59b0-py3-none-any.whl", hash = "sha256:ba9703e09d2c33c52fa798171f344c8123488fcd45017887981df088452d3c53"}, + {file = "opentelemetry_instrumentation_asgi-0.59b0.tar.gz", hash = "sha256:2509d6fe9fd829399ce3536e3a00426c7e3aa359fc1ed9ceee1628b56da40e7a"}, +] + +[package.dependencies] +asgiref = ">=3.0,<4.0" +opentelemetry-api = ">=1.12,<2.0" +opentelemetry-instrumentation = "0.59b0" +opentelemetry-semantic-conventions = "0.59b0" +opentelemetry-util-http = "0.59b0" + +[package.extras] +instruments = ["asgiref (>=3.0,<4.0)"] + +[[package]] +name = "opentelemetry-instrumentation-fastapi" +version = "0.59b0" +description = "OpenTelemetry FastAPI Instrumentation" +optional = true +python-versions = ">=3.9" +groups = ["main"] +markers = "extra == \"distributed-tracing\" or extra == \"all\"" +files = [ + {file = "opentelemetry_instrumentation_fastapi-0.59b0-py3-none-any.whl", hash = "sha256:0d8d00ff7d25cca40a4b2356d1d40a8f001e0668f60c102f5aa6bb721d660c4f"}, + {file = "opentelemetry_instrumentation_fastapi-0.59b0.tar.gz", hash = "sha256:e8fe620cfcca96a7d634003df1bc36a42369dedcdd6893e13fb5903aeeb89b2b"}, +] + +[package.dependencies] +opentelemetry-api = ">=1.12,<2.0" +opentelemetry-instrumentation = "0.59b0" +opentelemetry-instrumentation-asgi = "0.59b0" +opentelemetry-semantic-conventions = "0.59b0" +opentelemetry-util-http = "0.59b0" + +[package.extras] +instruments = ["fastapi (>=0.92,<1.0)"] + +[[package]] +name = "opentelemetry-instrumentation-httpx" +version = "0.59b0" +description = "OpenTelemetry HTTPX Instrumentation" +optional = true +python-versions = ">=3.9" +groups = ["main"] +markers = "extra == \"distributed-tracing\" or extra == \"all\"" +files = [ + {file = "opentelemetry_instrumentation_httpx-0.59b0-py3-none-any.whl", hash = "sha256:7dc9f66aef4ca3904d877f459a70c78eafd06131dc64d713b9b1b5a7d0a48f05"}, + {file = "opentelemetry_instrumentation_httpx-0.59b0.tar.gz", hash = "sha256:a1cb9b89d9f05a82701cc9ab9cfa3db54fd76932489449778b350bc1b9f0e872"}, +] + +[package.dependencies] +opentelemetry-api = ">=1.12,<2.0" +opentelemetry-instrumentation = "0.59b0" +opentelemetry-semantic-conventions = "0.59b0" +opentelemetry-util-http = "0.59b0" +wrapt = ">=1.0.0,<2.0.0" + +[package.extras] +instruments = ["httpx (>=0.18.0)"] + +[[package]] +name = "opentelemetry-instrumentation-requests" +version = "0.59b0" +description = "OpenTelemetry requests instrumentation" +optional = true +python-versions = ">=3.9" +groups = ["main"] +markers = "extra == \"distributed-tracing\" or extra == \"all\"" +files = [ + {file = "opentelemetry_instrumentation_requests-0.59b0-py3-none-any.whl", hash = "sha256:d43121532877e31a46c48649279cec2504ee1e0ceb3c87b80fe5ccd7eafc14c1"}, + {file = "opentelemetry_instrumentation_requests-0.59b0.tar.gz", hash = "sha256:9af2ffe3317f03074d7f865919139e89170b6763a0251b68c25e8e64e04b3400"}, +] + +[package.dependencies] +opentelemetry-api = ">=1.12,<2.0" +opentelemetry-instrumentation = "0.59b0" +opentelemetry-semantic-conventions = "0.59b0" +opentelemetry-util-http = "0.59b0" + +[package.extras] +instruments = ["requests (>=2.0,<3.0)"] + +[[package]] +name = "opentelemetry-instrumentation-starlette" +version = "0.59b0" +description = "OpenTelemetry Starlette Instrumentation" +optional = true +python-versions = ">=3.9" +groups = ["main"] +markers = "extra == \"distributed-tracing\" or extra == \"all\"" +files = [ + {file = "opentelemetry_instrumentation_starlette-0.59b0-py3-none-any.whl", hash = "sha256:a833b97d297e4b2aaf58041612663dbabb8380e3993c76a4e32b3e351470f321"}, + {file = "opentelemetry_instrumentation_starlette-0.59b0.tar.gz", hash = "sha256:3f033fd92d6a8e4122ebcb3d83afc5c64d6be7930e9094876eb02b8afbd08ba5"}, +] + +[package.dependencies] +opentelemetry-api = ">=1.12,<2.0" +opentelemetry-instrumentation = "0.59b0" +opentelemetry-instrumentation-asgi = "0.59b0" +opentelemetry-semantic-conventions = "0.59b0" +opentelemetry-util-http = "0.59b0" + +[package.extras] +instruments = ["starlette (>=0.13)"] + [[package]] name = "opentelemetry-proto" version = "1.38.0" @@ -3171,6 +3344,19 @@ files = [ opentelemetry-api = "1.38.0" typing-extensions = ">=4.5.0" +[[package]] +name = "opentelemetry-util-http" +version = "0.59b0" +description = "Web util for OpenTelemetry" +optional = true +python-versions = ">=3.9" +groups = ["main"] +markers = "extra == \"distributed-tracing\" or extra == \"all\"" +files = [ + {file = "opentelemetry_util_http-0.59b0-py3-none-any.whl", hash = "sha256:6d036a07563bce87bf521839c0671b507a02a0d39d7ea61b88efa14c6e25355d"}, + {file = "opentelemetry_util_http-0.59b0.tar.gz", hash = "sha256:ae66ee91be31938d832f3b4bc4eb8a911f6eddd38969c4a871b1230db2a0a560"}, +] + [[package]] name = "orjson" version = "3.11.9" @@ -3339,7 +3525,7 @@ files = [ {file = "packaging-24.2-py3-none-any.whl", hash = "sha256:09abb1bccd265c01f4a3aa3f7a7db064b36514d2cba19a2f694fe6150451a759"}, {file = "packaging-24.2.tar.gz", hash = "sha256:c228a6dc5e932d346bc5739379109d49e8853dd8223571c7c5b55260edc0b97f"}, ] -markers = {main = "extra == \"langchain\" or extra == \"all\" or extra == \"openai\" or python_version <= \"3.13\" and (extra == \"crewai\" or extra == \"all\" or extra == \"langchain\" or extra == \"openai\")"} +markers = {main = "extra == \"langchain\" or extra == \"all\" or extra == \"openai\" or extra == \"distributed-tracing\" or python_version <= \"3.13\" and (extra == \"crewai\" or extra == \"all\" or extra == \"langchain\" or extra == \"openai\" or extra == \"distributed-tracing\")"} [[package]] name = "pathspec" @@ -6510,8 +6696,9 @@ cffi = {version = ">=1.11", markers = "platform_python_implementation == \"PyPy\ cffi = ["cffi (>=1.11)"] [extras] -all = ["aiohttp", "crewai", "cryptography", "langchain", "langchain-core", "langsmith", "litellm", "mcp", "openai", "openai-agents", "packaging", "pdfminer-six", "starlette", "uv"] +all = ["aiohttp", "crewai", "cryptography", "langchain", "langchain-core", "langsmith", "litellm", "mcp", "openai", "openai-agents", "opentelemetry-instrumentation-aiohttp-client", "opentelemetry-instrumentation-fastapi", "opentelemetry-instrumentation-httpx", "opentelemetry-instrumentation-requests", "opentelemetry-instrumentation-starlette", "packaging", "pdfminer-six", "starlette", "uv"] crewai = ["aiohttp", "crewai", "cryptography", "litellm", "mcp", "pdfminer-six", "uv"] +distributed-tracing = ["opentelemetry-instrumentation-aiohttp-client", "opentelemetry-instrumentation-fastapi", "opentelemetry-instrumentation-httpx", "opentelemetry-instrumentation-requests", "opentelemetry-instrumentation-starlette"] langchain = ["langchain", "langchain-core", "langsmith"] middleware = ["starlette"] openai = ["cryptography", "mcp", "openai", "openai-agents", "packaging"] @@ -6519,4 +6706,4 @@ openai = ["cryptography", "mcp", "openai", "openai-agents", "packaging"] [metadata] lock-version = "2.1" python-versions = "^3.11,<3.15" -content-hash = "5472c5c6216450621d5aae8e3dbdb21a3a071c4c2e3fb383f3552b4dc365d63c" +content-hash = "3e3ae3afd6d9eacc3717e1075d0de4d90a25bdf3a9eaab25c61be0091a353b94" diff --git a/pyproject.toml b/pyproject.toml index c9e1aa6d..d342b4e5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -30,7 +30,8 @@ langchain = ["langchain-core", "langchain", "langsmith (>=0.8.0)"] openai = ["openai (>=2.8.0,<3.0.0)", "packaging (>=24.2,<25.0)", "openai-agents (>=0.4.0,<1.0.0)", "cryptography (>=50.0.0)", "mcp (>=1.27.2,<2)"] crewai = ["crewai (>=0.152.0,<2.0.0); python_version < '3.14'", "litellm (>=1.83.14,<2.0.0); python_version < '3.14'", "uv (>=0.9.6); python_version < '3.14'", "aiohttp (>=3.14.1,<4); python_version < '3.14'", "cryptography (>=50.0.0)", "mcp (>=1.27.2,<2)", "pdfminer-six (>=20251107)"] middleware = ["starlette"] -all = ["langchain-core", "langchain", "langsmith (>=0.8.0)", "openai (>=2.8.0,<3.0.0)", "packaging (>=24.2,<25.0)", "openai-agents (>=0.4.0,<1.0.0)", "crewai (>=0.152.0,<2.0.0); python_version < '3.14'", "starlette", "litellm (>=1.83.14,<2.0.0); python_version < '3.14'", "uv (>=0.9.6); python_version < '3.14'", "aiohttp (>=3.14.1,<4); python_version < '3.14'", "cryptography (>=50.0.0)", "mcp (>=1.27.2,<2)", "pdfminer-six (>=20251107)"] +distributed-tracing = ["opentelemetry-instrumentation-fastapi (==0.59b0)", "opentelemetry-instrumentation-starlette (==0.59b0)", "opentelemetry-instrumentation-requests (==0.59b0)", "opentelemetry-instrumentation-httpx (==0.59b0)", "opentelemetry-instrumentation-aiohttp-client (==0.59b0)"] +all = ["langchain-core", "langchain", "langsmith (>=0.8.0)", "openai (>=2.8.0,<3.0.0)", "packaging (>=24.2,<25.0)", "openai-agents (>=0.4.0,<1.0.0)", "crewai (>=0.152.0,<2.0.0); python_version < '3.14'", "starlette", "litellm (>=1.83.14,<2.0.0); python_version < '3.14'", "uv (>=0.9.6); python_version < '3.14'", "aiohttp (>=3.14.1,<4); python_version < '3.14'", "cryptography (>=50.0.0)", "mcp (>=1.27.2,<2)", "pdfminer-six (>=20251107)", "opentelemetry-instrumentation-fastapi (==0.59b0)", "opentelemetry-instrumentation-starlette (==0.59b0)", "opentelemetry-instrumentation-requests (==0.59b0)", "opentelemetry-instrumentation-httpx (==0.59b0)", "opentelemetry-instrumentation-aiohttp-client (==0.59b0)"] @@ -58,6 +59,11 @@ typing-extensions = { version = ">=4.5.0" } opentelemetry-sdk = "^1.38.0" opentelemetry-api = "^1.38.0" opentelemetry-exporter-otlp-proto-http = "^1.38.0" +opentelemetry-instrumentation-fastapi = { version = "0.59b0", optional = true } +opentelemetry-instrumentation-starlette = { version = "0.59b0", optional = true } +opentelemetry-instrumentation-requests = { version = "0.59b0", optional = true } +opentelemetry-instrumentation-httpx = { version = "0.59b0", optional = true } +opentelemetry-instrumentation-aiohttp-client = { version = "0.59b0", optional = true } filelock = ">=3.20.1" idna = ">=3.15,<4" python-dotenv = ">=1.2.2" diff --git a/src/splunk_ao/__init__.py b/src/splunk_ao/__init__.py index d850956d..2a3203cb 100644 --- a/src/splunk_ao/__init__.py +++ b/src/splunk_ao/__init__.py @@ -58,6 +58,7 @@ ) from splunk_ao.experiment import Experiment from splunk_ao.handlers.agent_control import SplunkAOAgentControlBridge, setup_agent_control_bridge +from splunk_ao.http_instrumentation import instrument_distributed_tracing from splunk_ao.integration import Integration from splunk_ao.logger import SplunkAOLogger from splunk_ao.logger.control import ControlAppliesTo, ControlCheckStage, ControlResult, ControlSpan @@ -172,6 +173,7 @@ "get_annotation_queue", "get_annotation_queue_records", "get_tracing_headers", + "instrument_distributed_tracing", "is_dependency_available", "list_annotation_queue_fields", "list_annotation_queue_users", diff --git a/src/splunk_ao/decorator.py b/src/splunk_ao/decorator.py index ab56f31e..354f3d84 100644 --- a/src/splunk_ao/decorator.py +++ b/src/splunk_ao/decorator.py @@ -65,6 +65,7 @@ def call_llm(prompt, temperature=0.7): from splunk_ao.schema.datasets import DatasetRecord from splunk_ao.schema.metrics import LocalMetricConfig from splunk_ao.schema.trace import SPAN_TYPE +from splunk_ao.session_context import _session_id_context, set_session_context from splunk_ao.shared.exceptions import ConfigurationError from splunk_ao.utils import _get_timestamp from splunk_ao.utils.env_helpers import _get_mode_or_default @@ -91,8 +92,6 @@ def call_llm(prompt, temperature=0.7): _experiment_id_context: ContextVar[str | None] = ContextVar("experiment_id_context", default=None) _span_stack_context: ContextVar[list[WorkflowSpan] | None] = ContextVar("span_stack_context", default=None) _mode_context: ContextVar[LoggerModeType | None] = ContextVar("mode_context", default=None) -_session_id_context: ContextVar[str | None] = ContextVar("session_id_context", default=None) - # Context variables for dataset fields (ground truth/reference output) # These allow setting ground truth data that will be attached to all spans # created within the context, enabling scorers that require reference output. @@ -183,7 +182,15 @@ def __exit__( _trace_context.set(_get_or_init_list(_trace_stack).pop()) _mode_context.set(_get_or_init_list(_mode_stack).pop()) _span_stack_context.set(_get_or_init_list(_span_stack_stack).pop()) - _session_id_context.set(_get_or_init_list(_session_id_stack).pop()) + restored_session_id = _get_or_init_list(_session_id_stack).pop() + _session_id_context.set(restored_session_id) + restored_logger = self.get_logger_instance( + project=_project_context.get(), + agent_stream=_agent_stream_context.get(), + experiment_id=_experiment_id_context.get(), + ) + if isinstance(restored_logger, SplunkAOLogger): + restored_logger._set_active_session_id(restored_session_id) def __call__( self, @@ -1433,6 +1440,7 @@ def start_session( def clear_session(self) -> None: """Clear the session in the active context logger instance.""" self.get_logger_instance().clear_session() + set_session_context(None) def set_session(self, session_id: str) -> None: """ @@ -1444,6 +1452,7 @@ def set_session(self, session_id: str) -> None: The id of the session to set. """ self.get_logger_instance().set_session(session_id) + set_session_context(session_id) splunk_ao_context = SplunkAODecorator() diff --git a/src/splunk_ao/exporter/sink.py b/src/splunk_ao/exporter/sink.py index b32a6088..d23b7505 100644 --- a/src/splunk_ao/exporter/sink.py +++ b/src/splunk_ao/exporter/sink.py @@ -54,14 +54,16 @@ def shutdown(self) -> None: def build_batch_processor(exporter: SpanExporter, config: BatchConfig | None = None) -> BatchSpanProcessor: - """Build a BatchSpanProcessor with explicit SDK defaults.""" - batch_config = config or BatchConfig() + """Build a BatchSpanProcessor with OTel defaults or explicit SDK configuration.""" + if config is None: + return BatchSpanProcessor(exporter) + return BatchSpanProcessor( exporter, - max_queue_size=batch_config.max_queue_size, - schedule_delay_millis=batch_config.schedule_delay_millis, - export_timeout_millis=batch_config.export_timeout_millis, - max_export_batch_size=batch_config.max_export_batch_size, + max_queue_size=config.max_queue_size, + schedule_delay_millis=config.schedule_delay_millis, + export_timeout_millis=config.export_timeout_millis, + max_export_batch_size=config.max_export_batch_size, ) diff --git a/src/splunk_ao/handlers/base_async_handler.py b/src/splunk_ao/handlers/base_async_handler.py index 68053a96..fa7b45ba 100644 --- a/src/splunk_ao/handlers/base_async_handler.py +++ b/src/splunk_ao/handlers/base_async_handler.py @@ -44,6 +44,12 @@ async def async_commit(self) -> None: _logger.warning("Unable to add nodes to trace: Root node does not exist") return + if not getattr(self._splunk_ao_logger, "_ingestion_hook", None): + self._finish_incremental_node(root_node) + if self._flush_on_chain_end: + await self._splunk_ao_logger.async_flush() + return + try: if self._owned_root is not None: self._update_owned_root(root_node) @@ -86,11 +92,7 @@ async def async_commit(self) -> None: self._conclude_owned_state_on_failure() _logger.warning("Failed to commit async handler telemetry", exc_info=True) finally: - self._nodes.clear() - self._root_node = None - self._owned_trace = None - self._owned_root = None - self._owned_parent = None + self._reset_handler_state() async def async_end_node(self, run_id: UUID, **kwargs: Any) -> None: """ @@ -115,6 +117,13 @@ async def async_end_node(self, run_id: UUID, **kwargs: Any) -> None: # Update node parameters node.span_params.update(**kwargs) + if not getattr(self._splunk_ao_logger, "_ingestion_hook", None): + is_root = self._root_node is node + self._finish_incremental_node(node) + if is_root and self._flush_on_chain_end: + await self._splunk_ao_logger.async_flush() + return + # Check if this is the root node and commit if so root = self._root_node if root and node.run_id == root.run_id: diff --git a/src/splunk_ao/handlers/base_handler.py b/src/splunk_ao/handlers/base_handler.py index 63c2d62c..1eca5887 100644 --- a/src/splunk_ao/handlers/base_handler.py +++ b/src/splunk_ao/handlers/base_handler.py @@ -6,6 +6,7 @@ from uuid import UUID from splunk_ao import splunk_ao_context +from splunk_ao.handlers.span_lifecycle import HandlerSpanState, build_handler_step, finalize_handler_step from splunk_ao.logger import SplunkAOLogger from splunk_ao.schema.handlers import INTEGRATION, NODE_TYPE, Node from splunk_ao.schema.logged import LoggedAgentSpan, LoggedWorkflowSpan @@ -62,6 +63,7 @@ def __init__( self._owned_trace: Any = None self._owned_root: Any = None self._owned_parent: Any = None + self._active_steps: dict[str, HandlerSpanState] = {} def commit(self) -> None: """Commit the nodes to the trace using the Splunk AO Logger. Optionally flush the trace.""" @@ -79,6 +81,12 @@ def commit(self) -> None: _logger.warning("Unable to add nodes to trace: Root node does not exist") return + if not getattr(self._splunk_ao_logger, "_ingestion_hook", None): + self._finish_incremental_node(root_node) + if self._flush_on_chain_end: + self._splunk_ao_logger.flush() + return + try: if self._owned_root is not None: self._update_owned_root(root_node) @@ -127,11 +135,82 @@ def commit(self) -> None: self._conclude_owned_state_on_failure() _logger.warning("Failed to commit handler telemetry", exc_info=True) finally: - self._nodes.clear() - self._root_node = None - self._owned_trace = None - self._owned_root = None - self._owned_parent = None + self._reset_handler_state() + + def _reset_handler_state(self) -> None: + """Release one completed callback tree without touching caller-owned state.""" + self._nodes.clear() + self._active_steps.clear() + self._root_node = None + self._owned_trace = None + self._owned_root = None + self._owned_parent = None + + def _start_incremental_step(self, node: Node) -> None: + """Create and activate one real callback span for normal OTLP egress.""" + node_id = str(node.run_id) + if node is self._root_node and self._owned_root is not None: + self._active_steps[node_id] = HandlerSpanState(step=self._owned_root, activation=None) + return + + if node.parent_run_id is not None: + parent_state = self._active_steps.get(str(node.parent_run_id)) + parent = parent_state.step if parent_state is not None else None + else: + parent = self._owned_parent or self._owned_trace + if parent is None: + raise RuntimeError(f"No active parent is available for handler node {node_id}") + + step = build_handler_step(node) + self._splunk_ao_logger._register_handler_step(step, parent) + activation = self._splunk_ao_logger._activate_handler_step(step) + self._active_steps[node_id] = HandlerSpanState(step=step, activation=activation) + + def _finish_incremental_node(self, node: Node) -> None: + """Finalize and enqueue one callback span when that callback ends.""" + node_id = str(node.run_id) + state = self._active_steps.get(node_id) + if state is None: + _logger.warning("Unable to complete handler node %s: no active span state", node_id) + if node is self._root_node: + self._reset_handler_state() + return + + is_root = node is self._root_node + try: + if is_root: + node.span_params["output"] = self._root_output(node) + if state.step is self._owned_root: + node.span_params["input"] = serialize_to_str(node.span_params.get("input", "")) + final = finalize_handler_step(node, state) + final = self._splunk_ao_logger._replace_handler_step(state.step, final) + state.step = final + if state.step is self._owned_root or node_id == str(getattr(self._root_node, "run_id", "")): + self._owned_root = final if self._owned_root is not None else self._owned_root + + self._splunk_ao_logger._restore_handler_step_context(state.activation) + state.activation = None + if self._splunk_ao_logger.current_parent() is final: + self._splunk_ao_logger._set_current_parent(final._parent) + self._splunk_ao_logger._complete_handler_step(final) + self._active_steps.pop(node_id, None) + + if is_root: + root_output = node.span_params.get("output", "") + if self._owned_trace is not None: + self._conclude_owned_trace( + self._owned_trace, + output=SplunkAOLogger._coerce_output(root_output), + status_code=node.span_params.get("status_code"), + ) + except Exception: + self._splunk_ao_logger._restore_handler_step_context(state.activation) + state.activation = None + self._conclude_owned_state_on_failure() + _logger.warning("Failed to complete handler telemetry for node %s", node_id, exc_info=True) + finally: + if is_root: + self._reset_handler_state() def _start_owned_root(self, root_node: Node) -> None: """Open the handler-owned envelope and real root at callback start.""" @@ -139,9 +218,9 @@ def _start_owned_root(self, root_node: Node) -> None: return if not self._start_new_trace: - # Caller-owned mode deliberately leaves the caller's real - # operation current for outbound W3C propagation. The buffered - # handler root is reconstructed beneath it at commit. + # Caller-owned mode preserves the caller's proprietary cursor. + # The incremental handler root is explicitly parented beneath it + # and becomes the active OTel operation during its callback. self._owned_parent = self._splunk_ao_logger.current_parent() return @@ -157,8 +236,8 @@ def _start_owned_root(self, root_node: Node) -> None: return self._owned_parent = parent - # A leaf root is still logged at commit. There is no open Path 1 leaf - # identity to keep active without changing its completion semantics. + # The legacy ingestion-hook path logs a leaf root at commit. Normal + # OTLP egress creates and activates it through _start_incremental_step. if root_node.node_type not in ("agent", "chain", "workflow"): return @@ -462,6 +541,9 @@ def start_node(self, node_type: NODE_TYPE, parent_run_id: UUID | None, run_id: U # chain to an agent immediately before this callback reached us. if self._root_node is not None: self._sync_owned_root_kind(self._root_node) + root_state = self._active_steps.get(str(self._root_node.run_id)) + if root_state is not None and self._owned_root is not None: + root_state.step = self._owned_root # Add to parent's children if parent exists if parent_run_id: @@ -471,6 +553,13 @@ def start_node(self, node_type: NODE_TYPE, parent_run_id: UUID | None, run_id: U else: _logger.debug(f"Parent node {parent_node_id} not found for {node_id}") + if not getattr(self._splunk_ao_logger, "_ingestion_hook", None): + try: + self._start_incremental_step(node) + except Exception: + self._conclude_owned_state_on_failure() + _logger.warning("Failed to start handler span telemetry for node %s", node_id, exc_info=True) + return node def end_node(self, run_id: UUID, **kwargs: Any) -> None: @@ -496,6 +585,13 @@ def end_node(self, run_id: UUID, **kwargs: Any) -> None: # Update node parameters node.span_params.update(**kwargs) + if not getattr(self._splunk_ao_logger, "_ingestion_hook", None): + is_root = self._root_node is node + self._finish_incremental_node(node) + if is_root and self._flush_on_chain_end: + self._splunk_ao_logger.flush() + return + # Check if this is the root node and commit if so root = self._root_node if root and node.run_id == root.run_id: diff --git a/src/splunk_ao/handlers/openai_agents/handler.py b/src/splunk_ao/handlers/openai_agents/handler.py index 9c46590a..d4236263 100644 --- a/src/splunk_ao/handlers/openai_agents/handler.py +++ b/src/splunk_ao/handlers/openai_agents/handler.py @@ -9,6 +9,7 @@ from galileo_core.schemas.logging.span import LlmMetrics, LlmSpan from galileo_core.schemas.logging.span import Span as SplunkAOSpan from splunk_ao import SplunkAOLogger, splunk_ao_context +from splunk_ao.handlers.span_lifecycle import HandlerSpanState, build_handler_step, finalize_handler_step from splunk_ao.schema.handlers import Node from splunk_ao.utils import _get_timestamp from splunk_ao.utils.openai_agents import ( @@ -61,6 +62,8 @@ def __init__(self, splunk_ao_logger: SplunkAOLogger | None = None, flush_on_trac self._owned_trace: Any = None self._owned_root: Any = None self._owned_root_node_id: str | None = None + self._caller_parent: Any = None + self._active_steps: dict[str, HandlerSpanState] = {} def on_trace_start(self, trace: Trace) -> None: """Called when an OpenAI Agent trace starts.""" @@ -75,7 +78,8 @@ def on_trace_start(self, trace: Trace) -> None: }, ) self._nodes[trace.trace_id] = node - if self._splunk_ao_logger.current_parent() is None: + self._caller_parent = self._splunk_ao_logger.current_parent() + if self._caller_parent is None: try: self._owned_trace = self._splunk_ao_logger.start_trace( input="Agent Workflow", @@ -103,6 +107,9 @@ def on_trace_end(self, trace: Trace) -> None: self._conclude_current_trace_on_failure() _logger.warning("Failed to commit OpenAI Agents telemetry", exc_info=True) finally: + for state in self._active_steps.values(): + self._splunk_ao_logger._restore_handler_step_context(state.activation) + self._splunk_ao_logger._release_otel_context(state.step) self._nodes = {} self._last_output = None self._last_status_code = None @@ -110,6 +117,8 @@ def on_trace_end(self, trace: Trace) -> None: self._owned_trace = None self._owned_root = None self._owned_root_node_id = None + self._caller_parent = None + self._active_steps = {} def _commit_trace(self, trace: Trace) -> None: if not self._nodes: @@ -121,6 +130,20 @@ def _commit_trace(self, trace: Trace) -> None: _logger.warning(f"Root node {trace.trace_id} not found") return + if not getattr(self._splunk_ao_logger, "_ingestion_hook", None): + if self._active_steps: + _logger.warning( + "OpenAI Agents trace %s ended with %d unfinished callback spans", + trace.trace_id, + len(self._active_steps), + ) + if self._owned_trace is not None: + self._owned_trace.input = self._first_input or "Agent Workflow" + self._owned_trace.metrics.duration_ns = root_node.span_params.get("duration_ns") + if self._splunk_ao_logger.current_parent() is self._owned_trace: + self._splunk_ao_logger.conclude(output=self._last_output, status_code=self._last_status_code) + return + live_root_node = self._nodes.get(self._owned_root_node_id) if self._owned_root_node_id else None if self._owned_trace is not None and self._owned_root is not None and live_root_node is not None: self._owned_trace.input = self._first_input or "Agent Workflow" @@ -309,6 +332,56 @@ def _update_owned_root(self, node: Node) -> None: if metadata is not None: self._owned_root.user_metadata = convert_to_string_dict(metadata) + def _start_incremental_span(self, node: Node) -> None: + """Create and activate one real OpenAI Agents callback span.""" + node_id = str(node.run_id) + if node_id == self._owned_root_node_id and self._owned_root is not None: + self._active_steps[node_id] = HandlerSpanState(step=self._owned_root, activation=None) + return + + if str(node.parent_run_id) in self._active_steps: + parent = self._active_steps[str(node.parent_run_id)].step + else: + parent = self._owned_trace or self._caller_parent + if parent is None: + raise RuntimeError(f"No active parent is available for OpenAI Agents span {node_id}") + + step = build_handler_step(node, openai_agents=True) + self._splunk_ao_logger._register_handler_step(step, parent) + activation = self._splunk_ao_logger._activate_handler_step(step) + self._active_steps[node_id] = HandlerSpanState(step=step, activation=activation) + + def _finish_incremental_span(self, node: Node) -> None: + """Finalize and enqueue an OpenAI Agents span at callback completion.""" + node_id = str(node.run_id) + state = self._active_steps.get(node_id) + if state is None: + _logger.warning("Unable to complete OpenAI Agents span %s: no active state", node_id) + return + try: + if node.node_type in ("agent", "chain", "workflow") and not node.span_params.get("output"): + last_child = self._nodes.get(node.children[-1]) if node.children else None + if last_child is not None: + node.span_params["output"] = last_child.span_params.get("output", "") + final = finalize_handler_step(node, state, openai_agents=True) + final = self._splunk_ao_logger._replace_handler_step(state.step, final) + state.step = final + if node_id == self._owned_root_node_id: + self._owned_root = final + + self._splunk_ao_logger._restore_handler_step_context(state.activation) + state.activation = None + if self._splunk_ao_logger.current_parent() is final: + self._splunk_ao_logger._set_current_parent(final._parent) + self._splunk_ao_logger._complete_handler_step(final) + self._active_steps.pop(node_id, None) + self._last_output = node.span_params.get("output") + self._last_status_code = node.span_params.get("status_code", 200) + except Exception: + self._splunk_ao_logger._restore_handler_step_context(state.activation) + state.activation = None + _logger.warning("Failed to complete OpenAI Agents span %s", node_id, exc_info=True) + def on_span_start(self, span: Span[Any]) -> None: """Called when an OpenAI Agent span starts.""" span_id = span.span_id @@ -395,6 +468,12 @@ def on_span_start(self, span: Span[Any]) -> None: self._owned_root = None self._owned_root_node_id = None _logger.warning("Failed to start OpenAI Agents root telemetry", exc_info=True) + if not getattr(self._splunk_ao_logger, "_ingestion_hook", None): + try: + self._start_incremental_span(node) + except Exception: + self._conclude_current_trace_on_failure() + _logger.warning("Failed to start OpenAI Agents span telemetry for %s", span_id, exc_info=True) def on_span_end(self, span: Span[Any]) -> None: """Called when an OpenAI Agent span ends.""" @@ -479,6 +558,8 @@ def on_span_end(self, span: Span[Any]) -> None: # Update the node's parameters node.span_params.update(end_params) + if not getattr(self._splunk_ao_logger, "_ingestion_hook", None): + self._finish_incremental_span(node) def shutdown(self) -> None: """Called when the application stops. Flushes any remaining logs.""" diff --git a/src/splunk_ao/handlers/span_lifecycle.py b/src/splunk_ao/handlers/span_lifecycle.py new file mode 100644 index 00000000..2e76ef80 --- /dev/null +++ b/src/splunk_ao/handlers/span_lifecycle.py @@ -0,0 +1,143 @@ +"""Shared callback-to-span construction for incremental handler telemetry.""" + +from __future__ import annotations + +from dataclasses import dataclass +from datetime import UTC, datetime +from uuid import UUID, uuid4 + +from galileo_core.schemas.logging.agent import AgentType +from galileo_core.schemas.logging.span import LlmMetrics, RetrieverSpan, StepWithChildSpans, ToolSpan +from galileo_core.schemas.logging.step import BaseStep, Metrics +from splunk_ao.logger.logger import HandlerStepContext +from splunk_ao.schema.handlers import Node +from splunk_ao.schema.logged import LoggedAgentSpan, LoggedLlmSpan, LoggedWorkflowSpan +from splunk_ao.utils.retrievers import convert_to_documents +from splunk_ao.utils.serialization import convert_to_string_dict, serialize_to_str + + +@dataclass +class HandlerSpanState: + """Mutable ownership state for one in-flight framework callback.""" + + step: BaseStep + activation: HandlerStepContext | None + + +def _created_at(node: Node) -> datetime: + created_at = node.span_params.get("created_at") + if isinstance(created_at, datetime): + return created_at + start_time_iso = node.span_params.get("start_time_iso") + if isinstance(start_time_iso, str) and start_time_iso: + return datetime.fromisoformat(start_time_iso) + return datetime.now(tz=UTC) + + +def _metadata(node: Node) -> dict[str, str] | None: + metadata = node.span_params.get("metadata") + return convert_to_string_dict(metadata) if metadata is not None else None + + +def _step_number(metadata: dict[str, str] | None) -> int | None: + if not metadata or not (value := metadata.get("langgraph_step")): + return None + try: + return int(value) + except (TypeError, ValueError): + return None + + +def build_handler_step( + node: Node, *, step_id: UUID | None = None, children: list[BaseStep] | None = None, openai_agents: bool = False +) -> BaseStep: + """Build a schema-validated provisional or final span from a callback node.""" + params = node.span_params + metadata = _metadata(node) + step_number = _step_number(metadata) + input_value = params.get("input", "") + output = params.get("output", "") + status_code = params.get("status_code") + if params.get("error"): + output = params["error"] + status_code = 500 + if openai_agents and node.node_type in ("agent", "chain", "workflow", "tool"): + input_value = input_value or node.node_type.capitalize() + " Step" + + name = params.get("name") + if openai_agents and node.node_type in ("llm", "chat"): + name = name or "LLM Span" + common = { + "name": name, + "created_at": _created_at(node), + "user_metadata": metadata, + "tags": params.get("tags"), + "status_code": status_code, + "id": step_id or uuid4(), + "step_number": step_number, + } + duration_ns = params.get("duration_ns") + + if node.node_type in ("chain", "workflow") or (openai_agents and node.node_type == "agent"): + return LoggedWorkflowSpan( + **common, + input=input_value, + output=serialize_to_str(output), + metrics=Metrics(duration_ns=duration_ns), + spans=children or [], + ) + if node.node_type == "agent": + return LoggedAgentSpan( + **common, + input=input_value, + output=serialize_to_str(output), + metrics=Metrics(duration_ns=duration_ns), + spans=children or [], + agent_type=params.get("agent_type", AgentType.default), + ) + if node.node_type in ("llm", "chat"): + return LoggedLlmSpan( + **common, + input=input_value, + output=output, + metrics=LlmMetrics.model_validate( + { + "duration_ns": duration_ns, + "num_input_tokens": params.get("num_input_tokens"), + "num_output_tokens": params.get("num_output_tokens"), + "num_total_tokens": params.get("num_total_tokens", params.get("total_tokens")), + "time_to_first_token_ns": params.get("time_to_first_token_ns"), + "num_reasoning_tokens": params.get("num_reasoning_tokens"), + "num_cached_input_tokens": params.get("num_cached_input_tokens"), + } + ), + tools=params.get("tools"), + events=params.get("events"), + model=params.get("model"), + temperature=params.get("temperature"), + ) + if node.node_type == "retriever": + retriever_common = {**common, "status_code": None} + return RetrieverSpan( + **retriever_common, + input=serialize_to_str(input_value), + output=convert_to_documents(output, "output"), + metrics=Metrics(duration_ns=duration_ns), + spans=children or [], + ) + if node.node_type == "tool": + return ToolSpan( + **common, + input=serialize_to_str(input_value), + output=serialize_to_str(output) if output is not None else None, + metrics=Metrics(duration_ns=duration_ns), + spans=children or [], + tool_call_id=params.get("tool_call_id"), + ) + raise ValueError(f"Unsupported handler node type: {node.node_type}") + + +def finalize_handler_step(node: Node, state: HandlerSpanState, *, openai_agents: bool = False) -> BaseStep: + """Build the final validated form while retaining the provisional UUID and children.""" + children = list(state.step.spans) if isinstance(state.step, StepWithChildSpans) else None + return build_handler_step(node, step_id=state.step.id, children=children, openai_agents=openai_agents) diff --git a/src/splunk_ao/http_instrumentation.py b/src/splunk_ao/http_instrumentation.py new file mode 100644 index 00000000..12d37382 --- /dev/null +++ b/src/splunk_ao/http_instrumentation.py @@ -0,0 +1,119 @@ +"""Opt-in setup for supported upstream OpenTelemetry HTTP instrumentors.""" + +from __future__ import annotations + +from typing import Any, cast + +from opentelemetry import trace + +from splunk_ao.otel import TracerProvider +from splunk_ao.session_context import install_session_propagator + +_INSTALL_MESSAGE = ( + "Automatic distributed tracing requires optional dependencies. " + "Install them with: pip install 'splunk-ao[distributed-tracing]'" +) + + +_client_provider_ids: dict[str, int] = {} +_instrumented_apps: set[tuple[int, int, str]] = set() + + +def _load_instrumentors() -> dict[str, type]: + try: + from opentelemetry.instrumentation.aiohttp_client import AioHttpClientInstrumentor # noqa: PLC0415 + from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor # noqa: PLC0415 + from opentelemetry.instrumentation.httpx import HTTPXClientInstrumentor # noqa: PLC0415 + from opentelemetry.instrumentation.requests import RequestsInstrumentor # noqa: PLC0415 + from opentelemetry.instrumentation.starlette import StarletteInstrumentor # noqa: PLC0415 + except ImportError as exc: + raise ImportError(_INSTALL_MESSAGE) from exc + + return { + "fastapi": FastAPIInstrumentor, + "starlette": StarletteInstrumentor, + "requests": RequestsInstrumentor, + "httpx": HTTPXClientInstrumentor, + "aiohttp-client": AioHttpClientInstrumentor, + } + + +def _validate_client_ownership(names: tuple[str, ...], provider_id: int) -> None: + """Reject provider conflicts before any instrumentation is changed.""" + for name in names: + configured_provider_id = _client_provider_ids.get(name) + if configured_provider_id is not None and configured_provider_id != provider_id: + raise RuntimeError( + f"{name} is already instrumented through Splunk AO with another tracer provider; " + "configure each process-wide client instrumentor through one owner" + ) + + +def _instrument_app(app: Any, instrumentors: dict[str, Any], tracer_provider: TracerProvider) -> None: + try: + from fastapi import FastAPI # noqa: PLC0415 + from starlette.applications import Starlette # noqa: PLC0415 + except ImportError as exc: + raise ImportError(_INSTALL_MESSAGE) from exc + + if isinstance(app, FastAPI): + framework = "fastapi" + elif isinstance(app, Starlette): + framework = "starlette" + else: + raise TypeError("app must be a FastAPI or Starlette application") + + key = (id(app), id(tracer_provider), framework) + if key in _instrumented_apps: + return + instrumentors[framework].instrument_app(app, tracer_provider=tracer_provider) + _instrumented_apps.add(key) + + +def instrument_distributed_tracing( + *, + tracer_provider: TracerProvider | None = None, + app: Any | None = None, + instrument_requests: bool = True, + instrument_httpx: bool = True, + instrument_aiohttp_client: bool = True, +) -> None: + """Enable supported upstream OTel server and client instrumentation. + + Parameters + ---------- + tracer_provider : TracerProvider | None + Caller-owned provider used by the upstream instrumentors. The current + global provider is used when omitted; this function never replaces it. + app : Any | None + Optional FastAPI or Starlette application to instrument. + instrument_requests : bool + Whether to instrument process-wide Requests clients. + instrument_httpx : bool + Whether to instrument process-wide HTTPX sync and async clients. + instrument_aiohttp_client : bool + Whether to instrument process-wide aiohttp clients. + + Notes + ----- + Client instrumentation is process-wide. Configure each supported client + through one owner. The caller also owns provider shutdown. + """ + instrumentors = _load_instrumentors() + resolved_provider = tracer_provider or cast(TracerProvider, trace.get_tracer_provider()) + + clients = {"requests": instrument_requests, "httpx": instrument_httpx, "aiohttp-client": instrument_aiohttp_client} + enabled_clients = tuple(name for name, enabled in clients.items() if enabled) + provider_id = id(resolved_provider) + _validate_client_ownership(enabled_clients, provider_id) + + if app is not None: + _instrument_app(app, instrumentors, resolved_provider) + + for name in enabled_clients: + if _client_provider_ids.get(name) == provider_id: + continue + instrumentors[name]().instrument(tracer_provider=resolved_provider) + _client_provider_ids[name] = provider_id + + install_session_propagator() diff --git a/src/splunk_ao/logger/logger.py b/src/splunk_ao/logger/logger.py index a079b04c..68575824 100644 --- a/src/splunk_ao/logger/logger.py +++ b/src/splunk_ao/logger/logger.py @@ -75,6 +75,7 @@ SessionCreateRequest, TracesIngestRequest, ) +from splunk_ao.session_context import get_effective_session_id, set_session_context from splunk_ao.traces import Traces from splunk_ao.utils.decorators import async_warn_catch_exception, nop_async, nop_sync, warn_catch_exception from splunk_ao.utils.env_helpers import _get_mode_or_default @@ -106,6 +107,7 @@ class OtelIds: span_context: SpanContext parent_span_context: SpanContext | None exportable: bool + session_id: str | None @dataclass(frozen=True) @@ -119,6 +121,15 @@ class ActiveOtelContext: token: Token +@dataclass(frozen=True) +class HandlerStepContext: + """OTel activation owned by one framework callback operation.""" + + prior_context: otel_context.Context + token: Token + span_context: SpanContext + + @dataclass(frozen=True) class OtelContextState: """Request-local state for all active proprietary logger contexts.""" @@ -408,6 +419,7 @@ def _record_otel_ids( span_context=self._assign_otel_context(otel_trace_id, trace_flags, trace_state), parent_span_context=parent_span_context, exportable=not isinstance(step, Trace), + session_id=get_effective_session_id(self.session_id), ) self._otel_ids[step.id] = ids return ids @@ -571,7 +583,7 @@ def _emit_and_release(self, finished_step: BaseStep) -> None: span=finished_step, span_context=ids.span_context, parent_span_context=self._export_parent_context(finished_step, ids), - session_id=self.session_id, + session_id=ids.session_id, resource=self._resource, ) ) @@ -618,6 +630,66 @@ def _mark_potentially_parentable(self, span: BaseStep) -> None: if not self._ingestion_hook and span.id in self._otel_ids: self._pending_otel_steps.add(span.id) + def _register_handler_step(self, step: BaseStep, parent: StepWithChildSpans) -> None: + """Register a callback-owned step under an explicit proprietary parent.""" + step._parent = parent + step.dataset_input = parent.dataset_input + step.dataset_output = parent.dataset_output + step.dataset_metadata = parent.dataset_metadata + parent.add_child_span(step) + if isinstance(step, LoggedWorkflowSpan | LoggedAgentSpan) and isinstance(parent, LoggedTrace): + step.conversation_root = True + self._record_otel_ids(step, parent_step=parent) + + def _activate_handler_step(self, step: BaseStep) -> HandlerStepContext | None: + """Make a callback-owned stable identity current for transport propagation.""" + ids = self._otel_ids.get(step.id) + 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) + return HandlerStepContext( + prior_context=prior_context, token=otel_context.attach(active_context), span_context=ids.span_context + ) + + def _restore_handler_step_context(self, activation: HandlerStepContext | None) -> None: + """Restore context owned by one callback without disturbing caller state.""" + if activation is None: + return + try: + otel_context.detach(activation.token) + except (RuntimeError, ValueError): + current = otel_trace.get_current_span().get_span_context() + if current == activation.span_context: + otel_context.attach(activation.prior_context) + self._logger.warning("Failed to detach handler OTel context; restored the prior context when safe.") + + def _replace_handler_step(self, provisional: BaseStep, final: BaseStep) -> BaseStep: + """Replace a provisional callback model while preserving identity and topology.""" + parent = provisional._parent + final._parent = parent + final.dataset_input = provisional.dataset_input + final.dataset_output = provisional.dataset_output + final.dataset_metadata = provisional.dataset_metadata + if isinstance(provisional, LoggedWorkflowSpan | LoggedAgentSpan) and isinstance( + final, LoggedWorkflowSpan | LoggedAgentSpan + ): + final.conversation_root = provisional.conversation_root + if isinstance(provisional, StepWithChildSpans) and isinstance(final, StepWithChildSpans): + final.spans = provisional.spans + final._last_child_created_at = provisional._last_child_created_at + for child in final.spans: + child._parent = final + if parent is not None: + parent.spans = [final if child is provisional else child for child in parent.spans] + if self.current_parent() is provisional and isinstance(final, StepWithChildSpans): + self._set_current_parent(final) + return final + + def _complete_handler_step(self, step: BaseStep) -> None: + """Enqueue one completed handler operation through the shared completion seam.""" + self._complete_step(step) + def _current_span_id(self) -> uuid.UUID: """Return the current proprietary parent ID for internal lifecycle tests.""" current_parent = self.current_parent() @@ -2027,7 +2099,7 @@ async def _start_or_get_session_async( ) -> str: self._session_external_id = external_id if self._ingestion_hook: - self.session_id = str(uuid.uuid4()) + self._set_active_session_id(str(uuid.uuid4())) self._logger.info("Session started: session_id=%s, external_id=%s", self.session_id, external_id) return self.session_id @@ -2052,7 +2124,7 @@ async def _start_or_get_session_async( if sessions and len(sessions["records"]) > 0: session_id = sessions["records"][0]["id"] self._logger.info(f"Session {session_id} with external ID {external_id} already exists; using it.") - self.session_id = session_id + self._set_active_session_id(session_id) return session_id except Exception: self._logger.error("Failed to search for session with external ID %s", external_id, exc_info=True) @@ -2066,9 +2138,14 @@ async def _start_or_get_session_async( ) self._logger.info("Session started with ID: %s", session["id"]) - self.session_id = str(session["id"]) + self._set_active_session_id(str(session["id"])) return self.session_id + def _set_active_session_id(self, session_id: str | None) -> None: + """Update compatibility and request-local session state together.""" + self.session_id = session_id + set_session_context(session_id) + @nop_async async def async_start_session( self, @@ -2140,11 +2217,15 @@ def start_session( str The ID of the session (existing or newly created). """ - return async_run( + session_id = async_run( self._start_or_get_session_async( name=name, previous_session_id=previous_session_id, external_id=external_id, metadata=metadata ) ) + # ``async_run`` may execute in another context; publish the resolved ID + # into the synchronous caller's request-local context as well. + self._set_active_session_id(session_id) + return session_id @nop_sync @warn_catch_exception(exceptions=(Exception,)) @@ -2162,14 +2243,14 @@ def set_session(self, session_id: str) -> None: None """ self._logger.info("Setting the current session to %s", session_id) - self.session_id = session_id + self._set_active_session_id(session_id) self._logger.info("Current session set to %s", session_id) @nop_sync @warn_catch_exception(exceptions=(Exception,)) def clear_session(self) -> None: self._logger.info("Clearing the current session from the logger...") - self.session_id = None + self._set_active_session_id(None) self._logger.info("Current session cleared.") @nop_async diff --git a/src/splunk_ao/otel.py b/src/splunk_ao/otel.py index 28bfba5b..d1155c14 100644 --- a/src/splunk_ao/otel.py +++ b/src/splunk_ao/otel.py @@ -18,13 +18,12 @@ from splunk_ao.config import SplunkAOConfig from splunk_ao.converter import build_span_attributes from splunk_ao.decorator import ( + _agent_stream_context, _dataset_input_context, _dataset_metadata_context, _dataset_output_context, _experiment_id_context, - _agent_stream_context, _project_context, - _session_id_context, ) from splunk_ao.deployment import DeploymentMode, O11yConfig, StandaloneConfig from splunk_ao.exporter import ( @@ -35,6 +34,7 @@ resolve_routing, ) from splunk_ao.exporter.diagnostics import get_export_health +from splunk_ao.session_context import get_effective_session_id logger = logging.getLogger(__name__) @@ -243,7 +243,7 @@ def __init__( def on_start(self, span: Span, parent_context: context.Context | None = None) -> None: """Handle span start events by delegating to the underlying processor.""" - session_id = _session_id_context.get(None) + session_id = get_effective_session_id(context=parent_context) if session_id: span.set_attribute("gen_ai.conversation.id", session_id) @@ -324,12 +324,13 @@ def start_splunk_ao_span(splunk_ao_span: SplunkAOSpan) -> Generator[trace.Span, is_conversation_root = not trace.get_current_span().get_span_context().is_valid and isinstance( splunk_ao_span, WorkflowSpan | AgentSpan ) + session_id = get_effective_session_id() with tracer.start_as_current_span(splunk_ao_span.name) as span: try: yield span finally: try: - attributes = build_span_attributes(splunk_ao_span, _session_id_context.get(None)) + attributes = build_span_attributes(splunk_ao_span, session_id) if is_conversation_root: attributes[GEN_AI_CONVERSATION_ROOT] = True for key, value in attributes.items(): diff --git a/src/splunk_ao/session_context.py b/src/splunk_ao/session_context.py new file mode 100644 index 00000000..bda49cbe --- /dev/null +++ b/src/splunk_ao/session_context.py @@ -0,0 +1,107 @@ +"""Request-local session context and W3C baggage propagation.""" + +from __future__ import annotations + +from contextvars import ContextVar +from typing import Any + +from opentelemetry import baggage, propagate +from opentelemetry import context as otel_context +from opentelemetry.baggage.propagation import W3CBaggagePropagator +from opentelemetry.context import Context +from opentelemetry.propagators import textmap + +GEN_AI_CONVERSATION_ID = "gen_ai.conversation.id" + +_session_id_context: ContextVar[str | None] = ContextVar("session_id_context", default=None) + + +def set_session_context(session_id: str | None) -> None: + """Set the request-local explicit session selection.""" + _session_id_context.set(session_id) + + +def _baggage_session_id(context: Context | None = None) -> str | None: + """Resolve an opaque conversation ID from baggage.""" + value = baggage.get_baggage(GEN_AI_CONVERSATION_ID, context=context) + return value.strip() if isinstance(value, str) and value.strip() else None + + +def get_effective_session_id(logger_session_id: str | None = None, context: Context | None = None) -> str | None: + """Resolve the request-local, inbound, or compatibility session ID.""" + local_session_id = _session_id_context.get(None) + if local_session_id: + return local_session_id + baggage_session_id = _baggage_session_id(context) + if baggage_session_id: + return baggage_session_id + return logger_session_id + + +def _context_with_conversation_id(context: Context, session_id: str | None) -> Context: + """Return a context with the conversation baggage normalized or removed.""" + normalized = baggage.remove_baggage(GEN_AI_CONVERSATION_ID, context=context) + if session_id: + normalized = baggage.set_baggage(GEN_AI_CONVERSATION_ID, session_id, context=normalized) + return normalized + + +class SplunkAOSessionPropagator(textmap.TextMapPropagator): + """Preserve the configured propagator while normalizing session baggage.""" + + def __init__(self, delegate: textmap.TextMapPropagator) -> None: + self._delegate = delegate + self._baggage = W3CBaggagePropagator() + + def extract( + self, + carrier: textmap.CarrierT, + context: Context | None = None, + getter: textmap.Getter[textmap.CarrierT] = textmap.default_getter, + ) -> Context: + extracted = self._delegate.extract(carrier, context=context, getter=getter) + extracted = self._baggage.extract(carrier, context=extracted, getter=getter) + return _context_with_conversation_id(extracted, _baggage_session_id(extracted)) + + def inject( + self, + carrier: textmap.CarrierT, + context: Context | None = None, + setter: textmap.Setter[textmap.CarrierT] = textmap.default_setter, + ) -> None: + active = context if context is not None else otel_context.get_current() + normalized = _context_with_conversation_id(active, get_effective_session_id(context=active)) + self._delegate.inject(carrier, context=normalized, setter=setter) + self._baggage.inject(carrier, context=normalized, setter=setter) + + @property + def fields(self) -> set[str]: + return set(self._delegate.fields) | set(self._baggage.fields) + + +def _session_propagator() -> SplunkAOSessionPropagator: + current = propagate.get_global_textmap() + if isinstance(current, SplunkAOSessionPropagator): + return current + return SplunkAOSessionPropagator(current) + + +def install_session_propagator() -> None: + """Install the session-aware adapter without replacing its delegate.""" + current = propagate.get_global_textmap() + if not isinstance(current, SplunkAOSessionPropagator): + propagate.set_global_textmap(SplunkAOSessionPropagator(current)) + + +def inject_session_context( + carrier: Any, context: Context | None = None, setter: textmap.Setter[Any] = textmap.default_setter +) -> None: + """Inject configured trace context and normalized session baggage.""" + _session_propagator().inject(carrier, context=context, setter=setter) + + +def extract_session_context( + carrier: Any, context: Context | None = None, getter: textmap.Getter[Any] = textmap.default_getter +) -> Context: + """Extract configured trace context and normalize session baggage.""" + return _session_propagator().extract(carrier, context=context, getter=getter) diff --git a/src/splunk_ao/tracing.py b/src/splunk_ao/tracing.py index 96f59c4c..69352b67 100644 --- a/src/splunk_ao/tracing.py +++ b/src/splunk_ao/tracing.py @@ -2,11 +2,11 @@ from collections.abc import Mapping, MutableMapping -from opentelemetry import propagate from opentelemetry.context import Context from splunk_ao.exceptions import SplunkAOLoggerException from splunk_ao.logger.logger import _has_active_exportable_span_context +from splunk_ao.session_context import extract_session_context, inject_session_context def get_tracing_headers(carrier: MutableMapping[str, str] | None = None) -> MutableMapping[str, str]: @@ -34,11 +34,11 @@ def get_tracing_headers(carrier: MutableMapping[str, str] | None = None) -> Muta if carrier is None: carrier = {} - propagate.inject(carrier) + inject_session_context(carrier) return carrier def extract_tracing_context(carrier: Mapping[str, str]) -> Context: """Extract W3C trace context from an incoming text-map carrier.""" normalized = {str(key).lower(): value for key, value in carrier.items()} - return propagate.extract(normalized) + return extract_session_context(normalized) diff --git a/tests/conftest.py b/tests/conftest.py index ba59c4bb..3f046681 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -62,7 +62,7 @@ from splunk_ao.collaborator import CollaboratorRole from splunk_ao.config import SplunkAOConfig from splunk_ao.configuration import _CONFIGURATION_KEYS, Configuration -from splunk_ao.decorator import _mode_context +from splunk_ao.decorator import _mode_context, _session_id_context from splunk_ao.resources.models import DatasetContent, DatasetRow, DatasetRowValuesDict from splunk_ao.resources.models.messages_list_item import MessagesListItem from splunk_ao.utils.singleton import SplunkAOLoggerSingleton @@ -158,8 +158,10 @@ def _clear_otel_test_context() -> None: @pytest.fixture(autouse=True) def reset_otel_test_context() -> Generator[None, None, None]: _clear_otel_test_context() + _session_id_context.set(None) yield _clear_otel_test_context() + _session_id_context.set(None) @pytest.fixture(autouse=True) diff --git a/tests/test_async_base_handler.py b/tests/test_async_base_handler.py index f2934e1d..0484c139 100644 --- a/tests/test_async_base_handler.py +++ b/tests/test_async_base_handler.py @@ -3,6 +3,7 @@ from unittest.mock import Mock, patch import pytest +from opentelemetry.sdk.trace import ReadableSpan from splunk_ao import get_tracing_headers from splunk_ao.handlers.base_async_handler import SplunkAOAsyncBaseHandler @@ -10,6 +11,55 @@ from tests.testutils.setup import setup_mock_logstreams_client, setup_mock_projects_client, setup_mock_traces_client +class RecordingSink: + def __init__(self) -> None: + self.spans: list[ReadableSpan] = [] + self.force_flush_calls = 0 + + def emit(self, span: ReadableSpan) -> None: + self.spans.append(span) + + def force_flush(self) -> bool: + self.force_flush_calls += 1 + return True + + def shutdown(self) -> None: + return None + + +@pytest.mark.asyncio +async def test_async_child_enqueues_when_its_callback_ends() -> None: + sink = RecordingSink() + logger = SplunkAOLogger(project_id="project-id", agent_stream_id="stream-id", _sink=sink) + handler = SplunkAOAsyncBaseHandler(splunk_ao_logger=logger, flush_on_chain_end=False) + root_id = uuid.uuid4() + child_id = uuid.uuid4() + try: + await handler.async_start_node( + node_type="chain", parent_run_id=None, run_id=root_id, name="root", input="request" + ) + await handler.async_start_node( + node_type="llm", parent_run_id=root_id, run_id=child_id, name="child", input="prompt", model="model" + ) + + await handler.async_end_node(child_id, output="answer", model="model") + + assert [(span.attributes or {}).get("gen_ai.operation.name") for span in sink.spans] == ["chat"] + assert str(root_id) in handler._active_steps + assert sink.force_flush_calls == 0 + + await handler.async_end_node(root_id, output="done") + + assert [(span.attributes or {}).get("gen_ai.operation.name") for span in sink.spans] == [ + "chat", + "invoke_workflow", + ] + assert sink.spans[0].parent == sink.spans[1].context + assert sink.force_flush_calls == 0 + finally: + logger.terminate() + + class TestSplunkAOAsyncBaseHandlerCallback: @pytest.fixture @patch("splunk_ao.logger.logger.AgentStreams") diff --git a/tests/test_base_handler.py b/tests/test_base_handler.py index 826a407d..3764101f 100644 --- a/tests/test_base_handler.py +++ b/tests/test_base_handler.py @@ -3,6 +3,7 @@ from unittest.mock import Mock, patch import pytest +from opentelemetry.sdk.trace import ReadableSpan from splunk_ao import get_tracing_headers from splunk_ao.handlers.base_handler import SplunkAOBaseHandler @@ -11,6 +12,89 @@ from tests.testutils.setup import setup_mock_logstreams_client, setup_mock_projects_client, setup_mock_traces_client +class RecordingSink: + def __init__(self) -> None: + self.spans: list[ReadableSpan] = [] + self.force_flush_calls = 0 + + def emit(self, span: ReadableSpan) -> None: + self.spans.append(span) + + def force_flush(self) -> bool: + self.force_flush_calls += 1 + return True + + def shutdown(self) -> None: + return None + + +def operation_names(spans: list[ReadableSpan]) -> list[str | None]: + return [(span.attributes or {}).get("gen_ai.operation.name") for span in spans] + + +def test_normal_otel_child_enqueues_at_callback_end_without_flush() -> None: + sink = RecordingSink() + logger = SplunkAOLogger(project_id="project-id", agent_stream_id="stream-id", _sink=sink) + handler = SplunkAOBaseHandler(splunk_ao_logger=logger, flush_on_chain_end=False) + root_id = uuid.uuid4() + child_id = uuid.uuid4() + try: + handler.start_node(node_type="chain", parent_run_id=None, run_id=root_id, name="root", input="request") + root_step = handler._active_steps[str(root_id)].step + handler.start_node( + node_type="llm", parent_run_id=root_id, run_id=child_id, name="child", input="prompt", model="model" + ) + + handler.end_node(child_id, output="answer", model="model") + + assert operation_names(sink.spans) == ["chat"] + assert str(root_id) in handler._active_steps + assert logger.current_parent() is root_step + assert sink.force_flush_calls == 0 + + handler.end_node(root_id, output="done") + + assert operation_names(sink.spans) == ["chat", "invoke_workflow"] + child_span, root_span = sink.spans + assert child_span.parent == root_span.context + assert root_span.parent is None + assert [child.name for child in root_step.spans] == ["child"] + assert sink.force_flush_calls == 0 + finally: + logger.terminate() + + +def test_incremental_handler_root_is_active_beneath_caller_owned_operation() -> None: + sink = RecordingSink() + logger = SplunkAOLogger(project_id="project-id", agent_stream_id="stream-id", _sink=sink) + logger.start_trace(input="caller trace") + caller = logger.add_workflow_span(input="caller operation", name="caller") + handler = SplunkAOBaseHandler(splunk_ao_logger=logger, start_new_trace=False, flush_on_chain_end=False) + root_id = uuid.uuid4() + child_id = uuid.uuid4() + try: + handler.start_node(node_type="chain", parent_run_id=None, run_id=root_id, name="handler", input="request") + handler_root = handler._active_steps[str(root_id)].step + handler_context = logger._otel_ids[handler_root.id].span_context + headers = get_tracing_headers() + handler.start_node( + node_type="llm", parent_run_id=root_id, run_id=child_id, name="child", input="prompt", model="model" + ) + handler.end_node(child_id, output="answer", model="model") + handler.end_node(root_id, output="done") + + assert headers["traceparent"].split("-")[2] == format(handler_context.span_id, "016x") + assert logger.current_parent() is caller + child_span, handler_span = sink.spans + assert child_span.parent == handler_span.context + assert handler_span.parent == logger._otel_ids[caller.id].span_context + assert [step.name for step in caller.spans] == ["handler"] + finally: + logger.conclude(output="caller done") + logger.conclude(output="trace done") + logger.terminate() + + class TestSplunkAOBaseHandler: @pytest.fixture @patch("splunk_ao.logger.logger.AgentStreams") diff --git a/tests/test_exporter_sink.py b/tests/test_exporter_sink.py index 6dd7cef0..8c5f6c92 100644 --- a/tests/test_exporter_sink.py +++ b/tests/test_exporter_sink.py @@ -2,6 +2,7 @@ from collections.abc import Generator, Sequence from typing import Any +from unittest.mock import patch import pytest from opentelemetry import trace @@ -58,6 +59,15 @@ def test_batch_processor_wraps_exporter(shutdown_workers: list[Any]) -> None: assert isinstance(processor, BatchSpanProcessor) +def test_default_batch_processor_delegates_configuration_to_otel() -> None: + exporter = RecordingExporter() + + with patch("splunk_ao.exporter.sink.BatchSpanProcessor") as processor_class: + build_batch_processor(exporter) + + processor_class.assert_called_once_with(exporter) + + def test_batch_processor_accepts_custom_config(shutdown_workers: list[Any]) -> None: exporter = RecordingExporter() processor = build_batch_processor(exporter, BatchConfig(max_export_batch_size=1)) diff --git a/tests/test_http_instrumentation.py b/tests/test_http_instrumentation.py new file mode 100644 index 00000000..dd186b2b --- /dev/null +++ b/tests/test_http_instrumentation.py @@ -0,0 +1,325 @@ +from __future__ import annotations + +from collections.abc import Generator +from dataclasses import dataclass +from typing import Any +from unittest.mock import MagicMock, patch + +import pytest +import requests +from fastapi import FastAPI +from opentelemetry import baggage, context, propagate, trace +from opentelemetry.context import Context +from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor +from opentelemetry.instrumentation.requests import RequestsInstrumentor +from opentelemetry.propagators import textmap +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.trace import NonRecordingSpan, SpanContext, TraceFlags +from starlette.applications import Starlette +from starlette.testclient import TestClient + +from splunk_ao import instrument_distributed_tracing +from splunk_ao.http_instrumentation import _client_provider_ids, _instrumented_apps, _load_instrumentors +from splunk_ao.session_context import GEN_AI_CONVERSATION_ID, SplunkAOSessionPropagator, _session_id_context + + +class FakeTracerProvider: + def add_span_processor(self, span_processor: Any) -> None: + return None + + def get_tracer( + self, + instrumenting_module_name: str, + instrumenting_library_version: str | None = None, + schema_url: str | None = None, + attributes: Any | None = None, + ) -> MagicMock: + return MagicMock() + + +class CustomPropagator(textmap.TextMapPropagator): + def extract( + self, + carrier: textmap.CarrierT, + context: context.Context | None = None, + getter: textmap.Getter[textmap.CarrierT] = textmap.default_getter, + ) -> context.Context: + return context if context is not None else Context() + + def inject( + self, + carrier: textmap.CarrierT, + context: context.Context | None = None, + setter: textmap.Setter[textmap.CarrierT] = textmap.default_setter, + ) -> None: + setter.set(carrier, "x-custom-context", "preserved") + + @property + def fields(self) -> set[str]: + return {"x-custom-context"} + + +@dataclass +class RecordingInstrumentor: + instrument_calls: list[dict[str, Any]] + app_calls: list[tuple[Any, dict[str, Any]]] + + def instrument(self, **kwargs: Any) -> None: + self.instrument_calls.append(kwargs) + + def instrument_app(self, app: Any, **kwargs: Any) -> None: + self.app_calls.append((app, kwargs)) + + +def instrumentor_type(recorder: RecordingInstrumentor) -> type: + class Instrumentor: + def __new__(cls) -> RecordingInstrumentor: + return recorder + + @classmethod + def instrument_app(cls, app: Any, **kwargs: Any) -> None: + recorder.instrument_app(app, **kwargs) + + return Instrumentor + + +@pytest.fixture(autouse=True) +def reset_instrumentation_state() -> Generator[None, None, None]: + previous_propagator = propagate.get_global_textmap() + session_token = _session_id_context.set(None) + _client_provider_ids.clear() + _instrumented_apps.clear() + try: + yield + finally: + _client_provider_ids.clear() + _instrumented_apps.clear() + propagate.set_global_textmap(previous_propagator) + _session_id_context.reset(session_token) + + +@pytest.fixture +def instrumentors() -> tuple[dict[str, type], dict[str, RecordingInstrumentor]]: + recorders = { + name: RecordingInstrumentor(instrument_calls=[], app_calls=[]) + for name in ("fastapi", "starlette", "requests", "httpx", "aiohttp_client") + } + types = { + "fastapi": instrumentor_type(recorders["fastapi"]), + "starlette": instrumentor_type(recorders["starlette"]), + "requests": instrumentor_type(recorders["requests"]), + "httpx": instrumentor_type(recorders["httpx"]), + "aiohttp-client": instrumentor_type(recorders["aiohttp_client"]), + } + return types, recorders + + +def test_missing_extra_fails_before_instrumentation_or_propagator_change() -> None: + previous_propagator = propagate.get_global_textmap() + + with ( + patch("splunk_ao.http_instrumentation._load_instrumentors", side_effect=ImportError("install the extra")), + pytest.raises(ImportError, match="install the extra"), + ): + instrument_distributed_tracing(tracer_provider=FakeTracerProvider()) + + assert _client_provider_ids == {} + assert _instrumented_apps == set() + assert propagate.get_global_textmap() is previous_propagator + + +def test_distributed_tracing_extra_loads_every_supported_upstream_instrumentor() -> None: + assert set(_load_instrumentors()) == {"fastapi", "starlette", "requests", "httpx", "aiohttp-client"} + + +@pytest.mark.parametrize(("app", "expected"), [(FastAPI(), "fastapi"), (Starlette(), "starlette")]) +def test_instruments_matching_app_and_every_enabled_client_once( + app: Any, expected: str, instrumentors: tuple[dict[str, type], dict[str, RecordingInstrumentor]] +) -> None: + types, recorders = instrumentors + provider = FakeTracerProvider() + + with patch("splunk_ao.http_instrumentation._load_instrumentors", return_value=types): + instrument_distributed_tracing(tracer_provider=provider, app=app) + instrument_distributed_tracing(tracer_provider=provider, app=app) + + assert recorders[expected].app_calls == [(app, {"tracer_provider": provider})] + other_framework = "starlette" if expected == "fastapi" else "fastapi" + assert recorders[other_framework].app_calls == [] + assert recorders["requests"].instrument_calls == [{"tracer_provider": provider}] + assert recorders["httpx"].instrument_calls == [{"tracer_provider": provider}] + assert recorders["aiohttp_client"].instrument_calls == [{"tracer_provider": provider}] + assert isinstance(propagate.get_global_textmap(), SplunkAOSessionPropagator) + + +def test_disabled_clients_are_not_instrumented( + instrumentors: tuple[dict[str, type], dict[str, RecordingInstrumentor]], +) -> None: + types, recorders = instrumentors + + with patch("splunk_ao.http_instrumentation._load_instrumentors", return_value=types): + instrument_distributed_tracing( + tracer_provider=FakeTracerProvider(), + instrument_requests=False, + instrument_httpx=False, + instrument_aiohttp_client=False, + ) + + assert all(not recorder.instrument_calls for recorder in recorders.values()) + + +def test_provider_conflict_fails_before_partial_instrumentation( + instrumentors: tuple[dict[str, type], dict[str, RecordingInstrumentor]], +) -> None: + types, recorders = instrumentors + first_provider = FakeTracerProvider() + second_provider = FakeTracerProvider() + + with patch("splunk_ao.http_instrumentation._load_instrumentors", return_value=types): + instrument_distributed_tracing(tracer_provider=first_provider) + with pytest.raises(RuntimeError, match="another tracer provider"): + instrument_distributed_tracing(tracer_provider=second_provider, app=FastAPI()) + + assert recorders["fastapi"].app_calls == [] + assert all(len(recorders[name].instrument_calls) == 1 for name in ("requests", "httpx", "aiohttp_client")) + + +def test_uses_current_provider_without_replacing_it( + instrumentors: tuple[dict[str, type], dict[str, RecordingInstrumentor]], +) -> None: + types, recorders = instrumentors + provider = FakeTracerProvider() + + with ( + patch("splunk_ao.http_instrumentation._load_instrumentors", return_value=types), + patch("splunk_ao.http_instrumentation.trace.get_tracer_provider", return_value=provider), + patch.object(trace, "set_tracer_provider") as set_provider, + ): + instrument_distributed_tracing() + + set_provider.assert_not_called() + assert recorders["requests"].instrument_calls == [{"tracer_provider": provider}] + + +def test_installed_global_propagator_injects_only_conversation_session( + instrumentors: tuple[dict[str, type], dict[str, RecordingInstrumentor]], +) -> None: + types, _ = instrumentors + span_context = SpanContext( + trace_id=0x4BF92F3577B34DA6A3CE929D0E0E4736, + span_id=0x00F067AA0BA902B7, + is_remote=False, + trace_flags=TraceFlags.SAMPLED, + ) + token = context.attach(trace.set_span_in_context(NonRecordingSpan(span_context))) + _session_id_context.set("conversation-123") + try: + with patch("splunk_ao.http_instrumentation._load_instrumentors", return_value=types): + instrument_distributed_tracing( + tracer_provider=FakeTracerProvider(), + instrument_requests=False, + instrument_httpx=False, + instrument_aiohttp_client=False, + ) + carrier: dict[str, str] = {} + propagate.inject(carrier) + finally: + context.detach(token) + + extracted = propagate.extract(carrier) + assert baggage.get_baggage(GEN_AI_CONVERSATION_ID, context=extracted) == "conversation-123" + assert "splunk_ao.session.id" not in carrier.get("baggage", "") + + +def test_session_adapter_preserves_application_configured_propagator( + instrumentors: tuple[dict[str, type], dict[str, RecordingInstrumentor]], +) -> None: + types, _ = instrumentors + custom = CustomPropagator() + propagate.set_global_textmap(custom) + + with patch("splunk_ao.http_instrumentation._load_instrumentors", return_value=types): + instrument_distributed_tracing( + tracer_provider=FakeTracerProvider(), + instrument_requests=False, + instrument_httpx=False, + instrument_aiohttp_client=False, + ) + + carrier: dict[str, str] = {} + propagate.inject(carrier) + installed = propagate.get_global_textmap() + assert carrier["x-custom-context"] == "preserved" + assert isinstance(installed, SplunkAOSessionPropagator) + assert installed._delegate is custom + + +def test_real_requests_instrumentor_injects_trace_and_conversation_without_manual_headers( + monkeypatch: pytest.MonkeyPatch, +) -> None: + captured_headers: dict[str, str] = {} + + def capture_send(session: requests.Session, request: requests.PreparedRequest, **kwargs: Any) -> requests.Response: + del session, kwargs + captured_headers.update(request.headers) + response = requests.Response() + response.status_code = 200 + response.request = request + response.url = request.url + return response + + monkeypatch.setattr(requests.Session, "send", capture_send) + provider = TracerProvider() + requests_instrumentor = RequestsInstrumentor() + try: + instrument_distributed_tracing( + tracer_provider=provider, instrument_httpx=False, instrument_aiohttp_client=False + ) + _session_id_context.set("conversation-outbound") + with provider.get_tracer(__name__).start_as_current_span("caller") as caller: + response = requests.get("https://example.test/automatic-context", timeout=1) + expected_trace_id = f"{caller.get_span_context().trace_id:032x}" + + assert response.status_code == 200 + assert captured_headers["traceparent"].split("-")[1] == expected_trace_id + assert captured_headers["baggage"] == "gen_ai.conversation.id=conversation-outbound" + assert "splunk_ao.session.id" not in captured_headers["baggage"] + finally: + requests_instrumentor.uninstrument() + provider.shutdown() + + +def test_real_fastapi_instrumentor_extracts_trace_and_conversation_without_sdk_middleware() -> None: + expected_trace_id = 0x4BF92F3577B34DA6A3CE929D0E0E4736 + observed: dict[str, Any] = {} + provider = TracerProvider() + app = FastAPI() + + @app.get("/automatic-context") + def automatic_context() -> dict[str, bool]: + observed["trace_id"] = trace.get_current_span().get_span_context().trace_id + observed["conversation_id"] = baggage.get_baggage(GEN_AI_CONVERSATION_ID) + return {"ok": True} + + try: + instrument_distributed_tracing( + tracer_provider=provider, + app=app, + instrument_requests=False, + instrument_httpx=False, + instrument_aiohttp_client=False, + ) + with TestClient(app) as client: + response = client.get( + "/automatic-context", + headers={ + "traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01", + "baggage": "gen_ai.conversation.id=conversation-inbound", + }, + ) + + assert response.json() == {"ok": True} + assert observed == {"trace_id": expected_trace_id, "conversation_id": "conversation-inbound"} + finally: + FastAPIInstrumentor.uninstrument_app(app) + provider.shutdown() diff --git a/tests/test_logger_otel_egress.py b/tests/test_logger_otel_egress.py index 2a3a976c..250dcc04 100644 --- a/tests/test_logger_otel_egress.py +++ b/tests/test_logger_otel_egress.py @@ -76,6 +76,23 @@ def test_complete_leaf_is_enqueued_before_flush(otlp_logger: SplunkAOLogger, rec assert recording_sink.force_flush_calls == 0 +def test_path1_session_is_captured_when_stable_span_identity_is_assigned( + otlp_logger: SplunkAOLogger, recording_sink: RecordingSink +) -> None: + otlp_logger.set_session("conversation-at-start") + otlp_logger.start_trace(input="question") + otlp_logger.add_workflow_span(input="work", name="workflow") + + otlp_logger.set_session("conversation-changed-later") + otlp_logger.conclude(output="done") + otlp_logger.conclude(output="trace done") + + [workflow_span] = recording_sink.spans + assert workflow_span.attributes["gen_ai.conversation.id"] == "conversation-at-start" + assert "splunk_ao.session.id" not in workflow_span.attributes + otlp_logger.clear_session() + + def test_single_llm_trace_emits_only_real_child(otlp_logger: SplunkAOLogger, recording_sink: RecordingSink) -> None: otlp_logger.add_single_llm_span_trace(input="question", output="answer", model="gpt-5") diff --git a/tests/test_openai_agents.py b/tests/test_openai_agents.py index 205bc6f9..4e00e485 100644 --- a/tests/test_openai_agents.py +++ b/tests/test_openai_agents.py @@ -16,6 +16,7 @@ set_trace_processors, ) from agents.tracing import ResponseSpanData +from opentelemetry.sdk.trace import ReadableSpan from pydantic import BaseModel from pytest import MonkeyPatch, mark @@ -28,6 +29,81 @@ from tests.testutils.setup import setup_mock_logstreams_client, setup_mock_projects_client, setup_mock_traces_client +class RecordingSink: + def __init__(self) -> None: + self.spans: list[ReadableSpan] = [] + self.force_flush_calls = 0 + + def emit(self, span: ReadableSpan) -> None: + self.spans.append(span) + + def force_flush(self) -> bool: + self.force_flush_calls += 1 + return True + + def shutdown(self) -> None: + return None + + +def test_openai_agents_children_enqueue_before_trace_end() -> None: + 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) + trace = MagicMock(trace_id="trace-id", name="Agent trace", metadata={}) + root = Node( + node_type="workflow", + run_id="root-span-id", + parent_run_id="trace-id", + span_params={ + "input": "question", + "name": "agent root", + "start_time_iso": "2025-01-01T00:00:00+00:00", + "status_code": 200, + }, + ) + child = Node( + node_type="llm", + run_id="child-span-id", + parent_run_id="root-span-id", + span_params={ + "input": "prompt", + "name": "model call", + "start_time_iso": "2025-01-01T00:00:01+00:00", + "model": "model", + "status_code": 200, + }, + ) + try: + processor.on_trace_start(trace) + processor._nodes[str(root.run_id)] = root + processor._nodes[trace.trace_id].children.append(str(root.run_id)) + processor._start_owned_root(root) + processor._start_incremental_span(root) + processor._nodes[str(child.run_id)] = child + root.children.append(str(child.run_id)) + processor._start_incremental_span(child) + + child.span_params.update(output="answer", duration_ns=10) + processor._finish_incremental_span(child) + + assert [(span.attributes or {}).get("gen_ai.operation.name") for span in sink.spans] == ["chat"] + assert str(root.run_id) in processor._active_steps + assert sink.force_flush_calls == 0 + + root.span_params.update(output="answer", duration_ns=20) + processor._finish_incremental_span(root) + processor.on_trace_end(trace) + + assert [(span.attributes or {}).get("gen_ai.operation.name") for span in sink.spans] == [ + "chat", + "invoke_workflow", + ] + assert sink.spans[0].parent == sink.spans[1].context + assert sink.force_flush_calls == 0 + finally: + logger.terminate() + + class HomeworkOutput(BaseModel): is_homework: bool reasoning: str diff --git a/tests/test_otel_native_paths.py b/tests/test_otel_native_paths.py index d2850c74..68fe5c9c 100644 --- a/tests/test_otel_native_paths.py +++ b/tests/test_otel_native_paths.py @@ -3,7 +3,7 @@ from unittest.mock import MagicMock, patch import pytest -from opentelemetry import trace +from opentelemetry import baggage, context, trace from opentelemetry.sdk.resources import Resource from opentelemetry.sdk.trace import Event, ReadableSpan from opentelemetry.sdk.trace import TracerProvider as SDKTracerProvider @@ -365,6 +365,20 @@ def test_processor_does_not_put_routing_on_span_attributes() -> None: processor.shutdown() +def test_processor_reads_inbound_standard_conversation_baggage() -> None: + exporter = RecordingExporter() + processor = SplunkAOSpanProcessor(SpanProcessor=RecordingSpanProcessor, _exporter=exporter) + parent_context = baggage.set_baggage("gen_ai.conversation.id", "inbound-conversation", context.Context()) + span = MagicMock() + + processor.on_start(span, parent_context=parent_context) + + calls = {args[0]: args[1] for args, _ in span.set_attribute.call_args_list} + assert calls["gen_ai.conversation.id"] == "inbound-conversation" + assert "splunk_ao.session.id" not in calls + processor.shutdown() + + def test_processor_forwards_complete_routing_to_immutable_exporter() -> None: exporter = RecordingExporter() exporter_factory = MagicMock() diff --git a/tests/test_tracing.py b/tests/test_tracing.py index 71ff3bd3..473ee7dc 100644 --- a/tests/test_tracing.py +++ b/tests/test_tracing.py @@ -1,7 +1,7 @@ from collections.abc import Generator import pytest -from opentelemetry import context, trace +from opentelemetry import baggage, context, trace from opentelemetry.sdk.trace import ReadableSpan from opentelemetry.trace import NonRecordingSpan, SpanContext, TraceFlags, TraceState @@ -9,6 +9,7 @@ from splunk_ao.exceptions import SplunkAOLoggerException from splunk_ao.logger import SplunkAOLogger from splunk_ao.logger.logger import _otel_context_state +from splunk_ao.session_context import GEN_AI_CONVERSATION_ID, _session_id_context class RecordingSink: @@ -28,11 +29,13 @@ def shutdown(self) -> None: @pytest.fixture(autouse=True) def isolated_context() -> Generator[None, None, None]: state_token = _otel_context_state.set(None) + session_token = _session_id_context.set(None) token = context.attach(context.Context()) try: yield finally: context.detach(token) + _session_id_context.reset(session_token) _otel_context_state.reset(state_token) @@ -102,6 +105,100 @@ def test_get_tracing_headers_uses_real_path1_operation_and_no_routing_headers() logger.terminate() +def test_get_tracing_headers_injects_only_standard_conversation_baggage() -> None: + logger, _ = make_logger() + try: + logger.set_session("conversation-123") + logger.start_trace(input="request") + logger.add_workflow_span(input="work", name="operation") + + headers = get_tracing_headers() + + extracted = extract_tracing_context(headers) + assert baggage.get_baggage(GEN_AI_CONVERSATION_ID, context=extracted) == "conversation-123" + assert "splunk_ao.session.id" not in headers.get("baggage", "") + assert all( + private_name not in headers.get("baggage", "") + for private_name in ( + "project", + "agent_stream", + "logstream", + "experiment", + "agent_name", + "workflow_name", + "model_name", + "token", + ) + ) + logger.conclude(output="work") + logger.conclude(output="done") + finally: + logger.terminate() + + +def test_extract_tracing_context_restores_conversation_baggage() -> None: + extracted = extract_tracing_context( + { + "traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01", + "baggage": "gen_ai.conversation.id=conversation-456,unrelated=value", + } + ) + + assert baggage.get_baggage(GEN_AI_CONVERSATION_ID, context=extracted) == "conversation-456" + assert baggage.get_baggage("unrelated", context=extracted) == "value" + + +@pytest.mark.parametrize("header", ["gen_ai.conversation.id=", "gen_ai.conversation.id"]) +def test_extract_tracing_context_ignores_malformed_or_empty_conversation_baggage(header: str) -> None: + extracted = extract_tracing_context( + {"traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01", "baggage": header} + ) + + assert baggage.get_baggage(GEN_AI_CONVERSATION_ID, context=extracted) is None + + +def test_local_session_overrides_inbound_conversation_when_injecting() -> None: + inbound = extract_tracing_context( + { + "traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01", + "baggage": "gen_ai.conversation.id=inbound-session,unrelated=value", + } + ) + token = context.attach(inbound) + _session_id_context.set("local-session") + try: + headers = get_tracing_headers() + finally: + context.detach(token) + + extracted = extract_tracing_context(headers) + assert baggage.get_baggage(GEN_AI_CONVERSATION_ID, context=extracted) == "local-session" + assert baggage.get_baggage("unrelated", context=extracted) == "value" + + +def test_inbound_conversation_overrides_compatibility_logger_field_for_path1() -> None: + remote = extract_tracing_context( + { + "traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01", + "baggage": "gen_ai.conversation.id=inbound-session", + } + ) + token = context.attach(remote) + logger, sink = make_logger() + logger.session_id = "stale-logger-session" + try: + logger.start_trace(input="request") + logger.add_workflow_span(input="work", name="operation") + logger.conclude(output="work") + logger.conclude(output="done") + + [operation] = sink.spans + assert operation.attributes[GEN_AI_CONVERSATION_ID] == "inbound-session" + finally: + logger.terminate() + context.detach(token) + + def test_extract_tracing_context_is_case_insensitive_and_preserves_tracestate() -> None: extracted = extract_tracing_context( {"TRACEPARENT": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01", "TRACESTATE": "vendor=value"} From c836026cb1cf6cf590d446f573524d6f5e8bd9b7 Mon Sep 17 00:00:00 2001 From: pradystar Date: Thu, 13 Aug 2026 11:40:07 -0700 Subject: [PATCH 03/13] test(tracing): complete DT.A and DT.B regression coverage --- splunk-ao-adk/tests/test_span_manager.py | 51 +++++ tests/test_async_base_handler.py | 68 +++++++ tests/test_base_handler.py | 125 ++++++++++++ tests/test_crewai_handler.py | 61 ++++++ tests/test_http_instrumentation.py | 233 ++++++++++++++++++++++- tests/test_langchain.py | 46 +++++ tests/test_otel_native_paths.py | 43 +++++ 7 files changed, 626 insertions(+), 1 deletion(-) diff --git a/splunk-ao-adk/tests/test_span_manager.py b/splunk-ao-adk/tests/test_span_manager.py index 27eda6c2..0c18f254 100644 --- a/splunk-ao-adk/tests/test_span_manager.py +++ b/splunk-ao-adk/tests/test_span_manager.py @@ -4,10 +4,61 @@ from uuid import uuid4 import pytest +from opentelemetry.sdk.trace import ReadableSpan +from splunk_ao.handlers.base_handler import SplunkAOBaseHandler +from splunk_ao.logger.logger import SplunkAOLogger from splunk_ao_adk.span_manager import INTEGRATION_TAG, SpanManager +class RecordingSink: + def __init__(self) -> None: + self.spans: list[ReadableSpan] = [] + self.force_flush_calls = 0 + + def emit(self, span: ReadableSpan) -> None: + self.spans.append(span) + + def force_flush(self) -> bool: + self.force_flush_calls += 1 + return True + + def shutdown(self) -> None: + return None + + +def test_adk_span_manager_enqueues_llm_at_callback_end_without_flush(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("SPLUNK_AO_API_KEY", "test-api-key") + monkeypatch.setenv("SPLUNK_AO_CONSOLE_URL", "https://console.test") + sink = RecordingSink() + logger = SplunkAOLogger(project_id="project-id", agent_stream_id="stream-id", _sink=sink) + handler = SplunkAOBaseHandler(splunk_ao_logger=logger, integration="google_adk", flush_on_chain_end=False) + manager = SpanManager(handler) + root_id = uuid4() + llm_id = uuid4() + try: + manager.start_run(root_id, "question", agent_name="researcher") + manager.start_llm(llm_id, root_id, "prompt", model="model") + + manager.end_llm(llm_id, "answer", num_input_tokens=2, num_output_tokens=1, total_tokens=3) + + assert [(span.attributes or {}).get("gen_ai.operation.name") for span in sink.spans] == ["chat"] + assert str(root_id) in handler._active_steps + assert sink.force_flush_calls == 0 + + manager.end_run(root_id, "done") + + child, root = sink.spans + assert child.parent == root.context + assert [span.name for span in sink.spans] == [ + "chat model", + "invoke_workflow invocation [researcher]", + ] + assert sink.force_flush_calls == 0 + finally: + logger.terminate() + + class TestSpanManagerRunSpans: """Tests for run (invocation) span lifecycle.""" diff --git a/tests/test_async_base_handler.py b/tests/test_async_base_handler.py index 0484c139..cdc6dd10 100644 --- a/tests/test_async_base_handler.py +++ b/tests/test_async_base_handler.py @@ -1,3 +1,4 @@ +import asyncio import uuid from collections.abc import Generator from unittest.mock import Mock, patch @@ -8,6 +9,7 @@ from splunk_ao import get_tracing_headers from splunk_ao.handlers.base_async_handler import SplunkAOAsyncBaseHandler from splunk_ao.logger.logger import SplunkAOLogger +from splunk_ao.session_context import set_session_context from tests.testutils.setup import setup_mock_logstreams_client, setup_mock_projects_client, setup_mock_traces_client @@ -60,6 +62,72 @@ async def test_async_child_enqueues_when_its_callback_ends() -> None: logger.terminate() +@pytest.mark.asyncio +async def test_concurrent_async_handlers_keep_parent_and_session_context_isolated() -> None: + sink = RecordingSink() + logger = SplunkAOLogger(project_id="project-id", agent_stream_id="stream-id", _sink=sink) + first_started = asyncio.Event() + second_started = asyncio.Event() + + async def run_handler(label: str, own_started: asyncio.Event, other_started: asyncio.Event) -> None: + set_session_context(label) + handler = SplunkAOAsyncBaseHandler(splunk_ao_logger=logger, flush_on_chain_end=False) + root_id = uuid.uuid4() + child_id = uuid.uuid4() + try: + await handler.async_start_node( + node_type="chain", parent_run_id=None, run_id=root_id, name=f"{label}-root", input="request" + ) + own_started.set() + await other_started.wait() + await handler.async_start_node( + node_type="llm", + parent_run_id=root_id, + run_id=child_id, + name=f"{label}-child", + input="prompt", + model="model", + ) + await asyncio.sleep(0) + await handler.async_end_node(child_id, output="answer", model="model") + await handler.async_end_node(root_id, output="done") + finally: + set_session_context(None) + + try: + await asyncio.gather( + run_handler("conversation-a", first_started, second_started), + run_handler("conversation-b", second_started, first_started), + ) + + assert len(sink.spans) == 4 + for label in ("conversation-a", "conversation-b"): + request_spans = [ + span for span in sink.spans if (span.attributes or {}).get("gen_ai.conversation.id") == label + ] + root = next( + span + for span in request_spans + if (span.attributes or {}).get("gen_ai.operation.name") == "invoke_workflow" + ) + child = next( + span for span in request_spans if (span.attributes or {}).get("gen_ai.operation.name") == "chat" + ) + assert child.parent == root.context + assert child.context.trace_id == root.context.trace_id + assert (child.attributes or {}).get("gen_ai.conversation.id") == label + assert (root.attributes or {}).get("gen_ai.conversation.id") == label + root_trace_ids = { + span.context.trace_id + for span in sink.spans + if (span.attributes or {}).get("gen_ai.operation.name") == "invoke_workflow" + } + assert len(root_trace_ids) == 2 + assert sink.force_flush_calls == 0 + finally: + logger.terminate() + + class TestSplunkAOAsyncBaseHandlerCallback: @pytest.fixture @patch("splunk_ao.logger.logger.AgentStreams") diff --git a/tests/test_base_handler.py b/tests/test_base_handler.py index 3764101f..1e53a11b 100644 --- a/tests/test_base_handler.py +++ b/tests/test_base_handler.py @@ -1,9 +1,14 @@ +import json import uuid from collections.abc import Generator from unittest.mock import Mock, patch import pytest +from opentelemetry import context as otel_context +from opentelemetry import propagate +from opentelemetry.context import Context from opentelemetry.sdk.trace import ReadableSpan +from opentelemetry.trace.status import StatusCode from splunk_ao import get_tracing_headers from splunk_ao.handlers.base_handler import SplunkAOBaseHandler @@ -95,6 +100,126 @@ def test_incremental_handler_root_is_active_beneath_caller_owned_operation() -> logger.terminate() +BASELINE_HANDLER_TOPOLOGY = { + "invoke_workflow handler-root": ("invoke_workflow", None), + "execute_tool search": ("execute_tool", "invoke_workflow handler-root"), + "chat model": ("chat", "execute_tool search"), +} + + +def _handler_topology(carrier: dict[str, str] | None = None) -> dict[str, tuple[str | None, str | None]]: + sink = RecordingSink() + logger = SplunkAOLogger(project_id="project-id", agent_stream_id="stream-id", _sink=sink) + handler = SplunkAOBaseHandler(splunk_ao_logger=logger, flush_on_chain_end=False) + root_id = uuid.uuid4() + tool_id = uuid.uuid4() + llm_id = uuid.uuid4() + token = otel_context.attach(propagate.extract(carrier) if carrier is not None else Context()) + try: + handler.start_node(node_type="chain", parent_run_id=None, run_id=root_id, name="handler-root", input="request") + handler.start_node(node_type="tool", parent_run_id=root_id, run_id=tool_id, name="search", input="query") + handler.start_node( + node_type="llm", parent_run_id=tool_id, run_id=llm_id, name="model-call", input="prompt", model="model" + ) + handler.end_node(llm_id, output="answer", model="model") + handler.end_node(tool_id, output="result") + handler.end_node(root_id, output="done") + + names_by_context = {(span.context.trace_id, span.context.span_id): span.name for span in sink.spans} + return { + span.name: ( + (span.attributes or {}).get("gen_ai.operation.name"), + names_by_context.get((span.parent.trace_id, span.parent.span_id)) if span.parent is not None else None, + ) + for span in sink.spans + } + finally: + logger.terminate() + otel_context.detach(token) + + +def test_distributed_and_non_distributed_handlers_preserve_baseline_local_topology() -> None: + remote_carrier = { + "traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01", + "tracestate": "vendor=value", + } + + assert _handler_topology() == BASELINE_HANDLER_TOPOLOGY + assert _handler_topology(remote_carrier) == BASELINE_HANDLER_TOPOLOGY + + +def test_leaf_only_handler_root_is_active_exportable_and_enqueued_without_flush() -> None: + sink = RecordingSink() + logger = SplunkAOLogger(project_id="project-id", agent_stream_id="stream-id", _sink=sink) + handler = SplunkAOBaseHandler(splunk_ao_logger=logger, flush_on_chain_end=False) + root_id = uuid.uuid4() + try: + handler.start_node( + node_type="llm", parent_run_id=None, run_id=root_id, name="leaf-root", input="prompt", model="model" + ) + active_span_id = get_tracing_headers()["traceparent"].split("-")[2] + + handler.end_node(root_id, output="answer", model="model") + + [span] = sink.spans + assert span.name == "chat model" + assert (span.attributes or {}).get("gen_ai.operation.name") == "chat" + assert format(span.context.span_id, "016x") == active_span_id + assert span.parent is None + assert sink.force_flush_calls == 0 + finally: + logger.terminate() + + +def test_unsampled_remote_parent_suppresses_incremental_handler_export() -> None: + remote_context = propagate.extract({"traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-00"}) + token = otel_context.attach(remote_context) + sink = RecordingSink() + logger = SplunkAOLogger(project_id="project-id", agent_stream_id="stream-id", _sink=sink) + handler = SplunkAOBaseHandler(splunk_ao_logger=logger, flush_on_chain_end=False) + root_id = uuid.uuid4() + child_id = uuid.uuid4() + try: + handler.start_node(node_type="chain", parent_run_id=None, run_id=root_id, name="root", input="request") + handler.start_node( + node_type="llm", parent_run_id=root_id, run_id=child_id, name="child", input="prompt", model="model" + ) + handler.end_node(child_id, output="answer", model="model") + handler.end_node(root_id, output="done") + + assert sink.spans == [] + assert sink.force_flush_calls == 0 + finally: + logger.terminate() + otel_context.detach(token) + + +def test_handler_error_span_enqueues_with_error_status_before_root_end() -> None: + sink = RecordingSink() + logger = SplunkAOLogger(project_id="project-id", agent_stream_id="stream-id", _sink=sink) + handler = SplunkAOBaseHandler(splunk_ao_logger=logger, flush_on_chain_end=False) + root_id = uuid.uuid4() + child_id = uuid.uuid4() + try: + handler.start_node(node_type="chain", parent_run_id=None, run_id=root_id, name="root", input="request") + handler.start_node( + node_type="tool", parent_run_id=root_id, run_id=child_id, name="failing-tool", input="arguments" + ) + + handler.end_node(child_id, error="tool failed") + + [span] = sink.spans + assert span.name == "execute_tool failing-tool" + assert span.status.status_code is StatusCode.ERROR + assert json.loads(str((span.attributes or {}).get("gen_ai.tool.call.result"))) == {"value": "tool failed"} + assert str(root_id) in handler._active_steps + assert sink.force_flush_calls == 0 + + handler.end_node(root_id, output="failed", status_code=500) + finally: + logger.terminate() + + class TestSplunkAOBaseHandler: @pytest.fixture @patch("splunk_ao.logger.logger.AgentStreams") diff --git a/tests/test_crewai_handler.py b/tests/test_crewai_handler.py index 61cd33e5..96d36e10 100644 --- a/tests/test_crewai_handler.py +++ b/tests/test_crewai_handler.py @@ -5,11 +5,13 @@ from unittest.mock import Mock, patch import pytest +from opentelemetry.sdk.trace import ReadableSpan # Skip all tests in this module on Python 3.14+ (crewai doesn't support it yet) pytestmark = pytest.mark.skipif(sys.version_info >= (3, 14), reason="crewai does not support Python 3.14+") from splunk_ao.handlers.crewai.handler import CrewAIEventListener # noqa: E402 +from splunk_ao.logger.logger import SplunkAOLogger # noqa: E402 from splunk_ao.schema.handlers import NodeType # noqa: E402 from tests.testutils.setup import ( # noqa: E402 setup_mock_logstreams_client, @@ -76,6 +78,65 @@ def __init__(self, raw="Test output"): self.raw = raw +class RecordingSink: + def __init__(self) -> None: + self.spans: list[ReadableSpan] = [] + self.force_flush_calls = 0 + + def emit(self, span: ReadableSpan) -> None: + self.spans.append(span) + + def force_flush(self) -> bool: + self.force_flush_calls += 1 + return True + + def shutdown(self) -> None: + return None + + +def test_crewai_adapter_enqueues_task_at_callback_end_without_flush() -> None: + sink = RecordingSink() + logger = SplunkAOLogger(project_id="project-id", agent_stream_id="stream-id", _sink=sink) + crew_id = uuid.uuid4() + task_id = uuid.uuid4() + crew = MockCrew(crew_id=crew_id) + agent = MockAgent(crew=crew) + task = MockTask(task_id=task_id, description="Research market trends", agent=agent) + try: + with ( + patch("splunk_ao.handlers.crewai.handler._crewai_imports_resolved", True), + patch("splunk_ao.handlers.crewai.handler.CREWAI_AVAILABLE", False), + patch("splunk_ao.handlers.crewai.handler.LITE_LLM_AVAILABLE", False), + ): + callback = CrewAIEventListener(splunk_ao_logger=logger, start_new_trace=True, flush_on_crew_completed=False) + + callback._handle_crew_kickoff_started( + MockSource(id=crew_id), MockEvent(crew_name="Test Crew", inputs={"question": "market trends"}) + ) + callback._handle_task_started(MockSource(id=task_id), MockEvent(task=task)) + + callback._handle_task_completed(MockSource(id=task_id), MockEvent(output=MockOutput("Done"))) + + assert len(sink.spans) == 1 + assert sink.spans[0].name == "invoke_workflow Research market trends" + assert str(crew_id) in callback._handler._active_steps + assert sink.force_flush_calls == 0 + + callback._handle_crew_kickoff_completed( + MockSource(id=crew_id), MockEvent(output=MockOutput("Crew completed successfully")) + ) + + child, root = sink.spans + assert child.parent == root.context + assert [span.name for span in sink.spans] == [ + "invoke_workflow Research market trends", + "invoke_workflow Test Crew", + ] + assert sink.force_flush_calls == 0 + finally: + logger.terminate() + + @pytest.fixture def mock_splunk_ao_logger(): """Creates a mock Galileo logger for testing.""" diff --git a/tests/test_http_instrumentation.py b/tests/test_http_instrumentation.py index dd186b2b..1c176328 100644 --- a/tests/test_http_instrumentation.py +++ b/tests/test_http_instrumentation.py @@ -5,21 +5,33 @@ from typing import Any from unittest.mock import MagicMock, patch +import aiohttp +import httpx import pytest import requests +from aiohttp import TraceRequestExceptionParams, TraceRequestStartParams from fastapi import FastAPI +from multidict import CIMultiDict from opentelemetry import baggage, context, propagate, trace from opentelemetry.context import Context +from opentelemetry.instrumentation.aiohttp_client import AioHttpClientInstrumentor from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor +from opentelemetry.instrumentation.httpx import HTTPXClientInstrumentor from opentelemetry.instrumentation.requests import RequestsInstrumentor +from opentelemetry.instrumentation.starlette import StarletteInstrumentor from opentelemetry.propagators import textmap -from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace import ReadableSpan, TracerProvider from opentelemetry.trace import NonRecordingSpan, SpanContext, TraceFlags from starlette.applications import Starlette +from starlette.requests import Request +from starlette.responses import JSONResponse +from starlette.routing import Route from starlette.testclient import TestClient +from yarl import URL from splunk_ao import instrument_distributed_tracing from splunk_ao.http_instrumentation import _client_provider_ids, _instrumented_apps, _load_instrumentors +from splunk_ao.logger.logger import SplunkAOLogger from splunk_ao.session_context import GEN_AI_CONVERSATION_ID, SplunkAOSessionPropagator, _session_id_context @@ -71,6 +83,20 @@ def instrument_app(self, app: Any, **kwargs: Any) -> None: self.app_calls.append((app, kwargs)) +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: + return None + + def instrumentor_type(recorder: RecordingInstrumentor) -> type: class Instrumentor: def __new__(cls) -> RecordingInstrumentor: @@ -323,3 +349,208 @@ def automatic_context() -> dict[str, bool]: finally: FastAPIInstrumentor.uninstrument_app(app) provider.shutdown() + + +def test_automatic_requests_to_fastapi_continues_active_path1_context_without_manual_headers( + monkeypatch: pytest.MonkeyPatch, +) -> None: + captured_headers: dict[str, str] = {} + observed: dict[str, Any] = {} + + def capture_send(session: requests.Session, request: requests.PreparedRequest, **kwargs: Any) -> requests.Response: + del session, kwargs + captured_headers.update(request.headers) + response = requests.Response() + response.status_code = 200 + response.request = request + response.url = request.url + return response + + monkeypatch.setattr(requests.Session, "send", capture_send) + provider = TracerProvider() + app = FastAPI() + logger = SplunkAOLogger(project_id="project-id", agent_stream_id="stream-id", _sink=RecordingSink()) + + @app.get("/downstream") + def downstream() -> dict[str, bool]: + observed["trace_id"] = trace.get_current_span().get_span_context().trace_id + observed["conversation_id"] = baggage.get_baggage(GEN_AI_CONVERSATION_ID) + return {"ok": True} + + try: + instrument_distributed_tracing( + tracer_provider=provider, app=app, instrument_httpx=False, instrument_aiohttp_client=False + ) + logger.set_session("conversation-cross-service") + logger.start_trace(input="request") + path1_operation = logger.add_workflow_span(input="agent work", name="agent") + path1_context = logger._otel_ids[path1_operation.id].span_context + + response = requests.get("https://example.test/downstream", timeout=1) + downstream_token = context.attach(Context()) + try: + with TestClient(app) as client: + downstream_response = client.get( + "/downstream", + headers={"traceparent": captured_headers["traceparent"], "baggage": captured_headers["baggage"]}, + ) + finally: + context.detach(downstream_token) + + assert response.status_code == 200 + assert downstream_response.json() == {"ok": True} + assert captured_headers["traceparent"].split("-")[1] == format(path1_context.trace_id, "032x") + assert captured_headers["baggage"] == "gen_ai.conversation.id=conversation-cross-service" + assert observed == {"trace_id": path1_context.trace_id, "conversation_id": "conversation-cross-service"} + finally: + logger.clear_session() + if logger.current_parent() is not None: + logger.conclude(output="done", conclude_all=True) + logger.terminate() + RequestsInstrumentor().uninstrument() + FastAPIInstrumentor.uninstrument_app(app) + provider.shutdown() + + +def test_real_httpx_sync_instrumentation_injects_active_path1_context(monkeypatch: pytest.MonkeyPatch) -> None: + captured_headers: dict[str, str] = {} + + def capture_sync(transport: httpx.HTTPTransport, request: httpx.Request) -> httpx.Response: + del transport + captured_headers.update(request.headers) + return httpx.Response(200, request=request) + + monkeypatch.setattr(httpx.HTTPTransport, "handle_request", capture_sync) + provider = TracerProvider() + logger = SplunkAOLogger(project_id="project-id", agent_stream_id="stream-id", _sink=RecordingSink()) + try: + instrument_distributed_tracing( + tracer_provider=provider, instrument_requests=False, instrument_aiohttp_client=False + ) + logger.set_session("conversation-httpx") + logger.start_trace(input="request") + operation = logger.add_workflow_span(input="agent work", name="agent") + operation_context = logger._otel_ids[operation.id].span_context + + with httpx.Client() as client: + assert client.get("https://example.test/sync").status_code == 200 + + assert captured_headers["traceparent"].split("-")[1] == format(operation_context.trace_id, "032x") + assert captured_headers["baggage"] == "gen_ai.conversation.id=conversation-httpx" + finally: + logger.clear_session() + if logger.current_parent() is not None: + logger.conclude(output="done", conclude_all=True) + logger.terminate() + HTTPXClientInstrumentor().uninstrument() + provider.shutdown() + + +@pytest.mark.asyncio +async def test_real_httpx_async_instrumentation_injects_active_path1_context(monkeypatch: pytest.MonkeyPatch) -> None: + captured_headers: dict[str, str] = {} + + async def capture_async(transport: httpx.AsyncHTTPTransport, request: httpx.Request) -> httpx.Response: + del transport + captured_headers.update(request.headers) + return httpx.Response(200, request=request) + + monkeypatch.setattr(httpx.AsyncHTTPTransport, "handle_async_request", capture_async) + provider = TracerProvider() + logger = SplunkAOLogger(project_id="project-id", agent_stream_id="stream-id", _sink=RecordingSink()) + try: + instrument_distributed_tracing( + tracer_provider=provider, instrument_requests=False, instrument_aiohttp_client=False + ) + logger.set_session("conversation-httpx") + logger.start_trace(input="request") + operation = logger.add_workflow_span(input="agent work", name="agent") + operation_context = logger._otel_ids[operation.id].span_context + + async with httpx.AsyncClient() as client: + assert (await client.get("https://example.test/async")).status_code == 200 + + assert captured_headers["traceparent"].split("-")[1] == format(operation_context.trace_id, "032x") + assert captured_headers["baggage"] == "gen_ai.conversation.id=conversation-httpx" + finally: + logger.clear_session() + if logger.current_parent() is not None: + logger.conclude(output="done", conclude_all=True) + logger.terminate() + HTTPXClientInstrumentor().uninstrument() + provider.shutdown() + + +@pytest.mark.asyncio +async def test_real_aiohttp_instrumentation_installs_hook_that_injects_active_path1_context() -> None: + provider = TracerProvider() + logger = SplunkAOLogger(project_id="project-id", agent_stream_id="stream-id", _sink=RecordingSink()) + try: + instrument_distributed_tracing(tracer_provider=provider, instrument_requests=False, instrument_httpx=False) + logger.set_session("conversation-aiohttp") + logger.start_trace(input="request") + operation = logger.add_workflow_span(input="agent work", name="agent") + operation_context = logger._otel_ids[operation.id].span_context + + async with aiohttp.ClientSession() as session: + trace_config = next( + config + for config in session._trace_configs + if getattr(config, "_is_instrumented_by_opentelemetry", False) + ) + trace_context = trace_config.trace_config_ctx() + headers: CIMultiDict[str] = CIMultiDict() + url = URL("https://example.test/aiohttp") + start_params = TraceRequestStartParams("GET", url, headers) + for callback in trace_config.on_request_start: + await callback(session, trace_context, start_params) + + assert headers["traceparent"].split("-")[1] == format(operation_context.trace_id, "032x") + assert headers["baggage"] == "gen_ai.conversation.id=conversation-aiohttp" + + exception_params = TraceRequestExceptionParams("GET", url, headers, RuntimeError("synthetic stop")) + for callback in trace_config.on_request_exception: + await callback(session, trace_context, exception_params) + finally: + logger.clear_session() + if logger.current_parent() is not None: + logger.conclude(output="done", conclude_all=True) + logger.terminate() + AioHttpClientInstrumentor().uninstrument() + provider.shutdown() + + +def test_real_starlette_instrumentor_extracts_remote_context_without_sdk_middleware() -> None: + expected_trace_id = 0x4BF92F3577B34DA6A3CE929D0E0E4736 + observed: dict[str, Any] = {} + provider = TracerProvider() + + async def automatic_context(request: Request) -> JSONResponse: + del request + observed["trace_id"] = trace.get_current_span().get_span_context().trace_id + observed["conversation_id"] = baggage.get_baggage(GEN_AI_CONVERSATION_ID) + return JSONResponse({"ok": True}) + + app = Starlette(routes=[Route("/automatic-context", automatic_context)]) + try: + instrument_distributed_tracing( + tracer_provider=provider, + app=app, + instrument_requests=False, + instrument_httpx=False, + instrument_aiohttp_client=False, + ) + with TestClient(app) as client: + response = client.get( + "/automatic-context", + headers={ + "traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01", + "baggage": "gen_ai.conversation.id=conversation-starlette", + }, + ) + + assert response.json() == {"ok": True} + assert observed == {"trace_id": expected_trace_id, "conversation_id": "conversation-starlette"} + finally: + StarletteInstrumentor.uninstrument_app(app) + provider.shutdown() diff --git a/tests/test_langchain.py b/tests/test_langchain.py index ca8b6144..f2bb85e9 100644 --- a/tests/test_langchain.py +++ b/tests/test_langchain.py @@ -9,6 +9,7 @@ from langchain_core.documents import Document from langchain_core.messages import AIMessage, HumanMessage, SystemMessage, ToolMessage from langchain_core.outputs import ChatGeneration, LLMResult +from opentelemetry.sdk.trace import ReadableSpan from galileo_core.schemas.shared.document import Document as GalileoDocument from splunk_ao import Message, MessageRole, splunk_ao_context @@ -22,6 +23,51 @@ from tests.testutils.setup import setup_mock_logstreams_client, setup_mock_projects_client, setup_mock_traces_client +class RecordingSink: + def __init__(self) -> None: + self.spans: list[ReadableSpan] = [] + self.force_flush_calls = 0 + + def emit(self, span: ReadableSpan) -> None: + self.spans.append(span) + + def force_flush(self) -> bool: + self.force_flush_calls += 1 + return True + + def shutdown(self) -> None: + return None + + +def test_langchain_adapter_enqueues_child_at_callback_end_without_flush() -> None: + 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 = uuid.uuid4() + child_id = uuid.uuid4() + try: + callback.on_chain_start(serialized={"name": "root"}, inputs={"query": "question"}, run_id=root_id) + callback.on_tool_start(serialized={"name": "search"}, input_str="query", run_id=child_id, parent_run_id=root_id) + + callback.on_tool_end(output="result", run_id=child_id, parent_run_id=root_id) + + assert [(span.attributes or {}).get("gen_ai.operation.name") for span in sink.spans] == ["execute_tool"] + assert str(root_id) in callback._handler._active_steps + assert sink.force_flush_calls == 0 + + callback.on_chain_end(outputs={"answer": "done"}, run_id=root_id) + + child, root = sink.spans + assert child.parent == root.context + assert [(span.attributes or {}).get("gen_ai.operation.name") for span in sink.spans] == [ + "execute_tool", + "invoke_workflow", + ] + assert sink.force_flush_calls == 0 + finally: + logger.terminate() + + class TestSplunkAOCallback: @pytest.fixture @patch("splunk_ao.logger.logger.AgentStreams") diff --git a/tests/test_otel_native_paths.py b/tests/test_otel_native_paths.py index 68fe5c9c..bb493f03 100644 --- a/tests/test_otel_native_paths.py +++ b/tests/test_otel_native_paths.py @@ -534,6 +534,49 @@ def test_caller_owned_otel_and_path1_logger_interoperate_bidirectionally() -> No assert by_name["otel-downstream"].parent == operation_context +def test_explicit_conversation_session_is_applied_across_paths_1_2_and_3() -> None: + exporter = RecordingExporter() + processor = SplunkAOSpanProcessor(SpanProcessor=RecordingSpanProcessor, _exporter=exporter) + provider = SDKTracerProvider() + provider.add_span_processor(processor) + provider_token = _TRACE_PROVIDER_CONTEXT_VAR.set(provider) + tracer = provider.get_tracer("caller-owned") + sink = MagicMock() + sink.force_flush.return_value = True + logger = SplunkAOLogger(project_id="project-id", agent_stream_id="stream-id", _sink=sink) + try: + logger.set_session("conversation-all-paths") + logger.start_trace(input="request") + path1_operation = logger.add_workflow_span(input="path 1", name="path-1") + + with start_splunk_ao_span(WorkflowSpan(input="path 2", name="path-2")): + pass + with tracer.start_as_current_span("path-3"): + pass + + logger.conclude(output="done") + logger.conclude(output="complete") + + [path1_span] = [call.args[0] for call in sink.emit.call_args_list] + assert path1_span.name == "invoke_workflow path-1" + assert path1_span.attributes["gen_ai.conversation.id"] == "conversation-all-paths" + assert "splunk_ao.session.id" not in path1_span.attributes + + recording_processor = processor.processor + assert isinstance(recording_processor, RecordingSpanProcessor) + started = {span.name: span for span, _ in recording_processor.started} + assert started["path-2"].attributes["gen_ai.conversation.id"] == "conversation-all-paths" + assert started["path-3"].attributes["gen_ai.conversation.id"] == "conversation-all-paths" + assert "splunk_ao.session.id" not in started["path-2"].attributes + assert "splunk_ao.session.id" not in started["path-3"].attributes + assert path1_operation.id not in logger._otel_ids + finally: + logger.clear_session() + logger.terminate() + _TRACE_PROVIDER_CONTEXT_VAR.reset(provider_token) + provider.shutdown() + + def test_sdk_native_and_caller_owned_spans_keep_topology_and_queue_without_flush() -> None: exporter = RecordingExporter() processor = SplunkAOSpanProcessor(_exporter=exporter) From be84378c9cc9633ce4674572a49aeca67b85d9d7 Mon Sep 17 00:00:00 2001 From: pradystar Date: Thu, 13 Aug 2026 14:02:32 -0700 Subject: [PATCH 04/13] feat(tracing): add unified automatic distributed tracing setup --- AGENTS.md | 3 + ARCHITECTURE.md | 14 ++- CHANGELOG.md | 8 +- README.md | 28 +++--- poetry.lock | 36 ++++---- pyproject.toml | 6 +- src/splunk_ao/__init__.py | 3 +- src/splunk_ao/http_instrumentation.py | 118 ++++++++++++++++++++++---- tests/test_http_instrumentation.py | 83 +++++++++++++++++- 9 files changed, 240 insertions(+), 59 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 7a6f7509..2e3831ea 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -134,6 +134,9 @@ CI supports Python 3.11–3.14; root CI also spans Linux, macOS, and Windows. - Handler/decorator/OpenAI/ADK telemetry uses the internal logged-step path and converts completed steps to immutable OTel spans. The internal trace envelope is never exported as a span. - `start_splunk_ao_span()` is SDK-native OTel. `add_splunk_ao_span_processor()` and A2A instrument caller-owned OTel. + `configure_distributed_tracing()` is the combined automatic HTTP setup: it creates or accepts an application-owned + provider, registers Splunk AO export once per provider, configures supported upstream instrumentors, and returns the + provider for application shutdown. `instrument_distributed_tracing()` remains transport-only. - Never replace the process-global tracer provider. Register processors on the provided provider; respect ownership. - Treat ended `ReadableSpan` objects as immutable. Normalize by copying at export, never by mutating private fields. - Completed spans enqueue immediately. `flush()` drains completed work without ending active work; `terminate()` drains, diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index 714ddbe1..bfc39502 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -89,10 +89,16 @@ provider and calls `shutdown()`; SDK-owned logger/export resources use their own ### Distributed context and HTTP transports Supported automatic propagation reuses upstream OpenTelemetry instrumentors rather than custom SDK HTTP wrappers. -`instrument_distributed_tracing()` configures FastAPI/Starlette inbound instrumentation and Requests, HTTPX sync/async, -and aiohttp-client outbound instrumentation with a caller-owned provider. It never sets or replaces the global tracer -provider. Manual `get_tracing_headers()` and `TracingMiddleware` remain fallbacks for unsupported transports; supported -automatic instrumentation does not require either one on individual requests. +`configure_distributed_tracing()` creates or accepts an OpenTelemetry SDK provider, registers Splunk AO export, and +configures FastAPI/Starlette inbound instrumentation plus Requests, HTTPX sync/async, and aiohttp-client outbound +instrumentation. It returns the application-owned provider for shutdown and never sets or replaces the global tracer +provider. `instrument_distributed_tracing()` remains the transport-only entry point for applications that manage +processor registration separately. Manual `get_tracing_headers()` and `TracingMiddleware` remain fallbacks for +unsupported transports; supported automatic instrumentation does not require either one on individual requests. + +High-level setup is idempotent per provider for Splunk AO processor registration. Each process-wide client instrumentor +and each server application must still have one instrumentation owner; do not combine the helper with direct upstream +instrumentation of the same component. An explicit SDK session propagates as the standard W3C baggage member `gen_ai.conversation.id`. Export normalization may derive the local compatibility attribute `splunk_ao.session.id`, but that attribute is not propagated as a second diff --git a/CHANGELOG.md b/CHANGELOG.md index 9bfae46a..4c37977b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,9 +9,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added -- Added the `distributed-tracing` extra and - `instrument_distributed_tracing()` startup helper for supported upstream - FastAPI/Starlette, Requests, HTTPX, and aiohttp-client instrumentation. +- Added the `distributed-tracing` extra and high-level + `configure_distributed_tracing()` setup for Splunk AO export plus supported + upstream FastAPI/Starlette, Requests, HTTPX, and aiohttp-client instrumentation. + The lower-level `instrument_distributed_tracing()` transport helper remains + available for applications that configure their provider separately. - Explicit SDK sessions now propagate across supported services as standard W3C `gen_ai.conversation.id` baggage. SDK routing, authentication, and application identity are not propagated. diff --git a/README.md b/README.md index 17210603..9be9538b 100644 --- a/README.md +++ b/README.md @@ -370,26 +370,28 @@ elsewhere owns a separate provider and should be terminated separately. The `distributed-tracing` extra provides a supported one-time setup for upstream OpenTelemetry FastAPI/Starlette, Requests, HTTPX sync/async, and -aiohttp-client instrumentation: +aiohttp-client instrumentation. The high-level helper creates an application-owned +provider, registers Splunk AO export, and configures those transports together: ```python from fastapi import FastAPI -from opentelemetry.sdk.trace import TracerProvider - -from splunk_ao import instrument_distributed_tracing -from splunk_ao.otel import add_splunk_ao_span_processor - -provider = TracerProvider() -add_splunk_ao_span_processor(provider) +from splunk_ao import configure_distributed_tracing app = FastAPI() -instrument_distributed_tracing(tracer_provider=provider, app=app) +provider = configure_distributed_tracing(app=app) ``` -The application owns `provider` and calls `provider.shutdown()` during process -teardown. The helper never replaces the process-global provider. Configure -other OpenTelemetry/OpenInference agent or model instrumentations with the -same provider when they should join the trace. +The application owns the returned `provider` and calls `provider.shutdown()` +during process teardown. The helper never replaces the process-global provider. +Pass an existing OpenTelemetry SDK `TracerProvider` through the +`tracer_provider` argument when other OpenTelemetry/OpenInference agent or model +instrumentations should use that same provider and join the trace. + +For advanced composition, `add_splunk_ao_span_processor()` remains the +export-only API and `instrument_distributed_tracing()` remains the supported +transport-only API. Do not also directly instrument a client or application +component that is owned by the high-level helper; each component should have one +instrumentation owner. After this startup call, supported inbound requests extract W3C context and supported outbound clients inject it automatically. Application code does not diff --git a/poetry.lock b/poetry.lock index 4cfb87f4..68a4dc3a 100644 --- a/poetry.lock +++ b/poetry.lock @@ -7,7 +7,7 @@ description = "Happy Eyeballs for asyncio" optional = true python-versions = ">=3.9" groups = ["main"] -markers = "python_version <= \"3.13\" and (extra == \"crewai\" or extra == \"all\")" +markers = "python_version <= \"3.13\" and (extra == \"crewai\" or extra == \"all\" or extra == \"distributed-tracing\") or python_version < \"3.14\" and (extra == \"distributed-tracing\" or extra == \"all\" or extra == \"crewai\") or extra == \"distributed-tracing\" or extra == \"all\"" files = [ {file = "aiohappyeyeballs-2.6.1-py3-none-any.whl", hash = "sha256:f349ba8f4b75cb25c99c5c2d84e997e485204d2902a9597802b0371f09331fb8"}, {file = "aiohappyeyeballs-2.6.1.tar.gz", hash = "sha256:c3f9d0113123803ccadfdf3f0faa505bc78e6a72d1cc4806cbd719826e943558"}, @@ -20,7 +20,7 @@ description = "Async http client/server framework (asyncio)" optional = true python-versions = ">=3.10" groups = ["main"] -markers = "python_version <= \"3.13\" and (extra == \"crewai\" or extra == \"all\")" +markers = "python_version <= \"3.13\" and (extra == \"crewai\" or extra == \"all\" or extra == \"distributed-tracing\") or python_version < \"3.14\" and (extra == \"distributed-tracing\" or extra == \"all\" or extra == \"crewai\") or extra == \"distributed-tracing\" or extra == \"all\"" files = [ {file = "aiohttp-3.14.3-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:eb0495d778817619273c108784292be161a924b9f5ae5cbbc70a2caa6838250b"}, {file = "aiohttp-3.14.3-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:c3c200cf9757edd785051dc699c7ecbec22110dbfcb3fefc7a9f9695eda8ea7a"}, @@ -163,7 +163,7 @@ description = "aiosignal: a list of registered asynchronous callbacks" optional = true python-versions = ">=3.9" groups = ["main"] -markers = "python_version <= \"3.13\" and (extra == \"crewai\" or extra == \"all\")" +markers = "python_version <= \"3.13\" and (extra == \"crewai\" or extra == \"all\" or extra == \"distributed-tracing\") or python_version < \"3.14\" and (extra == \"distributed-tracing\" or extra == \"all\" or extra == \"crewai\") or extra == \"distributed-tracing\" or extra == \"all\"" files = [ {file = "aiosignal-1.4.0-py3-none-any.whl", hash = "sha256:053243f8b92b990551949e63930a839ff0cf0b0ebbe0597b0f3fb19e1a0fe82e"}, {file = "aiosignal-1.4.0.tar.gz", hash = "sha256:f47eecd9468083c2029cc99945502cb7708b082c232f9aca65da147157b251c7"}, @@ -699,7 +699,7 @@ files = [ {file = "colorama-0.4.6-py2.py3-none-any.whl", hash = "sha256:4f1d9991f5acc0ca119f9d443620b77f9d6b33703e51011c16baf57afb285fc6"}, {file = "colorama-0.4.6.tar.gz", hash = "sha256:08695f5cb7ed6e0531a20572697297273c47b8cae5a63ffc6d6ed5c201be6e44"}, ] -markers = {main = "(platform_system == \"Windows\" or python_version <= \"3.13\") and (extra == \"crewai\" or extra == \"all\" or platform_system == \"Windows\") and (os_name == \"nt\" or platform_system == \"Windows\" or sys_platform == \"win32\")", dev = "platform_system == \"Windows\" or sys_platform == \"win32\"", test = "sys_platform == \"win32\""} +markers = {main = "(python_version <= \"3.13\" or platform_system == \"Windows\") and (extra == \"crewai\" or extra == \"all\" or platform_system == \"Windows\") and (os_name == \"nt\" or platform_system == \"Windows\" or sys_platform == \"win32\")", dev = "platform_system == \"Windows\" or sys_platform == \"win32\"", test = "sys_platform == \"win32\""} [[package]] name = "coloredlogs" @@ -965,7 +965,7 @@ files = [ {file = "distro-1.9.0-py3-none-any.whl", hash = "sha256:7bffd925d65168f85027d8da9af6bddab658135b840670a223589bc0c8ef02b2"}, {file = "distro-1.9.0.tar.gz", hash = "sha256:2fa77c6fd8940f116ee1d6b94a2f90b13b5ea8d019b98bc8bafdcabcdd9bdbed"}, ] -markers = {main = "extra == \"openai\" or extra == \"all\" or extra == \"langchain\" or python_version <= \"3.13\" and (extra == \"crewai\" or extra == \"all\" or extra == \"openai\" or extra == \"langchain\")"} +markers = {main = "python_version <= \"3.13\" and (extra == \"crewai\" or extra == \"all\" or extra == \"openai\" or extra == \"langchain\") or extra == \"openai\" or extra == \"all\" or extra == \"langchain\""} [[package]] name = "docstring-parser" @@ -1171,7 +1171,7 @@ description = "A list-like structure which implements collections.abc.MutableSeq optional = true python-versions = ">=3.9" groups = ["main"] -markers = "python_version <= \"3.13\" and (extra == \"crewai\" or extra == \"all\")" +markers = "python_version <= \"3.13\" and (extra == \"crewai\" or extra == \"all\" or extra == \"distributed-tracing\") or python_version < \"3.14\" and (extra == \"distributed-tracing\" or extra == \"all\" or extra == \"crewai\") or extra == \"distributed-tracing\" or extra == \"all\"" files = [ {file = "frozenlist-1.7.0-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:cc4df77d638aa2ed703b878dd093725b72a824c3c546c076e8fdf276f78ee84a"}, {file = "frozenlist-1.7.0-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:716a9973a2cc963160394f701964fe25012600f3d311f60c790400b00e568b61"}, @@ -1935,7 +1935,7 @@ files = [ {file = "jiter-0.10.0-cp39-cp39-win_amd64.whl", hash = "sha256:1b28302349dc65703a9e4ead16f163b1c339efffbe1049c30a44b001a2a4fff9"}, {file = "jiter-0.10.0.tar.gz", hash = "sha256:07a7142c38aacc85194391108dc91b5b57093c978a9932bd86a36862759d9500"}, ] -markers = {main = "extra == \"openai\" or extra == \"all\" or python_version <= \"3.13\" and (extra == \"crewai\" or extra == \"all\" or extra == \"openai\")"} +markers = {main = "python_version <= \"3.13\" and (extra == \"crewai\" or extra == \"all\" or extra == \"openai\") or extra == \"openai\" or extra == \"all\""} [[package]] name = "json-repair" @@ -2631,7 +2631,7 @@ description = "multidict implementation" optional = true python-versions = ">=3.9" groups = ["main"] -markers = "python_version <= \"3.13\" and (extra == \"crewai\" or extra == \"all\")" +markers = "python_version <= \"3.13\" and (extra == \"crewai\" or extra == \"all\" or extra == \"distributed-tracing\") or python_version < \"3.14\" and (extra == \"distributed-tracing\" or extra == \"all\" or extra == \"crewai\") or extra == \"distributed-tracing\" or extra == \"all\"" files = [ {file = "multidict-6.6.4-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:b8aa6f0bd8125ddd04a6593437bad6a7e70f300ff4180a531654aa2ab3f6d58f"}, {file = "multidict-6.6.4-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:b9e5853bbd7264baca42ffc53391b490d65fe62849bf2c690fa3f6273dbcd0cb"}, @@ -2961,7 +2961,7 @@ files = [ {file = "openai-2.32.0-py3-none-any.whl", hash = "sha256:4dcc9badeb4bf54ad0d187453742f290226d30150890b7890711bda4f32f192f"}, {file = "openai-2.32.0.tar.gz", hash = "sha256:c54b27a9e4cb8d51f0dd94972ffd1a04437efeb259a9e60d8922b8bd26fe55e0"}, ] -markers = {main = "extra == \"openai\" or extra == \"all\" or python_version <= \"3.13\" and (extra == \"crewai\" or extra == \"all\" or extra == \"openai\")"} +markers = {main = "python_version <= \"3.13\" and (extra == \"crewai\" or extra == \"all\" or extra == \"openai\") or extra == \"openai\" or extra == \"all\""} [package.dependencies] anyio = ">=3.5.0,<5" @@ -3440,7 +3440,7 @@ files = [ {file = "orjson-3.11.9-cp314-cp314-win_arm64.whl", hash = "sha256:eebdbdeef0094e4f5aefa20dcd4eb2368ab5e7a3b4edea27f1e7b2892e009cf9"}, {file = "orjson-3.11.9.tar.gz", hash = "sha256:4fef17e1f8722c11587a6ef18e35902450221da0028e65dbaaa543619e68e48f"}, ] -markers = {main = "extra == \"langchain\" or extra == \"all\" or python_version <= \"3.13\" and (extra == \"langchain\" or extra == \"all\" or extra == \"crewai\")", test = "platform_python_implementation != \"PyPy\""} +markers = {main = "python_version <= \"3.13\" and (extra == \"crewai\" or extra == \"all\" or extra == \"langchain\") or extra == \"langchain\" or extra == \"all\"", test = "platform_python_implementation != \"PyPy\""} [[package]] name = "ormsgpack" @@ -3525,7 +3525,7 @@ files = [ {file = "packaging-24.2-py3-none-any.whl", hash = "sha256:09abb1bccd265c01f4a3aa3f7a7db064b36514d2cba19a2f694fe6150451a759"}, {file = "packaging-24.2.tar.gz", hash = "sha256:c228a6dc5e932d346bc5739379109d49e8853dd8223571c7c5b55260edc0b97f"}, ] -markers = {main = "extra == \"langchain\" or extra == \"all\" or extra == \"openai\" or extra == \"distributed-tracing\" or python_version <= \"3.13\" and (extra == \"crewai\" or extra == \"all\" or extra == \"langchain\" or extra == \"openai\" or extra == \"distributed-tracing\")"} +markers = {main = "python_version <= \"3.13\" and (extra == \"crewai\" or extra == \"all\" or extra == \"langchain\" or extra == \"openai\" or extra == \"distributed-tracing\") or extra == \"langchain\" or extra == \"all\" or extra == \"openai\" or extra == \"distributed-tracing\""} [[package]] name = "pathspec" @@ -3788,7 +3788,7 @@ description = "Accelerated property cache" optional = true python-versions = ">=3.9" groups = ["main"] -markers = "python_version <= \"3.13\" and (extra == \"crewai\" or extra == \"all\")" +markers = "python_version <= \"3.13\" and (extra == \"crewai\" or extra == \"all\" or extra == \"distributed-tracing\") or python_version < \"3.14\" and (extra == \"distributed-tracing\" or extra == \"all\" or extra == \"crewai\") or extra == \"distributed-tracing\" or extra == \"all\"" files = [ {file = "propcache-0.3.2-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:22d9962a358aedbb7a2e36187ff273adeaab9743373a272976d2e348d08c7770"}, {file = "propcache-0.3.2-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:0d0fda578d1dc3f77b6b5a5dce3b9ad69a8250a891760a548df850a5e8da87f3"}, @@ -4169,7 +4169,7 @@ files = [ {file = "pycparser-2.22-py3-none-any.whl", hash = "sha256:c3702b6d3dd8c7abc1afa565d7e63d53a1d0bd86cdc24edd75470f4de499cfcc"}, {file = "pycparser-2.22.tar.gz", hash = "sha256:491c8be9c040f5390f5bf44a5b07752bd07f56edf992381b05c701439eec10f6"}, ] -markers = {main = "(python_version <= \"3.13\" or extra == \"langchain\" or extra == \"all\" or extra == \"openai\" or extra == \"crewai\") and (platform_python_implementation != \"PyPy\" or extra == \"langchain\" or extra == \"all\") and implementation_name != \"PyPy\" and (extra == \"langchain\" or extra == \"all\" or extra == \"crewai\" or extra == \"openai\") and (platform_python_implementation == \"PyPy\" or python_version <= \"3.13\" or extra == \"openai\" or extra == \"all\" or extra == \"crewai\") and (platform_python_implementation == \"PyPy\" or extra == \"crewai\" or extra == \"all\" or extra == \"openai\")", test = "platform_python_implementation == \"PyPy\" and implementation_name != \"PyPy\""} +markers = {main = "(python_version <= \"3.13\" or extra == \"langchain\" or extra == \"all\" or extra == \"openai\" or extra == \"crewai\") and (platform_python_implementation != \"PyPy\" or extra == \"langchain\" or extra == \"all\") and implementation_name != \"PyPy\" and (extra == \"langchain\" or extra == \"all\" or extra == \"crewai\" or extra == \"openai\") and (platform_python_implementation == \"PyPy\" or extra == \"openai\" or extra == \"all\" or extra == \"crewai\")", test = "platform_python_implementation == \"PyPy\" and implementation_name != \"PyPy\""} [[package]] name = "pydantic" @@ -4761,7 +4761,7 @@ files = [ {file = "PyYAML-6.0.2-cp39-cp39-win_amd64.whl", hash = "sha256:39693e1f8320ae4f43943590b49779ffb98acb81f788220ea932a6b6c51004d8"}, {file = "pyyaml-6.0.2.tar.gz", hash = "sha256:d584d9ec91ad65861cc08d42e834324ef890a082e591037abe114850ff7bbc3e"}, ] -markers = {main = "extra == \"langchain\" or extra == \"all\" or python_version <= \"3.13\" and (extra == \"crewai\" or extra == \"all\" or extra == \"langchain\")"} +markers = {main = "python_version <= \"3.13\" and (extra == \"crewai\" or extra == \"all\" or extra == \"langchain\") or extra == \"langchain\" or extra == \"all\""} [[package]] name = "referencing" @@ -5385,7 +5385,7 @@ files = [ {file = "tenacity-9.1.2-py3-none-any.whl", hash = "sha256:f77bf36710d8b73a50b2dd155c97b870017ad21afe6ab300326b0371b3b05138"}, {file = "tenacity-9.1.2.tar.gz", hash = "sha256:1169d376c297e7de388d18b4481760d478b0e99a777cad3a9c86e556f4b697cb"}, ] -markers = {main = "extra == \"langchain\" or extra == \"all\" or python_version <= \"3.13\" and (extra == \"crewai\" or extra == \"all\" or extra == \"langchain\")"} +markers = {main = "python_version <= \"3.13\" and (extra == \"crewai\" or extra == \"all\" or extra == \"langchain\") or extra == \"langchain\" or extra == \"all\""} [package.extras] doc = ["reno", "sphinx"] @@ -6427,7 +6427,7 @@ description = "Yet another URL library" optional = true python-versions = ">=3.9" groups = ["main"] -markers = "python_version <= \"3.13\" and (extra == \"crewai\" or extra == \"all\")" +markers = "python_version <= \"3.13\" and (extra == \"crewai\" or extra == \"all\" or extra == \"distributed-tracing\") or python_version < \"3.14\" and (extra == \"distributed-tracing\" or extra == \"all\" or extra == \"crewai\") or extra == \"distributed-tracing\" or extra == \"all\"" files = [ {file = "yarl-1.20.1-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:6032e6da6abd41e4acda34d75a816012717000fa6839f37124a47fcefc49bec4"}, {file = "yarl-1.20.1-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:2c7b34d804b8cf9b214f05015c4fee2ebe7ed05cf581e7192c06555c71f4446a"}, @@ -6698,7 +6698,7 @@ cffi = ["cffi (>=1.11)"] [extras] all = ["aiohttp", "crewai", "cryptography", "langchain", "langchain-core", "langsmith", "litellm", "mcp", "openai", "openai-agents", "opentelemetry-instrumentation-aiohttp-client", "opentelemetry-instrumentation-fastapi", "opentelemetry-instrumentation-httpx", "opentelemetry-instrumentation-requests", "opentelemetry-instrumentation-starlette", "packaging", "pdfminer-six", "starlette", "uv"] crewai = ["aiohttp", "crewai", "cryptography", "litellm", "mcp", "pdfminer-six", "uv"] -distributed-tracing = ["opentelemetry-instrumentation-aiohttp-client", "opentelemetry-instrumentation-fastapi", "opentelemetry-instrumentation-httpx", "opentelemetry-instrumentation-requests", "opentelemetry-instrumentation-starlette"] +distributed-tracing = ["aiohttp", "opentelemetry-instrumentation-aiohttp-client", "opentelemetry-instrumentation-fastapi", "opentelemetry-instrumentation-httpx", "opentelemetry-instrumentation-requests", "opentelemetry-instrumentation-starlette"] langchain = ["langchain", "langchain-core", "langsmith"] middleware = ["starlette"] openai = ["cryptography", "mcp", "openai", "openai-agents", "packaging"] @@ -6706,4 +6706,4 @@ openai = ["cryptography", "mcp", "openai", "openai-agents", "packaging"] [metadata] lock-version = "2.1" python-versions = "^3.11,<3.15" -content-hash = "3e3ae3afd6d9eacc3717e1075d0de4d90a25bdf3a9eaab25c61be0091a353b94" +content-hash = "8219ff5a4bc2c3be82d45473a92cb251482d8b3b5c02ca3a41ef68432f52deb1" diff --git a/pyproject.toml b/pyproject.toml index d342b4e5..fd32d71f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -30,14 +30,14 @@ langchain = ["langchain-core", "langchain", "langsmith (>=0.8.0)"] openai = ["openai (>=2.8.0,<3.0.0)", "packaging (>=24.2,<25.0)", "openai-agents (>=0.4.0,<1.0.0)", "cryptography (>=50.0.0)", "mcp (>=1.27.2,<2)"] crewai = ["crewai (>=0.152.0,<2.0.0); python_version < '3.14'", "litellm (>=1.83.14,<2.0.0); python_version < '3.14'", "uv (>=0.9.6); python_version < '3.14'", "aiohttp (>=3.14.1,<4); python_version < '3.14'", "cryptography (>=50.0.0)", "mcp (>=1.27.2,<2)", "pdfminer-six (>=20251107)"] middleware = ["starlette"] -distributed-tracing = ["opentelemetry-instrumentation-fastapi (==0.59b0)", "opentelemetry-instrumentation-starlette (==0.59b0)", "opentelemetry-instrumentation-requests (==0.59b0)", "opentelemetry-instrumentation-httpx (==0.59b0)", "opentelemetry-instrumentation-aiohttp-client (==0.59b0)"] -all = ["langchain-core", "langchain", "langsmith (>=0.8.0)", "openai (>=2.8.0,<3.0.0)", "packaging (>=24.2,<25.0)", "openai-agents (>=0.4.0,<1.0.0)", "crewai (>=0.152.0,<2.0.0); python_version < '3.14'", "starlette", "litellm (>=1.83.14,<2.0.0); python_version < '3.14'", "uv (>=0.9.6); python_version < '3.14'", "aiohttp (>=3.14.1,<4); python_version < '3.14'", "cryptography (>=50.0.0)", "mcp (>=1.27.2,<2)", "pdfminer-six (>=20251107)", "opentelemetry-instrumentation-fastapi (==0.59b0)", "opentelemetry-instrumentation-starlette (==0.59b0)", "opentelemetry-instrumentation-requests (==0.59b0)", "opentelemetry-instrumentation-httpx (==0.59b0)", "opentelemetry-instrumentation-aiohttp-client (==0.59b0)"] +distributed-tracing = ["aiohttp (>=3.14.1,<4)", "opentelemetry-instrumentation-fastapi (==0.59b0)", "opentelemetry-instrumentation-starlette (==0.59b0)", "opentelemetry-instrumentation-requests (==0.59b0)", "opentelemetry-instrumentation-httpx (==0.59b0)", "opentelemetry-instrumentation-aiohttp-client (==0.59b0)"] +all = ["langchain-core", "langchain", "langsmith (>=0.8.0)", "openai (>=2.8.0,<3.0.0)", "packaging (>=24.2,<25.0)", "openai-agents (>=0.4.0,<1.0.0)", "crewai (>=0.152.0,<2.0.0); python_version < '3.14'", "starlette", "litellm (>=1.83.14,<2.0.0); python_version < '3.14'", "uv (>=0.9.6); python_version < '3.14'", "aiohttp (>=3.14.1,<4)", "cryptography (>=50.0.0)", "mcp (>=1.27.2,<2)", "pdfminer-six (>=20251107)", "opentelemetry-instrumentation-fastapi (==0.59b0)", "opentelemetry-instrumentation-starlette (==0.59b0)", "opentelemetry-instrumentation-requests (==0.59b0)", "opentelemetry-instrumentation-httpx (==0.59b0)", "opentelemetry-instrumentation-aiohttp-client (==0.59b0)"] [tool.poetry.dependencies] python = "^3.11,<3.15" -aiohttp = { version = ">=3.14.1,<4", optional = true, python = ">=3.11,<3.14" } +aiohttp = { version = ">=3.14.1,<4", optional = true } pydantic = "^2.11.9" pyjwt = ">=2.13.0,<3.0.0" wrapt = "^1.14" diff --git a/src/splunk_ao/__init__.py b/src/splunk_ao/__init__.py index 2a3203cb..8b4648f3 100644 --- a/src/splunk_ao/__init__.py +++ b/src/splunk_ao/__init__.py @@ -58,7 +58,7 @@ ) from splunk_ao.experiment import Experiment from splunk_ao.handlers.agent_control import SplunkAOAgentControlBridge, setup_agent_control_bridge -from splunk_ao.http_instrumentation import instrument_distributed_tracing +from splunk_ao.http_instrumentation import configure_distributed_tracing, instrument_distributed_tracing from splunk_ao.integration import Integration from splunk_ao.logger import SplunkAOLogger from splunk_ao.logger.control import ControlAppliesTo, ControlCheckStage, ControlResult, ControlSpan @@ -161,6 +161,7 @@ "ValidationError", "WorkflowSpan", "add_records_to_annotation_queue", + "configure_distributed_tracing", "create_annotation_queue", "create_annotation_queue_field", "create_api_key", diff --git a/src/splunk_ao/http_instrumentation.py b/src/splunk_ao/http_instrumentation.py index 12d37382..07749a2c 100644 --- a/src/splunk_ao/http_instrumentation.py +++ b/src/splunk_ao/http_instrumentation.py @@ -1,12 +1,14 @@ -"""Opt-in setup for supported upstream OpenTelemetry HTTP instrumentors.""" +"""Supported automatic OpenTelemetry export and HTTP instrumentation setup.""" from __future__ import annotations from typing import Any, cast +from weakref import WeakSet from opentelemetry import trace +from opentelemetry.sdk.trace import TracerProvider as SDKTracerProvider -from splunk_ao.otel import TracerProvider +from splunk_ao.otel import TracerProvider, add_splunk_ao_span_processor from splunk_ao.session_context import install_session_propagator _INSTALL_MESSAGE = ( @@ -17,6 +19,7 @@ _client_provider_ids: dict[str, int] = {} _instrumented_apps: set[tuple[int, int, str]] = set() +_configured_providers: WeakSet[SDKTracerProvider] = WeakSet() def _load_instrumentors() -> dict[str, type]: @@ -49,7 +52,7 @@ def _validate_client_ownership(names: tuple[str, ...], provider_id: int) -> None ) -def _instrument_app(app: Any, instrumentors: dict[str, Any], tracer_provider: TracerProvider) -> None: +def _resolve_app_framework(app: Any) -> str: try: from fastapi import FastAPI # noqa: PLC0415 from starlette.applications import Starlette # noqa: PLC0415 @@ -62,6 +65,11 @@ def _instrument_app(app: Any, instrumentors: dict[str, Any], tracer_provider: Tr framework = "starlette" else: raise TypeError("app must be a FastAPI or Starlette application") + return framework + + +def _instrument_app(app: Any, instrumentors: dict[str, Any], tracer_provider: TracerProvider) -> None: + framework = _resolve_app_framework(app) key = (id(app), id(tracer_provider), framework) if key in _instrumented_apps: @@ -70,6 +78,32 @@ def _instrument_app(app: Any, instrumentors: dict[str, Any], tracer_provider: Tr _instrumented_apps.add(key) +def _instrument_supported_transports( + *, + instrumentors: dict[str, type], + tracer_provider: TracerProvider, + app: Any | None, + instrument_requests: bool, + instrument_httpx: bool, + instrument_aiohttp_client: bool, +) -> None: + clients = {"requests": instrument_requests, "httpx": instrument_httpx, "aiohttp-client": instrument_aiohttp_client} + enabled_clients = tuple(name for name, enabled in clients.items() if enabled) + provider_id = id(tracer_provider) + _validate_client_ownership(enabled_clients, provider_id) + + if app is not None: + _instrument_app(app, instrumentors, tracer_provider) + + for name in enabled_clients: + if _client_provider_ids.get(name) == provider_id: + continue + instrumentors[name]().instrument(tracer_provider=tracer_provider) + _client_provider_ids[name] = provider_id + + install_session_propagator() + + def instrument_distributed_tracing( *, tracer_provider: TracerProvider | None = None, @@ -101,19 +135,73 @@ def instrument_distributed_tracing( """ instrumentors = _load_instrumentors() resolved_provider = tracer_provider or cast(TracerProvider, trace.get_tracer_provider()) + _instrument_supported_transports( + instrumentors=instrumentors, + tracer_provider=resolved_provider, + app=app, + instrument_requests=instrument_requests, + instrument_httpx=instrument_httpx, + instrument_aiohttp_client=instrument_aiohttp_client, + ) - clients = {"requests": instrument_requests, "httpx": instrument_httpx, "aiohttp-client": instrument_aiohttp_client} - enabled_clients = tuple(name for name, enabled in clients.items() if enabled) - provider_id = id(resolved_provider) - _validate_client_ownership(enabled_clients, provider_id) - if app is not None: - _instrument_app(app, instrumentors, resolved_provider) +def configure_distributed_tracing( + *, + tracer_provider: SDKTracerProvider | None = None, + app: Any | None = None, + instrument_requests: bool = True, + instrument_httpx: bool = True, + instrument_aiohttp_client: bool = True, +) -> SDKTracerProvider: + """Configure Splunk AO export and supported automatic HTTP propagation. - for name in enabled_clients: - if _client_provider_ids.get(name) == provider_id: - continue - instrumentors[name]().instrument(tracer_provider=resolved_provider) - _client_provider_ids[name] = provider_id + Parameters + ---------- + tracer_provider : SDKTracerProvider | None + Optional application-owned OpenTelemetry SDK provider. A new provider + is created when omitted. This function never replaces the process-global + provider. + app : Any | None + Optional FastAPI or Starlette application to instrument. + instrument_requests : bool + Whether to instrument process-wide Requests clients. + instrument_httpx : bool + Whether to instrument process-wide HTTPX sync and async clients. + instrument_aiohttp_client : bool + Whether to instrument process-wide aiohttp clients. - install_session_propagator() + Returns + ------- + SDKTracerProvider + The configured provider. The caller owns it and must call ``shutdown()`` + during process teardown. + + Notes + ----- + Repeated calls with the same provider do not attach another Splunk AO span + processor. Use :func:`instrument_distributed_tracing` instead when export + has already been configured separately on the provider. + """ + # Resolve optional imports before constructing a processor and its worker so + # a missing extra cannot leave partially configured background resources. + instrumentors = _load_instrumentors() + resolved_provider = tracer_provider or SDKTracerProvider() + clients = {"requests": instrument_requests, "httpx": instrument_httpx, "aiohttp-client": instrument_aiohttp_client} + enabled_clients = tuple(name for name, enabled in clients.items() if enabled) + _validate_client_ownership(enabled_clients, id(resolved_provider)) + if app is not None: + _resolve_app_framework(app) + + if resolved_provider not in _configured_providers: + add_splunk_ao_span_processor(resolved_provider) + _configured_providers.add(resolved_provider) + + _instrument_supported_transports( + instrumentors=instrumentors, + tracer_provider=resolved_provider, + app=app, + instrument_requests=instrument_requests, + instrument_httpx=instrument_httpx, + instrument_aiohttp_client=instrument_aiohttp_client, + ) + return resolved_provider diff --git a/tests/test_http_instrumentation.py b/tests/test_http_instrumentation.py index 1c176328..0ecd7bfa 100644 --- a/tests/test_http_instrumentation.py +++ b/tests/test_http_instrumentation.py @@ -29,8 +29,13 @@ from starlette.testclient import TestClient from yarl import URL -from splunk_ao import instrument_distributed_tracing -from splunk_ao.http_instrumentation import _client_provider_ids, _instrumented_apps, _load_instrumentors +from splunk_ao import configure_distributed_tracing, instrument_distributed_tracing +from splunk_ao.http_instrumentation import ( + _client_provider_ids, + _configured_providers, + _instrumented_apps, + _load_instrumentors, +) from splunk_ao.logger.logger import SplunkAOLogger from splunk_ao.session_context import GEN_AI_CONVERSATION_ID, SplunkAOSessionPropagator, _session_id_context @@ -114,11 +119,13 @@ def reset_instrumentation_state() -> Generator[None, None, None]: previous_propagator = propagate.get_global_textmap() session_token = _session_id_context.set(None) _client_provider_ids.clear() + _configured_providers.clear() _instrumented_apps.clear() try: yield finally: _client_provider_ids.clear() + _configured_providers.clear() _instrumented_apps.clear() propagate.set_global_textmap(previous_propagator) _session_id_context.reset(session_token) @@ -154,6 +161,78 @@ def test_missing_extra_fails_before_instrumentation_or_propagator_change() -> No assert propagate.get_global_textmap() is previous_propagator +def test_configure_distributed_tracing_creates_and_returns_provider( + instrumentors: tuple[dict[str, type], dict[str, RecordingInstrumentor]], +) -> None: + types, recorders = instrumentors + provider = TracerProvider() + processor = MagicMock() + app = FastAPI() + + with ( + patch("splunk_ao.http_instrumentation.SDKTracerProvider", return_value=provider), + patch("splunk_ao.http_instrumentation.add_splunk_ao_span_processor", return_value=processor) as add_processor, + patch("splunk_ao.http_instrumentation._load_instrumentors", return_value=types), + ): + configured_provider = configure_distributed_tracing(app=app) + + assert configured_provider is provider + add_processor.assert_called_once_with(provider) + assert recorders["fastapi"].app_calls == [(app, {"tracer_provider": provider})] + assert recorders["requests"].instrument_calls == [{"tracer_provider": provider}] + + +def test_configure_distributed_tracing_registers_processor_once_per_provider( + instrumentors: tuple[dict[str, type], dict[str, RecordingInstrumentor]], +) -> None: + types, recorders = instrumentors + provider = TracerProvider() + processor = MagicMock() + + with ( + patch("splunk_ao.http_instrumentation.add_splunk_ao_span_processor", return_value=processor) as add_processor, + patch("splunk_ao.http_instrumentation._load_instrumentors", return_value=types), + ): + first_result = configure_distributed_tracing(tracer_provider=provider) + second_result = configure_distributed_tracing(tracer_provider=provider) + + assert first_result is provider + assert second_result is provider + add_processor.assert_called_once_with(provider) + assert recorders["requests"].instrument_calls == [{"tracer_provider": provider}] + + +def test_configure_distributed_tracing_missing_extra_does_not_register_processor() -> None: + provider = TracerProvider() + + with ( + patch("splunk_ao.http_instrumentation._load_instrumentors", side_effect=ImportError("install the extra")), + patch("splunk_ao.http_instrumentation.add_splunk_ao_span_processor") as add_processor, + pytest.raises(ImportError, match="install the extra"), + ): + configure_distributed_tracing(tracer_provider=provider) + + add_processor.assert_not_called() + assert provider not in _configured_providers + + +def test_configure_distributed_tracing_invalid_app_does_not_register_processor( + instrumentors: tuple[dict[str, type], dict[str, RecordingInstrumentor]], +) -> None: + types, _ = instrumentors + provider = TracerProvider() + + with ( + patch("splunk_ao.http_instrumentation._load_instrumentors", return_value=types), + patch("splunk_ao.http_instrumentation.add_splunk_ao_span_processor") as add_processor, + pytest.raises(TypeError, match="FastAPI or Starlette"), + ): + configure_distributed_tracing(tracer_provider=provider, app=object()) + + add_processor.assert_not_called() + assert provider not in _configured_providers + + def test_distributed_tracing_extra_loads_every_supported_upstream_instrumentor() -> None: assert set(_load_instrumentors()) == {"fastapi", "starlette", "requests", "httpx", "aiohttp-client"} From a36406a31198d8f64685a70ee196d488c7e7d8cc Mon Sep 17 00:00:00 2001 From: pradystar Date: Thu, 13 Aug 2026 14:51:24 -0700 Subject: [PATCH 05/13] fix(tracing): suppress SDK control-plane HTTP spans --- AGENTS.md | 3 +++ ARCHITECTURE.md | 5 ++++ CHANGELOG.md | 3 +++ README.md | 5 ++++ src/splunk_ao/config.py | 54 ++++++++++++++++++++++++++++++++++----- tests/test_o11y_config.py | 25 +++++++++++++++--- 6 files changed, 86 insertions(+), 9 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 2e3831ea..a3087621 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -138,6 +138,9 @@ CI supports Python 3.11–3.14; root CI also spans Linux, macOS, and Windows. provider, registers Splunk AO export once per provider, configures supported upstream instrumentors, and returns the provider for application shutdown. `instrument_distributed_tracing()` remains transport-only. - Never replace the process-global tracer provider. Register processors on the provided provider; respect ownership. +- Keep SDK-owned authentication, validation, CRUD, routing-resolution, and other control-plane HTTP calls out of + application traces with scoped OTel HTTP suppression. Never use broad backend URL exclusions that can suppress user + traffic, and always restore instrumentation immediately after the SDK request. - Treat ended `ReadableSpan` objects as immutable. Normalize by copying at export, never by mutating private fields. - Completed spans enqueue immediately. `flush()` drains completed work without ending active work; `terminate()` drains, shuts down SDK-owned resources, and discards unfinished state. Caller-owned providers use `shutdown()`. diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index bfc39502..808621b4 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -100,6 +100,11 @@ High-level setup is idempotent per provider for Splunk AO processor registration and each server application must still have one instrumentation owner; do not combine the helper with direct upstream instrumentation of the same component. +SDK-owned authentication, health-check, CRUD, routing-resolution, token-refresh, and streaming control-plane requests +execute inside OTel HTTP-instrumentation suppression. This boundary is scoped to the SDK request and must restore the +caller's context afterward; application traffic to the same host must remain instrumentable. Apply suppression at the +shared configuration/API-client boundary, never by excluding deployment URLs globally. + An explicit SDK session propagates as the standard W3C baggage member `gen_ai.conversation.id`. Export normalization may derive the local compatibility attribute `splunk_ao.session.id`, but that attribute is not propagated as a second baggage member. The SDK must not propagate authentication, deployment, Project, Agent Stream/log stream, experiment, diff --git a/CHANGELOG.md b/CHANGELOG.md index 4c37977b..71521672 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -21,6 +21,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 `extract_tracing_context()` helpers for standard W3C `traceparent` and `tracestate` propagation. `TracingMiddleware` now extracts and scopes that OpenTelemetry context for Starlette/FastAPI requests. +- SDK-owned authentication, validation, and CRUD HTTP requests are suppressed + from automatic application HTTP instrumentation, preventing control-plane + calls from appearing as unrelated application traces. ### Changed diff --git a/README.md b/README.md index 9be9538b..d687eda9 100644 --- a/README.md +++ b/README.md @@ -399,6 +399,11 @@ call `get_tracing_headers()` for each request and does not install `TracingMiddleware`. Those APIs remain supported for transports and frameworks outside the automatic support matrix. +SDK-owned authentication, health-check, Project, Agent Stream, and other +control-plane HTTP requests are scoped out of automatic client instrumentation. +They therefore do not appear as application traces, while HTTP requests made by +the application remain instrumented normally. + An explicit SDK session propagates as the standard `gen_ai.conversation.id` W3C baggage member. Project, Agent Stream, experiment, deployment, authentication, model, workflow, and agent identity diff --git a/src/splunk_ao/config.py b/src/splunk_ao/config.py index bf002819..8b2c3a3c 100644 --- a/src/splunk_ao/config.py +++ b/src/splunk_ao/config.py @@ -2,6 +2,7 @@ # We need to ignore syntax errors until https://github.com/python/mypy/issues/17535 is resolved. import os from collections.abc import Iterator +from contextlib import contextmanager from typing import Any, ClassVar, Optional from httpx import Response @@ -17,7 +18,33 @@ from splunk_ao.shared.exceptions import ConfigurationError, MissingConfigurationError -class O11yApiClient(ApiClient): +@contextmanager +def _suppress_control_plane_http() -> Iterator[None]: + """Suppress optional OTel HTTP instrumentation around SDK-owned requests.""" + try: + from opentelemetry.instrumentation.utils import suppress_http_instrumentation # noqa: PLC0415 + except ImportError: + yield + return + + with suppress_http_instrumentation(): + yield + + +class _ControlPlaneApiClient(ApiClient): + """Keep SDK-owned HTTP requests out of application telemetry.""" + + async def arequest(self, method: RequestMethod, path: str, *args: Any, **kwargs: Any) -> Any: + with _suppress_control_plane_http(): + return await super().arequest(method, path, *args, **kwargs) + + @contextmanager + def stream_request(self, method: RequestMethod, path: str, *args: Any, **kwargs: Any) -> Iterator[Response]: + with _suppress_control_plane_http(), super().stream_request(method, path, *args, **kwargs) as response: + yield response + + +class O11yApiClient(_ControlPlaneApiClient): """API client for Splunk Observability Cloud AO endpoints.""" o11y_token: SecretStr @@ -41,8 +68,10 @@ def _prefixed(self, path: str) -> str: async def arequest(self, method: RequestMethod, path: str, *args: Any, **kwargs: Any) -> Any: return await super().arequest(method, self._prefixed(path), *args, **kwargs) + @contextmanager def stream_request(self, method: RequestMethod, path: str, *args: Any, **kwargs: Any) -> Iterator[Response]: - return super().stream_request(method, self._prefixed(path), *args, **kwargs) + with super().stream_request(method, self._prefixed(path), *args, **kwargs) as response: + yield response # Mapping of SPLUNK_AO_* → GALILEO_* env var pairs used by the bridge. @@ -105,7 +134,8 @@ def set_api_url(cls, api_url: str | Url | None, info: ValidationInfo) -> Url: """Derive the O11y API URL from its realm and preserve standalone validation.""" if cls._is_o11y_env(): return Url(O11yConfig.from_env().require_api_url()) - return super().set_api_url(api_url, info) + with _suppress_control_plane_http(): + return super().set_api_url(api_url, info) @model_validator(mode="after") def set_jwt_token(self) -> "SplunkAOConfig": @@ -114,7 +144,8 @@ def set_jwt_token(self) -> "SplunkAOConfig": self.jwt_token = None self.refresh_token = None return self - super().set_jwt_token() + with _suppress_control_plane_http(): + super().set_jwt_token() return self @model_validator(mode="after") @@ -126,7 +157,17 @@ def set_validated_api_client(self) -> "SplunkAOConfig": host=o11y.api_root, o11y_token=o11y.crud_token, jwt_token=SecretStr(""), ssl_context=self.ssl_context ) return self - super().set_validated_api_client() + with _suppress_control_plane_http(): + super().set_validated_api_client() + assert self.validated_api_client is not None + client = self.validated_api_client + self.validated_api_client = _ControlPlaneApiClient( + host=client.host, + jwt_token=client.jwt_token, + raise_on_unexpected_status=client.raise_on_unexpected_status, + ssl_context=client.ssl_context, + thread_local=client.thread_local, + ) return self def _uses_o11y_api_client(self) -> bool: @@ -136,7 +177,8 @@ def refresh_jwt_token(self) -> None: """Skip JWT refresh when authenticating directly with an O11y token.""" if self._uses_o11y_api_client(): return - super().refresh_jwt_token() + with _suppress_control_plane_http(): + super().refresh_jwt_token() @classmethod def get(cls, **kwargs: Any) -> "SplunkAOConfig": diff --git a/tests/test_o11y_config.py b/tests/test_o11y_config.py index f933458b..3a49cdee 100644 --- a/tests/test_o11y_config.py +++ b/tests/test_o11y_config.py @@ -7,13 +7,14 @@ from uuid import uuid4 import pytest +from opentelemetry.instrumentation.utils import is_http_instrumentation_enabled from pydantic import SecretStr from galileo_core.constants.request_method import RequestMethod from galileo_core.constants.routes import Routes from galileo_core.helpers.api_client import ApiClient from galileo_core.schemas.base_config import GalileoConfig -from splunk_ao.config import O11yApiClient, SplunkAOConfig +from splunk_ao.config import O11yApiClient, SplunkAOConfig, _ControlPlaneApiClient from splunk_ao.shared.exceptions import AmbiguousConfigurationError, MissingConfigurationError _CONFIG_ENV_VARS = ( @@ -85,7 +86,12 @@ def test_o11y_api_client_sync_request_preserves_prefix_and_header(monkeypatch: p async def fake_make_request( request_method: RequestMethod, base_url: str, endpoint: str, headers: dict[str, str], **kwargs: object ) -> dict: - captured.update(method=request_method, url=urljoin(base_url, endpoint), headers=headers) + captured.update( + method=request_method, + url=urljoin(base_url, endpoint), + headers=headers, + instrumentation_enabled=is_http_instrumentation_enabled(), + ) return {} monkeypatch.setattr(ApiClient, "make_request", staticmethod(fake_make_request)) @@ -95,7 +101,9 @@ async def fake_make_request( "method": RequestMethod.GET, "url": "https://app.lab0.observability.splunkcloud.com/ao/api/projects", "headers": {"accept": "application/json", "Content-Type": "application/json", "X-SF-Token": "tok"}, + "instrumentation_enabled": False, } + assert is_http_instrumentation_enabled() @pytest.mark.asyncio @@ -116,12 +124,13 @@ async def fake_make_request( def test_o11y_api_client_prefixes_streaming_requests(monkeypatch: pytest.MonkeyPatch) -> None: - captured: dict[str, str] = {} + captured: dict[str, object] = {} def fake_stream_request( self: ApiClient, method: RequestMethod, path: str, *args: object, **kwargs: object ) -> contextlib.AbstractContextManager[str]: captured["path"] = path + captured["instrumentation_enabled"] = is_http_instrumentation_enabled() return nullcontext(path) monkeypatch.setattr(ApiClient, "stream_request", fake_stream_request) @@ -129,6 +138,8 @@ def fake_stream_request( pass assert captured["path"] == "/ao/api/projects" + assert captured["instrumentation_enabled"] is False + assert is_http_instrumentation_enabled() @pytest.mark.parametrize("token_var", ["SPLUNK_AO_O11Y_TOKEN", "SPLUNK_AO_O11Y_API_TOKEN"]) @@ -253,17 +264,21 @@ def test_ambiguous_environment_fails_before_config_construction(monkeypatch: pyt def test_standalone_constructor_flow_still_uses_parent_validation(monkeypatch: pytest.MonkeyPatch) -> None: calls: list[object] = [] + instrumentation_states: list[bool] = [] async def fake_make_request(request_method: RequestMethod, base_url: str, endpoint: str, **kwargs: object) -> dict: calls.append(endpoint) + instrumentation_states.append(is_http_instrumentation_enabled()) return {"status": "ok"} def fake_get_jwt_token(*args: object, **kwargs: object) -> tuple[SecretStr, None]: calls.append("jwt") + instrumentation_states.append(is_http_instrumentation_enabled()) return SecretStr("jwt-token"), None def fake_request(self: ApiClient, method: RequestMethod, path: str, **kwargs: object) -> dict[str, str]: calls.append(path) + instrumentation_states.append(is_http_instrumentation_enabled()) return {"id": str(uuid4()), "email": "user@example.com", "role": "user"} monkeypatch.setattr(SplunkAOConfig, "_instance", None) @@ -275,6 +290,10 @@ def fake_request(self: ApiClient, method: RequestMethod, path: str, **kwargs: ob cfg = SplunkAOConfig.get(console_url="https://app.galileo.ai", api_key="key") assert not isinstance(cfg.validated_api_client, O11yApiClient) + assert isinstance(cfg.validated_api_client, _ControlPlaneApiClient) assert Routes.healthcheck in calls assert "jwt" in calls assert Routes.current_user in calls + assert instrumentation_states + assert not any(instrumentation_states) + assert is_http_instrumentation_enabled() From 3f1db22f3aa0c9cb3d179ff1438ab6ba58a303bc Mon Sep 17 00:00:00 2001 From: pradystar Date: Thu, 13 Aug 2026 16:12:15 -0700 Subject: [PATCH 06/13] docs(examples): add explicit and automatic distributed tracing samples --- .../distributed-tracing-auto/.env.example | 23 +++++ .../DT_2.0/distributed-tracing-auto/README.md | 39 ++++++++ .../distributed-tracing-auto/main_run.py | 90 +++++++++++++++++++ .../distributed-tracing-auto/requirements.txt | 6 ++ .../retrieval_service.py | 56 ++++++++++++ .../DT_2.0/distributed-tracing/.env.example | 19 ++++ .../DT_2.0/distributed-tracing/README.md | 25 ++++++ .../DT_2.0/distributed-tracing/main_run.py | 77 ++++++++++++++++ .../distributed-tracing/requirements.txt | 6 ++ .../distributed-tracing/retrieval_service.py | 49 ++++++++++ 10 files changed, 390 insertions(+) create mode 100644 examples/logging-samples/DT_2.0/distributed-tracing-auto/.env.example create mode 100644 examples/logging-samples/DT_2.0/distributed-tracing-auto/README.md create mode 100644 examples/logging-samples/DT_2.0/distributed-tracing-auto/main_run.py create mode 100644 examples/logging-samples/DT_2.0/distributed-tracing-auto/requirements.txt create mode 100644 examples/logging-samples/DT_2.0/distributed-tracing-auto/retrieval_service.py create mode 100644 examples/logging-samples/DT_2.0/distributed-tracing/.env.example create mode 100644 examples/logging-samples/DT_2.0/distributed-tracing/README.md create mode 100644 examples/logging-samples/DT_2.0/distributed-tracing/main_run.py create mode 100644 examples/logging-samples/DT_2.0/distributed-tracing/requirements.txt create mode 100644 examples/logging-samples/DT_2.0/distributed-tracing/retrieval_service.py diff --git a/examples/logging-samples/DT_2.0/distributed-tracing-auto/.env.example b/examples/logging-samples/DT_2.0/distributed-tracing-auto/.env.example new file mode 100644 index 00000000..90a5f99f --- /dev/null +++ b/examples/logging-samples/DT_2.0/distributed-tracing-auto/.env.example @@ -0,0 +1,23 @@ +# Select exactly one deployment block. Do not mix o11y and standalone variables. +# O11y Cloud +SPLUNK_AO_REALM="your-o11y-realm" +SPLUNK_AO_O11Y_TOKEN="your-o11y-ingest-or-ingest+api-token" +# SPLUNK_AO_O11Y_API_TOKEN="your-o11y-api-token" + +# Standalone +SPLUNK_AO_API_KEY="your-splunk-ao-api-key" +SPLUNK_AO_CONSOLE_URL="https://app.galileo.ai" + +# Configure project and agentstream +SPLUNK_AO_PROJECT="your-splunk-ao-project" +SPLUNK_AO_AGENT_STREAM="distributed-tracing-example" + +# OpenAI configuration +OPENAI_API_KEY="your-openai-api-key" +OPENAI_MODEL=gpt-5-nano + +OTEL_SERVICE_NAME=distributed-tracing-manual + +# Optional delay to see partial traces. +OTEL_BSP_SCHEDULE_DELAY="1000" +BATCH_DEMO_SLEEP_SECONDS="20" diff --git a/examples/logging-samples/DT_2.0/distributed-tracing-auto/README.md b/examples/logging-samples/DT_2.0/distributed-tracing-auto/README.md new file mode 100644 index 00000000..08939358 --- /dev/null +++ b/examples/logging-samples/DT_2.0/distributed-tracing-auto/README.md @@ -0,0 +1,39 @@ +## Configure + +Copy `.env.example` to `.env` and set the credentials and project and agentstream. + +```bash +python -m venv venv +source venv/bin/activate +pip install -r requirements.txt +cp .env.example .env +``` + +Start the retrieval service: +```bash +uvicorn retrieval_service:app --port 8000 +``` + +Then run the orchestrator: +```bash +python main_run.py +``` + +### Observe batching before the root ends + +To demonstrate incremental export, add these values to an existing `.env`: + +```dotenv +OTEL_BSP_SCHEDULE_DELAY="1000" +BATCH_DEMO_SLEEP_SECONDS="20" +``` + +`OTEL_BSP_SCHEDULE_DELAY` is in milliseconds and is read when each batch span +processor is created. After the distributed retrieval call returns, the example +prints a batching-demo message and keeps the orchestrator root active for 20 +seconds. + +Set`BATCH_DEMO_SLEEP_SECONDS="0"` to disable the pause. + +The local tree remains the same as the baseline. Completed children enter the +batch processor at their own end callbacks, so no per-trace flush is needed. diff --git a/examples/logging-samples/DT_2.0/distributed-tracing-auto/main_run.py b/examples/logging-samples/DT_2.0/distributed-tracing-auto/main_run.py new file mode 100644 index 00000000..237a1326 --- /dev/null +++ b/examples/logging-samples/DT_2.0/distributed-tracing-auto/main_run.py @@ -0,0 +1,90 @@ +"""Orchestrator using automatic W3C HTTP client instrumentation.""" + +import asyncio +import os +from uuid import uuid4 + +import httpx +from dotenv import load_dotenv + +from splunk_ao import configure_distributed_tracing, log, openai, splunk_ao_context + +load_dotenv() + +provider = configure_distributed_tracing() + +openai_client = openai.OpenAI() +OPENAI_MODEL = os.environ["OPENAI_MODEL"] +RETRIEVAL_SERVICE_URL = "http://localhost:8000" +BATCH_DEMO_SLEEP_SECONDS = float(os.getenv("BATCH_DEMO_SLEEP_SECONDS", "0")) + + +@log +async def orchestrator_agent(question: str) -> str: + """Run one independently traced RAG request.""" + analysis = analyze_question(question) + + # HTTPX creates the client span and injects the active W3C context. + async with httpx.AsyncClient(base_url=RETRIEVAL_SERVICE_URL, timeout=100.0) as client: + try: + response = await client.post("/retrieve", json={"query": question}) + response.raise_for_status() + retrieved_docs = response.json()["results"] + except httpx.HTTPError: + retrieved_docs = [] + + if BATCH_DEMO_SLEEP_SECONDS > 0: + print( + "Batching demo: analysis and distributed retrieval spans have ended; " + f"the orchestrator root remains active for {BATCH_DEMO_SLEEP_SECONDS:g} seconds." + ) + await asyncio.sleep(BATCH_DEMO_SLEEP_SECONDS) + + context = format_context(analysis, retrieved_docs) + response = openai_client.chat.completions.create( + messages=[ + { + "role": "system", + "content": ( + "Answer the user's question using only the provided context. " + f"If the context is insufficient, say so.\n\nContext:\n{context}" + ), + }, + {"role": "user", "content": question}, + ], + model=OPENAI_MODEL, + ) + return response.choices[0].message.content or "" + + +@log +def analyze_question(question: str) -> dict[str, object]: + question_lower = question.lower() + return { + "needs_company_info": any(word in question_lower for word in ("company", "work", "employer")), + "needs_location_info": any(word in question_lower for word in ("location", "where", "city", "live")), + "question_type": "factual", + } + + +@log +def format_context(analysis: dict[str, object], documents: list[str]) -> str: + rendered = [f"Analysis: {analysis}", "", "Retrieved Documents:"] + rendered.extend(f"{index}. {document}" for index, document in enumerate(documents, 1)) + return "\n".join(rendered) + + +async def main() -> None: + # One app-owned session is propagated as gen_ai.conversation.id baggage. + # set_session() is local-only; unlike start_session(), it performs no CRUD request. + splunk_ao_context.set_session(str(uuid4())) + for question in ("What did Galileo Galilei research?", "Where did Galileo Galilei work?"): + answer = await orchestrator_agent(question) + print(f"Question: {question}\nAnswer: {answer}\n") + + +if __name__ == "__main__": + try: + asyncio.run(main()) + finally: + provider.shutdown() diff --git a/examples/logging-samples/DT_2.0/distributed-tracing-auto/requirements.txt b/examples/logging-samples/DT_2.0/distributed-tracing-auto/requirements.txt new file mode 100644 index 00000000..d351de56 --- /dev/null +++ b/examples/logging-samples/DT_2.0/distributed-tracing-auto/requirements.txt @@ -0,0 +1,6 @@ +python-dotenv==1.2.2 +pydantic==2.11.9 +httpx==0.28.1 +fastapi==0.117.1 +uvicorn==0.37.0 +-e ../../../..[openai,distributed-tracing] diff --git a/examples/logging-samples/DT_2.0/distributed-tracing-auto/retrieval_service.py b/examples/logging-samples/DT_2.0/distributed-tracing-auto/retrieval_service.py new file mode 100644 index 00000000..a0a276bd --- /dev/null +++ b/examples/logging-samples/DT_2.0/distributed-tracing-auto/retrieval_service.py @@ -0,0 +1,56 @@ +"""FastAPI service using automatic W3C server instrumentation.""" + +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager + +from dotenv import load_dotenv +from fastapi import FastAPI +from pydantic import BaseModel + +from splunk_ao import configure_distributed_tracing, log + +load_dotenv() + + +@asynccontextmanager +async def lifespan(app: FastAPI) -> AsyncIterator[None]: + yield + provider.shutdown() + + +app = FastAPI(title="Retrieval Service", lifespan=lifespan) +provider = configure_distributed_tracing(app=app) + + +class RetrievalRequest(BaseModel): + query: str + + +class RetrievalResponse(BaseModel): + results: list[str] + + +@log(span_type="retriever") +def retrieval_service(query: str) -> list[str]: + knowledge_base = { + "birthplace": ["Galileo Galilei was born in Pisa, Italy in 1564."], + "profession": ["Galileo taught geometry, mechanics, and astronomy at the University of Padua."], + "research": ["Galileo's telescopic observations transformed our understanding of the universe."], + } + query_lower = query.lower() + results: list[str] = [] + for category, facts in knowledge_base.items(): + if category in query_lower or any(word in query_lower for word in ("work", "location", "education")): + results.extend(facts) + return results[:3] + + +@app.post("/retrieve", response_model=RetrievalResponse) +@log +async def retrieve_endpoint(request: RetrievalRequest) -> RetrievalResponse: + return RetrievalResponse(results=retrieval_service(request.query)) + + +@app.get("/health") +async def health() -> dict[str, str]: + return {"status": "healthy"} diff --git a/examples/logging-samples/DT_2.0/distributed-tracing/.env.example b/examples/logging-samples/DT_2.0/distributed-tracing/.env.example new file mode 100644 index 00000000..7ddde639 --- /dev/null +++ b/examples/logging-samples/DT_2.0/distributed-tracing/.env.example @@ -0,0 +1,19 @@ +# Select exactly one deployment block. Do not mix o11y and standalone variables. +# O11y Cloud +SPLUNK_AO_REALM="your-o11y-realm" +SPLUNK_AO_O11Y_TOKEN="your-o11y-ingest-or-ingest+api-token" +# SPLUNK_AO_O11Y_API_TOKEN="your-o11y-api-token" + +# Standalone +SPLUNK_AO_API_KEY="your-splunk-ao-api-key" +SPLUNK_AO_CONSOLE_URL="https://app.galileo.ai" + +# Configure project and agentstream +SPLUNK_AO_PROJECT="your-splunk-ao-project" +SPLUNK_AO_AGENT_STREAM="distributed-tracing-example" + +# OpenAI configuration +OPENAI_API_KEY="your-openai-api-key" +OPENAI_MODEL=gpt-5-nano + +OTEL_SERVICE_NAME=distributed-tracing-manual diff --git a/examples/logging-samples/DT_2.0/distributed-tracing/README.md b/examples/logging-samples/DT_2.0/distributed-tracing/README.md new file mode 100644 index 00000000..cb93a933 --- /dev/null +++ b/examples/logging-samples/DT_2.0/distributed-tracing/README.md @@ -0,0 +1,25 @@ +## Configure + +Copy `.env.example` to `.env` and set the credentials and project and agentstream. + +```bash +python -m venv venv +source venv/bin/activate +pip install -r requirements.txt +``` + +## Run + +Start the retrieval service in one terminal: +```bash +uvicorn retrieval_service:app --reload --port 8000 +``` + +Run the orchestrator in another: +```bash +python main_run.py +``` + +Each top level `orchestrator_agent()` call creates an independent trace. Its +outbound request carries the active orchestrator workflow span ID, and the +retrieval endpoint's workflow becomes its direct downstream child. diff --git a/examples/logging-samples/DT_2.0/distributed-tracing/main_run.py b/examples/logging-samples/DT_2.0/distributed-tracing/main_run.py new file mode 100644 index 00000000..e2d44dd4 --- /dev/null +++ b/examples/logging-samples/DT_2.0/distributed-tracing/main_run.py @@ -0,0 +1,77 @@ +"""Orchestrator process for the W3C distributed-tracing example.""" + +import asyncio +from uuid import uuid4 + +import httpx +from dotenv import load_dotenv + +from splunk_ao import get_tracing_headers, log, openai, splunk_ao_context + +load_dotenv() + +openai_client = openai.OpenAI() +RETRIEVAL_SERVICE_URL = "http://localhost:8000" + + +@log +async def orchestrator_agent(question: str) -> str: + """Run one independently traced RAG request.""" + analysis = analyze_question(question) + + # The active @log workflow is a real exportable operation. The returned + # carrier contains W3C traceparent and optional tracestate fields. + headers = get_tracing_headers() + async with httpx.AsyncClient(base_url=RETRIEVAL_SERVICE_URL, timeout=100.0) as client: + try: + response = await client.post("/retrieve", json={"query": question}, headers=headers) + response.raise_for_status() + retrieved_docs = response.json()["results"] + except httpx.HTTPError: + retrieved_docs = [] + + context = format_context(analysis, retrieved_docs) + response = openai_client.chat.completions.create( + messages=[ + { + "role": "system", + "content": ( + "Answer the user's question using only the provided context. " + f"If the context is insufficient, say so.\n\nContext:\n{context}" + ), + }, + {"role": "user", "content": question}, + ], + model="gpt-5-mini", + ) + return response.choices[0].message.content or "" + + +@log +def analyze_question(question: str) -> dict[str, object]: + question_lower = question.lower() + return { + "needs_company_info": any(word in question_lower for word in ("company", "work", "employer")), + "needs_location_info": any(word in question_lower for word in ("location", "where", "city", "live")), + "question_type": "factual", + } + + +@log +def format_context(analysis: dict[str, object], documents: list[str]) -> str: + rendered = [f"Analysis: {analysis}", "", "Retrieved Documents:"] + rendered.extend(f"{index}. {document}" for index, document in enumerate(documents, 1)) + return "\n".join(rendered) + + +async def main() -> None: + # One app-owned session is propagated as gen_ai.conversation.id baggage. + # set_session() is local-only; unlike start_session(), it performs no CRUD request. + splunk_ao_context.set_session(str(uuid4())) + for question in ("What did Galileo Galilei research?", "Where did Galileo Galilei work?"): + answer = await orchestrator_agent(question) + print(f"Question: {question}\nAnswer: {answer}\n") + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/examples/logging-samples/DT_2.0/distributed-tracing/requirements.txt b/examples/logging-samples/DT_2.0/distributed-tracing/requirements.txt new file mode 100644 index 00000000..5e0a8f23 --- /dev/null +++ b/examples/logging-samples/DT_2.0/distributed-tracing/requirements.txt @@ -0,0 +1,6 @@ +python-dotenv==1.2.2 +pydantic==2.11.9 +httpx==0.28.1 +fastapi==0.117.1 +uvicorn==0.37.0 +-e ../../../..[openai,middleware] diff --git a/examples/logging-samples/DT_2.0/distributed-tracing/retrieval_service.py b/examples/logging-samples/DT_2.0/distributed-tracing/retrieval_service.py new file mode 100644 index 00000000..b402514e --- /dev/null +++ b/examples/logging-samples/DT_2.0/distributed-tracing/retrieval_service.py @@ -0,0 +1,49 @@ +"""FastAPI retrieval process that continues incoming W3C context.""" + +from dotenv import load_dotenv +from fastapi import FastAPI +from pydantic import BaseModel + +from splunk_ao import log +from splunk_ao.middleware import TracingMiddleware + +load_dotenv() + +app = FastAPI(title="Retrieval Service") +app.add_middleware(TracingMiddleware) + + +class RetrievalRequest(BaseModel): + query: str + + +class RetrievalResponse(BaseModel): + results: list[str] + + +@log(span_type="retriever") +def retrieval_service(query: str) -> list[str]: + knowledge_base = { + "birthplace": ["Galileo Galilei was born in Pisa, Italy in 1564."], + "profession": ["Galileo taught geometry, mechanics, and astronomy at the University of Padua."], + "research": ["Galileo's telescopic observations transformed our understanding of the universe."], + } + query_lower = query.lower() + results: list[str] = [] + for category, facts in knowledge_base.items(): + if category in query_lower or any(word in query_lower for word in ("work", "location", "education")): + results.extend(facts) + return results[:3] + + +@app.post("/retrieve", response_model=RetrievalResponse) +@log +async def retrieve_endpoint(request: RetrievalRequest) -> RetrievalResponse: + # @log creates the real downstream workflow beneath the middleware's + # extracted remote parent; the retriever is its local child. + return RetrievalResponse(results=retrieval_service(request.query)) + + +@app.get("/health") +async def health() -> dict[str, str]: + return {"status": "healthy"} From 7b30e67f91f30834abb61b359c468a47a4a69fae Mon Sep 17 00:00:00 2001 From: pradystar Date: Mon, 17 Aug 2026 09:28:58 -0700 Subject: [PATCH 07/13] fix(tracing): harden automatic instrumentation ownership --- AGENTS.md | 5 + ARCHITECTURE.md | 4 + CHANGELOG.md | 6 +- README.md | 5 + splunk-ao-migration-tool/README.md | 26 +++-- src/splunk_ao/http_instrumentation.py | 44 +++++---- src/splunk_ao/logger/logger.py | 7 +- src/splunk_ao/middleware/tracing.py | 3 +- src/splunk_ao/session_context.py | 6 +- tests/test_http_instrumentation.py | 132 +++++++++++++++++++++++++- tests/test_tracing.py | 24 +++++ 11 files changed, 225 insertions(+), 37 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index a3087621..c7506d24 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -137,6 +137,9 @@ CI supports Python 3.11–3.14; root CI also spans Linux, macOS, and Windows. `configure_distributed_tracing()` is the combined automatic HTTP setup: it creates or accepts an application-owned provider, registers Splunk AO export once per provider, configures supported upstream instrumentors, and returns the provider for application shutdown. `instrument_distributed_tracing()` remains transport-only. +- Track automatic server-app ownership through weak application references and process-wide client ownership through + actual provider identity. Never use raw `id()` values as persistent ownership keys or transfer installed client + instrumentation merely because its owner is garbage-collected. - Never replace the process-global tracer provider. Register processors on the provided provider; respect ownership. - Keep SDK-owned authentication, validation, CRUD, routing-resolution, and other control-plane HTTP calls out of application traces with scoped OTel HTTP suppression. Never use broad backend URL exclusions that can suppress user @@ -148,6 +151,8 @@ CI supports Python 3.11–3.14; root CI also spans Linux, macOS, and Windows. code; sanitize and rate-limit diagnostics. - Preserve standard `gen_ai.*` attributes. New SDK-owned attributes use `splunk_ao.*`; do not introduce new proprietary `galileo.*` wire attributes. +- Explicit session selection is ambient within the current thread or async execution context and is shared by logger + instances in that context. Independent simultaneous sessions require separate execution contexts. - Changes to propagation, IDs, parents, content schemas, or routing need coverage across every affected telemetry path. ## Code Style diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index 808621b4..6e5862c8 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -110,6 +110,10 @@ may derive the local compatibility attribute `splunk_ao.session.id`, but that at baggage member. The SDK must not propagate authentication, deployment, Project, Agent Stream/log stream, experiment, endpoint, model, workflow, or agent identity in baggage. +Explicit session selection is ambient within the current thread or async execution context and is shared by logger +instances in that context. The most recent explicit selection applies to subsequently started telemetry and outbound +propagation. Independent simultaneous sessions require separate execution contexts. + ## Span Lifecycle and Export Ownership `SpanSink` owns the SDK's private provider and `BatchSpanProcessor` for internal telemetry. A completed operation is diff --git a/CHANGELOG.md b/CHANGELOG.md index 71521672..f66f172a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -15,8 +15,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 The lower-level `instrument_distributed_tracing()` transport helper remains available for applications that configure their provider separately. - Explicit SDK sessions now propagate across supported services as standard - W3C `gen_ai.conversation.id` baggage. SDK routing, authentication, and - application identity are not propagated. + W3C `gen_ai.conversation.id` baggage. Session selection is ambient within the + current thread or async execution context and is shared by logger instances + in that context. SDK routing, authentication, and application identity are + not propagated. - Added module-level `get_tracing_headers()` and `extract_tracing_context()` helpers for standard W3C `traceparent` and `tracestate` propagation. `TracingMiddleware` now extracts and scopes that diff --git a/README.md b/README.md index d687eda9..45c6e743 100644 --- a/README.md +++ b/README.md @@ -409,6 +409,11 @@ An explicit SDK session propagates as the standard experiment, deployment, authentication, model, workflow, and agent identity remain local and are never added to baggage by the SDK. +Explicit session selection is ambient within the current thread or async +execution context. Setting a session through any logger in that context selects +the conversation for subsequently started telemetry and outbound propagation +from that context. Use separate execution contexts for independent sessions. + #### Export diagnostics The SDK does not change your application's global logging configuration. diff --git a/splunk-ao-migration-tool/README.md b/splunk-ao-migration-tool/README.md index 250c0f95..b3bde187 100644 --- a/splunk-ao-migration-tool/README.md +++ b/splunk-ao-migration-tool/README.md @@ -25,7 +25,7 @@ Additionally there are a handful of **removed features** (Protect, `GalileoScore ### 1.1 Package Availability -> **`splunk-ao` is not yet published to PyPI.** +> **`splunk-ao` is not yet published to PyPI.** > Use one of the two installation methods below until a public release is available. **Option A — Install directly from GitHub (recommended for most users)** @@ -81,7 +81,7 @@ splunk-ao = { path = "../splunk-ao-python", develop = true } ### 1.2 Optional Extra Groups -The extras keys are unchanged (`langchain`, `openai`, `crewai`, `middleware`, `otel`, `all`). +The extras keys are unchanged (`langchain`, `openai`, `crewai`, `middleware`, `otel`, `all`). One new dependency was added to the `otel` and `all` extras: | Extra | Change | @@ -354,14 +354,24 @@ from `galileo-python`, it can be deleted at leisure or simply ignored. ## 6. HTTP Tracing Headers -If your services use distributed tracing and propagate Galileo headers between services, update the header names: +Distributed tracing now uses standard W3C Trace Context instead of renamed +Galileo/Splunk AO trace headers. Remove code that constructs or reads +`X-Galileo-Trace-ID`, `X-Galileo-Parent-ID`, `Splunk-AO-Trace-ID`, or +`Splunk-AO-Parent-ID`. -| Old | New | -|-----|-----| -| `X-Galileo-Trace-ID` | `Splunk-AO-Trace-ID` | -| `X-Galileo-Parent-ID` | `Splunk-AO-Parent-ID` | +For explicit propagation, use the module-level helper and pass its output to +the downstream request: + +```python +from splunk_ao import get_tracing_headers + +response = client.post(url, headers=get_tracing_headers()) +``` -The `get_tracing_headers()` function return value now uses the new header names. +The returned carrier uses W3C `traceparent` and, when applicable, `tracestate` +and baggage. An explicit SDK session is propagated as +`gen_ai.conversation.id` baggage. Applications using supported automatic HTTP +instrumentation do not need to call this helper for each request. --- diff --git a/src/splunk_ao/http_instrumentation.py b/src/splunk_ao/http_instrumentation.py index 07749a2c..5ff35fd4 100644 --- a/src/splunk_ao/http_instrumentation.py +++ b/src/splunk_ao/http_instrumentation.py @@ -3,7 +3,7 @@ from __future__ import annotations from typing import Any, cast -from weakref import WeakSet +from weakref import WeakKeyDictionary, WeakSet from opentelemetry import trace from opentelemetry.sdk.trace import TracerProvider as SDKTracerProvider @@ -17,8 +17,10 @@ ) -_client_provider_ids: dict[str, int] = {} -_instrumented_apps: set[tuple[int, int, str]] = set() +# Supported client instrumentors are process-wide. Retain their bounded owner +# objects so garbage collection cannot silently transfer installed wrappers. +_client_providers: dict[str, TracerProvider] = {} +_instrumented_apps: WeakKeyDictionary[Any, tuple[TracerProvider, str]] = WeakKeyDictionary() _configured_providers: WeakSet[SDKTracerProvider] = WeakSet() @@ -41,11 +43,11 @@ def _load_instrumentors() -> dict[str, type]: } -def _validate_client_ownership(names: tuple[str, ...], provider_id: int) -> None: +def _validate_client_ownership(names: tuple[str, ...], tracer_provider: TracerProvider) -> None: """Reject provider conflicts before any instrumentation is changed.""" for name in names: - configured_provider_id = _client_provider_ids.get(name) - if configured_provider_id is not None and configured_provider_id != provider_id: + configured_provider = _client_providers.get(name) + if configured_provider is not None and configured_provider is not tracer_provider: raise RuntimeError( f"{name} is already instrumented through Splunk AO with another tracer provider; " "configure each process-wide client instrumentor through one owner" @@ -68,14 +70,25 @@ def _resolve_app_framework(app: Any) -> str: return framework -def _instrument_app(app: Any, instrumentors: dict[str, Any], tracer_provider: TracerProvider) -> None: +def _validate_app_ownership(app: Any, tracer_provider: TracerProvider) -> str: + """Resolve the framework and reject ownership conflicts without mutation.""" framework = _resolve_app_framework(app) + configured_owner = _instrumented_apps.get(app) + if configured_owner is not None: + configured_provider, configured_framework = configured_owner + if configured_provider is not tracer_provider or configured_framework != framework: + raise RuntimeError( + f"{framework} app is already instrumented through Splunk AO with another tracer provider" + ) + return framework + - key = (id(app), id(tracer_provider), framework) - if key in _instrumented_apps: +def _instrument_app(app: Any, instrumentors: dict[str, Any], tracer_provider: TracerProvider) -> None: + framework = _validate_app_ownership(app, tracer_provider) + if app in _instrumented_apps: return instrumentors[framework].instrument_app(app, tracer_provider=tracer_provider) - _instrumented_apps.add(key) + _instrumented_apps[app] = (tracer_provider, framework) def _instrument_supported_transports( @@ -89,17 +102,16 @@ def _instrument_supported_transports( ) -> None: clients = {"requests": instrument_requests, "httpx": instrument_httpx, "aiohttp-client": instrument_aiohttp_client} enabled_clients = tuple(name for name, enabled in clients.items() if enabled) - provider_id = id(tracer_provider) - _validate_client_ownership(enabled_clients, provider_id) + _validate_client_ownership(enabled_clients, tracer_provider) if app is not None: _instrument_app(app, instrumentors, tracer_provider) for name in enabled_clients: - if _client_provider_ids.get(name) == provider_id: + if _client_providers.get(name) is tracer_provider: continue instrumentors[name]().instrument(tracer_provider=tracer_provider) - _client_provider_ids[name] = provider_id + _client_providers[name] = tracer_provider install_session_propagator() @@ -188,9 +200,9 @@ def configure_distributed_tracing( resolved_provider = tracer_provider or SDKTracerProvider() clients = {"requests": instrument_requests, "httpx": instrument_httpx, "aiohttp-client": instrument_aiohttp_client} enabled_clients = tuple(name for name, enabled in clients.items() if enabled) - _validate_client_ownership(enabled_clients, id(resolved_provider)) + _validate_client_ownership(enabled_clients, resolved_provider) if app is not None: - _resolve_app_framework(app) + _validate_app_ownership(app, resolved_provider) if resolved_provider not in _configured_providers: add_splunk_ao_span_processor(resolved_provider) diff --git a/src/splunk_ao/logger/logger.py b/src/splunk_ao/logger/logger.py index 68575824..76ddbfe5 100644 --- a/src/splunk_ao/logger/logger.py +++ b/src/splunk_ao/logger/logger.py @@ -2231,7 +2231,11 @@ def start_session( @warn_catch_exception(exceptions=(Exception,)) def set_session(self, session_id: str) -> None: """ - Set the session ID for the logger. + Set the explicit session for this logger and execution context. + + The execution-context selection is ambient and applies to subsequently + started telemetry from other logger instances in the same thread or + async context. Use separate execution contexts for independent sessions. Parameters ---------- @@ -2249,6 +2253,7 @@ def set_session(self, session_id: str) -> None: @nop_sync @warn_catch_exception(exceptions=(Exception,)) def clear_session(self) -> None: + """Clear this logger's session and the current execution-context selection.""" self._logger.info("Clearing the current session from the logger...") self._set_active_session_id(None) self._logger.info("Current session cleared.") diff --git a/src/splunk_ao/middleware/tracing.py b/src/splunk_ao/middleware/tracing.py index ceb677db..4ee7645b 100644 --- a/src/splunk_ao/middleware/tracing.py +++ b/src/splunk_ao/middleware/tracing.py @@ -13,8 +13,7 @@ async def process_request(data: dict): logger.start_trace(input=str(data), name="request") logger.add_workflow_span(input=str(data), name="process") result = await process(data) - logger.conclude(output=str(result)) - logger.conclude(output=str(result)) + logger.conclude(output=str(result), conclude_all=True) return {"result": result} finally: logger.terminate() diff --git a/src/splunk_ao/session_context.py b/src/splunk_ao/session_context.py index bda49cbe..76c737ba 100644 --- a/src/splunk_ao/session_context.py +++ b/src/splunk_ao/session_context.py @@ -1,4 +1,4 @@ -"""Request-local session context and W3C baggage propagation.""" +"""Execution-context-local session selection and W3C baggage propagation.""" from __future__ import annotations @@ -17,7 +17,7 @@ def set_session_context(session_id: str | None) -> None: - """Set the request-local explicit session selection.""" + """Set the explicit session selection for the current execution context.""" _session_id_context.set(session_id) @@ -28,7 +28,7 @@ def _baggage_session_id(context: Context | None = None) -> str | None: def get_effective_session_id(logger_session_id: str | None = None, context: Context | None = None) -> str | None: - """Resolve the request-local, inbound, or compatibility session ID.""" + """Resolve the execution-local, inbound, or compatibility session ID.""" local_session_id = _session_id_context.get(None) if local_session_id: return local_session_id diff --git a/tests/test_http_instrumentation.py b/tests/test_http_instrumentation.py index 0ecd7bfa..15641ab4 100644 --- a/tests/test_http_instrumentation.py +++ b/tests/test_http_instrumentation.py @@ -1,9 +1,11 @@ from __future__ import annotations +import gc from collections.abc import Generator from dataclasses import dataclass from typing import Any from unittest.mock import MagicMock, patch +from weakref import ref import aiohttp import httpx @@ -31,7 +33,7 @@ from splunk_ao import configure_distributed_tracing, instrument_distributed_tracing from splunk_ao.http_instrumentation import ( - _client_provider_ids, + _client_providers, _configured_providers, _instrumented_apps, _load_instrumentors, @@ -118,13 +120,13 @@ def instrument_app(cls, app: Any, **kwargs: Any) -> None: def reset_instrumentation_state() -> Generator[None, None, None]: previous_propagator = propagate.get_global_textmap() session_token = _session_id_context.set(None) - _client_provider_ids.clear() + _client_providers.clear() _configured_providers.clear() _instrumented_apps.clear() try: yield finally: - _client_provider_ids.clear() + _client_providers.clear() _configured_providers.clear() _instrumented_apps.clear() propagate.set_global_textmap(previous_propagator) @@ -156,8 +158,8 @@ def test_missing_extra_fails_before_instrumentation_or_propagator_change() -> No ): instrument_distributed_tracing(tracer_provider=FakeTracerProvider()) - assert _client_provider_ids == {} - assert _instrumented_apps == set() + assert not _client_providers + assert not _instrumented_apps assert propagate.get_global_textmap() is previous_propagator @@ -180,6 +182,7 @@ def test_configure_distributed_tracing_creates_and_returns_provider( add_processor.assert_called_once_with(provider) assert recorders["fastapi"].app_calls == [(app, {"tracer_provider": provider})] assert recorders["requests"].instrument_calls == [{"tracer_provider": provider}] + assert _client_providers["requests"] is provider def test_configure_distributed_tracing_registers_processor_once_per_provider( @@ -233,6 +236,39 @@ def test_configure_distributed_tracing_invalid_app_does_not_register_processor( assert provider not in _configured_providers +def test_configure_distributed_tracing_app_conflict_does_not_register_second_processor( + instrumentors: tuple[dict[str, type], dict[str, RecordingInstrumentor]], +) -> None: + types, recorders = instrumentors + app = FastAPI() + first_provider = TracerProvider() + second_provider = TracerProvider() + + with ( + patch("splunk_ao.http_instrumentation.add_splunk_ao_span_processor") as add_processor, + patch("splunk_ao.http_instrumentation._load_instrumentors", return_value=types), + ): + configure_distributed_tracing( + tracer_provider=first_provider, + app=app, + instrument_requests=False, + instrument_httpx=False, + instrument_aiohttp_client=False, + ) + with pytest.raises(RuntimeError, match="app is already instrumented"): + configure_distributed_tracing( + tracer_provider=second_provider, + app=app, + instrument_requests=False, + instrument_httpx=False, + instrument_aiohttp_client=False, + ) + + add_processor.assert_called_once_with(first_provider) + assert recorders["fastapi"].app_calls == [(app, {"tracer_provider": first_provider})] + assert second_provider not in _configured_providers + + def test_distributed_tracing_extra_loads_every_supported_upstream_instrumentor() -> None: assert set(_load_instrumentors()) == {"fastapi", "starlette", "requests", "httpx", "aiohttp-client"} @@ -257,6 +293,92 @@ def test_instruments_matching_app_and_every_enabled_client_once( assert isinstance(propagate.get_global_textmap(), SplunkAOSessionPropagator) +def test_instrumented_app_ownership_is_released_when_app_is_collected( + instrumentors: tuple[dict[str, type], dict[str, RecordingInstrumentor]], +) -> None: + types, recorders = instrumentors + provider = FakeTracerProvider() + app = FastAPI() + app_reference = ref(app) + + with patch("splunk_ao.http_instrumentation._load_instrumentors", return_value=types): + instrument_distributed_tracing( + tracer_provider=provider, + app=app, + instrument_requests=False, + instrument_httpx=False, + instrument_aiohttp_client=False, + ) + + assert len(_instrumented_apps) == 1 + recorders["fastapi"].app_calls.clear() + del app + gc.collect() + + assert app_reference() is None + assert not _instrumented_apps + + +def test_new_app_is_instrumented_after_previous_app_is_collected( + instrumentors: tuple[dict[str, type], dict[str, RecordingInstrumentor]], +) -> None: + types, recorders = instrumentors + provider = FakeTracerProvider() + first_app = FastAPI() + + with patch("splunk_ao.http_instrumentation._load_instrumentors", return_value=types): + instrument_distributed_tracing( + tracer_provider=provider, + app=first_app, + instrument_requests=False, + instrument_httpx=False, + instrument_aiohttp_client=False, + ) + recorders["fastapi"].app_calls.clear() + del first_app + gc.collect() + + second_app = FastAPI() + instrument_distributed_tracing( + tracer_provider=provider, + app=second_app, + instrument_requests=False, + instrument_httpx=False, + instrument_aiohttp_client=False, + ) + + assert recorders["fastapi"].app_calls == [(second_app, {"tracer_provider": provider})] + assert _instrumented_apps[second_app] == (provider, "fastapi") + + +def test_app_provider_conflict_fails_without_reinstrumenting_app( + instrumentors: tuple[dict[str, type], dict[str, RecordingInstrumentor]], +) -> None: + types, recorders = instrumentors + app = FastAPI() + first_provider = FakeTracerProvider() + second_provider = FakeTracerProvider() + + with patch("splunk_ao.http_instrumentation._load_instrumentors", return_value=types): + instrument_distributed_tracing( + tracer_provider=first_provider, + app=app, + instrument_requests=False, + instrument_httpx=False, + instrument_aiohttp_client=False, + ) + with pytest.raises(RuntimeError, match="app is already instrumented"): + instrument_distributed_tracing( + tracer_provider=second_provider, + app=app, + instrument_requests=False, + instrument_httpx=False, + instrument_aiohttp_client=False, + ) + + assert recorders["fastapi"].app_calls == [(app, {"tracer_provider": first_provider})] + + def test_disabled_clients_are_not_instrumented( instrumentors: tuple[dict[str, type], dict[str, RecordingInstrumentor]], ) -> None: diff --git a/tests/test_tracing.py b/tests/test_tracing.py index 473ee7dc..cf5d30ac 100644 --- a/tests/test_tracing.py +++ b/tests/test_tracing.py @@ -136,6 +136,30 @@ def test_get_tracing_headers_injects_only_standard_conversation_baggage() -> Non logger.terminate() +def test_explicit_session_selection_is_shared_across_loggers_in_execution_context() -> None: + first_logger, first_sink = make_logger() + second_logger, _ = make_logger() + try: + first_logger.set_session("conversation-first") + second_logger.set_session("conversation-current") + + first_logger.start_trace(input="request") + first_logger.add_workflow_span(input="work", name="operation") + headers = get_tracing_headers() + first_logger.conclude(output="work") + first_logger.conclude(output="done") + + [operation] = first_sink.spans + extracted = extract_tracing_context(headers) + assert first_logger.session_id == "conversation-first" + assert operation.attributes[GEN_AI_CONVERSATION_ID] == "conversation-current" + assert baggage.get_baggage(GEN_AI_CONVERSATION_ID, context=extracted) == "conversation-current" + finally: + second_logger.clear_session() + first_logger.terminate() + second_logger.terminate() + + def test_extract_tracing_context_restores_conversation_baggage() -> None: extracted = extract_tracing_context( { From ecaa72309697cae2b18ab0c869f4726ac1652410 Mon Sep 17 00:00:00 2001 From: Pradeep Nair Date: Wed, 26 Aug 2026 12:03:28 -0700 Subject: [PATCH 08/13] Update src/splunk_ao/handlers/base_handler.py Co-authored-by: Fernando Correia --- src/splunk_ao/handlers/base_handler.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/splunk_ao/handlers/base_handler.py b/src/splunk_ao/handlers/base_handler.py index 1eca5887..e256f5e6 100644 --- a/src/splunk_ao/handlers/base_handler.py +++ b/src/splunk_ao/handlers/base_handler.py @@ -185,8 +185,8 @@ def _finish_incremental_node(self, node: Node) -> None: final = finalize_handler_step(node, state) final = self._splunk_ao_logger._replace_handler_step(state.step, final) state.step = final - if state.step is self._owned_root or node_id == str(getattr(self._root_node, "run_id", "")): - self._owned_root = final if self._owned_root is not None else self._owned_root + if self._owned_root is not None and (state.step is self._owned_root or is_root): + self._owned_root = final self._splunk_ao_logger._restore_handler_step_context(state.activation) state.activation = None From 184124046e6eb99fcf63261d8f8394a6851d93aa Mon Sep 17 00:00:00 2001 From: Pradeep Nair Date: Wed, 26 Aug 2026 12:06:04 -0700 Subject: [PATCH 09/13] Update src/splunk_ao/handlers/span_lifecycle.py Co-authored-by: Fernando Correia --- src/splunk_ao/handlers/span_lifecycle.py | 18 +----------------- 1 file changed, 1 insertion(+), 17 deletions(-) diff --git a/src/splunk_ao/handlers/span_lifecycle.py b/src/splunk_ao/handlers/span_lifecycle.py index 2e76ef80..a94bdb9e 100644 --- a/src/splunk_ao/handlers/span_lifecycle.py +++ b/src/splunk_ao/handlers/span_lifecycle.py @@ -99,23 +99,7 @@ def build_handler_step( return LoggedLlmSpan( **common, input=input_value, - output=output, - metrics=LlmMetrics.model_validate( - { - "duration_ns": duration_ns, - "num_input_tokens": params.get("num_input_tokens"), - "num_output_tokens": params.get("num_output_tokens"), - "num_total_tokens": params.get("num_total_tokens", params.get("total_tokens")), - "time_to_first_token_ns": params.get("time_to_first_token_ns"), - "num_reasoning_tokens": params.get("num_reasoning_tokens"), - "num_cached_input_tokens": params.get("num_cached_input_tokens"), - } - ), - tools=params.get("tools"), - events=params.get("events"), - model=params.get("model"), - temperature=params.get("temperature"), - ) + output=output if output is not None else "", if node.node_type == "retriever": retriever_common = {**common, "status_code": None} return RetrieverSpan( From 5c1102121b902c81af887ee996d2b9e915ad3575 Mon Sep 17 00:00:00 2001 From: Pradeep Nair Date: Wed, 26 Aug 2026 12:07:09 -0700 Subject: [PATCH 10/13] Update src/splunk_ao/logger/logger.py Co-authored-by: Fernando Correia --- src/splunk_ao/logger/logger.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/splunk_ao/logger/logger.py b/src/splunk_ao/logger/logger.py index 76ddbfe5..bb88b34b 100644 --- a/src/splunk_ao/logger/logger.py +++ b/src/splunk_ao/logger/logger.py @@ -657,12 +657,14 @@ def _restore_handler_step_context(self, activation: HandlerStepContext | None) - if activation is None: return try: - otel_context.detach(activation.token) + activation.token.var.reset(activation.token) except (RuntimeError, ValueError): current = otel_trace.get_current_span().get_span_context() if current == activation.span_context: otel_context.attach(activation.prior_context) - self._logger.warning("Failed to detach handler OTel context; restored the prior context when safe.") + self._logger.debug( + "Handler OTel context was created in another execution context; restored the prior context when safe." + ) def _replace_handler_step(self, provisional: BaseStep, final: BaseStep) -> BaseStep: """Replace a provisional callback model while preserving identity and topology.""" From 1b51c31fd0b26e40a417c21337b5aea6c78893f6 Mon Sep 17 00:00:00 2001 From: Pradeep Nair Date: Wed, 26 Aug 2026 12:07:51 -0700 Subject: [PATCH 11/13] Update src/splunk_ao/config.py Co-authored-by: Fernando Correia --- src/splunk_ao/config.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/splunk_ao/config.py b/src/splunk_ao/config.py index 8b2c3a3c..6649fdf9 100644 --- a/src/splunk_ao/config.py +++ b/src/splunk_ao/config.py @@ -159,7 +159,9 @@ def set_validated_api_client(self) -> "SplunkAOConfig": return self with _suppress_control_plane_http(): super().set_validated_api_client() - assert self.validated_api_client is not None + client = self.validated_api_client + if client is None: + raise ConfigurationError("Validated API client was not initialized.") client = self.validated_api_client self.validated_api_client = _ControlPlaneApiClient( host=client.host, From 478d1427382688ce274ea363b41895e1b5f08bae Mon Sep 17 00:00:00 2001 From: pradystar Date: Thu, 27 Aug 2026 10:14:04 -0700 Subject: [PATCH 12/13] fix(tracing): harden distributed tracing lifecycle and instrumentation --- AGENTS.md | 7 +- ARCHITECTURE.md | 10 +- CHANGELOG.md | 15 + README.md | 17 +- .../distributed-tracing-auto/.env.example | 2 +- .../distributed-tracing/.env.example | 4 +- .../distributed-tracing/README.md | 13 +- src/splunk_ao/decorator.py | 54 ++- src/splunk_ao/handlers/base_handler.py | 39 +- .../handlers/langchain/async_handler.py | 2 + src/splunk_ao/handlers/langchain/handler.py | 2 + .../handlers/openai_agents/handler.py | 396 +++++++++++------- src/splunk_ao/handlers/span_lifecycle.py | 30 +- src/splunk_ao/http_instrumentation.py | 195 ++++++--- src/splunk_ao/logger/logger.py | 7 +- src/splunk_ao/session_context.py | 60 ++- src/splunk_ao/tracing.py | 28 +- src/splunk_ao/utils/singleton.py | 16 + tests/test_base_handler.py | 150 +++++-- tests/test_decorator.py | 20 + tests/test_http_instrumentation.py | 45 ++ tests/test_langchain.py | 29 +- tests/test_middleware_tracing.py | 15 +- tests/test_openai_agents.py | 125 +++++- tests/test_tracing.py | 66 ++- 25 files changed, 1016 insertions(+), 331 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index c7506d24..8cdad7e6 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -152,7 +152,12 @@ CI supports Python 3.11–3.14; root CI also spans Linux, macOS, and Windows. - Preserve standard `gen_ai.*` attributes. New SDK-owned attributes use `splunk_ao.*`; do not introduce new proprietary `galileo.*` wire attributes. - Explicit session selection is ambient within the current thread or async execution context and is shared by logger - instances in that context. Independent simultaneous sessions require separate execution contexts. + instances in that context. An explicit clear masks inbound baggage in that context. Independent simultaneous sessions + require separate execution contexts. +- Automatic HTTP setup wraps the current process-global propagator after application propagator configuration; replacing + it later removes the session wrapper. Remote W3C sampling flags are authoritative for SDK telemetry export. +- Process-wide integration processors must key mutable lifecycle state by framework trace identity and release + unfinished callback activations in reverse nesting order. - Changes to propagation, IDs, parents, content schemas, or routing need coverage across every affected telemetry path. ## Code Style diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index 6e5862c8..53147734 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -100,6 +100,10 @@ High-level setup is idempotent per provider for Splunk AO processor registration and each server application must still have one instrumentation owner; do not combine the helper with direct upstream instrumentation of the same component. +Automatic setup wraps the current process-global text-map propagator with a session-aware delegate. Applications must +install custom global propagators before invoking the helper. Replacing the global propagator afterward removes the +session adapter until setup is invoked again. This propagator mutation is independent of tracer-provider ownership. + SDK-owned authentication, health-check, CRUD, routing-resolution, token-refresh, and streaming control-plane requests execute inside OTel HTTP-instrumentation suppression. This boundary is scoped to the SDK request and must restore the caller's context afterward; application traffic to the same host must remain instrumentable. Apply suppression at the @@ -112,7 +116,11 @@ endpoint, model, workflow, or agent identity in baggage. Explicit session selection is ambient within the current thread or async execution context and is shared by logger instances in that context. The most recent explicit selection applies to subsequently started telemetry and outbound -propagation. Independent simultaneous sessions require separate execution contexts. +propagation. An explicit clear masks inbound conversation baggage for the rest of that execution context. Independent +simultaneous sessions require separate execution contexts. + +Remote W3C sampling is authoritative. Descendants of a valid remote parent with `sampled=0` retain its IDs for +continuation but do not enter the Splunk AO exporter. ## Span Lifecycle and Export Ownership diff --git a/CHANGELOG.md b/CHANGELOG.md index eec1281f..1db99186 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -44,6 +44,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 each completed operation into the existing `BatchSpanProcessor` at that operation's end callback. The deprecated `ingestion_hook` retains its whole-tree compatibility behavior. +- Incoming W3C sampling decisions are honored. A remote parent with its sampled + flag unset continues to propagate its trace identity, but its Splunk AO + descendants are not exported. - Default logger-owned batching now honors standard OpenTelemetry `OTEL_BSP_*` configuration, matching caller-owned OTel paths. Explicit internal `BatchConfig` values remain authoritative when supplied. @@ -60,6 +63,18 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Fixed +- Automatic HTTP setup now imports only the requested optional instrumentors + and rolls back components installed by the current call when later setup + fails, avoiding partial ownership state. +- Explicit `clear_session()` now masks an inbound + `gen_ai.conversation.id` for subsequent local spans and outbound propagation + in the same execution context. +- Handler cleanup now releases unfinished OTel activations in nesting order, + isolates malformed child telemetry from the remaining handler trace, and + keeps async LangChain callback activation in the application's execution + context. +- OpenAI Agents processor state is isolated per framework trace, including + concurrent traces, and its public lifecycle path now has regression coverage. - Agent and workflow output conversion now removes confirmed repeated input history and keeps only the last message as the terminal output; intermediate tool-call and tool-response messages are no longer included in diff --git a/README.md b/README.md index 45c6e743..b418beea 100644 --- a/README.md +++ b/README.md @@ -393,6 +393,14 @@ transport-only API. Do not also directly instrument a client or application component that is owned by the high-level helper; each component should have one instrumentation owner. +Both automatic setup helpers wrap the process-global OpenTelemetry text-map +propagator so an explicit Splunk AO session can travel as W3C baggage. If the +application uses a custom propagator (for example, a composite that also supports +B3), install it before calling the Splunk AO helper. Replacing the global +propagator afterward removes the session adapter; call the helper again after +that replacement if session propagation is still required. This does not replace +the process-global tracer provider. + After this startup call, supported inbound requests extract W3C context and supported outbound clients inject it automatically. Application code does not call `get_tracing_headers()` for each request and does not install @@ -412,7 +420,14 @@ remain local and are never added to baggage by the SDK. Explicit session selection is ambient within the current thread or async execution context. Setting a session through any logger in that context selects the conversation for subsequently started telemetry and outbound propagation -from that context. Use separate execution contexts for independent sessions. +from that context. `clear_session()` explicitly masks an inbound conversation +for later telemetry and outbound propagation in the same execution context. Use +separate execution contexts for independent sessions. + +Incoming W3C sampling decisions are honored. If an inbound `traceparent` has +the sampled flag unset (`sampled=0`), its Splunk AO descendants retain the trace +identity for propagation but are not exported. Configure upstream head sampling +accordingly when Agent Observability telemetry must be retained for every request. #### Export diagnostics diff --git a/examples/logging-samples/DT_2.0/distributed-tracing-auto/.env.example b/examples/logging-samples/DT_2.0/distributed-tracing-auto/.env.example index 90a5f99f..38ccf775 100644 --- a/examples/logging-samples/DT_2.0/distributed-tracing-auto/.env.example +++ b/examples/logging-samples/DT_2.0/distributed-tracing-auto/.env.example @@ -16,7 +16,7 @@ SPLUNK_AO_AGENT_STREAM="distributed-tracing-example" OPENAI_API_KEY="your-openai-api-key" OPENAI_MODEL=gpt-5-nano -OTEL_SERVICE_NAME=distributed-tracing-manual +OTEL_SERVICE_NAME=distributed-tracing-auto # Optional delay to see partial traces. OTEL_BSP_SCHEDULE_DELAY="1000" diff --git a/examples/logging-samples/distributed-tracing/.env.example b/examples/logging-samples/distributed-tracing/.env.example index 673bf267..7322bb02 100644 --- a/examples/logging-samples/distributed-tracing/.env.example +++ b/examples/logging-samples/distributed-tracing/.env.example @@ -1,9 +1,7 @@ SPLUNK_AO_API_KEY="your-splunk-ao-api-key" OPENAI_API_KEY="your-openai-api-key" SPLUNK_AO_PROJECT="your-splunk-ao-project" -SPLUNK_AO_LOGSTREAM="distributed-tracing-example" - -SPLUNK_AO_MODE=distributed +SPLUNK_AO_AGENT_STREAM="distributed-tracing-example" # Provide the console url below if you are not using app.galileo.ai # SPLUNK_AO_CONSOLE_URL="your-splunk-ao-console-url" diff --git a/examples/logging-samples/distributed-tracing/README.md b/examples/logging-samples/distributed-tracing/README.md index 7018526e..ec4a7737 100644 --- a/examples/logging-samples/distributed-tracing/README.md +++ b/examples/logging-samples/distributed-tracing/README.md @@ -13,13 +13,13 @@ This example demonstrates distributed tracing across two services using Splunk A ```bash # Navigate to the distributed-tracing example directory -cd python/logging-samples/distributed-tracing +cd examples/logging-samples/distributed-tracing # Create virtual environment python -m venv venv # Activate virtual environment -source venv/bin/activate +source venv/bin/activate ``` ### 2. Install Dependencies @@ -41,9 +41,6 @@ SPLUNK_AO_API_KEY="your-splunk-ao-api-key" # Required: OpenAI API key OPENAI_API_KEY="your-openai-api-key" -# Required: Enable distributed mode for distributed tracing -SPLUNK_AO_MODE=distributed - # Optional: Splunk AO project and log stream names SPLUNK_AO_PROJECT="your-splunk-ao-project" SPLUNK_AO_AGENT_STREAM=distributed-tracing-example @@ -60,7 +57,7 @@ You need to run both services in **separate terminal processes**. ```bash # Make sure you're in the distributed-tracing example directory -# cd python/logging-samples/distributed-tracing +# cd examples/logging-samples/distributed-tracing # Activate venv (if not already activated) # source venv/bin/activate @@ -82,7 +79,7 @@ INFO: Application startup complete. ```bash # In a NEW terminal, navigate to the distributed-tracing example directory -cd python/logging-samples/distributed-tracing +cd examples/logging-samples/distributed-tracing # Activate venv source venv/bin/activate @@ -100,4 +97,4 @@ When you run the orchestrator, you should see: Question: What did Galileo Galilei research? ============================================================ Answer: Galileo Galilei made scientific observations that transformed our understanding of the universe. -``` \ No newline at end of file +``` diff --git a/src/splunk_ao/decorator.py b/src/splunk_ao/decorator.py index 354f3d84..0b6eb871 100644 --- a/src/splunk_ao/decorator.py +++ b/src/splunk_ao/decorator.py @@ -65,7 +65,14 @@ def call_llm(prompt, temperature=0.7): from splunk_ao.schema.datasets import DatasetRecord from splunk_ao.schema.metrics import LocalMetricConfig from splunk_ao.schema.trace import SPAN_TYPE -from splunk_ao.session_context import _session_id_context, set_session_context +from splunk_ao.session_context import ( + SessionSelection, + _session_id_context, # noqa: F401 # retained as a private compatibility re-export + explicit_session_id, + get_session_selection, + restore_session_selection, + set_session_context, +) from splunk_ao.shared.exceptions import ConfigurationError from splunk_ao.utils import _get_timestamp from splunk_ao.utils.env_helpers import _get_mode_or_default @@ -104,7 +111,7 @@ def call_llm(prompt, temperature=0.7): _agent_stream_stack: ContextVar[list[str | None] | None] = ContextVar("log_stream_stack", default=None) _trace_stack: ContextVar[list[Trace | None] | None] = ContextVar("trace_stack", default=None) _experiment_id_stack: ContextVar[list[str | None] | None] = ContextVar("experiment_id_stack", default=None) -_session_id_stack: ContextVar[list[str | None] | None] = ContextVar("session_id_stack", default=None) +_session_id_stack: ContextVar[list[SessionSelection] | None] = ContextVar("session_id_stack", default=None) _mode_stack: ContextVar[list[LoggerModeType] | None] = ContextVar("mode_stack", default=None) _span_stack_stack: ContextVar[list[list[WorkflowSpan]] | None] = ContextVar("span_stack_stack", default=None) @@ -173,7 +180,7 @@ def __exit__( experiment_id=_experiment_id_context.get(), ).flush() - _session_id_context.set(None) + set_session_context(None) # Pop values from the stacks and restore the previous context _project_context.set(_get_or_init_list(_project_stack).pop()) @@ -182,15 +189,22 @@ def __exit__( _trace_context.set(_get_or_init_list(_trace_stack).pop()) _mode_context.set(_get_or_init_list(_mode_stack).pop()) _span_stack_context.set(_get_or_init_list(_span_stack_stack).pop()) - restored_session_id = _get_or_init_list(_session_id_stack).pop() - _session_id_context.set(restored_session_id) - restored_logger = self.get_logger_instance( - project=_project_context.get(), - agent_stream=_agent_stream_context.get(), - experiment_id=_experiment_id_context.get(), - ) - if isinstance(restored_logger, SplunkAOLogger): - restored_logger._set_active_session_id(restored_session_id) + restored_selection = _get_or_init_list(_session_id_stack).pop() + restore_session_selection(restored_selection) + has_explicit_session, restored_session_id = explicit_session_id(restored_selection) + if has_explicit_session: + try: + restored_logger = SplunkAOLoggerSingleton().get_existing( + project=_project_context.get(), + agent_stream=_agent_stream_context.get(), + experiment_id=_experiment_id_context.get(), + mode=_mode_context.get(), + ) + except Exception: + _logger.debug("Could not resolve a cached logger while restoring session context", exc_info=True) + else: + if restored_logger is not None: + restored_logger._set_active_session_id(restored_session_id) def __call__( self, @@ -235,7 +249,7 @@ def __call__( _get_or_init_list(_trace_stack).append(_trace_context.get()) _get_or_init_list(_mode_stack).append(_mode_context.get()) _get_or_init_list(_span_stack_stack).append(_get_or_init_list(_span_stack_context).copy()) - _get_or_init_list(_session_id_stack).append(_session_id_context.get()) + _get_or_init_list(_session_id_stack).append(get_session_selection()) # Reset trace context values _span_stack_context.set([]) @@ -246,7 +260,7 @@ def __call__( _agent_stream_context.set(None) _experiment_id_context.set(None) _mode_context.set(_get_mode_or_default(None)) - _session_id_context.set(None) + set_session_context(None) # Override with explicitly provided values if project is not None: @@ -258,7 +272,6 @@ def __call__( if mode is not None: _mode_context.set(_get_mode_or_default(mode)) if session_id is not None: - _session_id_context.set(session_id) self.set_session(session_id) return self @@ -1341,7 +1354,7 @@ def reset(self) -> None: _mode_context.set(_get_mode_or_default(None)) _span_stack_context.set([]) _trace_context.set(None) - _session_id_context.set(None) + set_session_context(None) # Clear all stacks _get_or_init_list(_project_stack).clear() _get_or_init_list(_agent_stream_stack).clear() @@ -1402,7 +1415,7 @@ def init( _mode_context.set(_get_mode_or_default(mode)) _span_stack_context.set([]) _trace_context.set(None) - _session_id_context.set(None) + set_session_context(None) def start_session( self, @@ -1430,17 +1443,13 @@ def start_session( str The id of the newly created session. """ - session_id = self.get_logger_instance().start_session( + return self.get_logger_instance().start_session( name=name, previous_session_id=previous_session_id, external_id=external_id, metadata=metadata ) - _session_id_context.set(session_id) - return session_id - def clear_session(self) -> None: """Clear the session in the active context logger instance.""" self.get_logger_instance().clear_session() - set_session_context(None) def set_session(self, session_id: str) -> None: """ @@ -1452,7 +1461,6 @@ def set_session(self, session_id: str) -> None: The id of the session to set. """ self.get_logger_instance().set_session(session_id) - set_session_context(session_id) splunk_ao_context = SplunkAODecorator() diff --git a/src/splunk_ao/handlers/base_handler.py b/src/splunk_ao/handlers/base_handler.py index e256f5e6..619c5846 100644 --- a/src/splunk_ao/handlers/base_handler.py +++ b/src/splunk_ao/handlers/base_handler.py @@ -6,7 +6,13 @@ from uuid import UUID from splunk_ao import splunk_ao_context -from splunk_ao.handlers.span_lifecycle import HandlerSpanState, build_handler_step, finalize_handler_step +from splunk_ao.handlers.span_lifecycle import ( + HandlerSpanState, + build_handler_step, + discard_handler_step, + finalize_handler_step, + release_handler_steps, +) from splunk_ao.logger import SplunkAOLogger from splunk_ao.schema.handlers import INTEGRATION, NODE_TYPE, Node from splunk_ao.schema.logged import LoggedAgentSpan, LoggedWorkflowSpan @@ -139,8 +145,8 @@ def commit(self) -> None: def _reset_handler_state(self) -> None: """Release one completed callback tree without touching caller-owned state.""" + release_handler_steps(self._splunk_ao_logger, self._active_steps) self._nodes.clear() - self._active_steps.clear() self._root_node = None self._owned_trace = None self._owned_root = None @@ -162,9 +168,14 @@ def _start_incremental_step(self, node: Node) -> None: raise RuntimeError(f"No active parent is available for handler node {node_id}") step = build_handler_step(node) - self._splunk_ao_logger._register_handler_step(step, parent) - activation = self._splunk_ao_logger._activate_handler_step(step) - self._active_steps[node_id] = HandlerSpanState(step=step, activation=activation) + state = HandlerSpanState(step=step, activation=None) + self._active_steps[node_id] = state + try: + self._splunk_ao_logger._register_handler_step(step, parent) + state.activation = self._splunk_ao_logger._activate_handler_step(step) + except Exception: + discard_handler_step(self._splunk_ao_logger, self._active_steps, node_id) + raise def _finish_incremental_node(self, node: Node) -> None: """Finalize and enqueue one callback span when that callback ends.""" @@ -179,13 +190,16 @@ def _finish_incremental_node(self, node: Node) -> None: is_root = node is self._root_node try: if is_root: + for active_node_id in reversed(tuple(self._active_steps)): + if active_node_id != node_id: + discard_handler_step(self._splunk_ao_logger, self._active_steps, active_node_id) node.span_params["output"] = self._root_output(node) if state.step is self._owned_root: node.span_params["input"] = serialize_to_str(node.span_params.get("input", "")) final = finalize_handler_step(node, state) final = self._splunk_ao_logger._replace_handler_step(state.step, final) state.step = final - if self._owned_root is not None and (state.step is self._owned_root or is_root): + if self._owned_root is not None and is_root: self._owned_root = final self._splunk_ao_logger._restore_handler_step_context(state.activation) @@ -204,9 +218,9 @@ def _finish_incremental_node(self, node: Node) -> None: status_code=node.span_params.get("status_code"), ) except Exception: - self._splunk_ao_logger._restore_handler_step_context(state.activation) - state.activation = None - self._conclude_owned_state_on_failure() + discard_handler_step(self._splunk_ao_logger, self._active_steps, node_id) + if is_root: + self._conclude_owned_state_on_failure() _logger.warning("Failed to complete handler telemetry for node %s", node_id, exc_info=True) finally: if is_root: @@ -557,7 +571,12 @@ def start_node(self, node_type: NODE_TYPE, parent_run_id: UUID | None, run_id: U try: self._start_incremental_step(node) except Exception: - self._conclude_owned_state_on_failure() + discard_handler_step(self._splunk_ao_logger, self._active_steps, node_id) + if node is self._root_node: + self._conclude_owned_state_on_failure() + self._owned_trace = None + self._owned_root = None + self._owned_parent = None _logger.warning("Failed to start handler span telemetry for node %s", node_id, exc_info=True) return node diff --git a/src/splunk_ao/handlers/langchain/async_handler.py b/src/splunk_ao/handlers/langchain/async_handler.py index 6c21cd16..59421834 100644 --- a/src/splunk_ao/handlers/langchain/async_handler.py +++ b/src/splunk_ao/handlers/langchain/async_handler.py @@ -42,6 +42,8 @@ class SplunkAOAsyncCallback(AsyncCallbackHandler): The async handler for managing the trace. """ + run_inline = True + def __init__( self, splunk_ao_logger: SplunkAOLogger | None = None, diff --git a/src/splunk_ao/handlers/langchain/handler.py b/src/splunk_ao/handlers/langchain/handler.py index 5a957065..c6266a76 100644 --- a/src/splunk_ao/handlers/langchain/handler.py +++ b/src/splunk_ao/handlers/langchain/handler.py @@ -41,6 +41,8 @@ class SplunkAOCallback(BaseCallbackHandler): The handler for managing the trace. """ + run_inline = True + def __init__( self, splunk_ao_logger: SplunkAOLogger | None = None, diff --git a/src/splunk_ao/handlers/openai_agents/handler.py b/src/splunk_ao/handlers/openai_agents/handler.py index d4236263..51f56641 100644 --- a/src/splunk_ao/handlers/openai_agents/handler.py +++ b/src/splunk_ao/handlers/openai_agents/handler.py @@ -1,5 +1,7 @@ import logging +import threading import uuid +from dataclasses import dataclass, field from datetime import UTC, datetime from typing import Any, cast @@ -9,7 +11,13 @@ from galileo_core.schemas.logging.span import LlmMetrics, LlmSpan from galileo_core.schemas.logging.span import Span as SplunkAOSpan from splunk_ao import SplunkAOLogger, splunk_ao_context -from splunk_ao.handlers.span_lifecycle import HandlerSpanState, build_handler_step, finalize_handler_step +from splunk_ao.handlers.span_lifecycle import ( + HandlerSpanState, + build_handler_step, + discard_handler_step, + finalize_handler_step, + release_handler_steps, +) from splunk_ao.schema.handlers import Node from splunk_ao.utils import _get_timestamp from splunk_ao.utils.openai_agents import ( @@ -25,6 +33,21 @@ _logger = logging.getLogger(__name__) +@dataclass +class _OpenAITraceState: + """Mutable lifecycle state owned by one OpenAI Agents trace.""" + + nodes: dict[str, Node] = field(default_factory=dict) + last_output: Any = None + last_status_code: int | None = None + first_input: Any = None + owned_trace: Any = None + owned_root: Any = None + owned_root_node_id: str | None = None + caller_parent: Any = None + active_steps: dict[str, HandlerSpanState] = field(default_factory=dict) + + class SplunkAOTracingProcessor(TracingProcessor): """ OpenAI Agents TracingProcessor for logging traces to Splunk AO. @@ -38,8 +61,8 @@ class SplunkAOTracingProcessor(TracingProcessor): The Splunk AO logger instance. _flush_on_trace_end : bool Whether to automatically flush the log batch to Splunk AO when a trace ends. - _nodes : dict[str, Node] - Stores Node objects keyed by their OpenAI span_id or trace_id (for root). + _trace_states : dict[str, _OpenAITraceState] + Stores independent lifecycle state keyed by OpenAI trace ID. """ def __init__(self, splunk_ao_logger: SplunkAOLogger | None = None, flush_on_trace_end: bool = False): @@ -55,122 +78,130 @@ def __init__(self, splunk_ao_logger: SplunkAOLogger | None = None, flush_on_trac """ self._splunk_ao_logger: SplunkAOLogger = splunk_ao_logger or splunk_ao_context.get_logger_instance() self._flush_on_trace_end: bool = flush_on_trace_end - self._nodes: dict[str, Node] = {} - self._last_output: Any = None - self._last_status_code: int | None = None - self._first_input: Any = None - self._owned_trace: Any = None - self._owned_root: Any = None - self._owned_root_node_id: str | None = None - self._caller_parent: Any = None - self._active_steps: dict[str, HandlerSpanState] = {} + self._trace_states: dict[str, _OpenAITraceState] = {} + self._state_lock = threading.RLock() def on_trace_start(self, trace: Trace) -> None: """Called when an OpenAI Agent trace starts.""" - node = Node( - node_type="agent", - run_id=trace.trace_id, - span_params={ - "start_time": _get_timestamp(), - "start_time_iso": datetime.now(UTC).isoformat(), - "name": trace.name, - "metadata": convert_to_string_dict(trace.metadata), - }, - ) - self._nodes[trace.trace_id] = node - self._caller_parent = self._splunk_ao_logger.current_parent() - if self._caller_parent is None: - try: - self._owned_trace = self._splunk_ao_logger.start_trace( - input="Agent Workflow", - name="Trace", - created_at=datetime.fromisoformat(node.span_params["start_time_iso"]), - ) - except Exception: - self._owned_trace = None - _logger.warning("Failed to start OpenAI Agents trace telemetry", exc_info=True) + trace_id = str(trace.trace_id) + with self._state_lock: + previous = self._trace_states.pop(trace_id, None) + if previous is not None: + release_handler_steps(self._splunk_ao_logger, previous.active_steps) + self._conclude_current_trace_on_failure(previous) + _logger.warning("Replacing unfinished OpenAI Agents trace state for %s", trace_id) + + node = Node( + node_type="agent", + run_id=cast(uuid.UUID, trace_id), + span_params={ + "start_time": _get_timestamp(), + "start_time_iso": datetime.now(UTC).isoformat(), + "name": trace.name, + "metadata": convert_to_string_dict(trace.metadata), + }, + ) + state = _OpenAITraceState(nodes={trace_id: node}) + state.caller_parent = self._splunk_ao_logger.current_parent() + self._trace_states[trace_id] = state + if state.caller_parent is None: + try: + state.owned_trace = self._splunk_ao_logger.start_trace( + input="Agent Workflow", + name="Trace", + created_at=datetime.fromisoformat(node.span_params["start_time_iso"]), + ) + except Exception: + state.owned_trace = None + _logger.warning("Failed to start OpenAI Agents trace telemetry", exc_info=True) def on_trace_end(self, trace: Trace) -> None: """Called when an OpenAI Agent trace ends.""" - node = self._nodes.get(trace.trace_id) - if not node: - _logger.warning(f"End called for unknown trace_id {trace.trace_id}") - return - - node.span_params["duration_ns"] = convert_time_delta_to_ns(_get_timestamp() - node.span_params["start_time"]) + trace_id = str(trace.trace_id) + with self._state_lock: + state = self._trace_states.get(trace_id) + if state is None: + _logger.warning("End called for unknown trace_id %s", trace_id) + return + node = state.nodes.get(trace_id) + if node is None: + _logger.warning("Root node %s is missing at trace end", trace_id) + release_handler_steps(self._splunk_ao_logger, state.active_steps) + self._trace_states.pop(trace_id, None) + return + + node.span_params["duration_ns"] = convert_time_delta_to_ns( + _get_timestamp() - node.span_params["start_time"] + ) - try: - self._commit_trace(trace) - if self._flush_on_trace_end: - self._splunk_ao_logger.flush() - except Exception: - self._conclude_current_trace_on_failure() - _logger.warning("Failed to commit OpenAI Agents telemetry", exc_info=True) - finally: - for state in self._active_steps.values(): - self._splunk_ao_logger._restore_handler_step_context(state.activation) - self._splunk_ao_logger._release_otel_context(state.step) - self._nodes = {} - self._last_output = None - self._last_status_code = None - self._first_input = None - self._owned_trace = None - self._owned_root = None - self._owned_root_node_id = None - self._caller_parent = None - self._active_steps = {} - - def _commit_trace(self, trace: Trace) -> None: - if not self._nodes: + try: + if state.active_steps: + _logger.warning( + "OpenAI Agents trace %s ended with %d unfinished callback spans", + trace_id, + len(state.active_steps), + ) + root_was_unfinished = state.owned_root_node_id in state.active_steps + release_handler_steps(self._splunk_ao_logger, state.active_steps) + if root_was_unfinished: + state.owned_root = None + state.owned_root_node_id = None + self._commit_trace(trace, state) + if self._flush_on_trace_end: + self._splunk_ao_logger.flush() + except Exception: + self._conclude_current_trace_on_failure(state) + _logger.warning("Failed to commit OpenAI Agents telemetry", exc_info=True) + finally: + release_handler_steps(self._splunk_ao_logger, state.active_steps) + self._trace_states.pop(trace_id, None) + + def _commit_trace(self, trace: Trace, state: _OpenAITraceState) -> None: + if not state.nodes: _logger.warning("No nodes to commit") return - root_node = self._nodes.get(trace.trace_id) + trace_id = str(trace.trace_id) + root_node = state.nodes.get(trace_id) if root_node is None: - _logger.warning(f"Root node {trace.trace_id} not found") + _logger.warning("Root node %s not found", trace_id) return if not getattr(self._splunk_ao_logger, "_ingestion_hook", None): - if self._active_steps: - _logger.warning( - "OpenAI Agents trace %s ended with %d unfinished callback spans", - trace.trace_id, - len(self._active_steps), - ) - if self._owned_trace is not None: - self._owned_trace.input = self._first_input or "Agent Workflow" - self._owned_trace.metrics.duration_ns = root_node.span_params.get("duration_ns") - if self._splunk_ao_logger.current_parent() is self._owned_trace: - self._splunk_ao_logger.conclude(output=self._last_output, status_code=self._last_status_code) + if state.owned_trace is not None: + state.owned_trace.input = state.first_input or "Agent Workflow" + state.owned_trace.metrics.duration_ns = root_node.span_params.get("duration_ns") + if self._splunk_ao_logger.current_parent() is state.owned_trace: + self._splunk_ao_logger.conclude(output=state.last_output, status_code=state.last_status_code) return - live_root_node = self._nodes.get(self._owned_root_node_id) if self._owned_root_node_id else None - if self._owned_trace is not None and self._owned_root is not None and live_root_node is not None: - self._owned_trace.input = self._first_input or "Agent Workflow" - self._owned_trace.metrics.duration_ns = root_node.span_params.get("duration_ns") - self._update_owned_root(live_root_node) - self._log_node_tree(live_root_node, reuse_current=True) + live_root_node = state.nodes.get(state.owned_root_node_id) if state.owned_root_node_id else None + if state.owned_trace is not None and state.owned_root is not None and live_root_node is not None: + state.owned_trace.input = state.first_input or "Agent Workflow" + state.owned_trace.metrics.duration_ns = root_node.span_params.get("duration_ns") + self._update_owned_root(live_root_node, state) + self._log_node_tree(live_root_node, state, reuse_current=True) for child_id in root_node.children: - if child_id == self._owned_root_node_id: + if child_id == state.owned_root_node_id: continue - child = self._nodes.get(child_id) + child = state.nodes.get(child_id) if child is not None: - self._log_node_tree(child) - self._splunk_ao_logger.conclude(output=self._last_output, status_code=self._last_status_code) + self._log_node_tree(child, state) + self._splunk_ao_logger.conclude(output=state.last_output, status_code=state.last_status_code) return # A caller-owned operation remains current; log framework operations # beneath it without ever concluding the caller's parent. for child_id in root_node.children: - child = self._nodes.get(child_id) + child = state.nodes.get(child_id) if child is not None: - self._log_node_tree(child) + self._log_node_tree(child, state) - if self._owned_trace is not None: - self._splunk_ao_logger.conclude(output=self._last_output, status_code=self._last_status_code) + if state.owned_trace is not None: + self._splunk_ao_logger.conclude(output=state.last_output, status_code=state.last_status_code) - def _conclude_current_trace_on_failure(self) -> None: - if self._owned_trace is None: + def _conclude_current_trace_on_failure(self, state: _OpenAITraceState) -> None: + if state.owned_trace is None: return current_parent = self._splunk_ao_logger.current_parent() @@ -180,10 +211,10 @@ def _conclude_current_trace_on_failure(self) -> None: root = current_parent while root._parent is not None: root = root._parent - if root is self._owned_trace: + if root is state.owned_trace: self._splunk_ao_logger.conclude(output="", status_code=500, conclude_all=True) - def _log_node_tree(self, node: Node, reuse_current: bool = False) -> None: + def _log_node_tree(self, node: Node, state: _OpenAITraceState, reuse_current: bool = False) -> None: """ Log a node and its children recursively. @@ -276,9 +307,9 @@ def _log_node_tree(self, node: Node, reuse_current: bool = False) -> None: # Process all child nodes last_child = None for child_id in node.children: - child_node = self._nodes.get(child_id) + child_node = state.nodes.get(child_id) if child_node: - self._log_node_tree(child_node) + self._log_node_tree(child_node, state) last_child = child_node else: _logger.warning(f"Child node {child_id} not found") @@ -296,12 +327,12 @@ def _log_node_tree(self, node: Node, reuse_current: bool = False) -> None: duration_ns=node.span_params.get("duration_ns"), status_code=status_code, ) - self._last_status_code = status_code - self._last_output = output + state.last_status_code = status_code + state.last_output = output - def _start_owned_root(self, node: Node) -> None: + def _start_owned_root(self, node: Node, state: _OpenAITraceState) -> None: """Create the first top-level real operation while framework work is active.""" - if self._owned_trace is None or self._owned_root is not None: + if state.owned_trace is None or state.owned_root is not None: return if node.node_type not in ("agent", "chain", "workflow"): return @@ -309,86 +340,104 @@ def _start_owned_root(self, node: Node) -> None: metadata = node.span_params.get("metadata") if metadata is not None: metadata = convert_to_string_dict(metadata) - self._owned_root = self._splunk_ao_logger.add_workflow_span( + state.owned_root = self._splunk_ao_logger.add_workflow_span( input=node.span_params.get("input") or node.node_type.capitalize() + " Step", name=node.span_params.get("name"), metadata=metadata, tags=node.span_params.get("tags"), created_at=datetime.fromisoformat(node.span_params["start_time_iso"]), ) - if self._owned_root is not None: - self._owned_root_node_id = str(node.run_id) + if state.owned_root is not None: + state.owned_root_node_id = str(node.run_id) - def _update_owned_root(self, node: Node) -> None: + def _update_owned_root(self, node: Node, state: _OpenAITraceState) -> None: """Apply final OpenAI callback data to the already-live root.""" - if self._owned_root is None: + if state.owned_root is None: return - self._owned_root.input = node.span_params.get("input") or node.node_type.capitalize() + " Step" + state.owned_root.input = node.span_params.get("input") or node.node_type.capitalize() + " Step" if node.span_params.get("name") is not None: - self._owned_root.name = node.span_params["name"] + state.owned_root.name = node.span_params["name"] if node.span_params.get("tags") is not None: - self._owned_root.tags = node.span_params["tags"] + state.owned_root.tags = node.span_params["tags"] metadata = node.span_params.get("metadata") if metadata is not None: - self._owned_root.user_metadata = convert_to_string_dict(metadata) + state.owned_root.user_metadata = convert_to_string_dict(metadata) - def _start_incremental_span(self, node: Node) -> None: + def _start_incremental_span(self, node: Node, state: _OpenAITraceState) -> None: """Create and activate one real OpenAI Agents callback span.""" node_id = str(node.run_id) - if node_id == self._owned_root_node_id and self._owned_root is not None: - self._active_steps[node_id] = HandlerSpanState(step=self._owned_root, activation=None) + if node_id == state.owned_root_node_id and state.owned_root is not None: + state.active_steps[node_id] = HandlerSpanState(step=state.owned_root, activation=None) return - if str(node.parent_run_id) in self._active_steps: - parent = self._active_steps[str(node.parent_run_id)].step + if str(node.parent_run_id) in state.active_steps: + parent = state.active_steps[str(node.parent_run_id)].step else: - parent = self._owned_trace or self._caller_parent + parent = state.owned_trace or state.caller_parent if parent is None: raise RuntimeError(f"No active parent is available for OpenAI Agents span {node_id}") step = build_handler_step(node, openai_agents=True) - self._splunk_ao_logger._register_handler_step(step, parent) - activation = self._splunk_ao_logger._activate_handler_step(step) - self._active_steps[node_id] = HandlerSpanState(step=step, activation=activation) + span_state = HandlerSpanState(step=step, activation=None) + state.active_steps[node_id] = span_state + try: + self._splunk_ao_logger._register_handler_step(step, parent) + span_state.activation = self._splunk_ao_logger._activate_handler_step(step) + except Exception: + discard_handler_step(self._splunk_ao_logger, state.active_steps, node_id) + raise - def _finish_incremental_span(self, node: Node) -> None: + def _finish_incremental_span(self, node: Node, state: _OpenAITraceState) -> None: """Finalize and enqueue an OpenAI Agents span at callback completion.""" node_id = str(node.run_id) - state = self._active_steps.get(node_id) - if state is None: + span_state = state.active_steps.get(node_id) + if span_state is None: _logger.warning("Unable to complete OpenAI Agents span %s: no active state", node_id) return try: if node.node_type in ("agent", "chain", "workflow") and not node.span_params.get("output"): - last_child = self._nodes.get(node.children[-1]) if node.children else None + last_child = state.nodes.get(node.children[-1]) if node.children else None if last_child is not None: node.span_params["output"] = last_child.span_params.get("output", "") - final = finalize_handler_step(node, state, openai_agents=True) - final = self._splunk_ao_logger._replace_handler_step(state.step, final) - state.step = final - if node_id == self._owned_root_node_id: - self._owned_root = final - - self._splunk_ao_logger._restore_handler_step_context(state.activation) - state.activation = None + final = finalize_handler_step(node, span_state, openai_agents=True) + final = self._splunk_ao_logger._replace_handler_step(span_state.step, final) + span_state.step = final + if node_id == state.owned_root_node_id: + state.owned_root = final + + self._splunk_ao_logger._restore_handler_step_context(span_state.activation) + span_state.activation = None if self._splunk_ao_logger.current_parent() is final: self._splunk_ao_logger._set_current_parent(final._parent) self._splunk_ao_logger._complete_handler_step(final) - self._active_steps.pop(node_id, None) - self._last_output = node.span_params.get("output") - self._last_status_code = node.span_params.get("status_code", 200) + state.active_steps.pop(node_id, None) + state.last_output = node.span_params.get("output") + state.last_status_code = node.span_params.get("status_code", 200) except Exception: - self._splunk_ao_logger._restore_handler_step_context(state.activation) - state.activation = None + discard_handler_step(self._splunk_ao_logger, state.active_steps, node_id) + if node_id == state.owned_root_node_id: + self._conclude_current_trace_on_failure(state) + state.owned_trace = None + state.owned_root = None + state.owned_root_node_id = None _logger.warning("Failed to complete OpenAI Agents span %s", node_id, exc_info=True) def on_span_start(self, span: Span[Any]) -> None: """Called when an OpenAI Agent span starts.""" - span_id = span.span_id - trace_id = span.trace_id - parent_id = span.parent_id or span.trace_id # Parent is previous span or root trace + with self._state_lock: + self._on_span_start(span) + + def _on_span_start(self, span: Span[Any]) -> None: + """Start a span while the processor state lock is held.""" + span_id = str(span.span_id) + trace_id = str(span.trace_id) + parent_id = str(span.parent_id or span.trace_id) # Parent is previous span or root trace + state = self._trace_states.get(trace_id) + if state is None: + _logger.warning("Start called for span %s with unknown trace_id %s", span_id, trace_id) + return - if span_id in self._nodes: + if span_id in state.nodes: _logger.warning(f"Span node already exists for span_id {span_id}, overwriting...") # Determine span type and name @@ -413,8 +462,8 @@ def on_span_start(self, span: Span[Any]) -> None: "status_code": llm_data.get("status_code", 200), } ) - if not self._first_input and initial_params.get("input") != serialize_to_str(None): - self._first_input = initial_params.get("input") + if not state.first_input and initial_params.get("input") != serialize_to_str(None): + state.first_input = initial_params.get("input") elif splunk_ao_type == "tool": tool_data = _extract_tool_data(span.span_data) initial_params.update( @@ -450,35 +499,58 @@ def on_span_start(self, span: Span[Any]) -> None: splunk_ao_type = "workflow" # Create the node - node = Node(node_type=splunk_ao_type, span_params=initial_params, run_id=span_id, parent_run_id=parent_id) - self._nodes[span_id] = node + node = Node( + node_type=splunk_ao_type, + span_params=initial_params, + run_id=cast(uuid.UUID, span_id), + parent_run_id=cast(uuid.UUID, parent_id), + ) + state.nodes[span_id] = node # Add to parent's children list - parent_node = self._nodes.get(parent_id) + parent_node = state.nodes.get(parent_id) if not parent_node: _logger.warning(f"Parent node {parent_id} not found for span {span_id} in trace {trace_id}") return parent_node.children.append(span_id) if parent_id == trace_id: try: - self._start_owned_root(node) + self._start_owned_root(node, state) except Exception: - self._conclude_current_trace_on_failure() - self._owned_trace = None - self._owned_root = None - self._owned_root_node_id = None + self._conclude_current_trace_on_failure(state) + release_handler_steps(self._splunk_ao_logger, state.active_steps) + state.owned_trace = None + state.owned_root = None + state.owned_root_node_id = None _logger.warning("Failed to start OpenAI Agents root telemetry", exc_info=True) + return if not getattr(self._splunk_ao_logger, "_ingestion_hook", None): try: - self._start_incremental_span(node) + self._start_incremental_span(node, state) except Exception: - self._conclude_current_trace_on_failure() + discard_handler_step(self._splunk_ao_logger, state.active_steps, span_id) + if parent_id == trace_id: + self._conclude_current_trace_on_failure(state) + release_handler_steps(self._splunk_ao_logger, state.active_steps) + state.owned_trace = None + state.owned_root = None + state.owned_root_node_id = None _logger.warning("Failed to start OpenAI Agents span telemetry for %s", span_id, exc_info=True) def on_span_end(self, span: Span[Any]) -> None: """Called when an OpenAI Agent span ends.""" - span_id = span.span_id - node = self._nodes.get(span_id) + with self._state_lock: + self._on_span_end(span) + + def _on_span_end(self, span: Span[Any]) -> None: + """End a span while the processor state lock is held.""" + span_id = str(span.span_id) + trace_id = str(span.trace_id) + state = self._trace_states.get(trace_id) + if state is None: + _logger.warning("End called for span %s with unknown trace_id %s", span_id, trace_id) + return + node = state.nodes.get(span_id) if not node: _logger.warning(f"End called for unknown span_id {span_id}") return @@ -487,10 +559,14 @@ def on_span_end(self, span: Span[Any]) -> None: # Update node with final data splunk_ao_type = node.node_type - end_params: dict[str, Any] = {"end_time_iso": span.ended_at or datetime.now(UTC).isoformat()} - end_params["duration_ns"] = convert_time_delta_to_ns( - datetime.fromisoformat(span.ended_at) - datetime.fromisoformat(node.span_params["start_time_iso"]) - ) + end_time_iso = span.ended_at or datetime.now(UTC).isoformat() + end_params: dict[str, Any] = {"end_time_iso": end_time_iso} + try: + end_params["duration_ns"] = convert_time_delta_to_ns( + datetime.fromisoformat(end_time_iso) - datetime.fromisoformat(node.span_params["start_time_iso"]) + ) + except (TypeError, ValueError): + _logger.debug("OpenAI Agents span %s supplied an invalid timestamp", span_id) if splunk_ao_type == "llm": llm_data = _extract_llm_data(span.span_data) @@ -505,11 +581,11 @@ def on_span_end(self, span: Span[Any]) -> None: if node.span_params.get("input") is None: node.span_params["input"] = llm_data.get("input") if ( - not self._first_input + not state.first_input and node.span_params["input"] and node.span_params["input"] != serialize_to_str(None) ): - self._first_input = node.span_params["input"] + state.first_input = node.span_params["input"] # Extract embedded tool calls and merge with existing tool definitions if isinstance(span.span_data, ResponseSpanData) and span.span_data.response: @@ -559,7 +635,7 @@ def on_span_end(self, span: Span[Any]) -> None: # Update the node's parameters node.span_params.update(end_params) if not getattr(self._splunk_ao_logger, "_ingestion_hook", None): - self._finish_incremental_span(node) + self._finish_incremental_span(node, state) def shutdown(self) -> None: """Called when the application stops. Flushes any remaining logs.""" diff --git a/src/splunk_ao/handlers/span_lifecycle.py b/src/splunk_ao/handlers/span_lifecycle.py index 1e32e514..502c9005 100644 --- a/src/splunk_ao/handlers/span_lifecycle.py +++ b/src/splunk_ao/handlers/span_lifecycle.py @@ -4,6 +4,7 @@ from dataclasses import dataclass from datetime import UTC, datetime +from typing import TYPE_CHECKING from uuid import UUID, uuid4 from galileo_core.schemas.logging.agent import AgentType @@ -15,6 +16,9 @@ from splunk_ao.utils.retrievers import convert_to_documents from splunk_ao.utils.serialization import convert_to_string_dict, serialize_to_str +if TYPE_CHECKING: + from splunk_ao.logger import SplunkAOLogger + @dataclass class HandlerSpanState: @@ -30,10 +34,34 @@ def _created_at(node: Node) -> datetime: return created_at start_time_iso = node.span_params.get("start_time_iso") if isinstance(start_time_iso, str) and start_time_iso: - return datetime.fromisoformat(start_time_iso) + try: + return datetime.fromisoformat(start_time_iso) + except ValueError: + pass return datetime.now(tz=UTC) +def discard_handler_step(logger: SplunkAOLogger, active_steps: dict[str, HandlerSpanState], node_id: str) -> None: + """Discard one failed callback step without disturbing its siblings or owner.""" + state = active_steps.pop(node_id, None) + if state is None: + return + logger._restore_handler_step_context(state.activation) + state.activation = None + parent = state.step._parent + if logger.current_parent() is state.step: + logger._set_current_parent(parent) + if isinstance(parent, StepWithChildSpans): + parent.spans = [child for child in parent.spans if child is not state.step] + logger._release_otel_context(state.step) + + +def release_handler_steps(logger: SplunkAOLogger, active_steps: dict[str, HandlerSpanState]) -> None: + """Release unfinished callback activations and identities in stack order.""" + for node_id in reversed(tuple(active_steps)): + discard_handler_step(logger, active_steps, node_id) + + def _metadata(node: Node) -> dict[str, str] | None: metadata = node.span_params.get("metadata") return convert_to_string_dict(metadata) if metadata is not None else None diff --git a/src/splunk_ao/http_instrumentation.py b/src/splunk_ao/http_instrumentation.py index 5ff35fd4..75371eb4 100644 --- a/src/splunk_ao/http_instrumentation.py +++ b/src/splunk_ao/http_instrumentation.py @@ -2,6 +2,9 @@ from __future__ import annotations +import logging +from dataclasses import dataclass, field +from importlib import import_module from typing import Any, cast from weakref import WeakKeyDictionary, WeakSet @@ -11,6 +14,8 @@ from splunk_ao.otel import TracerProvider, add_splunk_ao_span_processor from splunk_ao.session_context import install_session_propagator +_logger = logging.getLogger(__name__) + _INSTALL_MESSAGE = ( "Automatic distributed tracing requires optional dependencies. " "Install them with: pip install 'splunk-ao[distributed-tracing]'" @@ -23,24 +28,26 @@ _instrumented_apps: WeakKeyDictionary[Any, tuple[TracerProvider, str]] = WeakKeyDictionary() _configured_providers: WeakSet[SDKTracerProvider] = WeakSet() +_INSTRUMENTOR_IMPORTS: dict[str, tuple[str, str]] = { + "fastapi": ("opentelemetry.instrumentation.fastapi", "FastAPIInstrumentor"), + "starlette": ("opentelemetry.instrumentation.starlette", "StarletteInstrumentor"), + "requests": ("opentelemetry.instrumentation.requests", "RequestsInstrumentor"), + "httpx": ("opentelemetry.instrumentation.httpx", "HTTPXClientInstrumentor"), + "aiohttp-client": ("opentelemetry.instrumentation.aiohttp_client", "AioHttpClientInstrumentor"), +} -def _load_instrumentors() -> dict[str, type]: - try: - from opentelemetry.instrumentation.aiohttp_client import AioHttpClientInstrumentor # noqa: PLC0415 - from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor # noqa: PLC0415 - from opentelemetry.instrumentation.httpx import HTTPXClientInstrumentor # noqa: PLC0415 - from opentelemetry.instrumentation.requests import RequestsInstrumentor # noqa: PLC0415 - from opentelemetry.instrumentation.starlette import StarletteInstrumentor # noqa: PLC0415 - except ImportError as exc: - raise ImportError(_INSTALL_MESSAGE) from exc - return { - "fastapi": FastAPIInstrumentor, - "starlette": StarletteInstrumentor, - "requests": RequestsInstrumentor, - "httpx": HTTPXClientInstrumentor, - "aiohttp-client": AioHttpClientInstrumentor, - } +def _load_instrumentors(names: tuple[str, ...] | None = None) -> dict[str, type[Any]]: + """Import only the instrumentors needed for the requested components.""" + requested = tuple(_INSTRUMENTOR_IMPORTS) if names is None else names + resolved: dict[str, type[Any]] = {} + for name in requested: + module_name, attribute = _INSTRUMENTOR_IMPORTS[name] + try: + resolved[name] = getattr(import_module(module_name), attribute) + except (AttributeError, ImportError) as exc: + raise ImportError(f"{_INSTALL_MESSAGE} (missing support for {name})") from exc + return resolved def _validate_client_ownership(names: tuple[str, ...], tracer_provider: TracerProvider) -> None: @@ -56,12 +63,19 @@ def _validate_client_ownership(names: tuple[str, ...], tracer_provider: TracerPr def _resolve_app_framework(app: Any) -> str: try: - from fastapi import FastAPI # noqa: PLC0415 from starlette.applications import Starlette # noqa: PLC0415 except ImportError as exc: raise ImportError(_INSTALL_MESSAGE) from exc - if isinstance(app, FastAPI): + fastapi_type: type[Any] | None + try: + from fastapi import FastAPI # noqa: PLC0415 + + fastapi_type = FastAPI + except ImportError: + fastapi_type = None + + if fastapi_type is not None and isinstance(app, fastapi_type): framework = "fastapi" elif isinstance(app, Starlette): framework = "starlette" @@ -83,12 +97,44 @@ def _validate_app_ownership(app: Any, tracer_provider: TracerProvider) -> str: return framework -def _instrument_app(app: Any, instrumentors: dict[str, Any], tracer_provider: TracerProvider) -> None: - framework = _validate_app_ownership(app, tracer_provider) +def _instrument_app(app: Any, framework: str, instrumentors: dict[str, Any], tracer_provider: TracerProvider) -> bool: if app in _instrumented_apps: - return + return False instrumentors[framework].instrument_app(app, tracer_provider=tracer_provider) _instrumented_apps[app] = (tracer_provider, framework) + return True + + +@dataclass +class _InstrumentationChanges: + """Components installed by one facade invocation and safe to roll back.""" + + tracer_provider: TracerProvider + clients: list[tuple[str, Any]] = field(default_factory=list) + app: Any | None = None + app_framework: str | None = None + app_instrumentor: Any | None = None + + def rollback(self) -> None: + """Undo only components installed by this invocation.""" + for name, instrumentor in reversed(self.clients): + try: + instrumentor.uninstrument() + except Exception: + _logger.warning("Failed to roll back %s instrumentation", name, exc_info=True) + finally: + if _client_providers.get(name) is self.tracer_provider: + _client_providers.pop(name, None) + + if self.app is not None and self.app_framework is not None and self.app_instrumentor is not None: + try: + self.app_instrumentor.uninstrument_app(self.app) + except Exception: + _logger.warning("Failed to roll back %s application instrumentation", self.app_framework, exc_info=True) + finally: + configured_owner = _instrumented_apps.get(self.app) + if configured_owner is not None and configured_owner[0] is self.tracer_provider: + _instrumented_apps.pop(self.app, None) def _instrument_supported_transports( @@ -96,24 +142,48 @@ def _instrument_supported_transports( instrumentors: dict[str, type], tracer_provider: TracerProvider, app: Any | None, + app_framework: str | None, instrument_requests: bool, instrument_httpx: bool, instrument_aiohttp_client: bool, -) -> None: +) -> _InstrumentationChanges: clients = {"requests": instrument_requests, "httpx": instrument_httpx, "aiohttp-client": instrument_aiohttp_client} enabled_clients = tuple(name for name, enabled in clients.items() if enabled) _validate_client_ownership(enabled_clients, tracer_provider) - - if app is not None: - _instrument_app(app, instrumentors, tracer_provider) - - for name in enabled_clients: - if _client_providers.get(name) is tracer_provider: - continue - instrumentors[name]().instrument(tracer_provider=tracer_provider) - _client_providers[name] = tracer_provider - - install_session_propagator() + changes = _InstrumentationChanges(tracer_provider=tracer_provider) + try: + if app is not None and app_framework is not None: + if _instrument_app(app, app_framework, instrumentors, tracer_provider): + changes.app = app + changes.app_framework = app_framework + changes.app_instrumentor = instrumentors[app_framework] + + for name in enabled_clients: + if _client_providers.get(name) is tracer_provider: + continue + instrumentor = instrumentors[name]() + instrumentor.instrument(tracer_provider=tracer_provider) + _client_providers[name] = tracer_provider + changes.clients.append((name, instrumentor)) + except Exception: + changes.rollback() + raise + return changes + + +def _requested_instrumentors( + *, app_framework: str | None, instrument_requests: bool, instrument_httpx: bool, instrument_aiohttp_client: bool +) -> tuple[str, ...]: + requested: list[str] = [] + if app_framework is not None: + requested.append(app_framework) + if instrument_requests: + requested.append("requests") + if instrument_httpx: + requested.append("httpx") + if instrument_aiohttp_client: + requested.append("aiohttp-client") + return tuple(requested) def instrument_distributed_tracing( @@ -143,18 +213,34 @@ def instrument_distributed_tracing( Notes ----- Client instrumentation is process-wide. Configure each supported client - through one owner. The caller also owns provider shutdown. + through one owner. The caller also owns provider shutdown. This function + wraps the current process-global text-map propagator so explicit Splunk AO + sessions can use W3C baggage. Install any custom global propagator before + calling this function; replacing it afterward removes the session adapter. """ - instrumentors = _load_instrumentors() resolved_provider = tracer_provider or cast(TracerProvider, trace.get_tracer_provider()) - _instrument_supported_transports( + app_framework = _validate_app_ownership(app, resolved_provider) if app is not None else None + requested = _requested_instrumentors( + app_framework=app_framework, + instrument_requests=instrument_requests, + instrument_httpx=instrument_httpx, + instrument_aiohttp_client=instrument_aiohttp_client, + ) + instrumentors = _load_instrumentors(requested) + changes = _instrument_supported_transports( instrumentors=instrumentors, tracer_provider=resolved_provider, app=app, + app_framework=app_framework, instrument_requests=instrument_requests, instrument_httpx=instrument_httpx, instrument_aiohttp_client=instrument_aiohttp_client, ) + try: + install_session_propagator() + except Exception: + changes.rollback() + raise def configure_distributed_tracing( @@ -192,28 +278,35 @@ def configure_distributed_tracing( ----- Repeated calls with the same provider do not attach another Splunk AO span processor. Use :func:`instrument_distributed_tracing` instead when export - has already been configured separately on the provider. + has already been configured separately on the provider. This function also + wraps the current process-global text-map propagator for session baggage. + Install a custom global propagator first; replacing it afterward removes + the session adapter. """ - # Resolve optional imports before constructing a processor and its worker so - # a missing extra cannot leave partially configured background resources. - instrumentors = _load_instrumentors() resolved_provider = tracer_provider or SDKTracerProvider() - clients = {"requests": instrument_requests, "httpx": instrument_httpx, "aiohttp-client": instrument_aiohttp_client} - enabled_clients = tuple(name for name, enabled in clients.items() if enabled) - _validate_client_ownership(enabled_clients, resolved_provider) - if app is not None: - _validate_app_ownership(app, resolved_provider) - - if resolved_provider not in _configured_providers: - add_splunk_ao_span_processor(resolved_provider) - _configured_providers.add(resolved_provider) - - _instrument_supported_transports( + app_framework = _validate_app_ownership(app, resolved_provider) if app is not None else None + requested = _requested_instrumentors( + app_framework=app_framework, + instrument_requests=instrument_requests, + instrument_httpx=instrument_httpx, + instrument_aiohttp_client=instrument_aiohttp_client, + ) + instrumentors = _load_instrumentors(requested) + changes = _instrument_supported_transports( instrumentors=instrumentors, tracer_provider=resolved_provider, app=app, + app_framework=app_framework, instrument_requests=instrument_requests, instrument_httpx=instrument_httpx, instrument_aiohttp_client=instrument_aiohttp_client, ) + try: + if resolved_provider not in _configured_providers: + add_splunk_ao_span_processor(resolved_provider) + _configured_providers.add(resolved_provider) + install_session_propagator() + except Exception: + changes.rollback() + raise return resolved_provider diff --git a/src/splunk_ao/logger/logger.py b/src/splunk_ao/logger/logger.py index 03ffa004..98449cb4 100644 --- a/src/splunk_ao/logger/logger.py +++ b/src/splunk_ao/logger/logger.py @@ -75,7 +75,7 @@ SessionCreateRequest, TracesIngestRequest, ) -from splunk_ao.session_context import get_effective_session_id, set_session_context +from splunk_ao.session_context import clear_session_context, get_effective_session_id, set_session_context from splunk_ao.traces import Traces from splunk_ao.utils.decorators import async_warn_catch_exception, nop_async, nop_sync, warn_catch_exception from splunk_ao.utils.env_helpers import _get_mode_or_default @@ -2145,7 +2145,10 @@ async def _start_or_get_session_async( def _set_active_session_id(self, session_id: str | None) -> None: """Update compatibility and request-local session state together.""" self.session_id = session_id - set_session_context(session_id) + if session_id is None: + clear_session_context() + else: + set_session_context(session_id) @nop_async async def async_start_session( diff --git a/src/splunk_ao/session_context.py b/src/splunk_ao/session_context.py index 76c737ba..2e1d5e5a 100644 --- a/src/splunk_ao/session_context.py +++ b/src/splunk_ao/session_context.py @@ -13,12 +13,51 @@ GEN_AI_CONVERSATION_ID = "gen_ai.conversation.id" -_session_id_context: ContextVar[str | None] = ContextVar("session_id_context", default=None) + +class _SessionSelection: + """Marker type for execution-local session selection state.""" + + +_SESSION_UNSET = _SessionSelection() +_SESSION_CLEARED = _SessionSelection() +SessionSelection = str | _SessionSelection | None + +_session_id_context: ContextVar[SessionSelection] = ContextVar("session_id_context", default=_SESSION_UNSET) def set_session_context(session_id: str | None) -> None: - """Set the explicit session selection for the current execution context.""" - _session_id_context.set(session_id) + """Set or reset the explicit session selection for this execution context. + + Passing ``None`` retains the historical reset behavior: inbound baggage and + the logger compatibility field may be consulted again. Use + :func:`clear_session_context` when an explicit clear must mask inbound + baggage for the remainder of the execution context. + """ + _session_id_context.set(_SESSION_UNSET if session_id is None else session_id) + + +def clear_session_context() -> None: + """Explicitly suppress local, inbound, and compatibility session selection.""" + _session_id_context.set(_SESSION_CLEARED) + + +def get_session_selection() -> SessionSelection: + """Return the raw execution-local selection for nested-context restoration.""" + return _session_id_context.get() + + +def restore_session_selection(selection: SessionSelection) -> None: + """Restore a previously captured execution-local selection.""" + _session_id_context.set(selection) + + +def explicit_session_id(selection: SessionSelection) -> tuple[bool, str | None]: + """Return whether a saved selection is explicit and its session value.""" + if selection is _SESSION_CLEARED: + return True, None + if isinstance(selection, str) and selection: + return True, selection + return False, None def _baggage_session_id(context: Context | None = None) -> str | None: @@ -29,9 +68,11 @@ def _baggage_session_id(context: Context | None = None) -> str | None: def get_effective_session_id(logger_session_id: str | None = None, context: Context | None = None) -> str | None: """Resolve the execution-local, inbound, or compatibility session ID.""" - local_session_id = _session_id_context.get(None) - if local_session_id: - return local_session_id + local_selection = _session_id_context.get() + if local_selection is _SESSION_CLEARED: + return None + if isinstance(local_selection, str) and local_selection: + return local_selection baggage_session_id = _baggage_session_id(context) if baggage_session_id: return baggage_session_id @@ -87,7 +128,12 @@ def _session_propagator() -> SplunkAOSessionPropagator: def install_session_propagator() -> None: - """Install the session-aware adapter without replacing its delegate.""" + """Wrap the process-global propagator with the session-aware adapter. + + The current propagator remains the delegate. Replacing the global + propagator later also removes this adapter until this function is called + again. + """ current = propagate.get_global_textmap() if not isinstance(current, SplunkAOSessionPropagator): propagate.set_global_textmap(SplunkAOSessionPropagator(current)) diff --git a/src/splunk_ao/tracing.py b/src/splunk_ao/tracing.py index 69352b67..164f6e83 100644 --- a/src/splunk_ao/tracing.py +++ b/src/splunk_ao/tracing.py @@ -1,14 +1,39 @@ """W3C distributed-tracing utilities for Splunk AO.""" from collections.abc import Mapping, MutableMapping +from typing import Any from opentelemetry.context import Context +from opentelemetry.propagators import textmap from splunk_ao.exceptions import SplunkAOLoggerException from splunk_ao.logger.logger import _has_active_exportable_span_context from splunk_ao.session_context import extract_session_context, inject_session_context +class _CaseInsensitiveGetter(textmap.Getter[Mapping[str, str]]): + """Read all case-insensitive carrier values without collapsing duplicates.""" + + def get(self, carrier: Mapping[str, str], key: str) -> list[str] | None: + getlist = getattr(carrier, "getlist", None) + if callable(getlist): + values = getlist(key) + normalized = [str(value) for value in values] + if key.lower() == "baggage" and len(normalized) > 1: + return [",".join(normalized)] + return normalized or None + values = [str(value) for candidate, value in carrier.items() if str(candidate).lower() == key.lower()] + if key.lower() == "baggage" and len(values) > 1: + return [",".join(values)] + return values or None + + def keys(self, carrier: Mapping[str, str]) -> list[str]: + return [str(key) for key in carrier] + + +_case_insensitive_getter: textmap.Getter[Any] = _CaseInsensitiveGetter() + + def get_tracing_headers(carrier: MutableMapping[str, str] | None = None) -> MutableMapping[str, str]: """Inject the active operation's W3C trace context into a carrier. @@ -40,5 +65,4 @@ def get_tracing_headers(carrier: MutableMapping[str, str] | None = None) -> Muta def extract_tracing_context(carrier: Mapping[str, str]) -> Context: """Extract W3C trace context from an incoming text-map carrier.""" - normalized = {str(key).lower(): value for key, value in carrier.items()} - return extract_session_context(normalized) + return extract_session_context(carrier, getter=_case_insensitive_getter) diff --git a/src/splunk_ao/utils/singleton.py b/src/splunk_ao/utils/singleton.py index eb985dec..294016bc 100644 --- a/src/splunk_ao/utils/singleton.py +++ b/src/splunk_ao/utils/singleton.py @@ -215,6 +215,22 @@ def get( self._splunk_ao_loggers[key] = logger return logger + def get_existing( + self, + *, + project: str | None = None, + project_id: str | None = None, + agent_stream: str | None = None, + agent_stream_id: str | None = None, + experiment_id: str | None = None, + mode: str | None = None, + ) -> SplunkAOLogger | None: + """Return a cached standard logger without constructing one.""" + key = SplunkAOLoggerSingleton._get_key( + project, project_id, agent_stream, agent_stream_id, _get_mode_or_default(mode), experiment_id + ) + return self._splunk_ao_loggers.get(key) + def reset( self, project: str | None = None, diff --git a/tests/test_base_handler.py b/tests/test_base_handler.py index 1e53a11b..46662cb1 100644 --- a/tests/test_base_handler.py +++ b/tests/test_base_handler.py @@ -11,9 +11,12 @@ from opentelemetry.trace.status import StatusCode from splunk_ao import get_tracing_headers +from splunk_ao.exceptions import SplunkAOLoggerException from splunk_ao.handlers.base_handler import SplunkAOBaseHandler +from splunk_ao.handlers.span_lifecycle import build_handler_step from splunk_ao.logger.logger import SplunkAOLogger -from splunk_ao.schema.logged import LoggedAgentSpan +from splunk_ao.schema.handlers import Node +from splunk_ao.schema.logged import LoggedAgentSpan, LoggedLlmSpan from tests.testutils.setup import setup_mock_logstreams_client, setup_mock_projects_client, setup_mock_traces_client @@ -37,6 +40,31 @@ def operation_names(spans: list[ReadableSpan]) -> list[str | None]: return [(span.attributes or {}).get("gen_ai.operation.name") for span in spans] +def test_handler_step_uses_current_time_for_malformed_framework_timestamp() -> None: + node = Node( + node_type="tool", + run_id=uuid.uuid4(), + span_params={"input": "arguments", "name": "tool", "start_time_iso": "not-an-iso-timestamp"}, + ) + + step = build_handler_step(node) + + assert step.created_at.tzinfo is not None + + +def test_handler_step_normalizes_missing_llm_output() -> None: + node = Node( + node_type="llm", + run_id=uuid.uuid4(), + span_params={"input": "prompt", "output": None, "name": "model", "model": "test-model"}, + ) + + step = build_handler_step(node) + + assert isinstance(step, LoggedLlmSpan) + assert step.output.content == "" + + def test_normal_otel_child_enqueues_at_callback_end_without_flush() -> None: sink = RecordingSink() logger = SplunkAOLogger(project_id="project-id", agent_stream_id="stream-id", _sink=sink) @@ -220,6 +248,74 @@ def test_handler_error_span_enqueues_with_error_status_before_root_end() -> None logger.terminate() +def test_failed_child_conversion_does_not_abort_remaining_handler_trace() -> None: + sink = RecordingSink() + logger = SplunkAOLogger(project_id="project-id", agent_stream_id="stream-id", _sink=sink) + handler = SplunkAOBaseHandler(splunk_ao_logger=logger, flush_on_chain_end=False) + root_id = uuid.uuid4() + failed_child_id = uuid.uuid4() + good_child_id = uuid.uuid4() + try: + handler.start_node(node_type="chain", parent_run_id=None, run_id=root_id, name="root", input="request") + root_step = handler._active_steps[str(root_id)].step + root_span_id = logger._otel_ids[root_step.id].span_context.span_id + handler.start_node( + node_type="llm", + parent_run_id=root_id, + run_id=failed_child_id, + name="bad child", + input="prompt", + model="model", + ) + + with patch("splunk_ao.handlers.base_handler.finalize_handler_step", side_effect=ValueError("bad child")): + handler.end_node(failed_child_id, output="invalid") + + assert str(failed_child_id) not in handler._active_steps + assert str(root_id) in handler._active_steps + assert get_tracing_headers()["traceparent"].split("-")[2] == format(root_span_id, "016x") + + handler.start_node( + node_type="llm", + parent_run_id=root_id, + run_id=good_child_id, + name="good child", + input="prompt", + model="model", + ) + handler.end_node(good_child_id, output="answer") + handler.end_node(root_id, output="done") + + assert operation_names(sink.spans) == ["chat", "invoke_workflow"] + assert all(span.status.status_code is not StatusCode.ERROR for span in sink.spans) + finally: + logger.terminate() + + +def test_root_completion_releases_unfinished_child_activation_and_identity() -> None: + sink = RecordingSink() + logger = SplunkAOLogger(project_id="project-id", agent_stream_id="stream-id", _sink=sink) + handler = SplunkAOBaseHandler(splunk_ao_logger=logger, flush_on_chain_end=False) + root_id = uuid.uuid4() + child_id = uuid.uuid4() + try: + handler.start_node(node_type="chain", parent_run_id=None, run_id=root_id, name="root", input="request") + handler.start_node( + node_type="llm", parent_run_id=root_id, run_id=child_id, name="unfinished", input="prompt", model="model" + ) + unfinished = handler._active_steps[str(child_id)].step + + handler.end_node(root_id, output="done") + + assert handler._active_steps == {} + assert unfinished.id not in logger._otel_ids + assert operation_names(sink.spans) == ["invoke_workflow"] + with pytest.raises(SplunkAOLoggerException, match="active exportable operation"): + get_tracing_headers() + finally: + logger.terminate() + + class TestSplunkAOBaseHandler: @pytest.fixture @patch("splunk_ao.logger.logger.AgentStreams") @@ -348,47 +444,35 @@ def test_end_node(self, handler: SplunkAOBaseHandler, splunk_ao_logger: SplunkAO def test_commit_calls_flush(self) -> None: """Test that commit() calls flush() when flush_on_chain_end=True.""" - # Given: a mock logger - mock_logger = Mock(spec=SplunkAOLogger) - mock_logger.start_trace = Mock() - mock_logger.conclude = Mock() - mock_logger.current_parent = Mock(return_value=None) - mock_logger.add_workflow_span = Mock() - mock_logger._set_current_parent = Mock() - handler = SplunkAOBaseHandler(splunk_ao_logger=mock_logger, flush_on_chain_end=True) - - # Setup a simple trace to commit + sink = RecordingSink() + logger = SplunkAOLogger(project_id="project-id", agent_stream_id="stream-id", _sink=sink) + handler = SplunkAOBaseHandler(splunk_ao_logger=logger, flush_on_chain_end=True) run_id = uuid.uuid4() - handler.start_node(node_type="chain", parent_run_id=None, run_id=run_id, name="Test", input="test") + try: + handler.start_node(node_type="chain", parent_run_id=None, run_id=run_id, name="Test", input="test") - # When: ending the node (which triggers commit) - handler.end_node(run_id, output="result") + handler.end_node(run_id, output="result") - # Then: flush is called - mock_logger.flush.assert_called_once() + assert operation_names(sink.spans) == ["invoke_workflow"] + assert sink.force_flush_calls == 1 + finally: + logger.terminate() def test_commit_no_flush_when_disabled(self) -> None: """Test that commit() doesn't call flush or terminate when flush_on_chain_end=False.""" - # Given: a mock logger with flush disabled - mock_logger = Mock(spec=SplunkAOLogger) - mock_logger.mode = "batch" - mock_logger.start_trace = Mock() - mock_logger.conclude = Mock() - mock_logger.current_parent = Mock(return_value=None) - mock_logger.add_workflow_span = Mock() - mock_logger._set_current_parent = Mock() - handler = SplunkAOBaseHandler(splunk_ao_logger=mock_logger, flush_on_chain_end=False) - - # Setup a simple trace to commit + sink = RecordingSink() + logger = SplunkAOLogger(project_id="project-id", agent_stream_id="stream-id", _sink=sink) + handler = SplunkAOBaseHandler(splunk_ao_logger=logger, flush_on_chain_end=False) run_id = uuid.uuid4() - handler.start_node(node_type="chain", parent_run_id=None, run_id=run_id, name="Test", input="test") + try: + handler.start_node(node_type="chain", parent_run_id=None, run_id=run_id, name="Test", input="test") - # When: ending the node (which triggers commit) - handler.end_node(run_id, output="result") + handler.end_node(run_id, output="result") - # Then: neither flush nor terminate is called - mock_logger.flush.assert_not_called() - mock_logger.terminate.assert_not_called() + assert operation_names(sink.spans) == ["invoke_workflow"] + assert sink.force_flush_calls == 0 + finally: + logger.terminate() def test_commit_failure_concludes_only_handler_owned_trace( self, handler: SplunkAOBaseHandler, splunk_ao_logger: SplunkAOLogger diff --git a/tests/test_decorator.py b/tests/test_decorator.py index f95167d3..b6470d3e 100644 --- a/tests/test_decorator.py +++ b/tests/test_decorator.py @@ -11,6 +11,7 @@ from splunk_ao import Message, MessageRole, log, splunk_ao_context, start_session from splunk_ao.decorator import _session_id_context from splunk_ao.schema.content_blocks import DataContentBlock, TextContentBlock +from splunk_ao.utils.singleton import SplunkAOLoggerSingleton from tests.testutils.setup import setup_mock_logstreams_client, setup_mock_projects_client, setup_mock_traces_client @@ -1505,6 +1506,25 @@ def test_session_id_nested_context_stacking( assert splunk_ao_context.get_logger_instance().session_id is None +@patch("splunk_ao.logger.logger.AgentStreams") +@patch("splunk_ao.logger.logger.Projects") +@patch("splunk_ao.logger.logger.Traces") +def test_context_exit_does_not_construct_logger_for_restored_outer_context( + mock_traces_client: Mock, mock_projects_client: Mock, mock_logstreams_client: Mock, reset_context +) -> None: + setup_mock_traces_client(mock_traces_client) + setup_mock_projects_client(mock_projects_client) + setup_mock_logstreams_client(mock_logstreams_client) + singleton = SplunkAOLoggerSingleton() + + with patch.object(singleton, "get", wraps=singleton.get) as get_logger: + with splunk_ao_context(project="inner-project", agent_stream="inner-stream"): + pass + + get_logger.assert_called_once() + assert len(singleton.get_all_loggers()) == 1 + + @patch("splunk_ao.logger.logger.AgentStreams") @patch("splunk_ao.logger.logger.Projects") @patch("splunk_ao.logger.logger.Traces") diff --git a/tests/test_http_instrumentation.py b/tests/test_http_instrumentation.py index 15641ab4..737a8f7f 100644 --- a/tests/test_http_instrumentation.py +++ b/tests/test_http_instrumentation.py @@ -82,6 +82,8 @@ def fields(self) -> set[str]: class RecordingInstrumentor: instrument_calls: list[dict[str, Any]] app_calls: list[tuple[Any, dict[str, Any]]] + uninstrument_calls: int = 0 + uninstrument_app_calls: list[Any] | None = None def instrument(self, **kwargs: Any) -> None: self.instrument_calls.append(kwargs) @@ -89,6 +91,14 @@ def instrument(self, **kwargs: Any) -> None: def instrument_app(self, app: Any, **kwargs: Any) -> None: self.app_calls.append((app, kwargs)) + def uninstrument(self) -> None: + self.uninstrument_calls += 1 + + def uninstrument_app(self, app: Any) -> None: + if self.uninstrument_app_calls is None: + self.uninstrument_app_calls = [] + self.uninstrument_app_calls.append(app) + class RecordingSink: def __init__(self) -> None: @@ -113,6 +123,10 @@ def __new__(cls) -> RecordingInstrumentor: def instrument_app(cls, app: Any, **kwargs: Any) -> None: recorder.instrument_app(app, **kwargs) + @classmethod + def uninstrument_app(cls, app: Any) -> None: + recorder.uninstrument_app(app) + return Instrumentor @@ -273,6 +287,17 @@ def test_distributed_tracing_extra_loads_every_supported_upstream_instrumentor() assert set(_load_instrumentors()) == {"fastapi", "starlette", "requests", "httpx", "aiohttp-client"} +def test_load_instrumentors_imports_only_requested_components() -> None: + requests_instrumentor = object() + with patch("splunk_ao.http_instrumentation.import_module") as import_module: + import_module.return_value.RequestsInstrumentor = requests_instrumentor + + resolved = _load_instrumentors(("requests",)) + + assert resolved == {"requests": requests_instrumentor} + import_module.assert_called_once_with("opentelemetry.instrumentation.requests") + + @pytest.mark.parametrize(("app", "expected"), [(FastAPI(), "fastapi"), (Starlette(), "starlette")]) def test_instruments_matching_app_and_every_enabled_client_once( app: Any, expected: str, instrumentors: tuple[dict[str, type], dict[str, RecordingInstrumentor]] @@ -411,6 +436,26 @@ def test_provider_conflict_fails_before_partial_instrumentation( assert all(len(recorders[name].instrument_calls) == 1 for name in ("requests", "httpx", "aiohttp_client")) +def test_client_instrumentation_failure_rolls_back_components_installed_by_call( + instrumentors: tuple[dict[str, type], dict[str, RecordingInstrumentor]], +) -> None: + types, recorders = instrumentors + provider = FakeTracerProvider() + app = FastAPI() + + with ( + patch("splunk_ao.http_instrumentation._load_instrumentors", return_value=types), + patch.object(recorders["httpx"], "instrument", side_effect=RuntimeError("httpx failed")), + pytest.raises(RuntimeError, match="httpx failed"), + ): + instrument_distributed_tracing(tracer_provider=provider, app=app, instrument_aiohttp_client=False) + + assert recorders["requests"].uninstrument_calls == 1 + assert recorders["fastapi"].uninstrument_app_calls == [app] + assert not _client_providers + assert not _instrumented_apps + + def test_uses_current_provider_without_replacing_it( instrumentors: tuple[dict[str, type], dict[str, RecordingInstrumentor]], ) -> None: diff --git a/tests/test_langchain.py b/tests/test_langchain.py index f2bb85e9..c43f87bc 100644 --- a/tests/test_langchain.py +++ b/tests/test_langchain.py @@ -1,18 +1,20 @@ import time import uuid from collections.abc import Generator +from typing import Any from unittest.mock import MagicMock, Mock, patch from uuid import UUID, uuid4 import pytest from langchain_core.agents import AgentFinish +from langchain_core.callbacks.manager import AsyncCallbackManager from langchain_core.documents import Document from langchain_core.messages import AIMessage, HumanMessage, SystemMessage, ToolMessage from langchain_core.outputs import ChatGeneration, LLMResult from opentelemetry.sdk.trace import ReadableSpan from galileo_core.schemas.shared.document import Document as GalileoDocument -from splunk_ao import Message, MessageRole, splunk_ao_context +from splunk_ao import Message, MessageRole, get_tracing_headers, splunk_ao_context from splunk_ao.config import SplunkAOConfig from splunk_ao.handlers.langchain import SplunkAOAsyncCallback, SplunkAOCallback from splunk_ao.handlers.langchain.utils import parse_llm_result, update_root_to_agent @@ -68,6 +70,31 @@ def test_langchain_adapter_enqueues_child_at_callback_end_without_flush() -> Non logger.terminate() +@pytest.mark.asyncio +@pytest.mark.parametrize("callback_type", [SplunkAOCallback, SplunkAOAsyncCallback], ids=["sync", "async"]) +async def test_async_langchain_dispatch_keeps_handler_root_active_for_application(callback_type: Any) -> None: + 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]) + root_id = uuid.uuid4() + try: + run_manager = await manager.on_chain_start( + serialized={"name": "root"}, inputs={"query": "question"}, run_id=root_id + ) + root_step = callback._handler._active_steps[str(root_id)].step + root_context = logger._otel_ids[root_step.id].span_context + + headers = get_tracing_headers() + + assert callback.run_inline is True + assert headers["traceparent"].split("-")[2] == format(root_context.span_id, "016x") + await run_manager.on_chain_end({"answer": "done"}) + assert [(span.attributes or {}).get("gen_ai.operation.name") for span in sink.spans] == ["invoke_workflow"] + finally: + logger.terminate() + + class TestSplunkAOCallback: @pytest.fixture @patch("splunk_ao.logger.logger.AgentStreams") diff --git a/tests/test_middleware_tracing.py b/tests/test_middleware_tracing.py index e6f5680a..db7a7e6d 100644 --- a/tests/test_middleware_tracing.py +++ b/tests/test_middleware_tracing.py @@ -123,12 +123,17 @@ def test_middleware_ignores_proprietary_headers(app_factory) -> None: assert exported.parent is None -def test_middleware_detaches_context_on_exception(app_factory) -> None: - make_app, _, _ = app_factory +@pytest.mark.asyncio +async def test_middleware_detaches_context_on_exception() -> None: + middleware = TracingMiddleware(Mock()) + request = Mock(headers={"traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"}) + + async def fail_in_same_task(_: Any) -> Any: + assert trace.get_current_span().get_span_context().trace_id == 0x4BF92F3577B34DA6A3CE929D0E0E4736 + raise RuntimeError("request failed") + with pytest.raises(RuntimeError, match="request failed"): - TestClient(make_app(fail=True)).get( - "/test", headers={"traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"} - ) + await middleware.dispatch(request, fail_in_same_task) assert not trace.get_current_span().get_span_context().is_valid diff --git a/tests/test_openai_agents.py b/tests/test_openai_agents.py index 4e00e485..d7b07d84 100644 --- a/tests/test_openai_agents.py +++ b/tests/test_openai_agents.py @@ -1,7 +1,10 @@ +import asyncio import os import uuid +from typing import Any from unittest.mock import MagicMock, Mock, patch +import agents.tracing as agents_tracing import httpx import pytest import respx @@ -15,7 +18,7 @@ Runner, set_trace_processors, ) -from agents.tracing import ResponseSpanData +from agents.tracing import ResponseSpanData, agent_span, generation_span from opentelemetry.sdk.trace import ReadableSpan from pydantic import BaseModel from pytest import MonkeyPatch, mark @@ -45,6 +48,10 @@ def shutdown(self) -> None: return None +def _trace_state(processor: SplunkAOTracingProcessor, trace_id: str) -> Any: + return processor._trace_states[trace_id] + + def test_openai_agents_children_enqueue_before_trace_end() -> None: sink = RecordingSink() logger = SplunkAOLogger(project_id="project-id", agent_stream_id="stream-id", _sink=sink) @@ -75,23 +82,24 @@ def test_openai_agents_children_enqueue_before_trace_end() -> None: ) try: processor.on_trace_start(trace) - processor._nodes[str(root.run_id)] = root - processor._nodes[trace.trace_id].children.append(str(root.run_id)) - processor._start_owned_root(root) - processor._start_incremental_span(root) - processor._nodes[str(child.run_id)] = child + state = _trace_state(processor, trace.trace_id) + state.nodes[str(root.run_id)] = root + state.nodes[trace.trace_id].children.append(str(root.run_id)) + processor._start_owned_root(root, state) + processor._start_incremental_span(root, state) + state.nodes[str(child.run_id)] = child root.children.append(str(child.run_id)) - processor._start_incremental_span(child) + processor._start_incremental_span(child, state) child.span_params.update(output="answer", duration_ns=10) - processor._finish_incremental_span(child) + processor._finish_incremental_span(child, state) assert [(span.attributes or {}).get("gen_ai.operation.name") for span in sink.spans] == ["chat"] - assert str(root.run_id) in processor._active_steps + assert str(root.run_id) in state.active_steps assert sink.force_flush_calls == 0 root.span_params.update(output="answer", duration_ns=20) - processor._finish_incremental_span(root) + processor._finish_incremental_span(root, state) processor.on_trace_end(trace) assert [(span.attributes or {}).get("gen_ai.operation.name") for span in sink.spans] == [ @@ -104,6 +112,79 @@ def test_openai_agents_children_enqueue_before_trace_end() -> None: logger.terminate() +def test_openai_agents_public_tracing_lifecycle_exports_and_cleans_state() -> None: + 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) + set_trace_processors([processor]) + try: + with agents_tracing.trace("Agent trace", trace_id="trace_public") as openai_trace: + with agent_span("Research agent", span_id="span_root", parent=openai_trace) as root_span: + with generation_span( + input=[{"role": "user", "content": "question"}], + output=[{"role": "assistant", "content": "answer"}], + model="test-model", + span_id="span_child", + parent=root_span, + ): + pass + + assert [(span.attributes or {}).get("gen_ai.operation.name") for span in sink.spans] == [ + "chat", + "invoke_workflow", + ] + assert sink.spans[0].parent == sink.spans[1].context + assert processor._trace_states == {} + finally: + set_trace_processors([]) + logger.terminate() + + +@pytest.mark.asyncio +async def test_openai_agents_concurrent_public_traces_keep_state_isolated() -> None: + 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) + first_started = asyncio.Event() + second_started = asyncio.Event() + + async def run_trace(label: str, own_started: asyncio.Event, other_started: asyncio.Event) -> None: + with agents_tracing.trace(f"{label} trace", trace_id=f"trace_{label}") as openai_trace: + with agent_span(f"{label} agent", span_id=f"span_{label}_root", parent=openai_trace) as root_span: + own_started.set() + await other_started.wait() + with generation_span( + input=[{"role": "user", "content": label}], + output=[{"role": "assistant", "content": f"{label} answer"}], + model="test-model", + span_id=f"span_{label}_child", + parent=root_span, + ): + pass + + set_trace_processors([processor]) + try: + await asyncio.gather( + run_trace("first", first_started, second_started), run_trace("second", second_started, first_started) + ) + + roots = [ + span for span in sink.spans if (span.attributes or {}).get("gen_ai.operation.name") == "invoke_workflow" + ] + children = [span for span in sink.spans if (span.attributes or {}).get("gen_ai.operation.name") == "chat"] + assert len(roots) == len(children) == 2 + child_parents = { + (child.parent.trace_id, child.parent.span_id) for child in children if child.parent is not None + } + root_contexts = {(root.context.trace_id, root.context.span_id) for root in roots} + assert child_parents == root_contexts + assert len({root.context.trace_id for root in roots}) == 2 + assert processor._trace_states == {} + finally: + set_trace_processors([]) + logger.terminate() + + class HomeworkOutput(BaseModel): is_homework: bool reasoning: str @@ -228,6 +309,9 @@ def test_processor_marks_direct_trace_child_agent( logger = SplunkAOLogger(project="test", agent_stream="test") processor = SplunkAOTracingProcessor(splunk_ao_logger=logger, flush_on_trace_end=False) logger.start_trace(input="input") + trace = MagicMock(trace_id="trace-id", name="Agent trace", metadata={}) + processor.on_trace_start(trace) + state = _trace_state(processor, trace.trace_id) processor._log_node_tree( Node( @@ -239,7 +323,8 @@ def test_processor_marks_direct_trace_child_agent( "name": "Agent step", "start_time_iso": "2025-01-01T00:00:00+00:00", }, - ) + ), + state, ) assert logger.traces[0].spans[0].conversation_root is True @@ -259,15 +344,15 @@ def test_commit_failure_concludes_handler_owned_trace( processor = SplunkAOTracingProcessor(splunk_ao_logger=logger) trace = MagicMock(trace_id="trace-id", name="Agent trace", metadata={}) processor.on_trace_start(trace) - owned_trace = processor._owned_trace + state = _trace_state(processor, trace.trace_id) + owned_trace = state.owned_trace with patch.object(processor, "_commit_trace", side_effect=RuntimeError("conversion failed")): processor.on_trace_end(trace) assert logger.current_parent() is None assert owned_trace.status_code == 500 - assert processor._nodes == {} - assert processor._owned_trace is None + assert processor._trace_states == {} @patch("splunk_ao.logger.logger.AgentStreams") @@ -289,8 +374,7 @@ def test_commit_failure_preserves_caller_owned_trace( processor.on_trace_end(trace) assert logger.current_parent() is caller_trace - assert processor._nodes == {} - assert processor._owned_trace is None + assert processor._trace_states == {} logger.conclude(output="done") @@ -307,6 +391,7 @@ def test_openai_agents_live_root_is_used_for_outbound_context_and_commit( processor = SplunkAOTracingProcessor(splunk_ao_logger=logger) trace = MagicMock(trace_id="trace-id", name="Agent trace", metadata={}) processor.on_trace_start(trace) + state = _trace_state(processor, trace.trace_id) node = Node( node_type="agent", run_id="root-span-id", @@ -320,10 +405,10 @@ def test_openai_agents_live_root_is_used_for_outbound_context_and_commit( "status_code": 200, }, ) - processor._nodes[str(node.run_id)] = node - processor._nodes[trace.trace_id].children.append(str(node.run_id)) - processor._start_owned_root(node) - live_root = processor._owned_root + state.nodes[str(node.run_id)] = node + state.nodes[trace.trace_id].children.append(str(node.run_id)) + processor._start_owned_root(node, state) + live_root = state.owned_root live_context = logger._otel_ids[live_root.id].span_context headers = get_tracing_headers() diff --git a/tests/test_tracing.py b/tests/test_tracing.py index cf5d30ac..1ae5d103 100644 --- a/tests/test_tracing.py +++ b/tests/test_tracing.py @@ -1,4 +1,4 @@ -from collections.abc import Generator +from collections.abc import Generator, Iterator, Mapping import pytest from opentelemetry import baggage, context, trace @@ -26,6 +26,28 @@ def shutdown(self) -> None: return None +class RepeatedHeaders(Mapping[str, str]): + """Minimal multi-value carrier used to verify the OTel Getter contract.""" + + def __init__(self, values: list[tuple[str, str]]) -> None: + self._values = values + + def __getitem__(self, key: str) -> str: + values = self.getlist(key) + if not values: + raise KeyError(key) + return values[-1] + + def __iter__(self) -> Iterator[str]: + return iter(dict(self._values)) + + def __len__(self) -> int: + return len(dict(self._values)) + + def getlist(self, key: str) -> list[str]: + return [value for candidate, value in self._values if candidate.lower() == key.lower()] + + @pytest.fixture(autouse=True) def isolated_context() -> Generator[None, None, None]: state_token = _otel_context_state.set(None) @@ -172,6 +194,21 @@ def test_extract_tracing_context_restores_conversation_baggage() -> None: assert baggage.get_baggage("unrelated", context=extracted) == "value" +def test_extract_tracing_context_preserves_repeated_baggage_headers() -> None: + extracted = extract_tracing_context( + RepeatedHeaders( + [ + ("traceparent", "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"), + ("baggage", "gen_ai.conversation.id=conversation-456"), + ("baggage", "unrelated=value"), + ] + ) + ) + + assert baggage.get_baggage(GEN_AI_CONVERSATION_ID, context=extracted) == "conversation-456" + assert baggage.get_baggage("unrelated", context=extracted) == "value" + + @pytest.mark.parametrize("header", ["gen_ai.conversation.id=", "gen_ai.conversation.id"]) def test_extract_tracing_context_ignores_malformed_or_empty_conversation_baggage(header: str) -> None: extracted = extract_tracing_context( @@ -223,6 +260,33 @@ def test_inbound_conversation_overrides_compatibility_logger_field_for_path1() - context.detach(token) +def test_clear_session_masks_inbound_conversation_for_later_spans_and_propagation() -> None: + remote = extract_tracing_context( + { + "traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01", + "baggage": "gen_ai.conversation.id=inbound-session,unrelated=value", + } + ) + token = context.attach(remote) + logger, sink = make_logger() + try: + logger.clear_session() + logger.start_trace(input="request") + logger.add_workflow_span(input="work", name="operation") + headers = get_tracing_headers() + logger.conclude(output="work") + logger.conclude(output="done") + + [operation] = sink.spans + assert GEN_AI_CONVERSATION_ID not in (operation.attributes or {}) + extracted = extract_tracing_context(headers) + assert baggage.get_baggage(GEN_AI_CONVERSATION_ID, context=extracted) is None + assert baggage.get_baggage("unrelated", context=extracted) == "value" + finally: + logger.terminate() + context.detach(token) + + def test_extract_tracing_context_is_case_insensitive_and_preserves_tracestate() -> None: extracted = extract_tracing_context( {"TRACEPARENT": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01", "TRACESTATE": "vendor=value"} From 04e9cf3540bbe2146e91668c011323b538dcb959 Mon Sep 17 00:00:00 2001 From: pradystar Date: Thu, 27 Aug 2026 14:29:10 -0700 Subject: [PATCH 13/13] docs(changelog): restore tracing changes to unreleased --- CHANGELOG.md | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index f4ca9739..01141d70 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,13 +7,6 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] -## [0.3.0] - -### Breaking Changes - -- Local configuration directory renamed from `~/.galileo` to `~/.splunk`. The override environment variable is now - `SPLUNK_AO_HOME_DIR` (previously `GALILEO_HOME_DIR`). - ### Added - Added the `distributed-tracing` extra and high-level @@ -77,6 +70,16 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 context. - OpenAI Agents processor state is isolated per framework trace, including concurrent traces, and its public lifecycle path now has regression coverage. + +## [0.3.0] + +### Breaking Changes + +- Local configuration directory renamed from `~/.galileo` to `~/.splunk`. The override environment variable is now + `SPLUNK_AO_HOME_DIR` (previously `GALILEO_HOME_DIR`). + +### Fixed + - Agent and workflow output conversion now removes confirmed repeated input history and keeps only the last message as the terminal output; intermediate tool-call and tool-response messages are no longer included in