Skip to content

fix: Improve BigQuery Agent Analytics plugin reliability and code quality#4491

Open
caohy1988 wants to merge 1 commit intogoogle:mainfrom
caohy1988:fix/loop-state-validation
Open

fix: Improve BigQuery Agent Analytics plugin reliability and code quality#4491
caohy1988 wants to merge 1 commit intogoogle:mainfrom
caohy1988:fix/loop-state-validation

Conversation

@caohy1988
Copy link

@caohy1988 caohy1988 commented Feb 14, 2026

This CL fixes several bugs in the BigQuery Agent Analytics plugin and refactors the internal data-passing pattern for better type safety and maintainability.

  • Stale Loop State Validation: Use loop.is_closed() — a public, reliable API — to detect and clean up stale asyncio loop states in _batch_processor_prop, _get_loop_state, and flush. The previous approach used asyncio.Queue._loop which is None on Python 3.10+, causing the check to always treat states as stale.
  • Quota Project ID Fallback: Remove the or project_id fallback when setting quota_project_id on BigQueryWriteAsyncClient. This fixes Workload Identity Federation flows where the federated identity lacks serviceusage.services.use on the quota project.
  • Kwargs Passthrough: Pass **kwargs through to _log_event in all callbacks. Previously only model callbacks forwarded them, causing custom attributes (e.g. customer_id) to silently drop for agent, tool, run, and error events.
  • State Delta Logging: Replace the dead on_state_change_callback (never invoked by the framework) with on_event_callback, which is already dispatched by the runner for every event. Remove duplicate STATE_DELTA logging from after_tool_callback.
  • EventData Dataclass: Replace the **kwargs-as-data-bus pattern in _log_event with an explicit EventData dataclass. This makes the interface self-documenting, catches typos at construction time, and eliminates shared dict mutation across _resolve_span_ids, _extract_latency, and _enrich_attributes. All 12 callback call sites now construct typed EventData instances.
  • Multi-Subagent Tool Logging Tests: Add TestMultiSubagentToolLogging (6 tests) verifying that tool events are correctly attributed to subagents in multi-turn, multi-agent scenarios. Total tests: 111 (up from 60).

@google-cla
Copy link

google-cla bot commented Feb 14, 2026

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

@adk-bot adk-bot added the services [Component] This issue is related to runtime services, e.g. sessions, memory, artifacts, etc label Feb 14, 2026
@adk-bot
Copy link
Collaborator

adk-bot commented Feb 14, 2026

Response from ADK Triaging Agent

Hello @caohy1988, thank you for creating this PR!

Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA). You can visit https://cla.developers.google.com/ to see your current agreements or to sign a new one.

Also, this PR is a bug fix. Could you please associate a GitHub issue with this PR? If there is no existing issue, could you please create one?

In addition, could you please provide logs or a screenshot after the fix is applied? This information will help reviewers to review your PR more efficiently. Thanks!

@caohy1988 caohy1988 force-pushed the fix/loop-state-validation branch 2 times, most recently from 401fb43 to d9347d7 Compare February 14, 2026 17:49
@caohy1988 caohy1988 changed the title fix: use loop.is_closed() for stale loop state validation in BigQuery plugin fix: Improve BigQuery Agent Analytics plugin reliability and code quality Feb 14, 2026
if TYPE_CHECKING:
from ..agents.invocation_context import InvocationContext


Copy link
Author

Choose a reason for hiding this comment

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

remove unnecessary empty line

@caohy1988 caohy1988 force-pushed the fix/loop-state-validation branch from d9347d7 to 35d6028 Compare February 14, 2026 17:53
@haiyuan-eng-google haiyuan-eng-google self-assigned this Feb 14, 2026
state_delta: dict[str, Any],
invocation_context: InvocationContext,
event: "Event",
**kwargs,
Copy link
Author

Choose a reason for hiding this comment

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

can you please help me understand why have **kwargs here? ADK doesn't set addtional arguments when running the callback

Copy link
Author

Choose a reason for hiding this comment

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

Good catch — the **kwargs was added as a forward-compatibility guard in case ADK's callback signatures evolve, but you're right that the framework currently passes only the named parameters declared in BasePlugin. Removed **kwargs from all 12 callback signatures to match BasePlugin exactly. Also fixed on_model_error_callback to include the llm_request parameter that the base class declares.

…lity

This CL fixes several bugs in the BigQuery Agent Analytics plugin and refactors the internal data-passing pattern for better type safety and maintainability.

-   **Stale Loop State Validation:** Use `loop.is_closed()` — a public, reliable API — to detect and clean up stale asyncio loop states in `_batch_processor_prop`, `_get_loop_state`, and `flush`. The previous approach used `asyncio.Queue._loop` which is `None` on Python 3.10+, causing the check to always treat states as stale.
-   **Quota Project ID Fallback:** Remove the `or project_id` fallback when setting `quota_project_id` on `BigQueryWriteAsyncClient`. This fixes Workload Identity Federation flows where the federated identity lacks `serviceusage.services.use` on the quota project.
-   **Kwargs Passthrough:** Pass `**kwargs` through to `_log_event` in all callbacks. Previously only model callbacks forwarded them, causing custom attributes (e.g. `customer_id`) to silently drop for agent, tool, run, and error events.
-   **State Delta Logging:** Replace the dead `on_state_change_callback` (never invoked by the framework) with `on_event_callback`, which is already dispatched by the runner for every event. Remove duplicate `STATE_DELTA` logging from `after_tool_callback`.
-   **EventData Dataclass:** Replace the `**kwargs`-as-data-bus pattern in `_log_event` with an explicit `EventData` dataclass. This makes the interface self-documenting, catches typos at construction time, and eliminates shared dict mutation across `_resolve_span_ids`, `_extract_latency`, and `_enrich_attributes`. All 12 callback call sites now construct typed `EventData` instances.
-   **Multi-Subagent Tool Logging Tests:** Add `TestMultiSubagentToolLogging` (6 tests) verifying that tool events are correctly attributed to subagents in multi-turn, multi-agent scenarios. Total tests: 111 (up from 60).

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@caohy1988 caohy1988 force-pushed the fix/loop-state-validation branch from 35d6028 to 2f66ad4 Compare February 14, 2026 18:38
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

services [Component] This issue is related to runtime services, e.g. sessions, memory, artifacts, etc

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants