feat(a2a): switch from callback-based to event-driven tool scheduler #17128
Conversation
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.
Summary of ChangesHello @abhipatel12, 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
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
|
![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 a significant architectural improvement by migrating the A2A server to an event-driven tool scheduler. The changes are well-structured, and the new logic for handling tool calls via the message bus is clearly implemented. The addition of the enableEventDrivenScheduler flag provides a safe way to manage this transition. I've found one area with redundant code that can be simplified for better maintainability.
| if (this.scheduler instanceof Scheduler) { | ||
| await this.scheduler.schedule(updatedRequests, abortSignal); | ||
| } else { | ||
| await this.scheduler.schedule(updatedRequests, abortSignal); | ||
| } |
There was a problem hiding this comment.
This if/else block contains identical code in both branches. Since the schedule method exists on both Scheduler and CoreToolScheduler with compatible signatures for this call site (the return value is not used), this conditional logic is redundant and can be simplified to a single line. This will improve code clarity and maintainability.
await this.scheduler.schedule(updatedRequests, abortSignal);
|
Size Change: +31.4 kB (+0.14%) Total Size: 23.2 MB
ℹ️ View Unchanged
|
gemini-cli
bot
added
the
status/need-issue
| return scheduler; | ||
| } | ||
|
|
||
| private _setupEventDrivenScheduler(): Scheduler { |
There was a problem hiding this comment.
What's the history behind the '_' prefix? Both of these are private instance functions. It seems like it's inconsistently used.
| messageBus.subscribe( | ||
| MessageBusType.TOOL_CALLS_UPDATE, | ||
| (message: unknown) => { | ||
| const event = message as ToolCallsUpdateMessage; |
There was a problem hiding this comment.
It looks like this cast is only necessary because you aren't passing the right type parameter to subscribe. I think you can instead do:
messageBus.subscribe<ToolCallsUpdateMessage>(
MessageBusType.TOOL_CALLS_UPDATE,
(message: ToolCallsUpdateMessage) => {
if (event.type !== MessageBusType.TOOL_CALLS_UPDATE) {
return;
}
...
}
)| // Inject a dummy onConfirm for legacy UI compatibility if needed, | ||
| // though A2A should use the correlationId-based path now. | ||
| onConfirm: async () => {}, | ||
| } as ToolCallConfirmationDetails); |
There was a problem hiding this comment.
Unneeded cast, please delete.
In general you should rarely need to cast in TS if you have the proper typings available.
Despite this Gemini CLI seems to add them all over. I'd recommend doing a quick Cmd+F for as before checking in each change until we fix it.
|
|
||
| messageBus.subscribe( | ||
| MessageBusType.TOOL_CALLS_UPDATE, | ||
| (message: unknown) => { |
There was a problem hiding this comment.
nit: consider extracting the bulk of this out when you have lambdas spanning more than a dozen or so lines.
|
|
||
| const toolCalls = event.toolCalls; | ||
|
|
||
| toolCalls.forEach((tc) => { |
There was a problem hiding this comment.
this lambda seems like another great candidate for a named function.
| if (tc.status === 'awaiting_approval' && tc.confirmationDetails) { | ||
| // Bridge the new serializable details back to the legacy shape for A2A UI | ||
| const details = | ||
| tc.confirmationDetails as SerializableConfirmationDetails; |
There was a problem hiding this comment.
Looks like another unnecessary cast.
| this.eventBus?.publish(statusUpdate); | ||
| } | ||
|
|
||
| // 5. Handle Auto-Execution (YOLO) |
There was a problem hiding this comment.
nit: another thing that might help with readability/maintainability here is if each of these numbered sections was its own function.
| this.config.getApprovalMode() === ApprovalMode.YOLO) | ||
| ) { | ||
| logger.info(`[Task] Auto-approving tool call ${callId}`); | ||
| void messageBus.publish({ |
There was a problem hiding this comment.
We're not awaiting this promise. Any risk of race condition in letting this code interleave?
| if (confirmationDetails.type === 'edit') { | ||
| const payload = part.data['newContent'] | ||
| ? ({ | ||
| newContent: part.data['newContent'] as string, |
There was a problem hiding this comment.
Should we be using typeof === 'string' instead so this doesn't cause a runtime failure if it's something else?
| const payload = part.data['newContent'] | ||
| ? ({ | ||
| newContent: part.data['newContent'] as string, | ||
| } as ToolConfirmationPayload) |
|
Hi there! Thank you for your contribution to Gemini CLI. To improve our contribution process and better track changes, we now require all pull requests to be associated with an existing issue, as announced in our recent discussion and as detailed in our CONTRIBUTING.md. This pull request is being closed because it is not currently linked to an issue. Once you have updated the description of this PR to link an issue (e.g., by adding How to link an issue: Thank you for your understanding and for being a part of our community! |
2 participants
Summary
Migrate
packages/a2a-serverto use the modern, event-drivenSchedulerby default. This change replaces the legacy direct callback mechanism with aMessageBuslistener in theTaskclass, aligning A2A with the architectural direction of Gemini CLI.Details
packages/a2a-serverto use theSchedulerinstead ofCoreToolScheduler.correlationId.enableEventDrivenScheduler(defaulting totruefor A2A) to provide an escape hatch while ensuring immediate migration.Schedulerclass frompackages/coreto enable its consumption in A2A.TOOL_CONFIRMATION_RESPONSEevents instead.Related Issues
Related to moving towards a default true state for
enableEventDrivenScheduler.How to Validate
npm test -w @google/gemini-cli-a2a-server -- src/agent/task.test.ts src/agent/task-event-driven.test.tsnpm run preflightPre-Merge Checklist