Fix: Prevent DefaultRequestHandler rejection by tracking per-agent task state

[ki1r0]

Description

This PR resolves #304 and also addresses a related issue discovered during the investigation and fix. The root cause has been identified in the comment of #304

✅ Summary of Fixes
🟢 Scenario 1 — Fix for issue #304
Problem: The host agent generates a new task_id, either regardless of the situation or when it fails to retrieve the task_id.
Fix: The host agent now:
Checks for an existing task for the target agent.
Reuses the task_id only if the task_id is successfully retrieved and task is not in a terminal state (TaskState.completed).
Otherwise, it sets task_id to None.
This allows the DefaultRequestHandler to generate the task_id and avoid TaskNotFound error.

🟠 Scenario 2 — New issue discovered and addressed in this PR
Problem: While addressing Scenario 1, it became clear that using a universal task_id across multiple remote agents introduces conflicts:
For example, remote_agent_1 generates a task_id, which is then reused by remote_agent_2.
Since remote_agent_2 has never seen that task ID before, the DefaultRequestHandler rejects it when there is a task_id an there is no task, throwing the same TaskNotFound error as in scenario 1
Fix: Introduced a per-agent tracking dictionary (state['agent_tasks']) that stores:
task_id, context_id, and status independently for each agent.
Ensures that each remote agent uses its own correct task state and ID, avoiding cross-agent contamination.

🧪 Scope
This PR applies the fix only in the demonstrated agent as a proof of concept.

Similar logic should be extended to other samples for consistent multi-agent compatibility.
Other samples with the same issue:

• a2a-mcp-without-framework
• azureaifoundry_sdk
• headless_agent_auth

Thank you for opening a Pull Request!
Before submitting your PR, there are a few things you can do to make sure it goes smoothly:

• Follow the CONTRIBUTING Guide.

Fixes #304 🦕

[gemini-code-assist]

Summary of Changes

Hello @ki1r0, 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 aims to prevent DefaultRequestHandler rejections by refining how task IDs are managed within a multi-agent system. It addresses two main scenarios: ensuring the host agent correctly reuses or generates task IDs based on task state, and resolving conflicts arising from a universal task ID by introducing per-agent task state tracking. This ensures each agent maintains its own correct task state, avoiding errors caused by shared or incorrectly reused task identifiers.

Highlights

• Fix for Issue #304: The host agent now intelligently reuses task_id only if the task is not in a terminal state, otherwise, it forces a new task_id generation by DefaultRequestHandler.
• Per-Agent Task State Tracking: Introduced a state['agent_tasks'] dictionary to store task_id, context_id, and status independently for each remote agent, preventing cross-agent contamination of task states.
• Context ID Management: context_id is now reset when a task is completed to ensure a fresh context for subsequent interactions with the same agent.
• Proof of Concept: The fix is currently applied only to the airbnb_planner_multiagent host agent as a proof of concept, with a note that similar logic should be extended to other samples.

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 in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

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 issue 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 or fill out our survey to provide feedback.

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 per-agent task state tracking mechanism to resolve issues with task and context ID management in a multi-agent scenario. The changes are logical and directly address the problems described. I've identified a critical issue regarding a missing import that would cause a runtime error, and a medium-severity suggestion to make the code more idiomatic and concise. Once these points are addressed, the changes look solid.

[holtskinner]

@rajeshvelicheti I think this was your sample, can you review?

[yedubi]

@holtskinner @rajeshvelicheti
also please see my comment, seems like this fix wouldn't work properly
#304 (comment)