Skip to content

Commit 33fcc18

Browse files
refactor this PR to only have merge mode
1 parent c0d9a57 commit 33fcc18

12 files changed

Lines changed: 171 additions & 219 deletions

File tree

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

Lines changed: 6 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -95,12 +95,8 @@ async def handle_input(user_message: str) -> str:
9595
sync_shutdown_observability,
9696
)
9797
from .telemetry import (
98-
clear_control_event_sink,
9998
clear_trace_context_provider,
100-
emit_control_events,
10199
get_trace_context_from_provider,
102-
has_control_event_sink,
103-
set_control_event_sink,
104100
set_trace_context_provider,
105101
)
106102
from .tracing import (
@@ -400,6 +396,7 @@ def init(
400396
observability_enabled: bool | None = None,
401397
log_config: dict[str, Any] | None = None,
402398
policy_refresh_interval_seconds: int = 60,
399+
merge_events: bool = False,
403400
**kwargs: object
404401
) -> Agent:
405402
"""
@@ -430,6 +427,9 @@ def init(
430427
{"enabled": True, "span_start": True, "span_end": True, "control_eval": True}
431428
policy_refresh_interval_seconds: Interval for background policy refresh loop.
432429
Defaults to 60 seconds. Set to 0 to disable background refresh.
430+
merge_events: Whether to merge local and server event creation in the
431+
SDK before enqueueing through the built-in observability path.
432+
Defaults to False.
433433
**kwargs: Additional metadata to store with the agent
434434
435435
Returns:
@@ -477,7 +477,6 @@ async def handle(message: str):
477477

478478
# Re-init behavior: always stop existing loop before mutating shared agent/session globals.
479479
_stop_policy_refresh_loop()
480-
clear_control_event_sink()
481480

482481
# Configure logging if provided (do this early before any logging happens)
483482
if log_config:
@@ -501,6 +500,7 @@ async def handle(message: str):
501500
state.current_agent = next_agent
502501
state.server_url = server_url or os.getenv('AGENT_CONTROL_URL') or 'http://localhost:8000'
503502
state.api_key = api_key
503+
state.merge_events = merge_events
504504

505505
# Merge auto-discovered steps from @control() decorators with explicit steps.
506506
# Explicit steps take precedence when (type, name) collides.
@@ -648,7 +648,7 @@ def _reset_state() -> None:
648648
state.server_controls = None
649649
state.server_url = None
650650
state.api_key = None
651-
clear_control_event_sink()
651+
state.merge_events = False
652652

653653

654654
async def ashutdown() -> None:
@@ -1313,10 +1313,6 @@ async def main():
13131313
"set_trace_context_provider",
13141314
"get_trace_context_from_provider",
13151315
"clear_trace_context_provider",
1316-
"set_control_event_sink",
1317-
"has_control_event_sink",
1318-
"emit_control_events",
1319-
"clear_control_event_sink",
13201316
# Observability
13211317
"init_observability",
13221318
"add_event",

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ def __init__(self) -> None:
2424
self.server_controls: list[dict[str, Any]] | None = None
2525
self.server_url: str | None = None
2626
self.api_key: str | None = None
27+
self.merge_events: bool = False
2728

2829

2930
# Singleton state instance

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

Lines changed: 36 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,6 @@
1919
from .client import AgentControlClient
2020
from .evaluation_events import build_control_execution_events, enqueue_observability_events
2121
from .observability import is_observability_enabled
22-
from .telemetry import emit_control_events, has_control_event_sink
2322
from .tracing import get_trace_and_span_ids
2423
from .validation import ensure_agent_name
2524

@@ -117,29 +116,40 @@ def _has_applicable_prefiltered_server_controls(
117116
)
118117

119118

120-
def _is_merged_event_mode_enabled(agent_name: str) -> bool:
121-
"""Return whether SDK-side merged event emission is safe for this request.
119+
def _is_merged_event_mode_enabled(
120+
agent_name: str,
121+
client: AgentControlClient | None = None,
122+
) -> bool:
123+
"""Return whether SDK-side merged event creation is safe for this request.
122124
123-
Merged reconstruction depends on initialized SDK session state:
124-
a registered sink, an initialized agent, and a cached server-control
125-
snapshot for the same agent. If any of those prerequisites are missing,
126-
evaluation falls back to the default behavior where the server remains the
127-
final emitter for server-side events.
125+
Merged event creation is a session-scoped option that depends on
126+
initialized SDK state: an initialized agent, cached server controls for the
127+
same agent, and observability enabled so the merged batch will actually be
128+
consumed by the built-in queue/Postgres path.
128129
129130
Args:
130131
agent_name: Normalized agent name for the current request.
132+
client: Optional client used for the current request. When provided,
133+
merged mode is allowed only if it targets the active initialized
134+
session server.
131135
132136
Returns:
133137
``True`` when the current SDK session has enough state to reconstruct
134-
and emit merged events safely.
138+
and enqueue merged events safely.
135139
"""
136-
if not has_control_event_sink():
140+
if not state.merge_events or not is_observability_enabled():
137141
return False
138142

139143
current_agent = state.current_agent
140144
if current_agent is None or current_agent.agent_name != agent_name:
141145
return False
142146

147+
if client is not None and state.server_url is not None:
148+
normalized_state_server_url = state.server_url.rstrip("/")
149+
normalized_client_server_url = client.base_url.rstrip("/")
150+
if normalized_client_server_url != normalized_state_server_url:
151+
return False
152+
143153
return state.server_controls is not None
144154

145155

@@ -202,11 +212,12 @@ async def check_evaluation(
202212
) -> EvaluationResult:
203213
"""Check if agent interaction is safe through the public SDK helper.
204214
205-
This helper preserves the default server-only evaluation path, but it also
206-
participates in the optional merged-event flow when a control event sink is
207-
registered. In that mode, the SDK asks the server to skip final event
208-
ingestion, reconstructs server events from the lightweight response, and
209-
emits them through the registered sink before returning the parsed result.
215+
This helper preserves the default server-only evaluation path, but it can
216+
also participate in merged event creation when the initialized SDK session
217+
has ``merge_events`` enabled. In that mode, the SDK asks the server to skip
218+
final event ingestion, reconstructs server events from the lightweight
219+
response, and enqueues them through the existing observability pipeline
220+
before returning the parsed result.
210221
211222
Args:
212223
client: Configured AgentControl client.
@@ -218,7 +229,7 @@ async def check_evaluation(
218229
The parsed evaluation result returned by the server.
219230
"""
220231
normalized_name = ensure_agent_name(agent_name)
221-
merged_emission_enabled = _is_merged_event_mode_enabled(normalized_name)
232+
merged_emission_enabled = _is_merged_event_mode_enabled(normalized_name, client)
222233
trace_id = None
223234
span_id = None
224235
headers: dict[str, str] | None = None
@@ -257,7 +268,7 @@ async def check_evaluation(
257268
span_id,
258269
normalized_name,
259270
)
260-
emit_control_events(server_events)
271+
enqueue_observability_events(server_events)
261272

262273
return cast(EvaluationResult, EvaluationResult.from_dict(evaluation_response.model_dump()))
263274

@@ -275,11 +286,11 @@ async def check_evaluation_with_local(
275286
"""Evaluate controls with local-first execution and configurable event flow.
276287
277288
This is the main decision boundary between the two supported event
278-
ingestion styles:
289+
creation styles:
279290
- default behavior: local events are reconstructed and queued immediately in
280291
the SDK, while server-side events are still emitted by the server
281292
- merged-event behavior: local and server events are reconstructed in the
282-
SDK and emitted once through a registered sink
293+
SDK and enqueued once through the built-in observability path
283294
284295
In both cases, the evaluation result itself stays lightweight and event
285296
reconstruction happens after evaluation completes.
@@ -377,8 +388,8 @@ def _with_parse_errors(result: EvaluationResult) -> EvaluationResult:
377388
combined_errors = (result.errors or []) + parse_errors
378389
return result.model_copy(update={"errors": combined_errors})
379390

380-
merged_emission_enabled = _is_merged_event_mode_enabled(normalized_name)
381-
should_reconstruct_local_events = merged_emission_enabled or is_observability_enabled()
391+
merged_emission_enabled = _is_merged_event_mode_enabled(normalized_name, client)
392+
should_reconstruct_local_events = is_observability_enabled()
382393

383394
local_result: EvaluationResponse | None = None
384395
local_events = []
@@ -409,7 +420,7 @@ def _with_parse_errors(result: EvaluationResult) -> EvaluationResult:
409420
if not local_result.is_safe:
410421
result = _with_parse_errors(EvaluationResult.model_validate(local_result.model_dump()))
411422
if merged_emission_enabled:
412-
emit_control_events(local_events)
423+
enqueue_observability_events(local_events)
413424
return result
414425

415426
if _has_applicable_prefiltered_server_controls(server_control_payloads, request):
@@ -444,18 +455,18 @@ def _with_parse_errors(result: EvaluationResult) -> EvaluationResult:
444455
if local_result is not None:
445456
result = _with_parse_errors(_merge_results(local_result, server_result))
446457
if merged_emission_enabled:
447-
emit_control_events(local_events + server_events)
458+
enqueue_observability_events(local_events + server_events)
448459
return result
449460

450461
result = _with_parse_errors(EvaluationResult.model_validate(server_result.model_dump()))
451462
if merged_emission_enabled:
452-
emit_control_events(server_events)
463+
enqueue_observability_events(server_events)
453464
return result
454465

455466
if local_result is not None:
456467
result = _with_parse_errors(EvaluationResult.model_validate(local_result.model_dump()))
457468
if merged_emission_enabled:
458-
emit_control_events(local_events)
469+
enqueue_observability_events(local_events)
459470
return result
460471

461472
return _with_parse_errors(EvaluationResult(is_safe=True, confidence=1.0))

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

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -127,12 +127,12 @@ def build_control_execution_events(
127127
) -> list[ControlExecutionEvent]:
128128
"""Reconstruct control execution events from an evaluation response.
129129
130-
This is the shared reconstruction step used by both supported ingestion
131-
styles:
130+
This is the shared reconstruction step used by both supported event
131+
creation styles:
132132
- the default SDK observability path, where reconstructed local events are
133133
queued into the existing SDK batcher
134134
- the merged-event path, where local and server events are reconstructed in
135-
the SDK and emitted together through a registered sink
135+
the SDK and queued together through the existing SDK batcher
136136
137137
Args:
138138
response: Evaluation response containing matches, errors, and
@@ -194,8 +194,8 @@ def build_control_execution_events(
194194
def enqueue_observability_events(events: list[ControlExecutionEvent]) -> None:
195195
"""Enqueue reconstructed events through the existing SDK observability path.
196196
197-
This preserves the default SDK behavior of forwarding local events through
198-
the existing observability batcher rather than a custom merged-event sink.
197+
This preserves the built-in SDK behavior of forwarding events through the
198+
existing observability batcher.
199199
200200
Args:
201201
events: Reconstructed control execution events to enqueue.
Lines changed: 1 addition & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,4 @@
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-
has_control_event_sink,
8-
set_control_event_sink,
9-
)
1+
"""Telemetry interfaces for provider-agnostic tracing."""
102
from .trace_context import (
113
TraceContext,
124
TraceContextProvider,
@@ -16,14 +8,9 @@
168
)
179

1810
__all__ = [
19-
"ControlEventSink",
2011
"TraceContext",
2112
"TraceContextProvider",
22-
"clear_control_event_sink",
2313
"clear_trace_context_provider",
24-
"emit_control_events",
2514
"get_trace_context_from_provider",
26-
"has_control_event_sink",
27-
"set_control_event_sink",
2815
"set_trace_context_provider",
2916
]

‎sdks/python/src/agent_control/telemetry/event_sink.py‎

Lines changed: 0 additions & 65 deletions
This file was deleted.

0 commit comments

Comments
 (0)