-
fix: isolate ADK tracer model output per invocation_id (#7266)
frostbyte_neo 发布于
2026-06-29 11:23:59 +00:00 - fix: isolate ADK tracer model output per invocation_id
A single OpikTracer instance is shared across all concurrent invocations
of an instrumented agent (the documented track_adk_agent_recursive
pattern). The most recent model output was stored on a plain instance
attribute, so two invocations in flight at the same time could overwrite
each other's output and a trace could be logged with another invocation's
answer. Trace/span identity was already contextvar-isolated; only the
output payload leaked.Key the last model output by ADK invocation_id in both OpikTracer and
LegacyOpikTracer. Sub-agent (span) completions read the entry without
removing it so the root (trace) completion still receives the output;
the root completion pops it.Adds a regression test covering concurrent isolation and the nested-agent
keep-for-root semantics.Closes #7265
- fix: drop cached output via open-agent count; skip modern tests on legacy ADK
Addresses review feedback on the PR:
- Under distributed_headers the root agent is a span (not a trace), so keying
the cached-output cleanup off the trace path leaked _last_model_output. Track
open agents per invocation_id (before_agent_callback increments,
after_agent_callback decrements) and drop the cached output once the
invocation's outermost agent finishes -- correct for nested sub-agents and
distributed root spans alike. Extract the shared cleanup into
helpers.drop_invocation_output_if_finished, used by both OpikTracer and
LegacyOpikTracer. - Skip the modern-tracer regression tests on ADK < 1.3.0 (where OpikTracer
resolves to LegacyOpikTracer) and add a distributed-root-as-span test.
- refactor: address /simplify review on the ADK tracer concurrency fix
- Move drop_invocation_output_if_finished into a finally block in both tracers
(hoisting invocation_id) so an exception during span/trace stamping cannot
skip cleanup and strand a cached model-output dict. - Reuse the autouse clear_context_storage fixture instead of a local one; add a
small tracer fixture in the concurrency tests. - Keep the span-vs-trace rationale only in the helper docstring; reuse the
invocation_id local in before_agent_callback.
- refactor: bounded per-invocation output cache (replace open-agent counter)
The open-agent depth counter could leak unboundedly: ADK skips
after_agent_callback on agent errors, early escalation/end_invocation, and
stream cancellation, so the decrement never ran and the per-invocation maps
grew forever on a long-lived shared tracer.Replace it with LastModelOutputCache: a thread-safe, size-bounded
(OrderedDict, evict-oldest) cache keyed by invocation_id, shared by both
tracers. Eviction caps memory regardless of whether after_agent_callback runs,
so there is no counter, no manual lifecycle, and no finally.Also:
- after_model_callback only caches a real (non-None) output, so a failed
conversion can't overwrite a previously cached output (matches legacy). - after_agent_callback is read-only again: invocation_id stays inside the try,
so a tracer-side error can't escape into the user's agent. - getstate excludes the cache (it holds a threading.Lock); setstate
recreates a fresh one. - Tests unit-test the cache (keying + bounded eviction) and the after_agent
read path instead of hand-seeding a counter.
Addresses code-review findings on unbounded growth, the getstate/pickle crash,
the None-clobber, the non-atomic counter, and tests that didn't exercise the fix.- fix: clear cached output up front so a failed model call can't stamp stale data
With output cached per invocation_id, an early-return or failed conversion in
after_model_callback used to leave the previous call's output in the cache, and
after_agent_callback would then stamp that stale value onto the trace/span.Clear the invocation's cache entry up front once we commit to processing a
final (non-partial) response, in both tracers, so any failure below (missing
span, conversion error, or a later error) leaves no stale value; it is re-set
only if conversion succeeds. Adds LastModelOutputCache.discard() plus tests for
discard and for after_agent_callback stamping None when nothing is cached.下载附件