Skip to content

Commit c532b45

Browse files
prepare-sync: OTelObserver sync core + ContextVar publish
Step 3 of PR-C.3. Renames ``_handle_started`` → ``_open_started_span`` and bakes idempotency into it: a short-circuit at the top returns early if a span already exists for the event's ``_StackKey``. That covers the common case where ``prepare_sync`` opened the span synchronously in the engine task and the async ``__call__`` later re-fires for the same event — the second call becomes a true no-op rather than opening a duplicate span. Observer-attached-late and test paths that bypass ``prepare_sync`` still get the span opened via ``__call__``'s fall-through. Adds the public ``prepare_sync(event)`` method. Routing-gated (only ``"started"``-phase non-LLM events qualify), it calls ``_open_started_span`` and then publishes the just-opened span via ``_set_active_observer_span``. The engine's ``innermost`` reads the ContextVar in step 4 to attach the span into the OTel context so logs emitted from inside the node body — even on the first line, before any ``await`` — pick up the right trace_id/span_id via the OTel ``LoggingHandler``. The Token returned by ``_set_active_observer_span`` is discarded on purpose: last-writer-wins is the documented contract — the next ``prepare_sync`` call overwrites, and the task-local context dies with the invocation task.
1 parent 4fb96cb commit c532b45

1 file changed

Lines changed: 62 additions & 6 deletions

File tree

src/openarmature/observability/otel/observer.py

Lines changed: 62 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -243,16 +243,67 @@ async def __call__(self, event: NodeEvent) -> None:
243243
self._emit_checkpoint_save_span(event)
244244
return
245245
if event.phase == "started":
246-
self._handle_started(event)
246+
# Idempotent — short-circuits inside ``_open_started_span``
247+
# if ``prepare_sync`` already opened the span synchronously
248+
# in the engine task. Falls through (and opens the span)
249+
# for observers attached after the engine started, or for
250+
# test paths that bypass ``prepare_sync``.
251+
self._open_started_span(event)
247252
elif event.phase == "completed":
248253
self._handle_completed(event)
249254

250255
# ------------------------------------------------------------------
251256
# Started / completed pairing
252257
# ------------------------------------------------------------------
253258

254-
def _handle_started(self, event: NodeEvent) -> None:
255-
"""Open a span for this attempt, push onto the in-flight map."""
259+
def prepare_sync(self, event: NodeEvent) -> None:
260+
"""Synchronous engine-task entry point: open the span for this
261+
attempt AND publish it via ``current_active_observer_span`` so
262+
the engine's ``innermost`` can attach it into the OTel context
263+
before the node body runs.
264+
265+
Called by ``_dispatch`` BEFORE ``queue.put_nowait`` for
266+
``"started"``-phase events. The async ``__call__`` later sees
267+
the span already in ``inv_state.open_spans`` and short-circuits.
268+
269+
Skipped for non-``"started"`` phases and for the LLM sentinel
270+
namespace — only graph-node started events participate in the
271+
engine-side attach. Errors don't leak: ``_dispatch`` wraps this
272+
call in try/except + ``warnings.warn`` matching the async path.
273+
"""
274+
if event.phase != "started" or event.namespace == _LLM_NAMESPACE:
275+
return
276+
from openarmature.observability.correlation import (
277+
_set_active_observer_span,
278+
current_invocation_id,
279+
)
280+
281+
self._open_started_span(event)
282+
invocation_id = current_invocation_id()
283+
if invocation_id is None:
284+
return
285+
inv_state = self._inv_states.get(invocation_id)
286+
if inv_state is None:
287+
return
288+
open_span = inv_state.open_spans.get(self._key_for(event))
289+
if open_span is None:
290+
return
291+
# Publish the span to the engine via the ContextVar. Discard
292+
# the Token — last-writer-wins is the documented contract
293+
# (next ``prepare_sync`` overwrites; task-local context dies
294+
# with the invocation task).
295+
_set_active_observer_span(open_span.span)
296+
297+
def _open_started_span(self, event: NodeEvent) -> None:
298+
"""Sync core: create the span + mutate ``inv_state.open_spans``.
299+
300+
Idempotent — short-circuits if a span already exists for this
301+
event's ``_StackKey``. That covers the common case where
302+
``prepare_sync`` opened the span synchronously in the engine
303+
task and the async ``__call__`` later re-fires for the same
304+
event; the second call becomes a true no-op rather than
305+
opening a duplicate span.
306+
"""
256307
from openarmature.observability.correlation import (
257308
current_correlation_id,
258309
current_invocation_id,
@@ -261,8 +312,13 @@ def _handle_started(self, event: NodeEvent) -> None:
261312
invocation_id = current_invocation_id()
262313
if invocation_id is None:
263314
return
264-
correlation_id = current_correlation_id()
265315
inv_state = self._inv_state_for(invocation_id)
316+
# Idempotency: a span already exists for this attempt — likely
317+
# opened by ``prepare_sync`` in the engine task. No-op to avoid
318+
# duplicates.
319+
if self._key_for(event) in inv_state.open_spans:
320+
return
321+
correlation_id = current_correlation_id()
266322

267323
# Lazily open the invocation span on the first event we see
268324
# for this invocation_id. Per-invocation_id scoping means
@@ -579,8 +635,8 @@ def _sync_subgraph_spans(
579635
any subgraph spans whose prefix is no longer an ancestor of
580636
the current event's namespace.
581637
582-
Called from ``_handle_started`` BEFORE opening the leaf node
583-
span. Detached-mode entries (subgraph or fan-out instance)
638+
Called from ``_open_started_span`` BEFORE opening the leaf
639+
node span. Detached-mode entries (subgraph or fan-out instance)
584640
are registered as detached roots so their inner spans live
585641
in a fresh trace.
586642
"""

0 commit comments

Comments
 (0)