feat: add retriever, trace and session

CHANGELOG.md

@@ -9,6 +9,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+- Agent Control bridges now preserve all supported `applies_to` values:
+  `llm_call`, `tool_call`, `retriever_call`, `trace_call`, and `session_call`.
 - `SplunkAOLogger.add_llm_span()` and `add_single_llm_span_trace()` accept
   `image_input_tokens`, `audio_input_tokens`, `audio_output_tokens`, and
   `image_output_tokens`: the image/audio share of the input and output token

openapi.yaml

@@ -15136,6 +15136,9 @@ components:
       enum:
         - llm_call
         - tool_call
+        - retriever_call
+        - trace_call
+        - session_call
       title: ControlAppliesTo
     ControlCheckStage:
       type: string

src/splunk_ao/handlers/agent_control/bridge.py

@@ -294,7 +294,7 @@ def _control_span_kwargs(event: Any) -> dict[str, Any]:
             "control_id": getattr(event, "control_id", None),
             "agent_name": getattr(event, "agent_name", None),
             "check_stage": ControlCheckStage(check_stage) if check_stage else None,
-            "applies_to": ControlAppliesTo(applies_to) if applies_to else None,
+            "applies_to": ControlAppliesTo(getattr(applies_to, "value", applies_to)) if applies_to else None,
             "evaluator_name": getattr(event, "evaluator_name", None),
             "selector_path": getattr(event, "selector_path", None),
         }

src/splunk_ao/logger/control.py

@@ -9,18 +9,29 @@
 
 from galileo_core.schemas.logging.step import Metrics
 
+
+class ControlAppliesTo(StrEnum):
+    llm_call = "llm_call"
+    tool_call = "tool_call"
+    retriever_call = "retriever_call"
+    trace_call = "trace_call"
+    session_call = "session_call"
+
+
 try:
-    from galileo_core.schemas.logging.control import ControlAppliesTo, ControlCheckStage, ControlResult
-    from galileo_core.schemas.logging.span import ControlSpan
+    from galileo_core.schemas.logging.control import ControlCheckStage, ControlResult
+    from galileo_core.schemas.logging.span import ControlSpan as _CoreControlSpan
 
     HAS_NATIVE_CONTROL_SPAN = True
+
+    class ControlSpan(_CoreControlSpan):
+        """ControlSpan whose applies_to field supports Agent Control event values."""
+
+        applies_to: ControlAppliesTo | None = Field(default=None, description="Parent execution type.")
+
 except ImportError:
     HAS_NATIVE_CONTROL_SPAN = False
 
-    class ControlAppliesTo(StrEnum):
-        llm_call = "llm_call"
-        tool_call = "tool_call"
-
     class ControlCheckStage(StrEnum):
         pre = "pre"
         post = "post"

src/splunk_ao/logger/logger.py

