[Bug]: After a streamed task is cancelled, its background producer never finishes #1322

if isinstance(event, _RequestCompleted):
    if request_id is not None and event.request_id == request_id:
        return
    continue            # never marked done
elif isinstance(event, _RequestStarted):
    continue            # never marked done
try:
    yield event
finally:
    tapped_queue.task_done()

import asyncio
import uuid

from a2a.server.agent_execution import AgentExecutor, RequestContext
from a2a.server.context import ServerCallContext
from a2a.server.events import EventQueue
from a2a.server.request_handlers import DefaultRequestHandler
from a2a.server.tasks import InMemoryTaskStore, TaskUpdater
from a2a.types.a2a_pb2 import (
    AgentCapabilities,
    AgentCard,
    CancelTaskRequest,
    Message,
    Part,
    Role,
    SendMessageRequest,
    Task,
    TaskState,
    TaskStatus,
)


class SlowAgent(AgentExecutor):
    """Starts a task, then works "forever" until it is cancelled."""

    async def execute(self, context: RequestContext, event_queue: EventQueue) -> None:
        await event_queue.enqueue_event(
            Task(
                id=context.task_id,
                context_id=context.context_id,
                status=TaskStatus(state=TaskState.TASK_STATE_SUBMITTED),
            )
        )
        await TaskUpdater(event_queue, context.task_id, context.context_id).start_work()
        await asyncio.Event().wait()  # long-running work

    async def cancel(self, context: RequestContext, event_queue: EventQueue) -> None:
        await TaskUpdater(event_queue, context.task_id, context.context_id).cancel()


def producers_still_running() -> list[str]:
    return [
        t.get_name()
        for t in asyncio.all_tasks()
        if t.get_name().startswith("producer:") and not t.done()
    ]


async def main() -> None:
    handler = DefaultRequestHandler(
        SlowAgent(),
        InMemoryTaskStore(),
        AgentCard(capabilities=AgentCapabilities(streaming=True)),
    )
    call = ServerCallContext()
    request = SendMessageRequest(
        message=Message(
            message_id=str(uuid.uuid4()), role=Role.ROLE_USER, parts=[Part(text="hi")]
        )
    )

    # 1. A client streams the task.
    stream = handler.on_message_send_stream(request, call)
    first = await anext(stream)
    task_id = first.id
    print("streaming task", task_id)

    # 2. While it streams, the client cancels it.
    async def cancel_soon() -> None:
        await asyncio.sleep(0.2)
        result = await handler.on_cancel_task(CancelTaskRequest(id=task_id), call)
        print("cancel answered:", TaskState.Name(result.status.state))

    canceller = asyncio.create_task(cancel_soon())

    # 3. The stream receives the cancel and ends normally.
    async for event in stream:
        print("stream got:", type(event).__name__, TaskState.Name(event.status.state))
    print("stream closed")
    await canceller

    # 4. Everything is done, yet the task's producer is still running.
    await asyncio.sleep(2)
    print("producers still running after 2 s:", producers_still_running())

    # 5. Only shutting down the whole handler frees it.
    await handler.aclose()
    print("after handler.aclose():", producers_still_running())


asyncio.run(main())

if isinstance(event, _RequestStarted | _RequestCompleted):
    tapped_queue.task_done()
if isinstance(event, _RequestCompleted):
    ...