[Bug]: Task state is not persisted to TaskStore after client disconnect

[azyobuzin]

What happened?

When the client-side connection is terminated, the EventConsumer stops processing. As a result, any changes to the task state after the disconnection are not persisted to the TaskStore. The task itself continues running in the background, but its updated state is no longer reflected in the TaskStore.

This issue was previously reported in this comment, but the parent issue was closed following a pull request merge.

Steps to Reproduce

To reproduce the issue, edit a2a-samples/samples/python/agents/helloworld as follows:

Update a2a-sdk to the latest commit of main branch (8c59545)

Implement HelloWorldAgentExecuter.execute:

from a2a.server.tasks import TaskUpdater
from a2a.utils import new_task

task = new_task(context.message)
updater = TaskUpdater(event_queue, task.id, task.context_id)
await updater.start_work()
await asyncio.sleep(5)
result = await self.agent.invoke()
await updater.complete(new_agent_text_message(result))

Add the following logger setup in __main__.py:

task_store_logger = logging.getLogger('a2a.server.tasks.inmemory_task_store')
task_store_logger.setLevel(logging.DEBUG)
task_store_logger.addHandler(logging.StreamHandler(sys.stdout))

Normal case

Send a request using the following curl command:

curl -X POST \
    -i -s \
    --header "Content-Type: application/json" \
    --data "{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"message/stream\",\"params\":{\"message\":{\"role\":\"user\",\"parts\":[{\"kind\":\"text\",\"text\":\"テストメッセージ\"}],\"messageId\":\"some-id\"}}}" \
    http://localhost:9999/

A normal run (without disconnection) produces logs like:

Attempting to get task with id: 42b2b321-37d9-4764-921c-2d2d0237be78
Task 42b2b321-37d9-4764-921c-2d2d0237be78 not found in store.
Task 42b2b321-37d9-4764-921c-2d2d0237be78 saved successfully.
Task 42b2b321-37d9-4764-921c-2d2d0237be78 saved successfully.
Task 42b2b321-37d9-4764-921c-2d2d0237be78 saved successfully.

Problematic case

Send the same request but cancel it by pressing Ctrl + C.

In this case, the last “saved successfully” log line is missing:

Attempting to get task with id: 42b2b321-37d9-4764-921c-2d2d0237be78
Task 42b2b321-37d9-4764-921c-2d2d0237be78 not found in store.
Task 42b2b321-37d9-4764-921c-2d2d0237be78 saved successfully.
Task 42b2b321-37d9-4764-921c-2d2d0237be78 saved successfully.

This shows that the call to updater.complete was not persisted.

Summary of the issue:

When the client disconnects, the EventConsumer stops, and subsequent state updates are no longer persisted by the TaskManager. This causes the actual progress of the task and the persisted state to diverge after a disconnect.

Relevant log output

Code of Conduct

• I agree to follow this project's Code of Conduct

[ovidiutaralesca]

Hi @azyobuzin , this was a sneaky issue. Great catch!

I've opened a small PR with a fix.