[Bug]: Cluster mode: SubscribeToTask on a replica holding a paused ActiveTask returns a stale INPUT_REQUIRED snapshot and closes #1323

"""Resubscribe on a replica that holds a paused ActiveTask returns a stale snapshot.

a2a-sdk 1.2.2, cluster mode (VersionedDatabaseTaskStore + DatabaseTaskEventStream),
three replicas sharing one SQLite database.
"""
import asyncio
import os
import tempfile

from sqlalchemy.ext.asyncio import create_async_engine

from a2a.helpers.proto_helpers import get_message_text
from a2a.server.agent_execution import AgentExecutor, RequestContext
from a2a.server.cluster import DatabaseTaskEventStream, VersionedDatabaseTaskStore
from a2a.server.context import ServerCallContext
from a2a.server.events.event_queue import EventQueue
from a2a.server.request_handlers import DefaultRequestHandler
from a2a.server.tasks.task_updater import TaskUpdater
from a2a.types import a2a_pb2 as pb

DB = os.path.join(tempfile.mkdtemp(), "cluster.db")


class TwoTurnAgent(AgentExecutor):
    """Turn 1 asks for input; turn 2 works for ~1.6s, then completes."""

    async def execute(self, context: RequestContext, queue: EventQueue) -> None:
        if context.current_task is None:
            await queue.enqueue_event(pb.Task(
                id=context.task_id, context_id=context.context_id,
                status=pb.TaskStatus(state=pb.TASK_STATE_SUBMITTED),
                history=[context.message]))
        updater = TaskUpdater(queue, context.task_id, context.context_id)
        if get_message_text(context.message) == "hi":
            await updater.requires_input()
            return
        for _ in range(4):
            await updater.start_work()
            await asyncio.sleep(0.4)
        await updater.complete()

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


CARD = pb.AgentCard(name="repro", capabilities=pb.AgentCapabilities(streaming=True))
STORES = []


def replica() -> DefaultRequestHandler:
    engine = create_async_engine(f"sqlite+aiosqlite:///{DB}")
    store, stream = VersionedDatabaseTaskStore(engine), DatabaseTaskEventStream(engine)
    STORES.extend([store, stream])
    return DefaultRequestHandler(agent_executor=TwoTurnAgent(), task_store=store,
                                 event_stream=stream, agent_card=CARD)


def send(text: str, task_id: str | None = None) -> pb.SendMessageRequest:
    msg = pb.Message(message_id=text, role=pb.ROLE_USER, context_id="ctx",
                     parts=[pb.Part(text=text)])
    if task_id:
        msg.task_id = task_id
    return pb.SendMessageRequest(
        message=msg, configuration=pb.SendMessageConfiguration(return_immediately=True))


def label(event) -> str:
    status = getattr(event, "status", None)
    return f"{type(event).__name__}({pb.TaskState.Name(status.state)})" if status else type(event).__name__


async def resubscribe(name: str, handler: DefaultRequestHandler, task_id: str) -> None:
    events = []

    async def tail() -> None:
        async for event in handler.on_subscribe_to_task(
            pb.SubscribeToTaskRequest(id=task_id), ServerCallContext()
        ):
            events.append(label(event))

    try:
        await asyncio.wait_for(tail(), timeout=5)
        events.append("<stream closed>")
    except asyncio.TimeoutError:
        events.append("<still open after 5s>")
    print(f"resubscribe on {name}: {' -> '.join(events)}")


async def main() -> None:
    a, b, c = replica(), replica(), replica()
    for s in STORES:
        await s.initialize()
    ctx = ServerCallContext()

    task = await a.on_message_send(send("hi"), ctx)          # turn 1 on A -> INPUT_REQUIRED
    await asyncio.sleep(0.3)
    await b.on_message_send(send("go", task.id), ctx)        # turn 2 runs on B
    await asyncio.sleep(0.1)

    await asyncio.gather(
        resubscribe("A (ran turn 1)", a, task.id),
        resubscribe("C (never saw the task)", c, task.id),
    )
    final = await c.on_get_task(pb.GetTaskRequest(id=task.id), ctx)
    print("GetTask afterwards:", pb.TaskState.Name(final.status.state))


asyncio.run(main())