Multi-server deployment

[mykytanetipa]

on_message_send / on_message_send_stream

Both handlers share _setup_active_task (https://github.com/a2aproject/a2a-python/blob/v1.1.3/src/a2a/server/request_handlers/default_request_handler_v2.py#L192), so they
share their multi-server problems. 3 issues:

Issue 1.1: Two replicas writing one task silently lose each other's writes

Issue: TaskStore.save() has no version parameter and no return value, so no
implementation can reject a stale write.
Concurrent turns on two replicas produce last-writer-wins with no error anywhere.

Current behaviour:

@abstractmethod
async def save(self, task: Task, context: ServerCallContext) -> None:
    """Saves or updates a task in the store."""

@abstractmethod
async def get(self, task_id: str, context: ServerCallContext) -> Task | None:
    """Retrieves a task from the store by ID."""

The shipped DB implementation (server/tasks/database_task_store.py):

async def save(self, task: Task, context: ServerCallContext) -> None:
    """Saves or updates a task in the database for the resolved owner."""
    await self._ensure_initialized()
    owner = self.owner_resolver(context)
    db_task = self._to_orm(task, owner)
    async with self.async_session_maker.begin() as session:
        await session.merge(db_task)

2 replicas do concurrent write on same task_id -> Both save calls return None. Both replicas believe they succeeded. (while 1 was silently overwritten)

What the Go SDK did.

Go put the version in the task store interface (https://github.com/a2aproject/a2a-go/blob/v0.3.15/a2asrv/tasks.go#L59):

type TaskStore interface {
	// Save stores a task. Implementations might choose to store event and use the previous known TaskVersion
	// for optimistic concurrency control during updates.
	Save(ctx context.Context, task *a2a.Task, event a2a.Event, prev *a2a.Task, prevVersion a2a.TaskVersion) (a2a.TaskVersion, error)

	// Get retrieves a task by ID. If a Task doesn't exist the method should return [a2a.ErrTaskNotFound].
	Get(ctx context.Context, taskID a2a.TaskID) (*a2a.Task, a2a.TaskVersion, error)

	List(ctx context.Context, req *a2a.ListTasksRequest) (*a2a.ListTasksResponse, error)
}

type TaskVersion int64

// TaskVersionMissing is a special value used to denote that task version is not being tracked.
var TaskVersionMissing TaskVersion = 0

func (v TaskVersion) After(another TaskVersion) bool {
	// Consider every state "latest" if the version is not tracked.
	if another == TaskVersionMissing {
		return true
	}
	if v == TaskVersionMissing {
		return false
	}
	return another < v
}

s.mu.Lock()
defer s.mu.Unlock()

version := a2a.TaskVersion(1)
if stored := s.tasks[task.ID]; stored != nil {
  if prevVersion != a2a.TaskVersionMissing && stored.version != prevVersion {
    return a2a.TaskVersionMissing, a2a.ErrConcurrentTaskModification
  }
  version = stored.version + 1
}

res, err := tx.ExecContext(ctx, `
  UPDATE task SET
    state = ?, last_updated = ?, task_json = ?
  WHERE id = ? AND last_updated = ?
`, task.Status.State, newVersion, string(taskJSON), task.ID, int64(prevVersion))
...
if rows == 0 {
  return a2a.TaskVersionMissing, fmt.Errorf("optimistic concurrency failure: task updated by another transaction")
}

Proposed fix (Python SDK):

A parallel implementation of VersionedTaskStore, no modification for exsisting to exsisting interface to avoid breaking changes.

only is_missing and is_after are used in code

# a2a/server/cluster/version.py
class TaskVersion:
    MISSING: 'Final[TaskVersion]'

    def __init__(self, value: int | str) -> None:
        self._value = value

    def __eq__(self, other: object) -> bool:
        return isinstance(other, TaskVersion) and self._value == other._value

    def __hash__(self) -> int:
        return hash(self._value)

    def __repr__(self) -> str:
        return 'TaskVersion.MISSING' if self.is_missing else f'TaskVersion({self._value!r})'

    @property
    def is_missing(self) -> bool:
        return self._value == 0

    def is_after(self, other: 'TaskVersion') -> bool:
        """True if self represents a strictly later state than other."""
        if other.is_missing:
            return True
        if self.is_missing:
            return False
        return other._value < self._value  # noqa: SLF001


TaskVersion.MISSING = TaskVersion(0)

# a2a/server/cluster/task_store.py
class ConcurrentTaskModificationError(Exception):
    """Raised by VersionedTaskStore.save when prev_version is stale.

    Callers treat this as "reload and decide", never as a fatal error.
    """


class VersionedTaskStore(ABC):
    """TaskStore with optimistic concurrency control.

    Distinct from TaskStore so existing implementations keep working untouched;
    wrap one in LegacyTaskStoreAdapter to use it here.
    """

    @abstractmethod
    async def save(
        self,
        task: Task,
        *,
        event: Event | None,
        prev: Task | None,
        prev_version: TaskVersion,
        context: ServerCallContext,
    ) -> TaskVersion:
        """Persist `task`, returning its new version.

        Args:
            event: the event that produced this state, or None for a direct write.
                Stores backed by an event log should persist it in the same
                transaction as the task row.
            prev: the task as previously read, for implementations that diff.
            prev_version: the version `task` was derived from. Implementations
                MUST raise ConcurrentTaskModificationError if the stored version
                differs, unless prev_version is TaskVersion.MISSING.

        Raises:
            ConcurrentTaskModificationError: prev_version is stale.
        """

    @abstractmethod
    async def get(
        self, task_id: str, context: ServerCallContext
    ) -> tuple[Task | None, TaskVersion]:
        """Return the task and its current version, or (None, MISSING)."""

    @abstractmethod
    async def list(
        self, params: ListTasksRequest, context: ServerCallContext
    ) -> ListTasksResponse: ...

    @abstractmethod
    async def delete(self, task_id: str, context: ServerCallContext) -> None: ...

# a2a/server/cluster/database_task_store.py  (sketch of the delta)
async def save(self, task, *, event, prev, prev_version, context) -> TaskVersion:
    owner = self.owner_resolver(context)
    new_version = time.time_ns()

    async with self.async_session_maker.begin() as session:
        if prev_version.is_missing:
            # First write for this task: INSERT, letting the PK collide loudly
            # rather than silently clobbering a concurrent creator.
            session.add(self._to_orm(task, owner, version=new_version))
            try:
                await session.flush()
            except IntegrityError as e:
                raise ConcurrentTaskModificationError(task.id) from e
        else:
            result = await session.execute(
                update(self.task_model)
                .where(
                    self.task_model.id == task.id,
                    self.task_model.owner == owner,
                    self.task_model.version == prev_version._value,   # the CAS
                )
                .values(**self._to_columns(task, version=new_version))
            )
            if result.rowcount == 0:
                raise ConcurrentTaskModificationError(task.id)

        if event is not None:
            # Same transaction as the task row: this is what makes the event log
            # a valid source for cross-replica replay (see §2).
            session.add(self._to_event_row(task.id, event, version=new_version))

    return TaskVersion(new_version)

---

Issue 1.2: Caching an ActiveTask skips the terminal-state check

Issue: The cached path returns an existing ActiveTask without calling
start() (contains the only terminal guard), so the guard that rejects work on a completed task never runs -> on a multi-replica
deployment a task can reach a terminal state on another replica while this one still holds a live
ActiveTask for it.

Current behaviour:

# server/agent_execution/active_task_registry.py (get_or_create() method)
async def get_or_create(
  ...
) -> ActiveTask:
    with self._lock:
        ...
        if task_id in self._active_tasks:
            return self._active_tasks[task_id]

        ...

    await active_task.start(
        call_context=call_context,
        create_task_if_missing=create_task_if_missing,
    )
    return active_task

# server/agent_execution/active_task.py (start() method)
async def start(
  ...
) -> None:
    ...
    task = await self._task_manager.get_task()

    if task:
        self._task_created.set()
        if task.status.state in TERMINAL_TASK_STATES:
            raise InvalidParamsError(
                message=f'Task {task.id} is in terminal state: {task.status.state}'
            )
    elif not create_task_if_missing:
        raise TaskNotFoundError
    ...

What the Go SDK did: NaN, the check is not attached to a cached object, store is re-read on every entry point.

and Resubscribe likewise consults the store rather than a registry
(distributed_manager.go:62). Because Go's distributed manager holds no per-task state, there is
no cache to go stale. This is the payoff of the "no in-process map" property, not a separate fix.

e.g. Task becomes COMPLETED on 1st replica, but 2nd can still wait in the INPUT_REQUIRED and new request on it would skip the terminal guard and rewrite the status (rewrite in concurrent requests case).

Proposed fix.

async def get_or_create(
    self, task_id, call_context, context_id=None,
    create_task_if_missing=False, initial_message=None,
) -> ActiveTask:
    with self._lock:
        if self._closed:
            raise RuntimeError('ActiveTaskRegistry is closed')
        existing = self._active_tasks.get(task_id)

    if existing is not None:
        # Reuse still has to re-check the store: another replica may have
        # advanced or finished this task while this ActiveTask sat idle
        # between turns. start() is not re-run, so do the check here.
        await existing.revalidate(call_context)
        return existing
    ...

# ActiveTask
async def revalidate(self, call_context: ServerCallContext) -> None:
    """Re-read the store before reusing this ActiveTask for a new request.

    Mirrors the checks start() performs on first use. Only call between
    requests: it drops the cached snapshot (see TaskManager.invalidate).
    """
    async with self._lock:
        if self._is_finished.is_set():
            raise InvalidParamsError(
                f'Task {self._task_id} is already completed. Cannot start it again.'
            )
        self._task_manager._call_context = call_context
        self._task_manager.invalidate()
        task = await self._task_manager.get_task()
        if task is None:
            raise TaskNotFoundError
        if task.status.state in TERMINAL_TASK_STATES:
            # Terminal elsewhere: retire the local ActiveTask so the next
            # request does not find it either.
            self._is_finished.set()
            self._request_queue.shutdown(immediate=True)
            raise InvalidParamsError(
                message=f'Task {task.id} is in terminal state: {task.status.state}'
            )

Issue 1.3: input_required / multi-turn resume across replicas (fixed by 1.2)

An ActiveTask survives input_required, so the next turn arriving after another replica advanced the task is computed against a stale snapshot.

on_subscribe_to_task

Issue 2.1: Resubscribe on the wrong replica hangs forever

Issue:

• Without a shared store the task does not exist on the other replica, so start() raises TaskNotFoundError -> the resubscribe fails outright.

• With a shared store the task is found, so start() succeeds, but on this replica it means creating a new ActiveTask with its own producer / consumer pair. That pair has no work to do (the agent runs on the other replica), so the client receives only the existing task snapshot and then the stream hangs with no further events.

+ resource leak

# default_request_handler_v2.py
async def on_subscribe_to_task(
    self,
    params: SubscribeToTaskRequest,
    context: ServerCallContext,
) -> AsyncGenerator[Event, None]:
    task_id = params.id

    active_task = await self._active_task_registry.get_or_create(
        task_id,
        call_context=context,
        create_task_if_missing=False,
    )

    async for event in active_task.subscribe(include_initial_task=True):
        yield event

# active_task.py (get_or_create() method)
existing = self._active_tasks.get(task_id) # separete process -> None

if existing is None:
    # create new task
    task_manager = TaskManager(...)

    active_task = ActiveTask(...)
    self._active_tasks[task_id] = active_task
...
await active_task.start(
     call_context=call_context,
     create_task_if_missing=create_task_if_missing,
)

# active_task.py (start() method)
# creates own producer that will wait for "await self._request_queue.get()"
self._producer_task = asyncio.create_task(
    self._run_producer(), name=f'producer:{self._task_id}'
)
self._consumer_task = asyncio.create_task(
    self._run_consumer(), name=f'consumer:{self._task_id}'
)

What the Go SDK did:

Go treats resubscribe as "read the snapshot, then tail a shared log", with the store as the source
of truth for existence. distributedManager.Resubscribe (internal/taskexec/distributed_manager.go:61-70):

func (m *distributedManager) Resubscribe(ctx context.Context, taskID a2a.TaskID) (Subscription, error) {
	if _, _, err := m.taskStore.Get(ctx, taskID); err != nil {
		return nil, a2a.ErrTaskNotFound
	}
	queue, err := m.queueManager.GetOrCreate(ctx, taskID)
	if err != nil {
		return nil, fmt.Errorf("failed to get event queue: %w", err)
	}
	return newRemoteSubscription(queue, m.taskStore, taskID), nil
}

No registry lookup, no producer spawned. The real work is in remoteSubscription.Events
(internal/taskexec/subscription.go):

2 main steps:

1. fetch versioned task snapshot from Task Store
2. stream remote event trough process shared db layer

func (s *remoteSubscription) Events(ctx context.Context) iter.Seq2[a2a.Event, error] {
	return func(yield func(a2a.Event, error) bool) {
		if s.consumed {
			yield(nil, fmt.Errorf("subscription already consumed"))
			return
		}
		s.consumed = true

		defer func() {
			if err := s.queue.Close(); err != nil {
				log.Warn(ctx, "queue close failed", "error", err)
			}
		}()

		task, snapshotVersion, err := s.store.Get(ctx, s.tid)
		if err != nil && !errors.Is(err, a2a.ErrTaskNotFound) {
			yield(nil, fmt.Errorf("task snapshot loading failed: %w", err))
			return
		}

		if task != nil {
			if !yield(task, nil) {
				return
			}
			if task.Status.State.Terminal() {
				return
			}
		}

		for {
			event, version, err := s.queue.Read(ctx)
			if version != a2a.TaskVersionMissing && !version.After(snapshotVersion) {
				log.Info(ctx, "skipping old event", "event", event, "version", version)
				continue
			}
			if err != nil {
				yield(nil, err)
				return
			}
			if !yield(event, nil) {
				return
			}
			if taskupdate.IsFinal(event) {
				return
			}
		}
	}
}

m.queueManager is the injected eventqueue.Manager from ClusterConfig (a2asrv/handler.go). In the reference cluster deployment it polls the shared event table (examples/clustermode/server/eventqueue.go):

rows, err := db.QueryContext(ctx, `
SELECT event_json, task_version, id
FROM task_event
WHERE task_id = ? AND id > ?
ORDER BY id ASC
LIMIT 10
`, taskID, pollFromID)

• Version-filtered replay — :135. Events already reflected in the snapshot are dropped, so a fan our consumer does not duplicate.

Proposed fix:

# a2a/server/cluster/event_bus.py
@dataclass(frozen=True)
class VersionedEvent:
    """An event plus the task version produced by applying it.

    `version` is what makes replay safe: a subscriber that already holds a
    snapshot at version V discards anything not strictly after V.
    """
    event: Event
    version: TaskVersion


class TaskEventBus(ABC):
    """Cross-replica delivery of task events."""

    @abstractmethod
    async def publish(self, task_id: str, event: VersionedEvent) -> None:
        """Publish one event for `task_id` to all replicas.

        Stores that persist events transactionally with the task (recommended)
        may implement this as a no-op and let `subscribe` read the log directly.
        """

    @abstractmethod
    def subscribe(
        self, task_id: str, *, after: TaskVersion
    ) -> AsyncIterator[VersionedEvent]:
        """Yield events for `task_id` newer than `after`.

        Must block rather than end when no events are available. The iterator
        ends only when closed by the consumer or the channel is destroyed.
        """

    @abstractmethod
    async def destroy(self, task_id: str) -> None:
        """Release resources for a task that reached a terminal state."""

@validate_request_params
@validate(
    lambda self: self._agent_card.capabilities.streaming,
    'Streaming is not supported by the agent',
)
async def on_subscribe_to_task(
    self,
    params: SubscribeToTaskRequest,
    context: ServerCallContext,
) -> AsyncGenerator[Event, None]:
    task_id = params.id

    # Fast path: this replica is running the agent, so tap its live queue.
    local = await self._active_task_registry.get(task_id)
    if local is not None:
        async for event in local.subscribe(include_initial_task=True):
            yield event
        return

    # Otherwise serve from shared state. Do NOT create an ActiveTask: that
    # spawns a producer that will never receive work and a subscriber queue
    # nothing will ever write to (the pre-fix hang).
    async for event in self._subscribe_remote(task_id, context):
        yield event

async def _subscribe_remote(
    self, task_id: str, context: ServerCallContext
) -> AsyncGenerator[Event, None]:
    """Snapshot from the store, then tail the shared bus.

    Mirrors a2a-go's remoteSubscription.Events.
    """
    task, snapshot_version = await self._task_store.get(task_id, context)
    if task is None:
        raise TaskNotFoundError

    yield task

    if task.status.state in TERMINAL_TASK_STATES:
        # Nothing further will be produced; end the stream cleanly instead
        # of tailing a channel that will never see another event.
        return

    subscription = self._event_bus.subscribe(task_id, after=snapshot_version)
    try:
        async for versioned in subscription:
            # Drop anything already reflected in the snapshot. Required
            # because a log-backed bus replays and a pub/sub bus may
            # redeliver.
            if not versioned.version.is_after(snapshot_version):
                continue
            yield versioned.event
            if _is_final(versioned.event):
                return
    finally:
        await subscription.aclose()

+ fix events fan out to subscribers:

async def _enqueue_to_subscribers(
    self, event: Event, updated_task: Task | None
) -> None:
    if updated_task is not None:
        updated_task_copy = Task()
        updated_task_copy.CopyFrom(updated_task)
        if event is updated_task:
            event = updated_task_copy
        updated_task = updated_task_copy

    await self.active_task._event_queue_subscribers.enqueue_event(
        cast('Any', (event, updated_task))
    )

    # Fan out to other replicas. Skipped for internal lifecycle sentinels
    # and exceptions, which are meaningless outside this process (Issue 7.4).
    bus = self.active_task._event_bus
    if bus is not None and _is_publishable(event):
        await bus.publish(
            self.active_task._task_id,
            VersionedEvent(
                event=event,
                version=self.active_task._task_manager.current_version,
            ),
        )

    self.active_task._event_queue_agent.task_done()

No breaking changes:
ActiveTask gains an optional _event_bus, defaulting to None, which preserves today's
single-process behaviour exactly.

on_cancel_task

Issue 3.1 — Cancel on the wrong replica doesn't cancel actual executing task

Issue: Same as subscribe: with shared Task Store new ActiveTask instance is created for another replica process -> cancel on the wrong object -> main replica still runs task

Current behaviour: (same get_or_create())

The handler (default_request_handler_v2.py:160-181):

async def on_cancel_task(
    self,
    params: CancelTaskRequest,
    context: ServerCallContext,
) -> Task | None:
    task_id = params.id

    try:
        active_task = await self._active_task_registry.get_or_create(
            task_id, call_context=context, create_task_if_missing=False
        )
        result = await active_task.cancel(context)
    except InvalidParamsError as e:
        raise TaskNotCancelableError from e

What the Go SDK did:

Go's distributed cancel never waits on a local object. It writes to the shared work queue and
waits on the shared event stream (internal/taskexec/distributed_manager.go):

func (m *distributedManager) Cancel(ctx context.Context, params *a2a.TaskIDParams) (*a2a.Task, error) {
	task, _, err := m.taskStore.Get(ctx, params.ID)
	if err != nil {
		return nil, fmt.Errorf("failed to load a task: %w", err)
	}

	if task.Status.State == a2a.TaskStateCanceled {
		return task, nil
	}

	if task.Status.State.Terminal() {
		return nil, fmt.Errorf("task in non-cancelable state %q: %w", task.Status.State, a2a.ErrTaskNotCancelable)
	}

	queue, err := m.queueManager.GetOrCreate(ctx, params.ID)
	if err != nil {
		return nil, fmt.Errorf("failed to get or create queue: %w", err)
	}

	if _, err := m.workQueue.Write(ctx, &workqueue.Payload{
		Type:         workqueue.PayloadTypeCancel,
		TaskID:       params.ID,
		CancelParams: params,
	}); err != nil {
		return nil, fmt.Errorf("failed to create work item: %w", err)
	}

	subscription := newRemoteSubscription(queue, m.taskStore, params.ID)
	var cancelationResult a2a.SendMessageResult
	var cancelationErr error
	for event, err := range subscription.Events(ctx) {
		if err != nil {
			cancelationErr = err
			break
		}
		if taskupdate.IsFinal(event) {
			if result, ok := event.(a2a.SendMessageResult); ok {
				cancelationResult = result
			}
			break
		}
	}
	if cancelationResult == nil && cancelationErr != nil {
		task, _, err := m.taskStore.Get(ctx, params.ID)
		if err != nil {
			cancelationErr = err
		} else {
			cancelationResult = task
		}
	}
	return convertToCancelationResult(ctx, cancelationResult, cancelationErr)
}

Proposed fix:

implement _cancel_remote that relies on Task Store:

@validate_request_params
async def on_cancel_task(
    self,
    params: CancelTaskRequest,
    context: ServerCallContext,
) -> Task | None:
    task_id = params.id

    task, version = await self._task_store.get(task_id, context)
    if task is None:
        raise TaskNotFoundError

    # Idempotent: already cancelled is success, not an error.
    if task.status.state == TaskState.TASK_STATE_CANCELED:
        return task
    if task.status.state in TERMINAL_TASK_STATES:
        raise TaskNotCancelableError

    local = await self._active_task_registry.get_for_read(task_id)
    if local is not None:
        # This replica owns the execution: stop it directly.
        try:
            return await local.cancel(context)
        except InvalidParamsError as e:
            raise TaskNotCancelableError from e

    # Executing elsewhere (or nowhere). Record the cancellation in shared
    # state and let the owning replica notice via CAS -- see (c).
    return await self._cancel_remote(task_id, task, version, context)

async def _cancel_remote(
    self,
    task_id: str,
    task: Task,
    version: TaskVersion,
    context: ServerCallContext,
) -> Task:
    """Cancel a task whose agent is running on another replica.

    No RPC to that replica. Writing the cancelled state bumps the task
    version, so the owning replica's next save raises
    ConcurrentTaskModificationError and it abandons the execution
    (see _abort_on_conflict). This is a2a-go's mechanism, and it is why
    Issue 1.1's CAS is a prerequisite for cancel working at all.
    """
    cancelled = Task()
    cancelled.CopyFrom(task)
    cancelled.status.state = TaskState.TASK_STATE_CANCELED
    cancelled.status.timestamp.GetCurrentTime()

    event = TaskStatusUpdateEvent(
        task_id=task_id,
        context_id=task.context_id,
        status=cancelled.status,
        final=True,
    )

    for attempt in range(_MAX_CANCEL_ATTEMPTS):
        try:
            new_version = await self._task_store.save(
                cancelled,
                event=event,
                prev=task,
                prev_version=version,
                context=context,
            )
        except ConcurrentTaskModificationError:
            # The owning replica wrote between our read and our write.
            # Re-read: it may have finished, or we may just need a newer
            # version to CAS against.
            task, version = await self._task_store.get(task_id, context)
            if task is None:
                raise TaskNotFoundError from None
            if task.status.state == TaskState.TASK_STATE_CANCELED:
                return task
            if task.status.state in TERMINAL_TASK_STATES:
                raise TaskNotCancelableError from None
            cancelled.CopyFrom(task)
            cancelled.status.state = TaskState.TASK_STATE_CANCELED
            cancelled.status.timestamp.GetCurrentTime()
            continue

        # Tell subscribers on every replica, including any tailing this task.
        await self._event_bus.publish(
            task_id, VersionedEvent(event=event, version=new_version)
        )
        return cancelled

    raise InternalError(
        message=f'Could not cancel task {task_id} after '
                f'{_MAX_CANCEL_ATTEMPTS} attempts due to concurrent updates'
    )

+ ensure that running task is actually canceled:

async def _process_event(self, event: Event) -> None:
    try:
        await self._process_event_inner(event)
    except ConcurrentTaskModificationError:
        self.active_task._task_manager.invalidate()
        task = await self.active_task._task_manager.get_task()
        if task is not None and task.status.state == TaskState.TASK_STATE_CANCELED:
            # Cancelled by another replica. Stop producing; do not overwrite
            # the cancelled state with whatever this execution was doing.
            logger.info(
                'Consumer[%s]: task cancelled elsewhere, aborting execution',
                self.active_task._task_id,
            )
            await self.active_task._abort_on_conflict(task)
            return
        if task is not None and task.status.state in TERMINAL_TASK_STATES:
            await self.active_task._handle_terminal_state(task)
            return
        await self._process_event_inner(event)