From 25252111ebc907fbb04f11abf2b915e117cb5065 Mon Sep 17 00:00:00 2001 From: Aditya Mehra Date: Mon, 12 Oct 2026 14:14:24 -0700 Subject: [PATCH 1/2] Address review comments --- src/splunk_ao/converter/attribute_mapping.py | 2 +- src/splunk_ao/converter/span_converter.py | 4 ++-- src/splunk_ao/otel.py | 3 +-- src/splunk_ao/session_context.py | 7 +++++-- src/splunk_ao/simulation_run_context.py | 6 +++--- tests/test_splunk_ao_context.py | 8 ++++++++ 6 files changed, 20 insertions(+), 10 deletions(-) diff --git a/src/splunk_ao/converter/attribute_mapping.py b/src/splunk_ao/converter/attribute_mapping.py index 77aa2efb..11cdc056 100644 --- a/src/splunk_ao/converter/attribute_mapping.py +++ b/src/splunk_ao/converter/attribute_mapping.py @@ -13,10 +13,10 @@ from galileo_core.schemas.logging.span import AgentSpan, LlmSpan, RetrieverSpan, ToolSpan, WorkflowSpan from galileo_core.schemas.logging.step import BaseStep, StepType from splunk_ao.logger.control import ControlSpan +from splunk_ao.simulation_run_context import SPLUNK_AO_SIMULATION_RUN_ID 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 33da32a7..6b748975 100644 --- a/src/splunk_ao/converter/span_converter.py +++ b/src/splunk_ao/converter/span_converter.py @@ -11,8 +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 SPLUNK_AO_SIMULATION_RUN_ID, build_span_attributes -from splunk_ao.simulation_run_context import validate_simulation_run_id +from splunk_ao.converter.attribute_mapping import build_span_attributes +from splunk_ao.simulation_run_context import SPLUNK_AO_SIMULATION_RUN_ID, validate_simulation_run_id from splunk_ao.utils.headers_data import get_package_version _EPOCH = datetime(1970, 1, 1, tzinfo=UTC) diff --git a/src/splunk_ao/otel.py b/src/splunk_ao/otel.py index 522e194f..d89f42bf 100644 --- a/src/splunk_ao/otel.py +++ b/src/splunk_ao/otel.py @@ -17,7 +17,6 @@ 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, @@ -36,7 +35,7 @@ ) from splunk_ao.exporter.diagnostics import get_export_health from splunk_ao.session_context import get_effective_session_id -from splunk_ao.simulation_run_context import get_simulation_run_id +from splunk_ao.simulation_run_context import SPLUNK_AO_SIMULATION_RUN_ID, get_simulation_run_id logger = logging.getLogger(__name__) diff --git a/src/splunk_ao/session_context.py b/src/splunk_ao/session_context.py index 98bd89f2..6717cf6f 100644 --- a/src/splunk_ao/session_context.py +++ b/src/splunk_ao/session_context.py @@ -11,8 +11,11 @@ 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 +from splunk_ao.simulation_run_context import ( + SPLUNK_AO_SIMULATION_RUN_ID, + get_baggage_simulation_run_id, + get_simulation_run_id, +) GEN_AI_CONVERSATION_ID = "gen_ai.conversation.id" diff --git a/src/splunk_ao/simulation_run_context.py b/src/splunk_ao/simulation_run_context.py index 9296d402..09735e5b 100644 --- a/src/splunk_ao/simulation_run_context.py +++ b/src/splunk_ao/simulation_run_context.py @@ -5,16 +5,16 @@ from opentelemetry import baggage from opentelemetry.context import Context -from splunk_ao.converter.attribute_mapping import SPLUNK_AO_SIMULATION_RUN_ID +SPLUNK_AO_SIMULATION_RUN_ID = "splunk_ao.simulation.run_id" _simulation_run_id_context: ContextVar[str | None] = ContextVar("simulation_run_id_context", default=None) def validate_simulation_run_id(simulation_run_id: object) -> str: - """Validate a public simulation run ID and return it unchanged.""" + """Validate a public simulation run ID and return its normalized value.""" 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 + return simulation_run_id.strip() def get_baggage_simulation_run_id(context: Context | None = None) -> str | None: diff --git a/tests/test_splunk_ao_context.py b/tests/test_splunk_ao_context.py index e0d70da7..c7372fbc 100644 --- a/tests/test_splunk_ao_context.py +++ b/tests/test_splunk_ao_context.py @@ -84,6 +84,14 @@ def test_simulation_run_scope_restores_the_previous_value() -> None: _simulation_run_id_context.reset(outer_token) +def test_simulation_run_id_is_normalized_at_public_context_boundaries() -> None: + with splunk_ao_context.simulation_run(" run-1\n"): + assert _simulation_run_id_context.get() == "run-1" + + with splunk_ao_context(simulation_run_id="\trun-2 "): + assert _simulation_run_id_context.get() == "run-2" + + def test_init_clears_active_simulation_run(monkeypatch: pytest.MonkeyPatch) -> None: logger = Mock() singleton = Mock() From ee20cc38df5a1717c46024cdd151501407c8cd4f Mon Sep 17 00:00:00 2001 From: Aditya Mehra Date: Mon, 12 Oct 2026 14:54:02 -0700 Subject: [PATCH 2/2] test(telemetry): mock context logger lookup --- tests/test_splunk_ao_context.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/test_splunk_ao_context.py b/tests/test_splunk_ao_context.py index c7372fbc..14af4177 100644 --- a/tests/test_splunk_ao_context.py +++ b/tests/test_splunk_ao_context.py @@ -84,7 +84,9 @@ def test_simulation_run_scope_restores_the_previous_value() -> None: _simulation_run_id_context.reset(outer_token) -def test_simulation_run_id_is_normalized_at_public_context_boundaries() -> None: +def test_simulation_run_id_is_normalized_at_public_context_boundaries(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(splunk_ao_context, "get_logger_instance", Mock()) + with splunk_ao_context.simulation_run(" run-1\n"): assert _simulation_run_id_context.get() == "run-1"