Skip to content

feat(plugins): improve BigQuery analytics observability and delivery - #12

Draft
caohy1988 wants to merge 8 commits into
mainfrom
agent/bqaa-observability-reliability
Draft

feat(plugins): improve BigQuery analytics observability and delivery#12
caohy1988 wants to merge 8 commits into
mainfrom
agent/bqaa-observability-reliability

Conversation

@caohy1988

@caohy1988 caohy1988 commented Aug 5, 2026

Copy link
Copy Markdown
Owner

Summary

This mono draft PR updates only the BigQuery Agent Analytics plugin and its unit tests. It addresses four related upstream reports while keeping each concern in a separate commit so the work can be split for upstream review later.

Related upstream issues:

Issue coverage

google#6465 — retry duplicate visibility and opt-in exactly-once delivery

  • Adds an unconditional, per-row event_id UUID to the physical schema, emitted row, and generated views. The same ID survives retries because it is assigned before enqueue/write processing.
  • Adds additive schema-upgrade coverage for existing tables.
  • Keeps BigQueryLoggerConfig.exactly_once_delivery = False, preserving default-stream behavior.
  • In opt-in mode, creates loop-local committed streams and appends batches with explicit monotonically increasing offsets.
  • Keeps ambiguous-send state sticky for the entire logical batch, even when a later attempt receives a definitive rejection.
  • Accepts ALREADY_EXISTS as confirmation only after the same batch previously had an ambiguous send at that offset. A first-attempt conflict poisons the local offset state instead of silently confirming the wrong batch.
  • Treats NOT_FOUND / OUT_OF_RANGE and exhausted ambiguous outcomes as desynchronization, accounts loss, and rotates before later writes.
  • Rotation does not wait for optional finalization of the old stream; the old stream is retained for bounded shutdown cleanup so a hung finalizer cannot stall the single writer.
  • Documents stream-creation quota exposure and the retry/rotation windows in which rows may still be dropped.

google#5644 — LLM response termination detail

  • Exposes finish_reason in v_llm_response and row attributes.
  • Emits finish_reason and error_message only on the final response, avoiding progressive-SSE double counting.
  • Routes response diagnostics through the existing sanitized error_message column without changing status semantics.
  • Supports enum and string-valued finish reasons and response-like objects that omit optional fields.

google#6529 — FunctionNode terminal event visibility (stage 1)

  • Emits NODE_OUTPUT for final workflow-node events carrying output, including legal content+output and error+output events.
  • Suppresses only message_as_output duplication and preserves Pydantic results as structured JSON.
  • Emits NODE_ERROR only for final node-execution failures.
  • Keeps model finish/block diagnostics such as MAX_TOKENS and MODEL_ARMOR in LLM_RESPONSE instead of misclassifying them as node failures.
  • Preserves state-delta rows as separate events and adds typed views with node identity fields.
  • Does not change framework event production.

google#5889 — tracing cleanup

  • Retains internal ID-only span tracking.
  • Removes the dead module-level OTel tracer allocation.
  • Keeps the real-provider exporter regression test and removes the vacuous mock-only tracer test.

Review findings addressed

  • B1: cross-batch exactly-once confirm-and-drop cascade.
  • B2: streamed finish_reason / error_message double counting.
  • M1: hot-path rotation blocked indefinitely on stream finalization.
  • M2: partial/model termination events emitted false NODE_ERROR rows.
  • M3: event.content is None and if/elif guards discarded valid node output.
  • Final ambiguity invariant: any terminal path after an ambiguous committed-stream send now poisons the stream before the next batch.
  • Progressive-SSE termination metadata cardinality is documented separately from legacy aggregator and mixed LiteLLM stream behavior.
  • Lower-priority findings: config guarantee boundaries, schema description, non-enum finish reasons, missing wiring/backoff/default-stream/Pydantic tests, inert tracer-test patches, and commit hygiene.

Live BigQuery validation

The review-fix head was validated against ADC project test-project-0728-467323 with temporary datasets that were deleted after the probes. Live committed-stream tests covered lost acknowledgements, retry exhaustion, rotation, and duplicate offsets with zero duplicate or silently lost rows; progressive SSE emitted exactly one finish-reason row per turn.

The final ambiguity branch added in this update is covered by a deterministic cross-batch regression. A final quick pass remains appropriate before undrafting.

Verification

  • pre-commit run --files src/google/adk/plugins/bigquery_agent_analytics_plugin.py tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py
  • PYTHONPATH=src <review-venv>/bin/pytest -q tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py
  • Result: 486 passed, 6 skipped, with 6 existing/experimental warnings.
  • Branch rebased onto upstream main at 0c79d1a8; fork main fast-forwarded to the same commit.

Commit map

  1. fix(plugins): Make BigQuery retry duplicates identifiable
  2. feat(plugins): Expose BigQuery LLM termination details
  3. fix(plugins): Capture workflow node results in BigQuery
  4. feat(plugins): add offset-aware BigQuery delivery
  5. fix(plugins): remove dead BigQuery analytics tracer
  6. refactor(plugins): harden model response metadata access
  7. fix(plugins): address BigQuery analytics review findings
  8. fix(plugins): close BigQuery ambiguity edge case

Attribution / CLA note

The committed-stream implementation was developed independently from the public design discussions in google#6465 and google#6466. It does not copy or cherry-pick PR google#6466 and intentionally has no Co-authored-by trailer while that contributor's CLA status is unresolved. Credit to @addenergyx for reporting the duplicate-delivery failure mode and advancing the public offset-aware design discussion.

@caohy1988

Copy link
Copy Markdown
Owner Author

CI follow-up after the first draft run:

  • All Mypy jobs pass on Python 3.10, 3.11, 3.12, and 3.13.
  • All unit-test jobs pass on Python 3.10 through 3.14.
  • All A2A v0.3 jobs pass on Python 3.10 through 3.14.
  • The sole red check is the repository's Pre-commit Linter, before any code-related failure: its all-files update-constraints hook invokes uv, but the pre-commit job environment has no uv binary (scripts/update_constraints.sh: uv: command not found) for every constraints file. Local pre-commit on both changed files passes. I have not modified workflow/config files because this PR is intentionally restricted to the BQ analytics plugin and its unit tests.

