Skip to content

fix(agentanalytics): preserve invocation spans for late callbacks - #1911

Open
jabrailkhalil wants to merge 1 commit into
google:mainfrom
jabrailkhalil:fix/analytics-late-callbacks
Open

jabrailkhalil wants to merge 1 commit into
google:mainfrom
jabrailkhalil:fix/analytics-late-callbacks

Conversation

@jabrailkhalil

Copy link
Copy Markdown
Contributor

Related: #1886, specifically callbacks arriving after AfterRun. The other bookkeeping items remain open.

Problem and solution

AfterRun removes both the active trace and its saved identity. A later callback, such as a deferred compaction summary, then creates another active trace with a root span that no starting row has. It also takes capacity away from live invocations.

I retain the completed root identity in the existing bounded tombstone cache. Late callbacks share a lazily created state there, so late agent/model/tool start/end pairs keep their span IDs and latency without returning the completed invocation to the active map. A new invocation-start callback with the same ID gets a fresh root; a capacity-evicted, unfinished invocation still restores its old identity.

Behavior change

Events arriving after AfterRun join the completed invocation's root span instead of inventing another root. The original invocation's span state is not copied into the late state. Completed entries share the existing tombstone capacity/pruning policy, with at most one late-state object per entry; this is not a strict byte or TTL guarantee.

This does not resolve overlapping runs with the same invocation ID or the other items in #1886. No public API or dependency versions change.

Testing Plan

Local Linux/amd64 Docker runtime, Go 1.26.7 and golangci-lint 2.3.1, with both modules in a workspace:

  • Root: 7,190 passing test/subtest results, 104 passing packages, two existing Vertex AI replay skips.
  • Analytics: 800 passing test/subtest results, no skips.
  • Both modules pass go build -mod=readonly ./..., go test -race -mod=readonly -count=1 -shuffle=on -timeout=180s -json ./..., golangci-lint run, golangci-lint fmt --diff and go mod tidy -diff.
  • New regressions pass 10 shuffled repetitions with the race detector. They cover a real Runner plugin chain before/after analytics AfterRun, late agent/model/tool pairing, terminal-row OTel adoption, 32 concurrent late callbacks while another invocation stays active, invocation-ID reuse, eviction restoration and cache capacity.
  • Reverting only trace_manager.go makes TestRunner_AfterRunEventKeepsInvocationSpan/after_end fail: derived event does not retain the completed invocation's span identity. Nine targeted mutations also fail their intended assertions, including disabling completed-state retention, state sharing and the new-run/eviction guards.

I also ran a separate real-compaction probe: the drained control passes on the original source, while stopping early fails. With this fix:

--- PASS: TestRunner_CompactionAfterRunKeepsInvocationSpan
    --- PASS: TestRunner_CompactionAfterRunKeepsInvocationSpan/drained
    --- PASS: TestRunner_CompactionAfterRunKeepsInvocationSpan/stopped_early

The probe uses the current workspace's compaction API. It is separate because committing that import would force the standalone analytics module's declared core dependency from v2.0.0 to v2.6.0 solely for this test. The shipping Runner regression uses the existing plugin API. Integration tests use the local fake BigQuery Storage Write server; no live Gemini/BigQuery or live-provider RunLive E2E was run.

Python parity

Python clears its context-local span state after the run, and its post-run compactor appends directly to the session service. Go deliberately routes those summaries through plugin callbacks (runner/runner.go); retaining completed identity handles that lifecycle difference.

Checklist

  • Read CONTRIBUTING and reviewed the full diff in a fresh context.
  • Documented the behavior change and cache boundary.
  • Added source-revert and targeted mutation evidence.
  • New and existing local tests pass.
  • Live-provider manual E2E; not run, as stated above.

Implementation and regression tests written with an OpenAI assistant under my supervision.

@jabrailkhalil

Copy link
Copy Markdown
Contributor Author

I investigated the red test (.) job. Its failure is TestA2AMultiHopInputRequiredCancellation at a2a_e2e_test.go:1321: CancelTask returns cancelation failed: queue is closed. The upstream analytics test and both lint jobs passed.

I compared exact head 0b7603ee with base 8cb74251: all changes are under plugin/agentanalytics; root source and dependency files are unchanged. The failing A2A package has the same 564-entry test dependency graph on both commits and does not import analytics.

Fresh Linux/Go 1.26.6 checks, without changing source or assertions:

  • Head/base diagnosis and complete suites: 200/200 isolated cancellation repetitions pass on each commit, with -race and the failing job's shuffle seed 1792439652092133875. The unchanged PR head also passes the complete root race/shuffle suite (7190 tests/subtests, 2 skips, 104 packages) and analytics suite (800 tests/subtests, no skips); both build and tidy checks pass.
  • Complete A2A package, 20 repetitions per commit: 7220 tests/subtests pass on each, including all 20 cancellation cases, using the same seed and race detector.

I could not reproduce the queue-closed error in these checks. This points to an intermittent A2A cancellation failure rather than an analytics regression, but I am not claiming to have fixed its underlying cause. Could a maintainer rerun the failed upstream test (.) job? I have left the PR code and CI requirements unchanged.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant