refactor: Separate task state building from persistence

[jmesnil]

This refactoring introduces TaskStateProcessor as a global singleton to
manage task state lifecycle and addresses a critical architectural issue
in task state management.

Problem:

• Task state was persisted multiple times during event propagation,
  causing redundant database writes and tight coupling between event
  processing and persistence

Solution:

• Created TaskStateProcessor as @ApplicationScoped singleton to maintain
  in-flight tasks globally
• Separated state building (TaskStateProcessor) from persistence (TaskStore)
• Tasks are now persisted once at appropriate lifecycle points instead of
  for each event
• Tasks are explicitly removed from TaskStateProcessor after final
  persistence, ensuring bounded memory usage

Changes:

• New: TaskStateProcessor - global singleton managing task state in memory
• Modified: DefaultRequestHandler - injects TaskStateProcessor and ensures
  cleanup after task completion
• Modified: ResultAggregator - uses TaskStateProcessor for state building,
  removes tasks after background consumption
• Modified: TaskManager - delegates state building to TaskStateProcessor
• Updated: All test files to inject TaskStateProcessor instances
• Fixed: Test isolation issues with defensive cleanup

Impact:

• Performance: Significantly reduced I/O operations (one write vs multiple)
• Memory: Bounded usage that scales with concurrent tasks, not total tasks
• Reliability: Improved test isolation and clearer task lifecycle

See doc/adr/0001_task_state_management_refactoring.md for detailed
architectural decision record.

• Follow the CONTRIBUTING Guide.
• Make your Pull Request title in the https://www.conventionalcommits.org/ specification.
• Ensure the tests pass
• Appropriate READMEs were updated (if necessary)

[gemini-code-assist]

Summary of Changes

Hello @jmesnil, I'm Gemini Code Assist¹! 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 introduces a significant refactoring to the task state management system. The primary goal is to optimize performance and improve reliability by decoupling the in-memory building of task state from its persistence to the database. By centralizing in-memory state management in a new TaskStateProcessor and ensuring that database writes occur only once per task lifecycle, the system reduces redundant I/O operations, manages memory more efficiently, and provides a clearer, more robust task lifecycle.

Highlights

• Separation of Concerns: Introduced a new TaskStateProcessor component to manage in-memory task state lifecycle, explicitly separating the concerns of state building from persistence. This addresses the architectural issue of redundant database writes during event propagation.
• Optimized Persistence: Task state is now persisted only once at appropriate lifecycle points (e.g., completion, cancellation) instead of multiple times for each event, significantly reducing I/O operations and database load.
• Bounded Memory Usage: Tasks are explicitly removed from the TaskStateProcessor after their final state is persisted, ensuring that memory usage scales with concurrent tasks rather than accumulating all processed tasks.
• Refactored Core Components: The DefaultRequestHandler, ResultAggregator, and TaskManager have been updated to leverage the new TaskStateProcessor for state management, streamlining their logic and improving maintainability.
• Improved Test Isolation: The new architecture naturally enhances test isolation, as each test can create its own TaskStateProcessor instance, leading to more reliable and reproducible test results.

Using Gemini Code Assist

The 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 /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

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 .gemini/ folder in the base of the repository. Detailed instructions can be found here.

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

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces a significant and beneficial refactoring by separating task state building from persistence through the new TaskStateProcessor. This change effectively reduces I/O operations and improves the task lifecycle management. The implementation is solid, and the accompanying Architectural Decision Record (ADR) is very helpful. I've provided a couple of suggestions: one to improve clarity in the ADR and another to simplify and fix a potentially buggy logic in DefaultRequestHandler. Overall, this is a great improvement to the architecture.

[kabir]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces a significant and valuable refactoring of task state management by introducing TaskStateProcessor. This change effectively separates in-memory state building from persistence, which should reduce I/O operations and improve performance as described in the ADR. The introduction of TaskStateProcessor and the delegation of state management from TaskManager and ResultAggregator are well-executed.

My review includes two main points:

1. An architectural inconsistency in onMessageSendStream where state is still persisted on every event, which seems to contradict the main goal of this refactoring.
2. A suggestion to improve the robustness of ResultAggregator by enforcing its dependency on TaskStateProcessor via a null check in the constructor.

Overall, this is a great improvement to the architecture. Addressing these points will enhance consistency and robustness.

[gemini-code-assist]

This change introduces task state persistence on every event for streaming responses. While the comment explains the rationale is to save state before sending an event to the client, this seems to contradict the primary goal of this refactoring, which is to reduce I/O and persist the task state only once at the end of its lifecycle.

The PR description states: "Tasks are now persisted once at appropriate lifecycle points instead of for each event". This change for streaming calls deviates from that principle.

If this per-event persistence is indeed necessary for streaming to handle client disconnects and allow for resubscription, this trade-off should be explicitly documented in the Architectural Decision Record (ADR) to clarify why streaming responses are treated differently from blocking/non-blocking calls. The current ADR gives the impression that all task processing now follows the "persist once" model.

[gemini-code-assist]

The TaskStateProcessor seems to be a required dependency for the ResultAggregator. However, the constructor allows it to be null, which forces a null check later in the code (lines 221-224).

To improve robustness and clarify the contract of this class, I suggest making stateProcessor a non-null dependency. You can enforce this by adding a null check in the constructor, similar to how it's done in TaskManager. This would also make the class more consistent with others that have required dependencies.

This change would also require updating ResultAggregatorTest to provide a mock or real TaskStateProcessor instance instead of null.

    public ResultAggregator(TaskManager taskManager, Message message, Executor executor, TaskStateProcessor stateProcessor) {
        io.a2a.util.Assert.checkNotNullParam("taskManager", taskManager);
        io.a2a.util.Assert.checkNotNullParam("executor", executor);
        io.a2a.util.Assert.checkNotNullParam("stateProcessor", stateProcessor);
        this.taskManager = taskManager;
        this.message = message;
        this.executor = executor;
        this.stateProcessor = stateProcessor;
    }

[jmesnil]

superseded by #515