Skip to content

Commit 58eb627

Browse files
Add provider agnostic traceing
1 parent b364cab commit 58eb627

8 files changed

Lines changed: 300 additions & 2 deletions

File tree

‎sdks/python/src/agent_control/__init__.py‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -107,6 +107,14 @@ async def handle_input(user_message: str) -> str:
107107
is_otel_available,
108108
with_trace,
109109
)
110+
from .telemetry import (
111+
clear_control_event_sink,
112+
clear_trace_context_provider,
113+
emit_control_events,
114+
get_trace_context_from_provider,
115+
set_control_event_sink,
116+
set_trace_context_provider,
117+
)
110118
from .validation import ensure_agent_name
111119

112120
# Module logger
@@ -1305,6 +1313,12 @@ async def main():
13051313
"get_current_span_id",
13061314
"with_trace",
13071315
"is_otel_available",
1316+
"set_trace_context_provider",
1317+
"get_trace_context_from_provider",
1318+
"clear_trace_context_provider",
1319+
"set_control_event_sink",
1320+
"emit_control_events",
1321+
"clear_control_event_sink",
13081322
# Observability
13091323
"init_observability",
13101324
"add_event",
Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
1+
"""Telemetry interfaces for provider-agnostic tracing and event emission."""
2+
3+
from .event_sink import (
4+
ControlEventSink,
5+
clear_control_event_sink,
6+
emit_control_events,
7+
set_control_event_sink,
8+
)
9+
from .trace_context import (
10+
TraceContext,
11+
TraceContextProvider,
12+
clear_trace_context_provider,
13+
get_trace_context_from_provider,
14+
set_trace_context_provider,
15+
)
16+
17+
__all__ = [
18+
"ControlEventSink",
19+
"TraceContext",
20+
"TraceContextProvider",
21+
"clear_control_event_sink",
22+
"clear_trace_context_provider",
23+
"emit_control_events",
24+
"get_trace_context_from_provider",
25+
"set_control_event_sink",
26+
"set_trace_context_provider",
27+
]
Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,33 @@
1+
"""Provider-agnostic sink for merged control execution events."""
2+
3+
from collections.abc import Callable
4+
5+
from agent_control_models import ControlExecutionEvent
6+
7+
ControlEventSink = Callable[[list[ControlExecutionEvent]], None]
8+
9+
_control_event_sink: ControlEventSink | None = None
10+
11+
12+
def set_control_event_sink(sink: ControlEventSink | None) -> None:
13+
"""Register a sink for merged control execution events."""
14+
global _control_event_sink
15+
_control_event_sink = sink
16+
17+
18+
def emit_control_events(events: list[ControlExecutionEvent]) -> None:
19+
"""Emit merged control execution events to the registered sink."""
20+
if not events or _control_event_sink is None:
21+
return
22+
23+
try:
24+
_control_event_sink(events)
25+
except Exception:
26+
# Sink failures should not break control evaluation.
27+
pass
28+
29+
30+
def clear_control_event_sink() -> None:
31+
"""Clear the registered control event sink."""
32+
global _control_event_sink
33+
_control_event_sink = None
Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
1+
"""Provider-agnostic trace context interface for external tracing systems."""
2+
3+
from collections.abc import Callable
4+
from typing import TypedDict
5+
6+
7+
class TraceContext(TypedDict):
8+
"""Resolved trace context for a control evaluation."""
9+
10+
trace_id: str
11+
span_id: str
12+
13+
14+
TraceContextProvider = Callable[[], TraceContext | None]
15+
16+
_trace_context_provider: TraceContextProvider | None = None
17+
18+
19+
def set_trace_context_provider(provider: TraceContextProvider | None) -> None:
20+
"""Register a provider that returns the current trace context."""
21+
global _trace_context_provider
22+
_trace_context_provider = provider
23+
24+
25+
def get_trace_context_from_provider() -> TraceContext | None:
26+
"""Return trace context from the registered provider, if any."""
27+
if _trace_context_provider is None:
28+
return None
29+
30+
try:
31+
trace_context = _trace_context_provider()
32+
except Exception:
33+
# Provider failures should not break control evaluation.
34+
return None
35+
36+
if trace_context is None:
37+
return None
38+
39+
trace_id = trace_context.get("trace_id")
40+
span_id = trace_context.get("span_id")
41+
if not isinstance(trace_id, str) or not isinstance(span_id, str):
42+
return None
43+
44+
return {
45+
"trace_id": trace_id,
46+
"span_id": span_id,
47+
}
48+
49+
50+
def clear_trace_context_provider() -> None:
51+
"""Clear the registered trace context provider."""
52+
global _trace_context_provider
53+
_trace_context_provider = None

