feat(a2a): switch from callback-based to event-driven tool scheduler#21467
feat(a2a): switch from callback-based to event-driven tool scheduler#21467cocosheng-g merged 12 commits intomainfrom
Conversation
|
Hi @cocosheng-g, thank you so much for your contribution to Gemini CLI! We really appreciate the time and effort you've put into this. We're making some updates to our contribution process to improve how we track and review changes. Please take a moment to review our recent discussion post: Improving Our Contribution Process & Introducing New Guidelines. Key Update: Starting January 26, 2026, the Gemini CLI project will require all pull requests to be associated with an existing issue. Any pull requests not linked to an issue by that date will be automatically closed. Thank you for your understanding and for being a part of our community! |
cocosheng-g
changed the title
|
Size Change: +264 B (0%) Total Size: 26.2 MB ℹ️ View Unchanged
|
|
/gemini review |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces an event-driven scheduler to the Task class, providing an alternative to the existing callback-based CoreToolScheduler. The Task constructor now conditionally initializes either the new Scheduler (if isEventDrivenSchedulerEnabled is true) or the legacy scheduler. Key changes include new methods (setupEventDrivenScheduler, handleEventDrivenToolCallsUpdate, handleEventDrivenToolCall, checkInputRequiredState) to subscribe to TOOL_CALLS_UPDATE messages from a MessageBus, process tool call statuses (including output, terminal states, and confirmation requests), stash correlation IDs, publish status updates to the A2A event bus, and manage auto-execution (YOLO mode). The tool confirmation flow (_handleToolConfirmationPart) has been updated to publish TOOL_CONFIRMATION_RESPONSE messages via the MessageBus when a correlation ID is present, while maintaining a fallback for legacy callback-based confirmations. A new enableEventDrivenScheduler configuration setting, defaulting to true, controls this feature, and a new test file (task-event-driven.test.ts) validates the new scheduler's functionality. A review comment points out an issue in the new tool confirmation logic, specifically that the confirmed property in the TOOL_CONFIRMATION_RESPONSE message is incorrectly set to true only for ProceedOnce outcomes, suggesting it should be true for all outcomes that lead to tool execution (i.e., not Cancel or ModifyWithEditor).
gemini-cli
bot
added
the
area/agent
cocosheng-g
force-pushed
the
pr-17128-takeover
branch
2 times, most recently
from
b0632cc to
abc05ff
Compare
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request introduces an event-driven scheduler for tool calls within the Task class, allowing for a more reactive and asynchronous handling of tool execution and confirmations. The changes include integrating a new Scheduler class, conditionally enabling it via a new configuration setting (isEventDrivenSchedulerEnabled), and refactoring tool call lifecycle management to use a MessageBus for status updates and confirmation responses. New methods were added to handle event-driven tool call updates, manage pending confirmations, and determine input-required states. The schedule method was updated to pre-register tools and use a fire-and-forget approach for the event-driven scheduler, while tool confirmation handling was enhanced to support both legacy callbacks and message bus-based responses, including auto-approval in 'YOLO' mode. Comprehensive unit tests for the new event-driven scheduler functionality were also added.
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request significantly refactors the tool scheduling mechanism within the Highlights
Changelog
Activity
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
This change transitions packages/a2a-server to use the event-driven Scheduler by default. It replaces the legacy direct callback mechanism with a MessageBus listener in the Task class to handle tool status updates, live output, and confirmations. - Added experimental.enableEventDrivenScheduler setting (defaults to true). - Refactored Task.ts to support both legacy and event-driven schedulers. - Implemented bus-based tool confirmation responses using correlationId. - Exported Scheduler from packages/core. - Added unit tests for the event-driven flow in A2A.
…ely closing executor loop
cocosheng-g
force-pushed
the
pr-17128-takeover
branch
from
abc05ff to
9a856df
Compare
There was a problem hiding this comment.
Code Review
The pull request introduces an event-driven scheduler to the Task class, allowing it to handle tool call updates and confirmations via a message bus, with a fallback to the legacy scheduler. This involves adding new methods for event handling, managing tool confirmation correlation IDs, and updating the task state based on tool execution and approval flows, including a "YOLO" mode for auto-execution. A critical issue was identified where the legacy scheduler's cancelAll method does not properly abort its AbortSignal, potentially leading to dangling processes.
…ension * Preserved continuous currentAgentMessageId during each turn to prevent separate UI bubbles * Added processedToolCallIds to prevent state manager from triggering multiple generation cycles * Added test coverage for MCP Server user journey operations
This commit adds tests for ProceedAlwaysServer, ProceedAlwaysTool, and ProceedAlwaysAndSave, and fixes a missing outcome mapping in the A2A task implementation to support all PR 21467 event-driven flows.
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request is a significant refactoring that migrates the A2A server to a modern, event-driven scheduler, effectively addressing critical issues like deadlocks and race conditions, thereby improving robustness and reliability. While the architectural change is sound and the implementation includes comprehensive test coverage, a critical security regression has been identified. The new asynchronous flow bypasses the mechanism intended to prevent sensitive GCP credentials from leaking into tool calls during tool confirmation. This vulnerability needs to be addressed to ensure environment variable protection. Additionally, a potential memory leak related to MessageBus subscriptions has been identified, requiring proper cleanup.
This commit addresses PR 21467 code review: (1) Adds missing await on the message bus publish in _handleToolConfirmationPart to prevent GCP variables from prematurely restoring and leaking into external tools. (2) Adds dispose() to Task.ts and explicitly unregisters the message bus listener from cache when tasks reach terminal states in executor.ts, preventing unbounded memory leaks over the session.
This commit addresses edge cases for PR 21467: * Fixes a bug in YOLO mode where auto-approved tool calls incorrectly triggered an `input-required` task state by adding the call ID to the confirmed tools set. * Adds comprehensive test coverage for YOLO mode auto-approvals, parallel tool calls, executing background tools, and unknown tool ID confirmations.
* Ensures pendingCorrelationIds is safely deleted alongside pendingToolConfirmationDetails to prevent memory leaks in long-running tasks. * Adds a regression test to verify that delayed or spurious confirmations are safely ignored once the ID is cleaned up.
Address stalls and task abortion regressions in YOLO mode by implementing unconditional core policy approval and primary execution loop isolation. Add comprehensive unit test coverage for PolicyEngine and CoderAgentExecutor.
|
/patch preview |
|
🚀 [Step 1/4] Patch workflow(s) waiting for approval! 📋 Details:
⏳ Status: The patch creation workflow has been triggered and is waiting for deployment approval. Please visit the specific workflow links below and approve the runs. 🔗 Track Progress: |
…21467) Co-authored-by: Abhi <abhipatel@google.com> Co-authored-by: Adam Weidman <adamfweidman@google.com> # Conflicts: # packages/a2a-server/src/config/config.ts
|
🚀 [Step 2/4] Patch PR Created! 📋 Patch Details:
📝 Next Steps:
🔗 Track Progress: |
|
🚀 [Step 3/4] Patch Release Waiting for Approval! 📋 Release Details:
⏳ Status: The patch release has been triggered and is waiting for deployment approval. Please visit the specific workflow run link below and approve the deployment. You'll receive another update when it completes. 🔗 Track Progress: |
|
✅ [Step 4/4] Patch Release Complete! 📦 Release Details:
🎉 Status: Your patch has been successfully released and published to npm! 📝 What's Available:
🔗 Links: |
…oogle-gemini#21467) Co-authored-by: Abhi <abhipatel@google.com> Co-authored-by: Adam Weidman <adamfweidman@google.com>
…oogle-gemini#21467) Co-authored-by: Abhi <abhipatel@google.com> Co-authored-by: Adam Weidman <adamfweidman@google.com>
…oogle-gemini#21467) Co-authored-by: Abhi <abhipatel@google.com> Co-authored-by: Adam Weidman <adamfweidman@google.com>
…oogle-gemini#21467) Co-authored-by: Abhi <abhipatel@google.com> Co-authored-by: Adam Weidman <adamfweidman@google.com>
3 participants
No description provided.