@@ -1841,6 +1841,7 @@ def add_control_span(
             "status_code": status_code,
             "metrics": Metrics(duration_ns=duration_ns),
             "id": span_id or uuid.uuid4(),
+            "session_id": self.session_id,
             "trace_id": trace_id,
             "parent_id": parent_id,
             "step_number": step_number,

tests/test_agent_control.py

@@ -7,6 +7,7 @@
 from splunk_ao import AgentControlTarget, AgentControlTargetUnresolvedError, get_agent_control_target
 from splunk_ao.constants import DEFAULT_AGENT_STREAM_NAME, DEFAULT_PROJECT_NAME
 from splunk_ao.decorator import splunk_ao_context
+from splunk_ao.logger.control import ControlAppliesTo
 from splunk_ao.utils.singleton import SplunkAOLoggerSingleton
 
 
@@ -42,6 +43,14 @@ def _stub_cached_loggers(monkeypatch, loggers: dict[tuple[str, ...], SimpleNames
     monkeypatch.setattr(SplunkAOLoggerSingleton, "get_all_loggers", lambda self: loggers)
 
 
+def test_control_applies_to_supports_agent_control_values() -> None:
+    # Given: the supported Agent Control parent execution types
+    applies_to = {value.value for value in ControlAppliesTo}
+
+    # Then: the SDK preserves all values without adding unsupported types
+    assert applies_to == {"llm_call", "tool_call", "retriever_call", "trace_call", "session_call"}
+
+
 def test_get_agent_control_target_uses_explicit_log_stream_id() -> None:
     # Given: an explicit Galileo log stream ID
     agent_stream_id = str(uuid4())

tests/test_agent_control_bridge.py

@@ -13,7 +13,7 @@
 
 from splunk_ao.exporter.span_transform import copy_span_for_export
 from splunk_ao.handlers.agent_control import setup_agent_control_bridge
-from splunk_ao.logger.control import ControlResult, ControlSpan
+from splunk_ao.logger.control import ControlAppliesTo, ControlResult, ControlSpan
 from splunk_ao.logger.logger import SplunkAOLogger
 from tests.testutils.setup import setup_mock_logstreams_client, setup_mock_projects_client, setup_mock_traces_client
 
@@ -320,8 +320,13 @@ def test_idle_new_logger_does_not_mask_active_logger_context(
 @patch("splunk_ao.logger.logger.AgentStreams")
 @patch("splunk_ao.logger.logger.Projects")
 @patch("splunk_ao.logger.logger.Traces")
+@pytest.mark.parametrize("applies_to", [value.value for value in ControlAppliesTo])
 def test_agent_control_event_converts_to_control_span_in_batch_mode(
-    mock_traces_client: Mock, mock_projects_client: Mock, mock_logstreams_client: Mock, fake_agent_control_modules
+    mock_traces_client: Mock,
+    mock_projects_client: Mock,
+    mock_logstreams_client: Mock,
+    fake_agent_control_modules,
+    applies_to: str,
 ) -> None:
     # Given: a batch logger with an active parent and a matching Agent Control event
     mock_traces_client_instance = setup_mock_traces_client(mock_traces_client)
@@ -330,8 +335,10 @@ def test_agent_control_event_converts_to_control_span_in_batch_mode(
     logger = SplunkAOLogger(project="my_project", agent_stream="my_log_stream")
     logger.start_trace(input="trace input")
     workflow = logger.add_workflow_span(input="workflow input", name="workflow")
+    session_id = str(uuid.uuid4())
+    logger.set_session(session_id)
     bridge = setup_agent_control_bridge(logger)
-    event = _make_event(logger)
+    event = _make_event(logger, applies_to=applies_to)
 
     # When: the bridge receives the event through the public sink contract
     result = bridge.write_events([event])
@@ -345,7 +352,10 @@ def test_agent_control_event_converts_to_control_span_in_batch_mode(
     assert isinstance(control_span, ControlSpan)
     assert control_span.id == uuid.UUID(event.control_execution_id)
     assert control_span.trace_id == logger.traces[0].id
+    assert control_span.session_id == uuid.UUID(session_id)
     assert control_span.parent_id == workflow.id
+    assert control_span.applies_to == ControlAppliesTo(applies_to)
+    assert control_span.model_dump(mode="json")["applies_to"] == applies_to
     assert control_span.name == "toxicity-guardrail"
     assert control_span.input == "selected text"
     assert control_span.output == ControlResult(action="observe", matched=True, confidence=0.91)
@@ -380,20 +390,22 @@ def test_agent_control_event_converts_to_control_span_in_batch_mode(
         "agent_control.control_name": "toxicity-guardrail",
         "agent_control.agent_name": "assistant",
         "agent_control.check_stage": "pre",
-        "agent_control.applies_to": "llm_call",
+        "agent_control.applies_to": applies_to,
         "agent_control.evaluator_name": "regex",
         "agent_control.selector_path": "input",
         "agent_control.action": "observe",
         "agent_control.matched": True,
         "agent_control.confidence": 0.91,
     }
     assert exported_attrs["gen_ai.operation.name"] == "control"
+    assert exported_attrs["gen_ai.conversation.id"] == session_id
     assert "splunk_ao.operation.name" not in exported_attrs
     assert json.loads(exported_attrs["gen_ai.input.messages"]) == [
         {"parts": [{"content": "selected text", "type": "text"}], "role": "user"}
     ]
     assert exported.context == emitted.context
     assert exported.parent == logger._otel_ids[workflow.id].span_context
+    assert exported.context.trace_id == logger._otel_ids[workflow.id].span_context.trace_id
     epoch = datetime.datetime(1970, 1, 1, tzinfo=datetime.UTC)
     elapsed = event.timestamp - epoch
     expected_start_ns = ((elapsed.days * 86_400 + elapsed.seconds) * 1_000_000_000) + elapsed.microseconds * 1_000

tests/test_attribute_mapping.py

@@ -852,15 +852,16 @@ def test_orchestration_keeps_non_json_strings_as_text_messages() -> None:
     ]
 
 
-def test_control_mapping_exports_fully_populated_backend_contract() -> None:
+@pytest.mark.parametrize("applies_to", [value.value for value in ControlAppliesTo])
+def test_control_mapping_exports_fully_populated_backend_contract(applies_to: str) -> None:
     span = LoggedControlSpan(
         name="PII Guard",
         input="question",
         output=ControlResult(action="deny", matched=True, confidence=0.97),
         control_id=42,
         agent_name="planner",
         check_stage=ControlCheckStage.pre,
-        applies_to=ControlAppliesTo.llm_call,
+        applies_to=ControlAppliesTo(applies_to),
         evaluator_name="pii-check",
         selector_path="$.input",
         tags=["agent_control", "control"],
@@ -877,7 +878,7 @@ def test_control_mapping_exports_fully_populated_backend_contract() -> None:
         "agent_control.control_name": "PII Guard",
         "agent_control.agent_name": "planner",
         "agent_control.check_stage": "pre",
-        "agent_control.applies_to": "llm_call",
+        "agent_control.applies_to": applies_to,
         "agent_control.evaluator_name": "pii-check",
         "agent_control.selector_path": "$.input",
         "agent_control.action": "deny",

tests/test_span_converter.py

@@ -165,7 +165,8 @@ def test_converter_leaves_final_export_normalization_to_the_sink() -> None:
     assert "splunk_ao.system" not in attributes
 
 
-def test_fully_populated_logged_control_span_preserves_otlp_envelope_and_attributes() -> None:
+@pytest.mark.parametrize("applies_to", [value.value for value in ControlAppliesTo])
+def test_fully_populated_logged_control_span_preserves_otlp_envelope_and_attributes(applies_to: str) -> None:
     parent_context = make_context(span_id=PARENT_SPAN_ID)
     span = LoggedControlSpan(
         name="PII Guard",
@@ -181,7 +182,7 @@ def test_fully_populated_logged_control_span_preserves_otlp_envelope_and_attribu
         control_id=42,
         agent_name="planner",
         check_stage=ControlCheckStage.pre,
-        applies_to=ControlAppliesTo.llm_call,
+        applies_to=ControlAppliesTo(applies_to),
         evaluator_name="pii-check",
         selector_path="$.input",
     )
@@ -197,6 +198,7 @@ def test_fully_populated_logged_control_span_preserves_otlp_envelope_and_attribu
     assert result.status.status_code is StatusCode.ERROR
     assert attrs["galileo.span.kind"] == "control"
     assert attrs["agent_control.control_id"] == 42
+    assert attrs["agent_control.applies_to"] == applies_to
     assert attrs["agent_control.action"] == "deny"
     assert attrs["agent_control.matched"] is False
     assert attrs["agent_control.error_message"] == "blocked"