‎sdks/python/src/agent_control/tracing.py‎

Lines changed: 20 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,8 @@
3131
from contextlib import contextmanager
3232
from contextvars import ContextVar, Token
3333

34+
from .telemetry.trace_context import get_trace_context_from_provider
35+
3436
# Context variables for trace/span propagation
3537
_trace_id_var: ContextVar[str | None] = ContextVar("trace_id", default=None)
3638
_span_id_var: ContextVar[str | None] = ContextVar("span_id", default=None)
@@ -94,8 +96,9 @@ def get_trace_and_span_ids() -> tuple[str, str]:
9496
9597
Priority:
9698
1. Context variable (set by with_trace or explicitly)
97-
2. OpenTelemetry context (if OTEL is installed and active)
98-
3. Generate new OTEL-compatible IDs
99+
2. External provider
100+
3. OpenTelemetry context (if OTEL is installed and active)
101+
4. Generate new OTEL-compatible IDs
99102
100103
Returns:
101104
Tuple of (trace_id, span_id) - both are hex strings
@@ -114,6 +117,11 @@ def get_trace_and_span_ids() -> tuple[str, str]:
114117
if trace_id is not None and span_id is not None:
115118
return trace_id, span_id
116119

120+
# Try external provider
121+
trace_context = get_trace_context_from_provider()
122+
if trace_context:
123+
return trace_context["trace_id"], trace_context["span_id"]
124+
117125
# Try OpenTelemetry context
118126
otel_trace_id, otel_span_id = _get_otel_ids()
119127

@@ -136,6 +144,11 @@ def get_current_trace_id() -> str | None:
136144
if trace_id is not None:
137145
return trace_id
138146

147+
# Try external provider
148+
trace_context = get_trace_context_from_provider()
149+
if trace_context:
150+
return trace_context["trace_id"]
151+
139152
# Try OpenTelemetry
140153
otel_trace_id, _ = _get_otel_ids()
141154
return otel_trace_id
@@ -153,6 +166,11 @@ def get_current_span_id() -> str | None:
153166
if span_id is not None:
154167
return span_id
155168

169+
# Try external provider
170+
trace_context = get_trace_context_from_provider()
171+
if trace_context:
172+
return trace_context["span_id"]
173+
156174
# Try OpenTelemetry
157175
_, otel_span_id = _get_otel_ids()
158176
return otel_span_id
Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
1+
"""Tests for the telemetry merged control event sink interface."""
2+
3+
from datetime import UTC, datetime
4+
5+
from agent_control.telemetry.event_sink import (
6+
clear_control_event_sink,
7+
emit_control_events,
8+
set_control_event_sink,
9+
)
10+
from agent_control_models import ControlExecutionEvent
11+
12+
13+
def _event() -> ControlExecutionEvent:
14+
return ControlExecutionEvent(
15+
control_execution_id="ce-1",
16+
trace_id="a" * 32,
17+
span_id="b" * 16,
18+
agent_name="test-agent",
19+
control_id=1,
20+
control_name="pii_check",
21+
check_stage="pre",
22+
applies_to="llm_call",
23+
action="allow",
24+
matched=False,
25+
confidence=0.95,
26+
timestamp=datetime.now(UTC),
27+
metadata={},
28+
)
29+
30+
31+
def teardown_function() -> None:
32+
clear_control_event_sink()
33+
34+
35+
def test_emit_control_events_calls_registered_sink() -> None:
36+
seen: list[list[ControlExecutionEvent]] = []
37+
38+
def _sink(events: list[ControlExecutionEvent]) -> None:
39+
seen.append(events)
40+
41+
event = _event()
42+
set_control_event_sink(_sink)
43+
44+
emit_control_events([event])
45+
46+
assert seen == [[event]]
47+
48+
49+
def test_emit_control_events_noops_without_sink() -> None:
50+
emit_control_events([_event()])
51+
52+
53+
def test_emit_control_events_swallows_sink_failures() -> None:
54+
def _sink(_events: list[ControlExecutionEvent]) -> None:
55+
raise RuntimeError("boom")
56+
57+
set_control_event_sink(_sink)
58+
59+
emit_control_events([_event()])
Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,48 @@
1+
"""Tests for the telemetry trace context provider interface."""
2+
3+
from agent_control.telemetry.trace_context import (
4+
clear_trace_context_provider,
5+
get_trace_context_from_provider,
6+
set_trace_context_provider,
7+
)
8+
9+
10+
def teardown_function() -> None:
11+
clear_trace_context_provider()
12+
13+
14+
def test_get_trace_context_from_provider_returns_registered_context() -> None:
15+
set_trace_context_provider(
16+
lambda: {
17+
"trace_id": "a" * 32,
18+
"span_id": "b" * 16,
19+
}
20+
)
21+
22+
assert get_trace_context_from_provider() == {
23+
"trace_id": "a" * 32,
24+
"span_id": "b" * 16,
25+
}
26+
27+
28+
def test_get_trace_context_from_provider_returns_none_when_unset() -> None:
29+
assert get_trace_context_from_provider() is None
30+
31+
32+
def test_get_trace_context_from_provider_swallows_provider_failures() -> None:
33+
def _raising_provider():
34+
raise RuntimeError("boom")
35+
36+
set_trace_context_provider(_raising_provider)
37+
38+
assert get_trace_context_from_provider() is None
39+
40+
41+
def test_get_trace_context_from_provider_returns_none_for_invalid_shape() -> None:
42+
set_trace_context_provider( # type: ignore[arg-type]
43+
lambda: {
44+
"trace_id": "a" * 32,
45+
}
46+
)
47+
48+
assert get_trace_context_from_provider() is None

