diff --git a/AGENTS.md b/AGENTS.md index 1a20ae00..73aebaa4 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -152,7 +152,9 @@ 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. - `splunk_ao_context.simulation_run()` scopes a run ID for logger-created and native OpenTelemetry spans. Capture it at - span start as the direct `splunk_ao.simulation_run.id` span attribute. Keep it out of Resource routing. + span start as the direct `splunk_ao.simulation.run_id` span attribute. Propagate it in W3C baggage for supported + cross-service calls. Nested `splunk_ao_context(...)` scopes start without the outer run ID unless explicitly set and + restore the outer ID on exit; `init()` clears it. Keep it out of Resource routing. - Explicit session selection is ambient within the current thread or async execution context and is shared by logger instances in that context. An explicit clear masks inbound baggage in that context. Independent simultaneous sessions require separate execution contexts. diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index 4b0413ec..f24d2f41 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -101,9 +101,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. +Automatic setup wraps the current process-global text-map propagator with a delegate for session and simulation-run +baggage. Applications must install custom global propagators before invoking the helper. Replacing the global +propagator afterward removes the SDK 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 @@ -112,8 +113,9 @@ shared configuration/API-client boundary, never by excluding deployment URLs glo 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. +baggage member. An active simulation run propagates as `splunk_ao.simulation.run_id` baggage and becomes the direct +`splunk_ao.simulation.run_id` span attribute on exported spans. 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 @@ -121,8 +123,11 @@ propagation. An explicit clear masks inbound conversation baggage for the rest o simultaneous sessions require separate execution contexts. `splunk_ao_context.simulation_run(run_id)` scopes a run ID to the current execution context. Logger-created and native -OpenTelemetry spans capture it when each span starts. They export it as the direct `splunk_ao.simulation_run.id` span -attribute. This context does not change Project or Agent Stream routing, Resource attributes, or experiment routing. +OpenTelemetry spans capture it when each span starts. They export it as the direct `splunk_ao.simulation.run_id` span +attribute. Supported W3C propagation carries the run ID across service boundaries. Nested `splunk_ao_context(...)` +scopes start without the outer run ID unless an ID is explicitly supplied, then restore the outer ID on exit. Calling +`splunk_ao_context.init()` clears the active run ID. This context does not change Project or Agent Stream routing, +Resource attributes, or experiment routing. 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. diff --git a/CHANGELOG.md b/CHANGELOG.md index da943b8f..8356d587 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,8 +10,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added - Added `splunk_ao_context.simulation_run()` to scope a run ID across separate - conversations. Logger-created and native OpenTelemetry spans capture it as - the direct `splunk_ao.simulation_run.id` span attribute without changing routing. + conversations. The ID can also be selected with + `splunk_ao_context(simulation_run_id=...)`. Logger-created and native + OpenTelemetry spans capture it as the direct `splunk_ao.simulation.run_id` span attribute. Supported W3C baggage + carries the run ID across service boundaries without changing routing. Nested + routing contexts start a fresh run-ID scope and restore the outer ID on exit; + `init()` clears the active run ID. ## [0.4.0] diff --git a/README.md b/README.md index fdf23975..49e4d987 100644 --- a/README.md +++ b/README.md @@ -394,12 +394,12 @@ 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. +propagator. This lets explicit Splunk AO sessions and simulation run IDs travel +as W3C baggage. If the application uses a custom propagator, install it before +calling the Splunk AO helper. Replacing the global propagator afterward removes +the SDK adapter. Call the helper again after that replacement if session or run +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 @@ -413,9 +413,11 @@ 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 -remain local and are never added to baggage by the SDK. +`gen_ai.conversation.id` W3C baggage member. An active simulation run also +propagates as `splunk_ao.simulation.run_id` baggage. The SDK adds that run ID to +each exported span as the `splunk_ao.simulation.run_id` span attribute. Project, +Agent Stream, 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 @@ -426,7 +428,11 @@ separate execution contexts for independent sessions. Use `splunk_ao_context.simulation_run()` to add one run ID to spans across separate conversations. The SDK captures the value at span start as the direct -`splunk_ao.simulation_run.id` span attribute. This scope does not change routing. +`splunk_ao.simulation.run_id` span attribute. Supported W3C propagation carries +the value to downstream services. A nested `splunk_ao_context(...)` starts a +fresh run-ID scope unless it supplies `simulation_run_id`; exiting restores the +outer run ID. Calling `splunk_ao_context.init()` clears the active run ID. The +run ID does not change routing, experiment selection, or Resource attributes. ```python from splunk_ao import splunk_ao_context @@ -435,6 +441,13 @@ with splunk_ao_context.simulation_run("simulation-run-123"): run_conversation() ``` +You can also set the ID on a routing context: + +```python +with splunk_ao_context(project="my_project", agent_stream="my_stream", simulation_run_id="simulation-run-123"): + run_conversation() +``` + 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 diff --git a/src/splunk_ao/converter/attribute_mapping.py b/src/splunk_ao/converter/attribute_mapping.py index 39b5dde9..77aa2efb 100644 --- a/src/splunk_ao/converter/attribute_mapping.py +++ b/src/splunk_ao/converter/attribute_mapping.py @@ -16,6 +16,7 @@ SPLUNK_AO_SYSTEM = "splunk_ao.system" SPLUNK_AO_SYSTEM_VALUE = "splunk_ao_python" +SPLUNK_AO_SIMULATION_RUN_ID = "splunk_ao.simulation.run_id" CONTENT_ALIAS_BY_GEN_AI: Mapping[str, str] = { "gen_ai.input.messages": "splunk_ao.input.messages", diff --git a/src/splunk_ao/converter/span_converter.py b/src/splunk_ao/converter/span_converter.py index 6314809b..33da32a7 100644 --- a/src/splunk_ao/converter/span_converter.py +++ b/src/splunk_ao/converter/span_converter.py @@ -11,7 +11,8 @@ from opentelemetry.trace import SpanContext, SpanKind, Status, StatusCode from galileo_core.schemas.logging.step import BaseStep, StepType -from splunk_ao.converter.attribute_mapping import build_span_attributes +from splunk_ao.converter.attribute_mapping import SPLUNK_AO_SIMULATION_RUN_ID, build_span_attributes +from splunk_ao.simulation_run_context import validate_simulation_run_id from splunk_ao.utils.headers_data import get_package_version _EPOCH = datetime(1970, 1, 1, tzinfo=UTC) @@ -92,7 +93,8 @@ def convert_span( start_time_ns = _to_unix_ns(span.created_at) attributes = build_span_attributes(span, session_id) if simulation_run_id is not None: - attributes["splunk_ao.simulation_run.id"] = simulation_run_id + simulation_run_id = validate_simulation_run_id(simulation_run_id) + attributes[SPLUNK_AO_SIMULATION_RUN_ID] = simulation_run_id return ReadableSpan( name=_span_name(span, step_type), diff --git a/src/splunk_ao/decorator.py b/src/splunk_ao/decorator.py index 82266194..3c13deb4 100644 --- a/src/splunk_ao/decorator.py +++ b/src/splunk_ao/decorator.py @@ -74,7 +74,7 @@ def call_llm(prompt, temperature=0.7): set_session_context, ) from splunk_ao.shared.exceptions import ConfigurationError -from splunk_ao.simulation_run_context import _simulation_run_id_context +from splunk_ao.simulation_run_context import _simulation_run_id_context, validate_simulation_run_id from splunk_ao.utils import _get_timestamp from splunk_ao.utils.env_helpers import _get_mode_or_default from splunk_ao.utils.serialization import EventSerializer, convert_time_delta_to_ns, serialize_to_str @@ -249,6 +249,9 @@ def __call__( SplunkAODecorator The decorator instance for use in a with statement """ + if simulation_run_id is not None: + simulation_run_id = validate_simulation_run_id(simulation_run_id) + # Push current values onto the stacks _get_or_init_list(_project_stack).append(_project_context.get()) _get_or_init_list(_agent_stream_stack).append(_agent_stream_context.get()) @@ -1479,6 +1482,7 @@ def set_session(self, session_id: str) -> None: @contextmanager def simulation_run(self, simulation_run_id: str) -> Generator[None, None, None]: """Bind emitted spans to one simulation run without changing logger routing.""" + simulation_run_id = validate_simulation_run_id(simulation_run_id) token = _simulation_run_id_context.set(simulation_run_id) try: yield diff --git a/src/splunk_ao/http_instrumentation.py b/src/splunk_ao/http_instrumentation.py index 75371eb4..332dadf6 100644 --- a/src/splunk_ao/http_instrumentation.py +++ b/src/splunk_ao/http_instrumentation.py @@ -215,8 +215,9 @@ def instrument_distributed_tracing( Client instrumentation is process-wide. Configure each supported client 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. + sessions and simulation runs can use W3C baggage. Install custom global + propagators before calling this function. Replacing one later removes the + SDK adapter. """ resolved_provider = tracer_provider or cast(TracerProvider, trace.get_tracer_provider()) app_framework = _validate_app_ownership(app, resolved_provider) if app is not None else None @@ -279,9 +280,9 @@ 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. 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. + wraps the current process-global text-map propagator for session and run + baggage. Install a custom global propagator first. Replacing it later + removes the SDK adapter. """ resolved_provider = tracer_provider or SDKTracerProvider() app_framework = _validate_app_ownership(app, resolved_provider) if app is not None else None diff --git a/src/splunk_ao/otel.py b/src/splunk_ao/otel.py index 880e6d49..522e194f 100644 --- a/src/splunk_ao/otel.py +++ b/src/splunk_ao/otel.py @@ -17,6 +17,7 @@ from galileo_core.schemas.logging.span import Span as SplunkAOSpan from splunk_ao.config import SplunkAOConfig from splunk_ao.converter import build_span_attributes +from splunk_ao.converter.attribute_mapping import SPLUNK_AO_SIMULATION_RUN_ID from splunk_ao.decorator import ( _agent_stream_context, _dataset_input_context, @@ -245,12 +246,12 @@ 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 = get_effective_session_id(context=parent_context) - simulation_run_id = get_simulation_run_id() + simulation_run_id = get_simulation_run_id(context=parent_context) if session_id: span.set_attribute("gen_ai.conversation.id", session_id) - if simulation_run_id: - span.set_attribute("splunk_ao.simulation_run.id", simulation_run_id) + if simulation_run_id is not None: + span.set_attribute(SPLUNK_AO_SIMULATION_RUN_ID, simulation_run_id) # Set dataset attributes for ground truth/reference output support _apply_dataset_attributes( diff --git a/src/splunk_ao/session_context.py b/src/splunk_ao/session_context.py index 2e1d5e5a..98bd89f2 100644 --- a/src/splunk_ao/session_context.py +++ b/src/splunk_ao/session_context.py @@ -1,4 +1,4 @@ -"""Execution-context-local session selection and W3C baggage propagation.""" +"""Execution-context-local session selection and Splunk AO baggage propagation.""" from __future__ import annotations @@ -11,6 +11,9 @@ from opentelemetry.context import Context from opentelemetry.propagators import textmap +from splunk_ao.converter.attribute_mapping import SPLUNK_AO_SIMULATION_RUN_ID +from splunk_ao.simulation_run_context import get_baggage_simulation_run_id, get_simulation_run_id + GEN_AI_CONVERSATION_ID = "gen_ai.conversation.id" @@ -87,8 +90,16 @@ def _context_with_conversation_id(context: Context, session_id: str | None) -> C return normalized +def _context_with_simulation_run_id(context: Context, simulation_run_id: str | None) -> Context: + """Return a context with the simulation run baggage normalized or removed.""" + normalized = baggage.remove_baggage(SPLUNK_AO_SIMULATION_RUN_ID, context=context) + if simulation_run_id: + normalized = baggage.set_baggage(SPLUNK_AO_SIMULATION_RUN_ID, simulation_run_id, context=normalized) + return normalized + + class SplunkAOSessionPropagator(textmap.TextMapPropagator): - """Preserve the configured propagator while normalizing session baggage.""" + """Preserve the configured propagator while normalizing session and run baggage.""" def __init__(self, delegate: textmap.TextMapPropagator) -> None: self._delegate = delegate @@ -102,7 +113,8 @@ def extract( ) -> 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)) + extracted = _context_with_conversation_id(extracted, _baggage_session_id(extracted)) + return _context_with_simulation_run_id(extracted, get_baggage_simulation_run_id(extracted)) def inject( self, @@ -112,6 +124,7 @@ def inject( ) -> 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)) + normalized = _context_with_simulation_run_id(normalized, get_simulation_run_id(context=active)) self._delegate.inject(carrier, context=normalized, setter=setter) self._baggage.inject(carrier, context=normalized, setter=setter) @@ -128,7 +141,7 @@ def _session_propagator() -> SplunkAOSessionPropagator: def install_session_propagator() -> None: - """Wrap the process-global propagator with the session-aware adapter. + """Wrap the process-global propagator with the session and run adapter. The current propagator remains the delegate. Replacing the global propagator later also removes this adapter until this function is called @@ -142,12 +155,12 @@ def install_session_propagator() -> None: 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.""" + """Inject configured trace context and normalized session and run 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.""" + """Extract configured trace context and normalize session and run baggage.""" return _session_propagator().extract(carrier, context=context, getter=getter) diff --git a/src/splunk_ao/simulation_run_context.py b/src/splunk_ao/simulation_run_context.py index 67515d08..9296d402 100644 --- a/src/splunk_ao/simulation_run_context.py +++ b/src/splunk_ao/simulation_run_context.py @@ -2,9 +2,28 @@ from contextvars import ContextVar +from opentelemetry import baggage +from opentelemetry.context import Context + +from splunk_ao.converter.attribute_mapping import SPLUNK_AO_SIMULATION_RUN_ID + _simulation_run_id_context: ContextVar[str | None] = ContextVar("simulation_run_id_context", default=None) -def get_simulation_run_id() -> str | None: - """Return the simulation run ID active in the current execution context.""" - return _simulation_run_id_context.get() +def validate_simulation_run_id(simulation_run_id: object) -> str: + """Validate a public simulation run ID and return it unchanged.""" + if not isinstance(simulation_run_id, str) or not simulation_run_id.strip(): + raise ValueError("simulation_run_id must be a non-empty string") + return simulation_run_id + + +def get_baggage_simulation_run_id(context: Context | None = None) -> str | None: + """Return a non-empty simulation run ID from W3C baggage.""" + value = baggage.get_baggage(SPLUNK_AO_SIMULATION_RUN_ID, context=context) + return value.strip() if isinstance(value, str) and value.strip() else None + + +def get_simulation_run_id(context: Context | None = None) -> str | None: + """Resolve the local simulation run ID, then the inbound baggage value.""" + local_value = _simulation_run_id_context.get() + return local_value if local_value is not None else get_baggage_simulation_run_id(context) diff --git a/tests/conftest.py b/tests/conftest.py index 3f046681..78cb28e2 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, _session_id_context +from splunk_ao.decorator import _mode_context, _session_id_context, _simulation_run_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 @@ -159,9 +159,11 @@ def _clear_otel_test_context() -> None: def reset_otel_test_context() -> Generator[None, None, None]: _clear_otel_test_context() _session_id_context.set(None) + _simulation_run_id_context.set(None) yield _clear_otel_test_context() _session_id_context.set(None) + _simulation_run_id_context.set(None) @pytest.fixture(autouse=True) diff --git a/tests/test_http_instrumentation.py b/tests/test_http_instrumentation.py index 737a8f7f..6f10e030 100644 --- a/tests/test_http_instrumentation.py +++ b/tests/test_http_instrumentation.py @@ -31,7 +31,8 @@ from starlette.testclient import TestClient from yarl import URL -from splunk_ao import configure_distributed_tracing, instrument_distributed_tracing +from splunk_ao import configure_distributed_tracing, instrument_distributed_tracing, splunk_ao_context +from splunk_ao.converter.attribute_mapping import SPLUNK_AO_SIMULATION_RUN_ID from splunk_ao.http_instrumentation import ( _client_providers, _configured_providers, @@ -473,7 +474,7 @@ def test_uses_current_provider_without_replacing_it( assert recorders["requests"].instrument_calls == [{"tracer_provider": provider}] -def test_installed_global_propagator_injects_only_conversation_session( +def test_installed_global_propagator_injects_session_and_simulation_run_baggage( instrumentors: tuple[dict[str, type], dict[str, RecordingInstrumentor]], ) -> None: types, _ = instrumentors @@ -494,12 +495,14 @@ def test_installed_global_propagator_injects_only_conversation_session( instrument_aiohttp_client=False, ) carrier: dict[str, str] = {} - propagate.inject(carrier) + with splunk_ao_context.simulation_run("run-123"): + propagate.inject(carrier) finally: context.detach(token) extracted = propagate.extract(carrier) assert baggage.get_baggage(GEN_AI_CONVERSATION_ID, context=extracted) == "conversation-123" + assert baggage.get_baggage(SPLUNK_AO_SIMULATION_RUN_ID, context=extracted) == "run-123" assert "splunk_ao.session.id" not in carrier.get("baggage", "") diff --git a/tests/test_logger_otel_egress.py b/tests/test_logger_otel_egress.py index 91a63610..eb59070a 100644 --- a/tests/test_logger_otel_egress.py +++ b/tests/test_logger_otel_egress.py @@ -7,6 +7,7 @@ from opentelemetry import context, propagate, trace from opentelemetry.sdk.trace import ReadableSpan +from splunk_ao.converter.attribute_mapping import SPLUNK_AO_SIMULATION_RUN_ID from splunk_ao.deployment import DeploymentMode from splunk_ao.exceptions import SplunkAOLoggerException from splunk_ao.exporter.span_transform import copy_span_for_export @@ -109,7 +110,7 @@ def test_simulation_run_id_is_captured_when_stable_span_identity_is_assigned( _simulation_run_id_context.reset(token) [workflow_span] = recording_sink.spans - assert workflow_span.attributes["splunk_ao.simulation_run.id"] == "run-at-start" + assert workflow_span.attributes[SPLUNK_AO_SIMULATION_RUN_ID] == "run-at-start" assert "splunk_ao.experiment.id" not in workflow_span.attributes diff --git a/tests/test_otel_native_paths.py b/tests/test_otel_native_paths.py index 422faed7..01949d88 100644 --- a/tests/test_otel_native_paths.py +++ b/tests/test_otel_native_paths.py @@ -14,6 +14,7 @@ from galileo_core.schemas.logging.span import WorkflowSpan from splunk_ao import get_tracing_headers +from splunk_ao.converter.attribute_mapping import SPLUNK_AO_SIMULATION_RUN_ID from splunk_ao.decorator import ( _agent_stream_context, _dataset_input_context, @@ -364,7 +365,7 @@ def test_processor_does_not_put_routing_on_span_attributes() -> None: calls = {args[0]: args[1] for args, _ in span.set_attribute.call_args_list} assert not ROUTING_KEYS.intersection(calls) assert calls["gen_ai.conversation.id"] == "session-id" - assert calls["splunk_ao.simulation_run.id"] == "run-id" + assert calls[SPLUNK_AO_SIMULATION_RUN_ID] == "run-id" assert "splunk_ao.experiment.id" not in calls assert calls["splunk_ao.dataset.input"] == "question" processor.shutdown() @@ -374,16 +375,31 @@ 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()) + parent_context = baggage.set_baggage(SPLUNK_AO_SIMULATION_RUN_ID, "inbound-run", context=parent_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 calls[SPLUNK_AO_SIMULATION_RUN_ID] == "inbound-run" assert "splunk_ao.session.id" not in calls processor.shutdown() +def test_processor_ignores_empty_inbound_simulation_run_baggage() -> None: + exporter = RecordingExporter() + processor = SplunkAOSpanProcessor(SpanProcessor=RecordingSpanProcessor, _exporter=exporter) + parent_context = baggage.set_baggage(SPLUNK_AO_SIMULATION_RUN_ID, "", 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 SPLUNK_AO_SIMULATION_RUN_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_span_converter.py b/tests/test_span_converter.py index ccf99f0d..ab825e36 100644 --- a/tests/test_span_converter.py +++ b/tests/test_span_converter.py @@ -1,7 +1,7 @@ from __future__ import annotations from datetime import UTC, datetime -from typing import Any +from typing import Any, cast import pytest from opentelemetry.sdk.resources import Resource @@ -13,7 +13,7 @@ from galileo_core.schemas.logging.trace import Trace from galileo_core.schemas.shared.document import Document from splunk_ao.converter import SpanConverter, span_converter -from splunk_ao.converter.attribute_mapping import build_span_attributes +from splunk_ao.converter.attribute_mapping import SPLUNK_AO_SIMULATION_RUN_ID, build_span_attributes from splunk_ao.logger.control import ControlAppliesTo, ControlCheckStage, ControlResult, ControlSpan from splunk_ao.schema.logged import LoggedControlSpan from splunk_ao.utils.headers_data import get_package_version @@ -167,15 +167,28 @@ def test_converter_leaves_final_export_normalization_to_the_sink() -> None: assert "splunk_ao.system" not in attributes +def test_simulation_run_attribute_name_is_stable() -> None: + assert SPLUNK_AO_SIMULATION_RUN_ID == "splunk_ao.simulation.run_id" + + +@pytest.mark.parametrize("simulation_run_id", ["", " ", 123]) +def test_converter_rejects_invalid_simulation_run_ids(simulation_run_id: object) -> None: + with pytest.raises(ValueError, match="simulation_run_id must be a non-empty string"): + convert( + WorkflowSpan(input="work", output="done", name="workflow", created_at=CREATED_AT), + simulation_run_id=cast(str, simulation_run_id), + ) + + @pytest.mark.parametrize(("span", "expected_name", "expected_kind"), supported_spans()) def test_simulation_run_id_is_a_direct_attribute_for_every_span_type( span: BaseStep, expected_name: str, expected_kind: SpanKind ) -> None: result = convert(span, simulation_run_id="run-1") - assert (result.attributes or {})["splunk_ao.simulation_run.id"] == "run-1" + assert (result.attributes or {})[SPLUNK_AO_SIMULATION_RUN_ID] == "run-1" assert (result.attributes or {})["gen_ai.conversation.id"] == "session-1" - assert "splunk_ao.simulation_run.id" not in result.resource.attributes + assert SPLUNK_AO_SIMULATION_RUN_ID not in result.resource.attributes assert "splunk_ao.experiment.id" not in (result.attributes or {}) diff --git a/tests/test_splunk_ao_context.py b/tests/test_splunk_ao_context.py index b14320dd..e0d70da7 100644 --- a/tests/test_splunk_ao_context.py +++ b/tests/test_splunk_ao_context.py @@ -1,3 +1,4 @@ +from typing import cast from unittest.mock import Mock, patch import pytest @@ -37,17 +38,6 @@ def test_nested_context_restoration( assert _experiment_id_context.get() is None assert _simulation_run_id_context.get() is None - -def test_simulation_run_scope_restores_the_previous_value(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr(splunk_ao_context, "get_logger_instance", Mock()) - outer_token = _simulation_run_id_context.set("outer-run") - try: - with splunk_ao_context.simulation_run("inner-run"): - assert _simulation_run_id_context.get() == "inner-run" - assert _simulation_run_id_context.get() == "outer-run" - finally: - _simulation_run_id_context.reset(outer_token) - # First level context with splunk_ao_context(project="project1", agent_stream="log_stream1", simulation_run_id="run-1"): assert _project_context.get() == "project1" @@ -55,39 +45,69 @@ def test_simulation_run_scope_restores_the_previous_value(monkeypatch: pytest.Mo assert _experiment_id_context.get() is None assert _simulation_run_id_context.get() == "run-1" - # Second level context + # A nested routing context starts without the outer simulation run. with splunk_ao_context(project="project2", agent_stream="log_stream2"): assert _project_context.get() == "project2" assert _agent_stream_context.get() == "log_stream2" assert _experiment_id_context.get() is None assert _simulation_run_id_context.get() is None - # Third level context with splunk_ao_context(project="project3", agent_stream="log_stream3", simulation_run_id="run-3"): assert _project_context.get() == "project3" assert _agent_stream_context.get() == "log_stream3" assert _experiment_id_context.get() is None assert _simulation_run_id_context.get() == "run-3" - # After exiting third level, should be back to second level assert _project_context.get() == "project2" assert _agent_stream_context.get() == "log_stream2" assert _experiment_id_context.get() is None assert _simulation_run_id_context.get() is None - # After exiting second level, should be back to first level assert _project_context.get() == "project1" assert _agent_stream_context.get() == "log_stream1" assert _experiment_id_context.get() is None assert _simulation_run_id_context.get() == "run-1" - # After exiting first level, should be back to None assert _project_context.get() is None assert _agent_stream_context.get() is None assert _experiment_id_context.get() is None assert _simulation_run_id_context.get() is None +def test_simulation_run_scope_restores_the_previous_value() -> None: + outer_token = _simulation_run_id_context.set("outer-run") + try: + with splunk_ao_context.simulation_run("inner-run"): + assert _simulation_run_id_context.get() == "inner-run" + assert _simulation_run_id_context.get() == "outer-run" + finally: + _simulation_run_id_context.reset(outer_token) + + +def test_init_clears_active_simulation_run(monkeypatch: pytest.MonkeyPatch) -> None: + logger = Mock() + singleton = Mock() + singleton.get.return_value = logger + monkeypatch.setattr("splunk_ao.decorator.SplunkAOLoggerSingleton", lambda: singleton) + + with splunk_ao_context.simulation_run("active-run"): + splunk_ao_context.init(project="project", agent_stream="stream") + assert _simulation_run_id_context.get() is None + + +@pytest.mark.parametrize("simulation_run_id", [None, "", " \t", 123]) +def test_simulation_run_rejects_invalid_ids(simulation_run_id: object) -> None: + with pytest.raises(ValueError, match="simulation_run_id must be a non-empty string"): + with splunk_ao_context.simulation_run(cast(str, simulation_run_id)): + pytest.fail("invalid simulation run ID entered the scope") + + if simulation_run_id is not None: + with pytest.raises(ValueError, match="simulation_run_id must be a non-empty string"): + splunk_ao_context(simulation_run_id=cast(str, simulation_run_id)) + + assert _simulation_run_id_context.get() is None + + @patch("splunk_ao.logger.logger.AgentStreams") @patch("splunk_ao.logger.logger.Projects") @patch("splunk_ao.logger.logger.Traces") diff --git a/tests/test_tracing.py b/tests/test_tracing.py index 1ae5d103..99685474 100644 --- a/tests/test_tracing.py +++ b/tests/test_tracing.py @@ -5,11 +5,13 @@ 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 import extract_tracing_context, get_tracing_headers, splunk_ao_context +from splunk_ao.converter.attribute_mapping import SPLUNK_AO_SIMULATION_RUN_ID 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 +from splunk_ao.simulation_run_context import _simulation_run_id_context class RecordingSink: @@ -52,11 +54,13 @@ def getlist(self, key: str) -> list[str]: def isolated_context() -> Generator[None, None, None]: state_token = _otel_context_state.set(None) session_token = _session_id_context.set(None) + run_token = _simulation_run_id_context.set(None) token = context.attach(context.Context()) try: yield finally: context.detach(token) + _simulation_run_id_context.reset(run_token) _session_id_context.reset(session_token) _otel_context_state.reset(state_token) @@ -158,6 +162,25 @@ def test_get_tracing_headers_injects_only_standard_conversation_baggage() -> Non logger.terminate() +def test_get_tracing_headers_injects_simulation_run_baggage() -> None: + logger, sink = make_logger() + try: + with splunk_ao_context.simulation_run("run-123"): + 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") + + extracted = extract_tracing_context(headers) + assert baggage.get_baggage(SPLUNK_AO_SIMULATION_RUN_ID, context=extracted) == "run-123" + assert any( + (span.attributes or {}).get(SPLUNK_AO_SIMULATION_RUN_ID) == "run-123" for span in sink.spans + ) + finally: + 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() @@ -194,6 +217,18 @@ def test_extract_tracing_context_restores_conversation_baggage() -> None: assert baggage.get_baggage("unrelated", context=extracted) == "value" +def test_extract_tracing_context_restores_simulation_run_baggage() -> None: + extracted = extract_tracing_context( + { + "traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01", + "baggage": f"{SPLUNK_AO_SIMULATION_RUN_ID}=run-456,unrelated=value", + } + ) + + assert baggage.get_baggage(SPLUNK_AO_SIMULATION_RUN_ID, context=extracted) == "run-456" + assert baggage.get_baggage("unrelated", context=extracted) == "value" + + def test_extract_tracing_context_preserves_repeated_baggage_headers() -> None: extracted = extract_tracing_context( RepeatedHeaders(