Skip to content

Commit f1ca27c

Browse files
feat(sdk): [Enterprise Integration]: Add provider agnostic traceing (#145)
**Summary** Added a new provider-agnostic telemetry package to the AgentControl Python SDK for external trace context resolution and merged control event emission. Updated tracing to consult a registered external trace context provider before falling back to OTEL context. Exported the new telemetry APIs from the top-level agent_control package. Added focused tests to ensure provider/sink failures do not affect existing behavior. **Scope** - User-facing/API changes: - New SDK APIs: 1. set_trace_context_provider(...) 2. get_trace_context_from_provider() 3. clear_trace_context_provider() 4. set_control_event_sink(...) 5. emit_control_events(...) 6. clear_control_event_sink() - Internal changes: - Added: 1. agent_control/telemetry/trace_context.py 2. agent_control/telemetry/event_sink.py 3. agent_control/telemetry/__init__.py 4. Updated tracing.py to use external provider before OTEL fallback. 5. Added tests for provider behavior, sink behavior, and tracing precedence. **Risk and Rollout** **Risk level:** low **Rollback plan:** Revert the new telemetry package and the small tracing/export changes in the SDK. Since the change is additive and inactive unless a provider or sink is explicitly registered, rollback is straightforward. **Testing** - Added or updated automated tests - Ran make check (or explained why not) - Ran targeted SDK tests only; did not run the full project check suite. - Manually verified behavior **Checklist** - Linked issue/spec (if applicable) - Updated docs/examples for user-facing changes - Included any required follow-up tasks
1 parent b76014f commit f1ca27c

8 files changed

Lines changed: 294 additions & 17 deletions

File tree

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

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -78,11 +78,7 @@ async def handle_input(user_message: str) -> str:
7878
from ._control_registry import (
7979
clear as clear_step_registry,
8080
)
81-
82-
# Import client and operations modules
8381
from .client import AgentControlClient
84-
85-
# Import control decorator
8682
from .control_decorators import ControlSteerError, ControlViolationError, control
8783
from .evaluation import check_evaluation_with_local, evaluate_controls
8884
from .observability import (
@@ -98,8 +94,11 @@ async def handle_input(user_message: str) -> str:
9894
shutdown_observability,
9995
sync_shutdown_observability,
10096
)
101-
102-
# Import tracing and observability
97+
from .telemetry import (
98+
clear_trace_context_provider,
99+
get_trace_context_from_provider,
100+
set_trace_context_provider,
101+
)
103102
from .tracing import (
104103
get_current_span_id,
105104
get_current_trace_id,
@@ -1305,6 +1304,9 @@ async def main():
13051304
"get_current_span_id",
13061305
"with_trace",
13071306
"is_otel_available",
1307+
"set_trace_context_provider",
1308+
"get_trace_context_from_provider",
1309+
"clear_trace_context_provider",
13081310
# Observability
13091311
"init_observability",
13101312
"add_event",

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

Lines changed: 14 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
from ._state import state
2121
from .client import AgentControlClient
2222
from .observability import add_event, get_logger, is_observability_enabled
23+
from .tracing import get_trace_and_span_ids
2324
from .validation import ensure_agent_name
2425

2526
_logger = get_logger(__name__)
@@ -291,6 +292,13 @@ async def check_evaluation_with_local(
291292
httpx.HTTPError: If server request fails
292293
"""
293294
normalized_name = ensure_agent_name(agent_name)
295+
resolved_trace_id = trace_id
296+
resolved_span_id = span_id
297+
if trace_id is None or span_id is None:
298+
current_trace_id, current_span_id = get_trace_and_span_ids()
299+
resolved_trace_id = trace_id or current_trace_id
300+
resolved_span_id = span_id or current_span_id
301+
294302
# Partition controls by local flag
295303
local_controls: list[_ControlAdapter] = []
296304
parse_errors: list[ControlMatch] = []
@@ -389,8 +397,8 @@ def _with_parse_errors(result: EvaluationResult) -> EvaluationResult:
389397
local_result,
390398
request,
391399
applicable_local_controls,
392-
trace_id,
393-
span_id,
400+
resolved_trace_id,
401+
resolved_span_id,
394402
agent_name=event_agent_name,
395403
)
396404

@@ -409,10 +417,10 @@ def _with_parse_errors(result: EvaluationResult) -> EvaluationResult:
409417
if _has_applicable_prefiltered_server_controls(server_control_payloads, request):
410418
request_payload = request.model_dump(mode="json", exclude_none=True)
411419
headers: dict[str, str] = {}
412-
if trace_id:
413-
headers["X-Trace-Id"] = trace_id
414-
if span_id:
415-
headers["X-Span-Id"] = span_id
420+
if resolved_trace_id:
421+
headers["X-Trace-Id"] = resolved_trace_id
422+
if resolved_span_id:
423+
headers["X-Span-Id"] = resolved_span_id
416424

417425
response = await client.http_client.post(
418426
"/api/v1/evaluation",
Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
"""Telemetry interfaces for provider-agnostic tracing."""
2+
3+
from .trace_context import (
4+
TraceContext,
5+
TraceContextProvider,
6+
clear_trace_context_provider,
7+
get_trace_context_from_provider,
8+
set_trace_context_provider,
9+
)
10+
11+
__all__ = [
12+
"TraceContext",
13+
"TraceContextProvider",
14+
"clear_trace_context_provider",
15+
"get_trace_context_from_provider",
16+
"set_trace_context_provider",
17+
]
Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,55 @@
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+
if not trace_id or not span_id:
44+
return None
45+
46+
return {
47+
"trace_id": trace_id,
48+
"span_id": span_id,
49+
}
50+
51+
52+
def clear_trace_context_provider() -> None:
53+
"""Clear the registered trace context provider."""
54+
global _trace_context_provider
55+
_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

‎sdks/python/tests/test_observability_updates.py‎

Lines changed: 69 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,10 @@
1010
_map_applies_to,
1111
_merge_results,
1212
)
13+
from agent_control.telemetry.trace_context import (
14+
clear_trace_context_provider,
15+
set_trace_context_provider,
16+
)
1317
from agent_control_models import ControlDefinition
1418

1519
# =============================================================================
@@ -326,6 +330,9 @@ def test_fallback_warning_logged_only_once(self):
326330
class TestCheckEvaluationWithLocal:
327331
"""Tests for check_evaluation_with_local event emission and non_matches."""
328332

333+
def teardown_method(self) -> None:
334+
clear_trace_context_provider()
335+
329336
@pytest.mark.asyncio
330337
async def test_emits_events_when_trace_context_provided(self):
331338
"""Should emit observability events when trace_id and span_id are passed."""
@@ -398,7 +405,7 @@ async def test_emits_events_when_trace_context_provided(self):
398405

399406
@pytest.mark.asyncio
400407
async def test_emits_events_without_trace_context(self):
401-
"""Should still emit events when trace_id/span_id not provided (fallback IDs)."""
408+
"""Should resolve trace context from the provider when IDs are omitted."""
402409
from agent_control_models import EvaluationResponse, Step
403410

404411
mock_response = EvaluationResponse(
@@ -424,6 +431,12 @@ async def test_emits_events_without_trace_context(self):
424431
client = MagicMock()
425432
client.http_client = AsyncMock()
426433
step = Step(type="llm", name="test-step", input="hello")
434+
set_trace_context_provider(
435+
lambda: {
436+
"trace_id": "a" * 32,
437+
"span_id": "b" * 16,
438+
}
439+
)
427440

428441
with patch("agent_control.evaluation.ControlEngine", return_value=mock_engine), \
429442
patch("agent_control.evaluation.list_evaluators", return_value=["regex"]), \
@@ -438,8 +451,8 @@ async def test_emits_events_without_trace_context(self):
438451
)
439452
mock_emit.assert_called_once()
440453
call_args = mock_emit.call_args
441-
assert call_args[0][3] is None # trace_id passed as None
442-
assert call_args[0][4] is None # span_id passed as None
454+
assert call_args[0][3] == "a" * 32
455+
assert call_args[0][4] == "b" * 16
443456

444457
@pytest.mark.asyncio
445458
async def test_forwards_trace_headers_to_server(self):
@@ -492,6 +505,59 @@ async def test_forwards_trace_headers_to_server(self):
492505
assert headers["X-Trace-Id"] == "aaaa1111bbbb2222cccc3333dddd4444"
493506
assert headers["X-Span-Id"] == "eeee5555ffff6666"
494507

508+
@pytest.mark.asyncio
509+
async def test_forwards_provider_trace_headers_to_server_when_ids_omitted(self):
510+
"""Server POST should resolve trace headers from the provider when omitted."""
511+
from agent_control_models import Step
512+
513+
controls = [{
514+
"id": 1,
515+
"name": "server-ctrl",
516+
"control": {
517+
"condition": {
518+
"evaluator": {"name": "regex", "config": {"pattern": "test"}},
519+
"selector": {"path": "input"},
520+
},
521+
"action": {"decision": "deny"},
522+
"execution": "server",
523+
},
524+
}]
525+
526+
mock_http_response = MagicMock()
527+
mock_http_response.json.return_value = {
528+
"is_safe": True,
529+
"confidence": 1.0,
530+
"matches": None,
531+
"errors": None,
532+
"non_matches": None,
533+
}
534+
mock_http_response.raise_for_status = MagicMock()
535+
536+
client = MagicMock()
537+
client.http_client = AsyncMock()
538+
client.http_client.post = AsyncMock(return_value=mock_http_response)
539+
step = Step(type="llm", name="test-step", input="hello")
540+
set_trace_context_provider(
541+
lambda: {
542+
"trace_id": "c" * 32,
543+
"span_id": "d" * 16,
544+
}
545+
)
546+
547+
with patch("agent_control.evaluation.list_evaluators", return_value=["regex"]):
548+
await evaluation.check_evaluation_with_local(
549+
client=client,
550+
agent_name="agent-000000000001",
551+
step=step,
552+
stage="pre",
553+
controls=controls,
554+
)
555+
556+
call_kwargs = client.http_client.post.call_args
557+
headers = call_kwargs.kwargs.get("headers", {})
558+
assert headers["X-Trace-Id"] == "c" * 32
559+
assert headers["X-Span-Id"] == "d" * 16
560+
495561

496562
# =============================================================================
497563
# control_decorators non_matches dict conversion
Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,65 @@
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_returns_none_when_provider_returns_none() -> None:
33+
set_trace_context_provider(lambda: None)
34+
35+
assert get_trace_context_from_provider() is None
36+
37+
38+
def test_get_trace_context_from_provider_swallows_provider_failures() -> None:
39+
def _raising_provider():
40+
raise RuntimeError("boom")
41+
42+
set_trace_context_provider(_raising_provider)
43+
44+
assert get_trace_context_from_provider() is None
45+
46+
47+
def test_get_trace_context_from_provider_returns_none_for_invalid_shape() -> None:
48+
set_trace_context_provider( # type: ignore[arg-type]
49+
lambda: {
50+
"trace_id": "a" * 32,
51+
}
52+
)
53+
54+
assert get_trace_context_from_provider() is None
55+
56+
57+
def test_get_trace_context_from_provider_returns_none_for_empty_ids() -> None:
58+
set_trace_context_provider(
59+
lambda: {
60+
"trace_id": "",
61+
"span_id": "",
62+
}
63+
)
64+
65+
assert get_trace_context_from_provider() is None

0 commit comments

Comments
 (0)