‎sdks/python/tests/test_tracing.py‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
import pytest
44

5+
from agent_control.telemetry.trace_context import clear_trace_context_provider, set_trace_context_provider
56
from agent_control.tracing import (
67
_generate_span_id,
78
_generate_trace_id,
@@ -17,6 +18,10 @@
1718
)
1819

1920

21+
def teardown_function() -> None:
22+
clear_trace_context_provider()
23+
24+
2025
class TestIdGeneration:
2126
"""Tests for trace and span ID generation."""
2227

@@ -132,6 +137,30 @@ def test_get_current_ids_without_context(self):
132137
assert trace_id is None or isinstance(trace_id, str)
133138
assert span_id is None or isinstance(span_id, str)
134139

140+
def test_get_current_trace_id_uses_provider(self):
141+
"""Test that get_current_trace_id uses external provider before OTEL fallback."""
142+
expected_trace = "a" * 32
143+
set_trace_context_provider(
144+
lambda: {
145+
"trace_id": expected_trace,
146+
"span_id": "b" * 16,
147+
}
148+
)
149+
150+
assert get_current_trace_id() == expected_trace
151+
152+
def test_get_current_span_id_uses_provider(self):
153+
"""Test that get_current_span_id uses external provider before OTEL fallback."""
154+
expected_span = "b" * 16
155+
set_trace_context_provider(
156+
lambda: {
157+
"trace_id": "a" * 32,
158+
"span_id": expected_span,
159+
}
160+
)
161+
162+
assert get_current_span_id() == expected_span
163+
135164

136165
class TestWithTraceContextManager:
137166
"""Tests for the with_trace context manager."""
@@ -237,6 +266,23 @@ def test_get_trace_and_span_ids_uses_context(self):
237266
assert trace_id == expected_trace
238267
assert span_id == expected_span
239268

269+
def test_get_trace_and_span_ids_uses_provider_before_otel(self):
270+
"""Test that an external provider is checked before OTEL fallback."""
271+
expected_trace = "c" * 32
272+
expected_span = "d" * 16
273+
274+
set_trace_context_provider(
275+
lambda: {
276+
"trace_id": expected_trace,
277+
"span_id": expected_span,
278+
}
279+
)
280+
281+
trace_id, span_id = get_trace_and_span_ids()
282+
283+
assert trace_id == expected_trace
284+
assert span_id == expected_span
285+
240286

241287
class TestOtelAvailability:
242288
"""Tests for OpenTelemetry availability detection."""

0 commit comments

Comments
 (0)