@caohy1988 caohy1988 left a comment

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Full review at 8de2adb4 — approve with one requested change

Review surface: the 5 plugin commits (+876 on the plugin + its test file). The other 8 branch commits are upstream sync noise — verified none touch the plugin. Local run with the worktree's own src: 472 passed, matching CI. mypy matrices green; the pre-commit failure is the documented unrelated uv workflow issue.

Per-feature verification

  • event_id (google#6465 dedup key) — Solid. NULLABLE STRING column, _SCHEMA_VERSION 1→2, in _VIEW_COMMON_COLUMNS, stamped at row construction so transport retries preserve it; test_bigquery_retry_reuses_the_same_event_id proves the same ID survives an ack-lost retry. Existing tables pick it up because _maybe_upgrade_schema is diff-based, not label-gated.
  • Exactly-once opt-in (google#6465) — Traced every branch; the protocol is sound: frozen per-batch offset; in-band and raised ALREADY_EXISTS confirm (idempotent assignment, not +=); NOT_FOUND/OUT_OF_RANGE desync → accounted under offset_conflict → finalize old stream → rotate with 30s backoff; the request_sent/definitive_rejection pair correctly separates "server explicitly rejected" (no desync, offset safely reused) from "ack lost" (poison + rotate); offsets reset to 0 on a new stream; finalize is idempotent and retried across shutdown/close via the pending set. Test coverage is genuinely good: both ALREADY_EXISTS forms, all four desync forms, ambiguous-exhaustion rotation, finalize retry.
  • finish_reason (google#5644) — Attributes + typed view column; partial chunks omit the key (tested); .name on the str-enum is fine. error_message routing is proven sanitized by test (Bearer MODEL-SECRET[REDACTED], status stays "OK", secret absent from the whole row).
  • NODE_OUTPUT/NODE_ERROR (google#6529) — Guards are right: node_path + non-partial + content is None + output is not None prevents AGENT_RESPONSE double-logging; the message_as_output case is tested; STATE_DELTA coexists correctly; views expose node identity; dict/list/str payloads round-trip with a default=str fallback for exotic outputs.
  • Dead tracer cleanup (google#5889) — Removed, and the ~20 test patch sites correctly moved to create=True.

Findings

Three inline (1×P2, 2×P3) plus one nit: squash 8de2adb4 into cb5b14cb when splitting for upstream — it only adds the asserts and swaps the exception type, so it isn't a standalone change.

Verdict: approve with finding 1 (the default-mode retry change) addressed or explicitly documented. The engineering is careful and the test coverage is real; 2–4 are upstream-review polish.

Comment thread src/google/adk/plugins/bigquery_agent_analytics_plugin.py Outdated
Comment thread src/google/adk/plugins/bigquery_agent_analytics_plugin.py Outdated
Comment thread src/google/adk/plugins/bigquery_agent_analytics_plugin.py Outdated
@caohy1988

Copy link
Copy Markdown
Owner Author

Live e2e verification (real Vertex Gemini + real BigQuery)

Ran a fresh end-to-end suite against this branch (8de2adb) with gemini-2.5-flash on Vertex AI and live BigQuery Storage Write API — a normal LlmAgent turn, a forced-MAX_TOKENS turn, a two-FunctionNode workflow, and a failing FunctionNode, executed twice: default delivery and exactly_once_delivery=True. Temp datasets were deleted after verification.

Check Result
finish_reason on LLM_RESPONSE (google#5644) STOP and MAX_TOKENS in attributes.finish_reason and in v_llm_response.finish_reason, both modes
event_id coverage (google#6465) ✅ 100% of rows in both modes
NODE_OUTPUT rows (google#6529) pr12_wf_ok@1/step_one@1 / step_two@1 with full payloads, exposed via v_node_output with node_path/node_run_id
NODE_ERROR rows (google#6529) pr12_wf_err@1/boom_node@1 with RuntimeError + message via v_node_error
Span pollution (google#5889) ✅ global in-memory SpanExporter captured 0 plugin-scoped spans and 0 bare agent/llm_request/tool spans; only gcp.vertex.agent framework spans exported
Duplicate behavior Default mode caught a real transport-level duplicate in the wild — two byte-identical USER_MESSAGE_RECEIVED rows sharing one event_id (identical timestamps/invocation), i.e. exactly the google#6465 phenomenon occurring naturally, made identifiable by the new column. Exactly-once mode: zero duplicates across all runs.
Schema/views adk_schema_version:2 label, v_node_output/v_node_error created, event_id present in views
Unit suite (this branch, py3.13) 466 passed, 6 skipped (the "472" figure counts the 6 env-gated skips)

The MAX_TOKENS scenario also surfaced a behavior issue (a plain, non-workflow agent produced a NODE_ERROR row with error_code=MAX_TOKENS) — detailed in the review comment below.

@caohy1988

Copy link
Copy Markdown
Owner Author

Full review — 2 blockers, 5 majors; keep as draft until addressed

Line references are against 8de2adb (plugin = src/google/adk/plugins/bigquery_agent_analytics_plugin.py, tests = tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py).

Blockers

B1 — exactly-once mode can silently confirm-and-drop batches (cross-batch cascade).
request_sent / definitive_rejection are reset at the top of every retry attempt, and the exhaustion-desync decision reads only the final attempt's flags (plugin:2559-2561, 2671-2676). definitive_rejection is set for any in-band error, including retryable 4/13/14 (plugin:2596). Failure scenario:

  1. Batch A, attempt 1: client timeout, but the server committed the append (ambiguous outcome; the request_sent=True is then forgotten).
  2. Batch A, final attempt: in-band UNAVAILABLE → definitive_rejection=True → exhaustion with no desync; _next_offset not advanced, A's rows are in the table.
  3. Batch B appends at the same offset → ALREADY_EXISTS on its first attempt → unconditionally treated as confirmation (plugin:2598-2601, 2641-2645) → B's rows are silently lost with no drop counter, and if len(A) > len(B) subsequent batches keep hitting ALREADY_EXISTS and are dropped in a cascade until an OUT_OF_RANGE finally forces rotation.

Fix shape: make "an attempt with unknown outcome was sent" sticky per batch; desync on exhaustion whenever it is set; honor ALREADY_EXISTS as confirmation only when this batch previously sent at that offset — otherwise treat it as desync. The existing test at tests:9300 always precedes ALREADY_EXISTS with a send, which masks the first-attempt hole; B1 needs its own cross-batch test.

B2 — partial streaming rows carry finish_reason, double-counting the headline metric (google#5644).
EventData.finish_reason is set with no partial gate (plugin:7052-7061). In streaming mode the aggregator marks all chunks partial=True, the final chunk carries finish_reason, and the aggregated partial=False response carries it again — so every streamed turn writes finish_reason on ≥2 LLM_RESPONSE rows, and LLM_RESPONSE rows have no partial marker to filter on. COUNT(*) GROUP BY finish_reason — the exact query promised on google#5644 — over-counts. The streaming test at tests:1894 only covers a partial chunk whose finish_reason is None, so it can't catch this. Fix: gate finish_reason (and the new error_message pass-through, which lite_llm sets on non-STOP chunks) on not llm_response.partial.

Majors

M1 — NODE_ERROR has no partial guard → duplicate error rows. The NODE_OUTPUT branch checks event.partial is not True; the NODE_ERROR branch does not (plugin:6497). A streamed SAFETY/MAX_TOKENS failure under a workflow node emits NODE_ERROR twice (error-bearing partial chunk + aggregated event).

M2 — LLM finish anomalies masquerade as node failures. LlmResponse.create copies non-STOP finish_reason into error_code (llm_response.py:218, block_reason at :231), and every agent runs as a node in ADK 2 — so any MAX_TOKENS/SAFETY response now emits a status=ERROR NODE_ERROR row alongside its LLM_RESPONSE row, for all users, workflow or not (reproduced live: plain agent, node_path=maxtok_agent@1, error_code=MAX_TOKENS). Anyone alerting on v_node_error will page on ordinary MAX_TOKENS finishes, and "node failed" vs "model finished abnormally" become indistinguishable. Suggest excluding LLM-finish-reason error codes from the NODE_ERROR branch (they're already covered by finish_reason on LLM_RESPONSE) or tagging them distinctly.

M3 — event.content is None guard silently drops output. A user-yielded Event(content=..., output=...) from a FunctionNode is legal and validated, but produces no NODE_OUTPUT (plugin:6511); similarly the if/elif means an event with both error_code and output loses its output row. The message_as_output dedup doesn't require this guard — the runner already clears output for message-as-output nodes, and checking message_as_output directly would suffice.

M4 — stream finalize on shutdown has no timeout. shutdown(timeout=t) bounds only the worker; finally: await asyncio.shield(self._finalize_stream()) (plugin:2733-2738, 2815-2820) can hang for gapic's multi-minute FinalizeWriteStream retry deadline on a network blackhole, blowing the shutdown_timeout contract. Wrap in asyncio.wait_for and fall back to _pending_finalize_streams.

M5 — committed docs missing. No user-facing prose documents (a) CreateWriteStream quota exposure — one stream per event loop, plus one per rotation, plus re-creation after failures — or (b) that exactly-once degrades to at-most-once during desync/rotation windows (batches dropped under offset_conflict). Both were part of the google#6465 plan; the config docstring is the natural home.

Minors

  • Commit hygiene: the Feature request: BigQueryLoggerConfig.export_internal_spans flag to disable plugin span export google/adk-python#5889 dead-tracer removal and a finish_reason refactor are bundled into cb5b14c ("offset-aware delivery", empty body). Split them out or drop the tracer change; write bodies for cb5b14c/8de2adb. Also test_push_pop_does_not_call_tracer_start_span (tests:2917) is now vacuous — it patches a nonexistent attribute with create=True; the real-provider exporter test at tests:2954 carries the actual coverage, so either delete the vacuous test or restore its meaning.
  • Attribution: no Co-authored-by trailers on any commit. roanny gets prose credit in 701f946; the exactly-once commits credit addenergyx nowhere despite following the fix(plugins): add opt-in exactly-once delivery to BigQueryAgentAnalyticsPlugin google/adk-python#6466 design — add prose credit at minimum.
  • Default-path behavior change not called out: an empty response stream now raises and retries (plugin:2635-2636) where upstream treated it as success — defensible, but it is a new duplicate source in default mode and deserves a commit-message/docstring mention.
  • finish_reason.name has no non-enum fallback (plugin:7053): a string-valued finish_reason (duck-typed/model_constructed responses) raises AttributeError and @_safe_callback silently drops the whole LLM_RESPONSE row. Use getattr(fr, "name", str(fr)).
  • Rotation backoff is a hardcoded 30.0s (plugin:2498), no jitter, ignores retry_config; all batches during the window are dropped — sustained CreateWriteStream failure is 100% loss.
  • Definitively-not-committed batches (OUT_OF_RANGE/NOT_FOUND) are dropped rather than replayed on the replacement stream (plugin:2602-2605, 2649-2658) — consistent with the declared design but avoidable loss; document it (M5).
  • Test gaps: partial-chunk-with-finish_reason (catches B2); partial NODE_ERROR (M1); B1 cross-batch scenario; plugin-level exactly_once_delivery=True wiring test (only the helper is unit-tested, tests:9436-9451); fork/pickle child-not-reusing-parent-stream; _default-never-finalized; rotation-failure/backoff-drop; pydantic-model NODE_OUTPUT; non-enum finish_reason. Also test_ambiguous_exhaustion_poison_stream_and_rotates (tests:9382) uses a plain MagicMock for finalize_write_stream, so the await TypeErrors and the test silently exercises the finalize-failure branch — make it an AsyncMock (and add a deliberate failure-branch test).
  • assert offset_for_batch is not None (plugin:2599, 2632, 2643) vanishes under python -O; prefer explicit checks.

Verified good

  • event_id: stamped pre-enqueue (plugin:6322→6356), retry-stable (proven live — a wild transport duplicate shared one id), schema v2 additive upgrade works on real tables, exposed in views, not deniable, no PARTITION-BY-NULL trap (no dedup view ships).
  • Both raised and in-band forms of ALREADY_EXISTS / NOT_FOUND / OUT_OF_RANGE handled (plugin:2598-2605, 2641-2661), matching live evidence that duplicates surface in-band (code 6).
  • Default-off is genuinely inert: no offset field on _default requests, _default never finalized, new except branches map to upstream's drop semantics.
  • Offset lifecycle correct: fresh stream starts at 0, rotation resets, offset frozen per batch, advance-by-len(rows) matches all-or-nothing Arrow batch prep.
  • Single-writer invariant, fork/pickle reset (children create fresh streams), transport close on creation failure, schema-upgrade-before-stream-creation ordering: all check out.
  • finish_reason mechanics: enum-name string, None → key omitted, snake_case, view column exactly as promised on feat: project finish_reason from LlmResponse to attributes in BigQueryAgentAnalyticsPlugin LLM_RESPONSE events google/adk-python#5644; genai FinishReason's dynamic _missing_ keeps .name safe for unknown enum values.
  • error_message routing: dedicated column, real sanitization with fail-closed [REDACTED_SENSITIVE_TEXT], status semantics unchanged, no collision with LLM_ERROR.
  • Node identity plumbing end-to-end (adk.node.{path,run_id,parent_run_id}v_node_*), zero-regression guard for non-node events (node_path short-circuit), deliberate+tested STATE_DELTA/NODE_OUTPUT dual emission, crash-safe arbitrary-output handling via the fail-closed parse boundary.
  • Denylist grace: views omit finish_reason/node content columns when attributes/content are denied.

Suggested order: B1, B2, then M1-M3 (cheap row-semantics fixes), M4, then M5 + hygiene/attribution, then the test gaps. With those addressed this is a merge-quality change — the live run shows all four issues genuinely fixed under production conditions.

@caohy1988
caohy1988 force-pushed the agent/bqaa-observability-reliability branch from 8de2adb to 5fec54b Compare August 5, 2026 22:51
@caohy1988

Copy link
Copy Markdown
Owner Author

Fresh review — current head 5fec54bc

Verdict: changes requested; keep this PR in draft.

I reviewed the current head, rather than the older 8de2adb covered by the previous comments. Both original blockers still reproduce. I also found one additional committed-stream writer-stall defect. The earlier retry-scope, shutdown-timeout, and production-assert findings have been addressed.

Blockers

B1. Exactly-once retry can silently discard the next batch

The retry loop resets its ambiguity state on every attempt. If batch A may have committed but its acknowledgement was lost, a later definitive error clears that ambiguity. After retries exhaust, batch B reuses A's offset. An ALREADY_EXISTS response is then interpreted as confirmation of B, even though it may only confirm A.

Current-head reproduction:

  1. A sends at offset 0, then loses the acknowledgement.
  2. A retries at offset 0 and receives UNAVAILABLE.
  3. The writer remains at offset 0 without entering desynchronized state.
  4. B sends at offset 0 and receives ALREADY_EXISTS.
  5. B is counted as successful and the offset advances, although B was never written.

Code: retry and offset handling.

ALREADY_EXISTS only says data already exists at that offset; it does not identify which local batch wrote it. See the official AppendRows response semantics and offset requirements.

Required fix:

  • Keep a sticky had_ambiguous_send state for the entire logical batch, outside the attempt loop.
  • A later definitive rejection must not erase ambiguity from an earlier attempt.
  • Only treat ALREADY_EXISTS as confirmation when this same batch previously had an ambiguous send at this offset.
  • A first-attempt ALREADY_EXISTS should mark the local offset state invalid and rotate/desynchronize.

B2. Streaming responses still record finish_reason twice

The progressive-stream aggregator emits a partial terminal response containing finish_reason, followed by a final aggregate containing the same value. The plugin records the field unconditionally in after_model_callback.

Current-head reproduction:

STOP with text:
(partial=True, finish_reason=STOP)
(partial=False, finish_reason=STOP)

MAX_TOKENS without text:
(partial=True, finish_reason=MAX_TOKENS)
(partial=False, finish_reason=MAX_TOKENS)

The behavior originates in the streaming aggregator, which later emits the final aggregate.

Gate both finish_reason and error_message to final responses and add a test using the real streaming aggregator. The existing partial-response test constructs a partial response without a finish reason, so it cannot catch this regression.

Major findings

M1. Committed-stream rotation can stall the writer indefinitely

_ensure_writable_stream awaits finalization of the old stream without a timeout before creating its replacement.

I reproduced a non-returning finalizer: the writer remained blocked and no replacement stream was created. Because this is the single batch writer, the queue eventually fills and analytics events begin dropping.

The shutdown/close timeout fix is good, but it does not protect this hot-path rotation. Finalization is optional for committed streams according to the BigQuery streaming documentation. Bound or decouple finalization, retain failed streams for later cleanup, and create the replacement immediately.

M2. NODE_ERROR conflates model termination with node failure

The plugin emits NODE_ERROR for any event with node_path and error_code in this branch. Normal LLM events can receive error codes from finish reasons such as MAX_TOKENS through LlmResponse.

This produces false node-failure telemetry. With progressive streaming, the missing partial guard can also emit the false NODE_ERROR twice.

Within plugin scope, suppress known model finish/block codes from NODE_ERROR and reject partial events. Tests should prove:

  • MAX_TOKENS produces an LLM_RESPONSE, not a NODE_ERROR.
  • A real NodeRunner exception produces exactly one NODE_ERROR.

M3. Valid node output can be silently omitted

The NODE_OUTPUT condition requires event.content is None and is an elif after the error branch.

ADK permits an event to contain both content and output. NodeRunner specifically clears output only for message_as_output=True. Therefore, legal combined content/output events and error/output events lose their node-output record.

Gate on message_as_output, not on content is None, and make output/error handling independent branches.

Lower-priority findings

  • The exactly_once_delivery documentation does not disclose stream rotation, loss after retry exhaustion, rotation backoff, or CreateWriteStream quota implications. The config documentation should state the actual guarantee boundaries.
  • The physical schema still describes error_message as populated only when status=ERROR, although LLM_RESPONSE rows may now have status=OK plus an error message.
  • finish_reason.name assumes an enum. A response-like object containing a string can raise and silently lose the entire LLM response row. Use getattr(value, "name", str(value)).
  • The exact-once commit also contains dead-tracer cleanup and create=True test changes. One resulting test patches a nonexistent tracer and proves only that the nonexistent object was not called. Move the real cleanup into the Feature request: BigQueryLoggerConfig.export_internal_spans flag to disable plugin span export google/adk-python#5889 commit and remove the vacuous test before splitting upstream.

Verified good

  • event_id is additive, generated before enqueue, stable across retries, and exposed in views.
  • Default mode retains legacy append behavior and does not use offsets or committed-stream finalization.
  • error_message is now captured.
  • Raised and in-band Storage Write API errors are both recognized.
  • Shutdown/close finalization has a bounded deadline.
  • Production assert statements from the previous review were replaced.
  • Changes remain confined to the BQAA plugin and its tests.
  • Upstream main is two commits ahead, but neither commit overlaps these files.

Validation completed:

  • 477 plugin unit tests passed locally.
  • All changed-file pre-commit hooks passed locally.
  • PR unit, A2A, and mypy jobs pass across their configured Python versions.
  • The CI pre-commit job's remaining failure is environmental: uv is absent from that runner during update-constraints.
  • Worktree is clean.

The posted e2e report remains useful, but it was run against 8de2adb and tests the successful exactly-once path—not the cross-batch ambiguity sequence or progressive-stream double-counting reproduced above.

Assign each analytics row a stable event_id before enqueue so Storage Write API retries preserve a queryable deduplication key. Existing tables receive the nullable column through additive schema upgrade.

Refs google#6465
Project finish_reason into LLM response attributes and typed views, and route model response diagnostics through the existing sanitized error_message column without changing status semantics.

The finish_reason projection follows the approach proposed by @roanny in google#5704.

Refs google#5644
Emit NODE_OUTPUT and NODE_ERROR rows from workflow events that already reach the plugin, preserving node identity while avoiding message-as-output duplication.

Fixes google#6529
Add an opt-in committed-stream writer with explicit per-batch offsets, retry-stable requests, stream rotation after ambiguous outcomes, bounded finalization, and loss accounting.

The design follows the public discussion in google#6465 and @addenergyx's proposal in google#6466. This implementation was developed independently to avoid CLA-sensitive source derivation.

Refs google#6465
@caohy1988
caohy1988 force-pushed the agent/bqaa-observability-reliability branch from 5fec54b to 97c1a41 Compare August 6, 2026 03:01
@caohy1988

Copy link
Copy Markdown
Owner Author

Re-review at head 97c1a415 — all prior blockers fixed and live-verified; 1 new major, then ready to undraft

Two-part verification: a fresh adversarial code review of all 7 commits, and a live run against real Vertex Gemini + real BigQuery including SSE streaming and duplicate-offset probes against real committed streams. Line refs are against 97c1a415.

Live verification (real Vertex + BigQuery, both delivery modes)

Duplicate-offset probes — drove BatchProcessor directly against real COMMITTED streams with an ack-suppressing wrapper client (the google#6465 harness pattern):

Scenario Observed Table result
Suppress 1 ack, then retry Retry at same offset hit real in-band code 6 ("expected offset 1, received 0") → confirmed via sticky-ambiguity rule; next batch landed at offset 1 Each row exactly once
Suppress ALL acks (exhaustion) retry_exhausted=1 counted honestly, stream desynced, rotation created a replacement stream, next batch landed on it Each row exactly once

Zero duplicates, zero silent loss, in both scenarios. The B1 state machine holds against the real Storage Write API.

Agent e2e (plain + SSE runs of STOP and MAX_TOKENS agents, a 3-FunctionNode workflow incl. a pydantic-returning node, and a failing node — run twice: default and exactly_once_delivery=True):

Check Result
B2: SSE finish_reason count ✅ 2 LLM_RESPONSE rows per SSE turn, exactly 1 carries finish_reason (STOP and MAX_TOKENS, both modes)
NODE_ERROR conflation ✅ MAX_TOKENS agents (plain + SSE) produced zero NODE_ERROR rows; only the genuine RuntimeError node produced one, exactly once
NODE_OUTPUT structured output ✅ pydantic node serialized as structured JSON ({"patient":...,"score":99,"tags":["a","b"]})
Span export (google#5889) ✅ 0 plugin-scoped spans through a global in-memory SpanProcessor, both runs
event_id duplicates ✅ 0 duplicate groups in both datasets
Unit suite ✅ 485 passed, 6 skipped (env-gated), reproduced in a clean venv

Prior findings — all verified fixed in code (not just claimed)

  • B1 confirm-and-drop: sticky had_ambiguous_send per batch (plugin:2578, set 2726/2756, never reset); exhaustion desyncs iff sticky (2729-2732); _handle_already_exists (2510-2527) confirms only after a prior ambiguous send, else warns + desyncs + counts offset_conflict without advancing the offset. Full path×state matrix traced — every terminal path advances the offset only via _confirm_committed_delivery. Cross-batch test coverage pins the cascade (tests:9430, 9470).
  • Rotation stall: finalize decoupled — old stream goes to _pending_finalize_streams (2542), replacement created immediately (2544); finalize runs only at shutdown/close under the deadline budget (2795-2823). Pinned by tests:9548 (hung finalizer, bounded write) and 9506.
  • B2 partial double-count: is_partial gate at 7064; finish_reason (7146-7156) and error_message (7157-7161) stamped only on final responses; real-StreamingResponseAggregator test at tests:1922. Verified live (table above).
  • NODE_ERROR conflation: _LLM_RESPONSE_ERROR_CODES = all FinishReason + BlockedReason values (220-224), excluded at 6588; robust for str and enum forms (pydantic coerces to str; str-enum hashing matches); LiteLLM maps to canonical enums. Parametrized tests incl. BlockedReason.
  • NODE_OUTPUT guard: now message_as_output is not True instead of content is None (6600-6603), and error/output are independent ifs — content+output, error+output, and message-as-output cases all tested.
  • Shutdown finalize budget intact (2787-2823, pinned by tests:9658); python -O asserts replaced with an explicit desync-on-None check (2501-2507, tested); default mode reverted to byte-identical upstream behavior (empty response stream = success again, 2691; offset field never set — HasField asserted) plus event_id.
  • M5 docs: config docstring (1720-1727) now documents one stream per event loop, CreateWriteStream quota exposure, the 30s rotation backoff, and every drop boundary — text checked against the implementation, accurate.
  • Feature request: BigQueryLoggerConfig.export_internal_spans flag to disable plugin span export google/adk-python#5889 hygiene: 027847e contains only the tracer removal; the vacuous patch-a-nonexistent-attribute test was removed; the real-provider exporter regression test is retained.
  • Attribution/CLA: all 7 commit bodies non-empty; @roanny prose credit + Refs #5644 on 64baec5; @addenergyx prose credit with explicit independent-implementation statements on 8dcc72f/97c1a41; no Co-authored-by trailers (deliberate, CLA-safe); Fixes #6529/Fixes #5889 refs correct.

New findings

MAJOR-1 — non-retryable in-band error after an ambiguous send does not desync (plugin:2667-2684). The fall-through non-retryable branch counts non_retryable and returns without consulting had_ambiguous_send — the only terminal path that skips the check. Scenario: batch A's first attempt times out but committed (sticky flag set); a later attempt gets an in-band code outside {4,5,6,11,13,14} (e.g. 8 RESOURCE_EXHAUSTED, 7 PERMISSION_DENIED) → A dropped as non_retryable, stream NOT desynced, offset still points at A's rows. If batch B then also has an ambiguous first send, B's retry sees ALREADY_EXISTS (caused by A) and wrongly confirms B — silent, uncounted loss, with offset drift if len(A) != len(B). One-line fix: desync in that branch when exactly_once_delivery and had_ambiguous_send, plus a test. This is the last hole in the ambiguity invariant; every other terminal path honors it.

MAJOR-2 (residual, pre-existing ADK-core semantics — document rather than fix here). The "exactly one finish_reason row per call" guarantee holds for the progressive-SSE default (verified live). With PROGRESSIVE_SSE_STREAMING disabled, the legacy aggregator emits up to 3 non-partial responses per call, each stamped (reproduced with a probe); LiteLLM mixed text+tool streams emit 2. Not a plugin regression — strictly better than pre-fix — but nothing tests or documents the non-progressive modes. Suggest a docstring note ("one row per final response; legacy/LiteLLM streams can emit multiple finals") and, ideally, a legacy-mode test.

Minors:

  • 18 inert mock.patch(..., "tracer", ..., create=True) sites remain in tests, patching a nonexistent attribute — dead scaffolding worth deleting (the surrounding tests are otherwise meaningful).
  • _LLM_RESPONSE_ERROR_CODES is frozen at import: a future server-side finish reason produces a spurious NODE_ERROR; a custom node error code equal to a member (e.g. "OTHER") silently loses its NODE_ERROR row.
  • _pending_finalize_streams grows unboundedly in a long-lived process under stream churn (drained only at shutdown); an opportunistic fire-and-forget finalize during rotation would cap it.
  • create_stream has no plugin-level timeout (rotation 2544, initial 4538) — a hung CreateWriteStream still stalls the single writer; bounded only by transport defaults.
  • Rotation backoff is a hardcoded 30.0s, no jitter, ignores retry_config (2548) — documented, but inconsistent with the configurable policy.
  • No fork/child test pinning that a forked child creates a fresh committed stream (behavior is correct by construction via _reset_runtime_state, 5522-5523, but unpinned).
  • output_for multi-node attribution is not surfaced on NODE_OUTPUT rows (only path/run_id/parent_run_id).
  • 676bc5d is broken in isolation for string finish_reason (fixed one commit later) — bisect hazard only; commit timestamps are non-monotonic from the rewrite. Cosmetic.

Recommendation

Fix MAJOR-1 (one line + one test), add the MAJOR-2 docstring note, optionally sweep the inert tracer patches — then this is ready to undraft. Everything else on the list is quality-of-life and can trail. The live evidence at this head is comprehensive: all four issues are solved under production conditions, the exactly-once state machine survives real lost-ack and exhaustion-rotation scenarios without duplicating or silently losing a row, and default-mode behavior is upstream-identical plus a working dedup key.

@caohy1988 caohy1988 left a comment

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fresh review at 97c1a415 — approve

Rebased history verified: 7 clean per-issue commits, +379 plugin / +1249 tests vs merge-base. Local run with the worktree's own src: 491 passed (consistent with CI's 485 passed + 6 skipped).

Last round's findings — all fixed, each pinned by a test

  1. Empty-stream retry (P2) — now scoped to exactly_once_delivery (:2700-2702); default mode keeps the old assume-success behavior, pinned by test_default_mode_keeps_empty_response_as_success (exactly one call, zero drops).
  2. Unbounded finalize (P3)_finalize_stream_before(deadline) implements the remaining-budget contract with clean timeout/cancel handling, pinned by test_finalization_respects_remaining_close_budget (hanging finalize cancelled inside budget).
  3. Production asserts (P3) — replaced by _confirm_committed_delivery, which poisons-and-rotates on a missing offset instead of crashing; strictly safer, pinned by test_missing_committed_offset_desynchronizes_without_assertion.

The additional hardening — verified on close tracing

  • had_ambiguous_send/_handle_already_exists: first-attempt ALREADY_EXISTS desyncs instead of falsely confirming a batch that never had an ambiguous send — the confirm-and-drop cascade is closed, pinned in both in-band and raised forms.
  • Rotation no longer awaits finalization before creating the replacement; the old stream is parked in _pending_finalize_streams and finalized at shutdown. Rows stay visible, writer never blocks. Pinned.
  • finish_reason/error_message extraction gated on non-partial (streaming double-count fixed); getattr(finish_reason, "name", str(...)) covers plain-string reasons.
  • _LLM_RESPONSE_ERROR_CODES (FinishReason + BlockedReason values) keeps model-termination codes out of NODE_ERROR — no more double classification, pinned.
  • message_as_output is not True directly encodes the framework's dedup semantics (verified at runners.py:942-946 — the runner clears output on content-carrying copies before the plugin sees them).
  • Config docs now state the loss boundaries and CreateWriteStream quota exposure.
  • 027847e5 gives google#5889's dead-tracer removal its own commit; history splits cleanly per issue for upstream.

Remaining — nothing blocking

One informational inline note, plus a trivial nit: 027847e5's title slightly overstates (spans were already not created; the commit removes dead code) — irrelevant once split.

Verdict: approve. Every prior finding is fixed with a test behind it, and the branch is in clean per-issue split shape. The live BigQuery e2e run (SSE streaming + duplicate-offset probes) is the only remaining gate before un-drafting.

)
if (
event.output is not None
and getattr(node_info, "message_as_output", None) is not True

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P3, informational] With the content is None guard dropped, task-mode agents (finish_task FC args as output, no message_as_output) can now emit both AGENT_RESPONSE and NODE_OUTPUT for related events — the runner only clears output on content-carrying copies when message_as_output is set (runners.py:942-946). The telemetry is distinct (args vs response text) and test_content_and_output_event_preserves_node_output pins this as deliberate, so no change needed — just be ready to articulate that story when upstream review asks.

Remove the unused plugin tracer so analytics stack bookkeeping cannot create plugin-owned OpenTelemetry spans beside the framework spans. Keep existing trace tests compatible with the intentionally absent module tracer.

Fixes google#5889
Read optional finish-reason and diagnostic fields through guarded attribute access so response-like implementations remain compatible with the analytics callback.
Preserve ambiguous committed-stream outcomes across retries, reject first-attempt offset conflicts, and rotate without waiting on optional finalization.

Record streaming termination metadata only on final responses, distinguish model termination from node failures, and preserve structured node outputs.

Document stream quota and loss boundaries and add focused regression coverage. The committed-stream design follows the public discussions in google#6465 and @addenergyx's proposal in google#6466; this implementation was developed independently.
Desynchronize a committed stream when a terminal in-band rejection follows an ambiguous send so the next batch cannot be falsely confirmed at an occupied offset.

Document terminal-response multiplicity outside progressive SSE, add cross-batch regression coverage, and remove inert tracer patches from the analytics tests.
@caohy1988
caohy1988 force-pushed the agent/bqaa-observability-reliability branch from 97c1a41 to d8fac0b Compare August 6, 2026 17:08
@caohy1988

Copy link
Copy Markdown
Owner Author

Final review at d8fac0be — approved; ready to undraft

Delta since the previously reviewed head (97c1a415) is minimal and exactly what was requested: +5 plugin lines, one new test, 18 inert tracer patches removed (create=True count in the test file is now 0), history retitled so the google#5889 commit is fix(plugins): remove dead BigQuery analytics tracer, and the new d8fac0be commit carries a proper body.

Verified in code

  • MAJOR-1 closed — the non-retryable in-band branch now runs if self.exactly_once_delivery and had_ambiguous_send: self._desync_stream() before counting the drop, making it consistent with every other terminal path. The ambiguity invariant ("no terminal path may leave a possibly-occupied offset live") now holds everywhere.
  • Regression test is the right onetest_non_retryable_rejection_after_ambiguity_rotates_stream models the exact hole (timeout → in-band code 7 → next batch), and its mock returns ALREADY_EXISTS if the next batch ever touches the old stream, so a regression to false confirmation fails on three separate assertions (stream sequence, _next_offset, drop stats).
  • MAJOR-2 documentedafter_model_callback docstring now states termination metadata is recorded once per non-partial response, and that legacy aggregators / mixed LiteLLM streams can emit multiple terminal responses (progressive SSE emits one).

Verified by execution

  • Full plugin suite at this head: 486 passed, 6 skipped (env-gated), clean venv.
  • Live committed-stream probes re-run against real BigQuery at this head (ack suppressed once → in-band ALREADY_EXISTS confirmation; all acks suppressed → exhaustion → desync → rotation to a replacement stream): every row exactly once, zero duplicates, zero silent loss, honest drop accounting in both scenarios.
  • The observability paths (finish_reason single-stamping under SSE, NODE_OUTPUT/NODE_ERROR semantics, structured pydantic output, zero plugin span export, event_id dedup) were live-verified at the previous head and received only a docstring change since — results carry over.

Remaining items — all minor, all disclosed, none blocking

Frozen _LLM_RESPONSE_ERROR_CODES set ages both ways; _pending_finalize_streams unbounded growth under stream churn; no plugin-level timeout on create_stream; hardcoded 30s rotation backoff; no fork/child committed-stream test; output_for not surfaced on NODE_OUTPUT rows. Fine to trail as follow-ups.

Verdict: undraft when ready. For upstream submission, consider splitting as discussed: the small commits (google#5644 finish_reason, google#5889 tracer, google#6529 node rows, google#6465 event_id) review quickly and shouldn't wait on the exactly-once feature review.

@caohy1988 caohy1988 left a comment

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fresh review at d8fac0be — approve, no new findings

Delta vs 97c1a415: one 5-line logic change, a docstring, and a large test cleanup across 8 clean per-issue commits. Local run with the worktree's own src: 492 passed (consistent with CI's 486 + 6 skipped).

The new logic change — verified correct, hole closed

Desync when a terminal in-band non-retryable rejection follows an ambiguous send (:2683-2684). I re-walked the full ambiguity matrix; this was the last open branch:

  • Ambiguous send → in-band ALREADY_EXISTS → confirm ✓
  • Ambiguous send → in-band non-retryable → drop + desync ← this fix; previously the batch dropped but the stream stayed trusted, so the next batch reused the occupied offset and could be falsely confirmed at it
  • Ambiguous send → raised exception / retry exhaustion → desync ✓
  • No ambiguity → in-band ALREADY_EXISTS / NOT_FOUND / OUT_OF_RANGE → desync ✓
  • No ambiguity → in-band non-retryable → drop, no desync — correct: the offset was never applied and is safely reused

The cross-batch regression (test_non_retryable_rejection_after_ambiguity_rotates_stream) has exactly the right shape: batch A times out then gets PERMISSION_DENIED; batch B must land on the replacement stream at offset 0 — asserting the full stream sequence [committed-1, committed-1, committed-2], _next_offset == 1, and a single rotation. Without the fix, batch B writes to the poisoned stream.

The cleanup — verified

  • 18 inert tracer patches removed: correct — since 37abcef1 removed the module tracer, mock.patch(...tracer, create=True) was scaffolding around nothing. test_push_pop_does_not_export_spans_through_real_provider now asserts zero exported spans against the real module state — a stronger guarantee than the patched version.
  • SSE/LiteLLM cardinality docstring: accurate — progressive SSE yields one terminal response per call; legacy aggregators and mixed LiteLLM streams can emit several, which is exactly why the non-partial gating from the previous round matters.
  • 37abcef1 title corrected — last round's nit is closed.
  • Task-mode AGENT_RESPONSE + NODE_OUTPUT co-emission left unchanged — already recorded as deliberate in the earlier inline comment, so upstream has the context.

Verdict: approve. The ambiguity matrix is fully closed with a regression test on every branch, and the commit map splits cleanly per issue. Remaining gates are the ones you know: the queued CI matrix going green and the live BigQuery e2e run (SSE streaming + duplicate-offset probes) before un-drafting.

@caohy1988

Copy link
Copy Markdown
Owner Author

Live e2e addendum — full real-agent suite re-run at final head d8fac0be

Closes the three verification gaps noted in the final review: the agent e2e now has fresh evidence at this exact commit (not carried over from 97c1a415), the v1→v2 schema upgrade was exercised against a live populated table, and a SAFETY finish-reason attempt was made. Real gemini-2.5-flash on Vertex → plugin → real BigQuery, run twice (default and exactly_once_delivery=True); temp datasets deleted after verification.

Check Result
Live schema upgrade v1→v2 ✅ Pre-staged a populated v1 agent_events table (16 fields, adk_schema_version:1, one legacy row) — the plugin upgraded it in place: event_id added (17 fields), label bumped to 2, legacy row keeps NULL event_id, all 51 new rows populated
finish_reason single-stamping ✅ Plain turns: 1 row / 1 stamped; SSE turns: 2 rows / exactly 1 stamped (STOP and MAX_TOKENS, both modes)
NODE_OUTPUT / NODE_ERROR ✅ 3 outputs (incl. structured pydantic JSON), exactly 1 NODE_ERROR (RuntimeError), zero NODE_ERROR from MAX_TOKENS agents
Span export ✅ 0 plugin-scoped spans through a real global SpanProcessor, both modes
Duplicates Default mode caught a second wild at-least-once transport duplicate (byte-identical USER_MESSAGE_RECEIVED pair sharing one event_id, same timestamp/trace/content) — the documented _default behavior, now identifiable and dedupable. Exactly-once mode: zero duplicates, again
SAFETY finish reason Not reproduced — BLOCK_LOW_AND_ABOVE across three harm categories on a crime-narration prompt still yielded a benign STOP completion from gemini-2.5-flash. SAFETY/MALFORMED remain covered by parametrized unit tests only; the enum-name projection path is identical for all members

Two independent runs have now each caught a naturally occurring transport duplicate in default mode, and exactly-once mode has produced zero across every run — strong field evidence for both halves of the google#6465 fix at the exact commit under review.

@caohy1988

Copy link
Copy Markdown
Owner Author

Live BigQuery e2e evidence at d8fac0be — all green

Scratch dataset test-project-0728-467323.adk_pr12_e2e_20260806_184931 (7-day default table expiration, labeled), plugin loaded from this branch, gemini-2.5-flash on Vertex via ADC. Three live runs against real BigQuery.

Part A — LlmAgent + tool, default mode. Forced tool call. 17 rows, every event type present (LLM_REQUEST/LLM_RESPONSE ×2, TOOL_STARTING/TOOL_COMPLETED, AGENT_*, INVOCATION_*, AGENT_RESPONSE), event_id on 100% of rows, 17/17 distinct — no duplicates. Both LLM_RESPONSE rows carry attributes.finish_reason = "STOP" (google#5644 projection verified against real Vertex traffic).

Part B — graph Workflow with two FunctionNodes (one returns a dict, one raises). Exactly the rows google#6529 asked for:

  • NODE_OUTPUTattributes.adk.node.path = e2e_wf@1/step_one@1, output payload {"patient_id":"p-42","status":"found"} intact
  • NODE_ERRORnode_path = e2e_wf@1/step_two@1, status = ERROR, sanitized error_message, content.error_code = "ValueError"

Part C — exactly_once_delivery=True (separate table agent_events_xo). 8/8 rows landed via the committed-stream path: live create_write_stream + offset appends work, clean shutdown/finalize, 8/8 distinct event_id.

Views. Tables labeled adk_schema_version:2; all typed views created including the new v_node_output/v_node_error; v_llm_response exposes finish_reason end-to-end.

One observation (pre-existing, not a regression from this PR): views are named per-dataset (v_*, no table qualifier), so the Part C plugin re-created all views over agent_events_xo, replacing Part A's views over agent_events. The one-table-per-dataset assumption silently breaks when two plugins share a dataset with different table_ids — the new NODE view SQL itself is correct (verified manually against agent_events). Worth a known-limitation note if upstream ever asks about multi-table datasets.

Not covered here: the SSE streaming probe (this run was unary calls only). Terminal-metadata non-partial gating is unit-tested, but a live SSE pass is still worth doing before un-drafting.

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