From e1ec71a845c310210f38ba6f4da1f1770e5d327b Mon Sep 17 00:00:00 2001 From: Mykyta Netipa Date: Mon, 14 Sep 2026 22:17:04 +0000 Subject: [PATCH] docs(samples): add runnable multi-replica cluster-mode sample samples/clustermode runs several DefaultRequestHandlerV2 replicas sharing one database-backed versioned store and event stream, with a driver that exercises send / resubscribe / cancel routed across replicas. Serves as end-to-end documentation for deploying the server without task affinity. --- .github/actions/spelling/allow.txt | 2 + samples/clustermode/README.md | 74 +++++++++ samples/clustermode/__init__.py | 0 samples/clustermode/cluster_common.py | 171 ++++++++++++++++++++ samples/clustermode/exercise.py | 219 ++++++++++++++++++++++++++ samples/clustermode/run_cluster.py | 88 +++++++++++ samples/clustermode/server.py | 87 ++++++++++ 7 files changed, 641 insertions(+) create mode 100644 samples/clustermode/README.md create mode 100644 samples/clustermode/__init__.py create mode 100644 samples/clustermode/cluster_common.py create mode 100644 samples/clustermode/exercise.py create mode 100644 samples/clustermode/run_cluster.py create mode 100644 samples/clustermode/server.py diff --git a/.github/actions/spelling/allow.txt b/.github/actions/spelling/allow.txt index 8bb3d70dd..8c0448695 100644 --- a/.github/actions/spelling/allow.txt +++ b/.github/actions/spelling/allow.txt @@ -45,6 +45,7 @@ deepwiki denormal denormals drivername +dsn DSNs dunders ES256 @@ -119,6 +120,7 @@ poolclass postgres POSTGRES postgresql +procs proot proto protobuf diff --git a/samples/clustermode/README.md b/samples/clustermode/README.md new file mode 100644 index 000000000..1ec16a567 --- /dev/null +++ b/samples/clustermode/README.md @@ -0,0 +1,74 @@ +# Cluster-mode sample (multi-replica A2A) + +Runs several A2A server replicas that share one durable task store and one event +stream, so any replica can serve any request for any task. This demonstrates the +multi-server support in `a2a.server.cluster`: send, resubscribe, and cancel all +work regardless of which replica a request lands on — no sticky routing needed. + +## What it shows + +- **Shared, versioned task store** (`VersionedDatabaseTaskStore`) — concurrent + writes across replicas are serialized by optimistic concurrency (CAS); no lost + updates. +- **Shared event stream** (`DatabaseTaskEventStream`) — a subscription on one replica + streams events produced by an agent running on another. +- **Cancel via CAS** — a cancel on replica C stops an agent running on replica A + (A's next save fails the compare-and-swap and aborts). + +## Requirements + +Nothing extra by default: it uses a file-backed SQLite database +(`/tmp/a2a_cluster_demo.db`) shared by all replica processes. + +For a realistic setup, point every replica at a shared Postgres/MySQL: + +```bash +export A2A_CLUSTER_DSN='postgresql+asyncpg://user:pass@localhost/a2a' +``` + +## Run + +Start a 3-replica cluster (ports 41241, 41242, 41243): + +```bash +python -m samples.clustermode.run_cluster --replicas 3 --base-port 41241 +``` + +In another terminal, exercise it (one task, spread across replicas): + +```bash +python -m samples.clustermode.exercise --host 127.0.0.1 --ports 41241,41242,41243 +``` + +Expected output ends with: + +``` +OK: one task was started, observed, and cancelled across three different replicas. +``` + +## Files + +| File | Purpose | +|------|---------| +| `cluster_common.py` | Shared store/stream wiring, agent card, and the demo agent | +| `server.py` | One replica (a FastAPI app on one port) | +| `run_cluster.py` | Launches N replicas as subprocesses | +| `exercise.py` | Client that sends/gets/cancels one task across replicas | + +## How it maps to the API + +Every replica builds its handler the same way — the only multi-server-specific +part is passing a shared `event_stream` and a shared `VersionedTaskStore`: + +```python +handler = DefaultRequestHandler( + agent_executor=SlowEchoAgent(...), + task_store=VersionedDatabaseTaskStore(engine=..., create_table=False), + agent_card=agent_card, + event_stream=DatabaseTaskEventStream(engine=..., create_table=False), +) +``` + +Omit `event_stream` (and use a plain `InMemoryTaskStore`) and you get the ordinary +single-process behaviour — the multi-server path is fully opt-in. +``` diff --git a/samples/clustermode/__init__.py b/samples/clustermode/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/samples/clustermode/cluster_common.py b/samples/clustermode/cluster_common.py new file mode 100644 index 000000000..6d13093bc --- /dev/null +++ b/samples/clustermode/cluster_common.py @@ -0,0 +1,171 @@ +"""Shared wiring for the multi-replica (cluster mode) sample. + +Every replica in this sample builds its handler from the SAME durable task store +and the SAME event stream (pointed at one shared database). That is what lets any +replica serve any request for any task - send, resubscribe, and cancel all work +regardless of which replica the load balancer picks. + +By default this uses a file-backed SQLite database so the sample runs with no +external services. Set A2A_CLUSTER_DSN to a Postgres/MySQL async DSN for a more +realistic setup, e.g.: + + export A2A_CLUSTER_DSN='postgresql+asyncpg://user:pass@localhost/a2a' +""" + +import asyncio +import logging +import os +import tempfile + +from pathlib import Path + +from sqlalchemy.ext.asyncio import create_async_engine + +from a2a.helpers.proto_helpers import new_task_from_user_message +from a2a.server.agent_execution.agent_executor import AgentExecutor +from a2a.server.agent_execution.context import RequestContext +from a2a.server.cluster import VersionedDatabaseTaskStore +from a2a.server.cluster.database_event_stream import DatabaseTaskEventStream +from a2a.server.events.event_queue import EventQueue +from a2a.server.models import Base +from a2a.server.tasks.task_updater import TaskUpdater +from a2a.types import ( + AgentCapabilities, + AgentCard, + AgentInterface, + AgentSkill, + Part, + TaskState, +) + + +logger = logging.getLogger(__name__) + + +def default_sqlite_path() -> str: + """Path to the shared SQLite file used when no DSN is configured.""" + return os.environ.get( + 'A2A_CLUSTER_SQLITE', + str(Path(tempfile.gettempdir()) / 'a2a_cluster_demo.db'), + ) + + +def default_dsn() -> str: + """The shared-database DSN. SQLite file by default; override via env.""" + dsn = os.environ.get('A2A_CLUSTER_DSN') + if dsn: + return dsn + # A file (not :memory:) so separate replica processes share one database. + return f'sqlite+aiosqlite:///{default_sqlite_path()}' + + +async def init_schema(dsn: str) -> None: + """Creates the shared tables (task + task_events) once, up front.""" + engine = create_async_engine(dsn) + try: + async with engine.begin() as conn: + await conn.run_sync(Base.metadata.create_all) + finally: + await engine.dispose() + + +def build_cluster_backends( + dsn: str, +) -> tuple[VersionedDatabaseTaskStore, DatabaseTaskEventStream]: + """Builds a versioned store + event stream over the shared database. + + Each replica calls this with the same DSN. `create_table=False` because + `init_schema` owns table creation. + """ + store_engine = create_async_engine(dsn) + stream_engine = create_async_engine(dsn) + store = VersionedDatabaseTaskStore(engine=store_engine, create_table=False) + stream = DatabaseTaskEventStream( + engine=stream_engine, create_table=False, poll_interval_s=0.2 + ) + return store, stream + + +def build_agent_card(base_url: str) -> AgentCard: + """The agent card advertised by every replica (same logical agent).""" + return AgentCard( + name='Cluster Demo Agent', + description='A slow agent used to demonstrate multi-replica A2A.', + version='1.0.0', + capabilities=AgentCapabilities( + streaming=True, push_notifications=False + ), + default_input_modes=['text'], + default_output_modes=['text', 'task-status'], + skills=[ + AgentSkill( + id='cluster_demo', + name='Cluster Demo', + description='Echoes slowly so you can observe it across replicas.', + tags=['sample', 'cluster'], + examples=['hello'], + input_modes=['text'], + output_modes=['text', 'task-status'], + ) + ], + supported_interfaces=[ + AgentInterface( + protocol_binding='JSONRPC', + protocol_version='1.0', + url=f'{base_url}/a2a/jsonrpc', + ), + ], + ) + + +class SlowEchoAgent(AgentExecutor): + """Goes WORKING, emits ticks slowly, then completes. + + The deliberate slowness makes the multi-replica behaviour observable: you + can subscribe or cancel from a different replica while a task is mid-flight. + Resumability across replicas relies on this agent reading state only from + the task (it does not keep anything in process memory between turns). + """ + + def __init__(self, replica_id: str, ticks: int = 10) -> None: + self._replica_id = replica_id + self._ticks = ticks + + async def execute( + self, context: RequestContext, event_queue: EventQueue + ) -> None: + """Runs the slow echo: WORKING, ticks, artifact, complete.""" + updater = TaskUpdater( + event_queue, + str(context.task_id or ''), + str(context.context_id or ''), + ) + if context.current_task is None: + await event_queue.enqueue_event( + new_task_from_user_message(context.message) + ) + await updater.start_work( + message=updater.new_agent_message( + [Part(text=f'[{self._replica_id}] starting')] + ) + ) + for i in range(self._ticks): + await asyncio.sleep(1.0) + await updater.update_status( + TaskState.TASK_STATE_WORKING, + message=updater.new_agent_message( + [Part(text=f'[{self._replica_id}] tick {i + 1}')] + ), + ) + await updater.add_artifact( + [Part(text=f'[{self._replica_id}] done')], + name='response', + last_chunk=True, + ) + await updater.complete() + + async def cancel( + self, context: RequestContext, event_queue: EventQueue + ) -> None: + """No-op: the stop happens via the store CAS when CANCELED is recorded.""" + return diff --git a/samples/clustermode/exercise.py b/samples/clustermode/exercise.py new file mode 100644 index 000000000..210a25414 --- /dev/null +++ b/samples/clustermode/exercise.py @@ -0,0 +1,219 @@ +"""Exercises a running replica cluster across different replicas. + +Deliberately spreads one task's requests across replicas to show that, with a +shared store + event stream, a task started on replica 0 can be observed from +replica 1 and cancelled from replica 2 - consistently, because no per-task +state is pinned to a single replica. Uses raw JSON-RPC over httpx with protojson +bodies, so it depends only on the SDK's proto types. + +Usage (after run_cluster.py is up): + python -m samples.clustermode.exercise --host 127.0.0.1 --ports 41241,41242 +""" + +import argparse +import asyncio +import contextlib +import json +import uuid + +from typing import Any + +import httpx + +from google.protobuf.json_format import MessageToDict, ParseDict +from google.protobuf.message import Message as ProtoMessage + +from a2a.types import ( + CancelTaskRequest, + GetTaskRequest, + Message, + Part, + Role, + SendMessageRequest, + StreamResponse, + Task, + TaskState, +) + + +MIN_REPLICAS = 2 +_WORKING = TaskState.TASK_STATE_WORKING +_CANCELED = TaskState.TASK_STATE_CANCELED + + +def _endpoint(host: str, port: int) -> str: + return f'http://{host}:{port}/a2a/jsonrpc' + + +async def _rpc( + client: httpx.AsyncClient, + url: str, + method: str, + params_msg: ProtoMessage, +) -> dict[str, Any]: + """Sends one JSON-RPC call with a protojson params body; returns result.""" + body = { + 'jsonrpc': '2.0', + 'id': str(uuid.uuid4()), + 'method': method, + 'params': MessageToDict(params_msg), + } + # This sample speaks the v1.0 protocol; the header selects the v1 handler. + resp = await client.post( + url, json=body, headers={'A2A-Version': '1.0'}, timeout=30.0 + ) + resp.raise_for_status() + data = resp.json() + if 'error' in data: + raise RuntimeError(f'{method} error: {data["error"]}') + return data['result'] + + +async def _get_task(client: httpx.AsyncClient, url: str, task_id: str) -> Task: + """Fetches a task via GetTask from the given replica.""" + result = await _rpc(client, url, 'GetTask', GetTaskRequest(id=task_id)) + task = Task() + ParseDict(result, task) + return task + + +async def _start_streaming( + client: httpx.AsyncClient, url: str, send: SendMessageRequest +) -> tuple[str, asyncio.Task[None]]: + """Starts a streaming send; returns (task_id, drainer task). + + The drainer keeps the SSE stream open (so the agent keeps running) until it + is cancelled. The task id is taken from the first streamed Task event. + """ + loop = asyncio.get_event_loop() + task_id_future: asyncio.Future[str] = loop.create_future() + + async def _drain() -> None: + body = { + 'jsonrpc': '2.0', + 'id': str(uuid.uuid4()), + 'method': 'SendStreamingMessage', + 'params': MessageToDict(send), + } + headers = {'A2A-Version': '1.0', 'Accept': 'text/event-stream'} + with contextlib.suppress(Exception): + async with client.stream( + 'POST', url, json=body, headers=headers, timeout=120.0 + ) as resp: + async for line in resp.aiter_lines(): + if not line.startswith('data:'): + continue + payload = json.loads(line[len('data:') :].strip()) + result = payload.get('result') + if result is None: + continue + sr = StreamResponse() + ParseDict(result, sr) + if not task_id_future.done() and sr.HasField('task'): + task_id_future.set_result(sr.task.id) + + drainer = asyncio.create_task(_drain()) + try: + task_id = await asyncio.wait_for(task_id_future, timeout=30) + except TimeoutError: + drainer.cancel() + raise RuntimeError('no task event received from stream') from None + return task_id, drainer + + +async def run(host: str, ports: list[int]) -> None: + """Runs the cross-replica send / observe / cancel demonstration.""" + if len(ports) < MIN_REPLICAS: + raise SystemExit('Need at least 2 replicas to demonstrate the point.') + + urls = [_endpoint(host, p) for p in ports] + + async with httpx.AsyncClient() as client: + # 1) Start a streaming task on replica 0. The stream stays open (agent + # runs) while we observe and cancel from other replicas. + print(f'[send] -> replica 0 ({urls[0]}) (streaming)') + send = SendMessageRequest( + message=Message( + role=Role.ROLE_USER, + message_id=str(uuid.uuid4()), + parts=[Part(text='hello cluster')], + ) + ) + task_id, send_call = await _start_streaming(client, urls[0], send) + print(f' task {task_id} started') + + # 2) Poll GetTask from replica 1 until it observes WORKING. + print(f'[get] -> replica 1 ({urls[1]}) while it runs') + seen_working = False + for _ in range(8): + await asyncio.sleep(1.0) + try: + task = await _get_task(client, urls[1], task_id) + except RuntimeError: + continue # not persisted yet + print( + f' replica 1 sees state=' + f'{TaskState.Name(task.status.state)}' + ) + if task.status.state == _WORKING: + seen_working = True + break + if not seen_working: + raise RuntimeError('replica 1 never observed the task WORKING') + + # 3) Cancel from a different replica than the one running the agent. + cancel_idx = 2 if len(urls) > MIN_REPLICAS else 1 + print(f'[cancel] -> replica {cancel_idx} ({urls[cancel_idx]})') + result = await _rpc( + client, + urls[cancel_idx], + 'CancelTask', + CancelTaskRequest(id=task_id), + ) + cancelled = Task() + ParseDict(result, cancelled) + print( + f' cancel returned state=' + f'{TaskState.Name(cancelled.status.state)}' + ) + + # Stop draining the (now-cancelled) stream. CancelledError is a + # BaseException, so suppress it explicitly. + send_call.cancel() + with contextlib.suppress(asyncio.CancelledError, Exception): + await send_call + + # 4) Confirm from replica 0 that the cancel is visible everywhere. + await asyncio.sleep(1.0) + final = await _get_task(client, urls[0], task_id) + print( + f'[verify] replica 0 sees final state=' + f'{TaskState.Name(final.status.state)}' + ) + if final.status.state != _CANCELED: + raise RuntimeError( + 'expected the task to be CANCELED across all replicas, got ' + f'{TaskState.Name(final.status.state)}' + ) + print( + '\nOK: one task was started on replica 0, observed on replica 1, ' + 'and cancelled from another replica -- consistently.' + ) + + +def main() -> None: + """Parses args and runs the demonstration.""" + parser = argparse.ArgumentParser(description='Exercise an A2A cluster') + parser.add_argument('--host', default='127.0.0.1') + parser.add_argument( + '--ports', + default='41241,41242,41243', + help='Comma-separated replica ports', + ) + args = parser.parse_args() + ports = [int(p) for p in args.ports.split(',')] + asyncio.run(run(args.host, ports)) + + +if __name__ == '__main__': + main() diff --git a/samples/clustermode/run_cluster.py b/samples/clustermode/run_cluster.py new file mode 100644 index 000000000..b17b45dab --- /dev/null +++ b/samples/clustermode/run_cluster.py @@ -0,0 +1,88 @@ +"""Launches N cluster-mode replicas as subprocesses on consecutive ports. + +Usage: + python -m samples.clustermode.run_cluster --replicas 3 --base-port 41241 + +Each replica shares one database (SQLite file by default). Point any HTTP client +at any replica's /a2a/jsonrpc endpoint, or use exercise.py which deliberately +spreads requests for one task across replicas. +""" + +import argparse +import contextlib +import os +import signal +import subprocess +import sys +import time + +from pathlib import Path + +from samples.clustermode.cluster_common import default_sqlite_path + + +def main() -> None: + """Launches N replica subprocesses on consecutive ports.""" + parser = argparse.ArgumentParser(description='Run an A2A replica cluster') + parser.add_argument('--replicas', type=int, default=3) + parser.add_argument('--host', default='127.0.0.1') + parser.add_argument('--base-port', type=int, default=41241) + parser.add_argument('--ticks', type=int, default=10) + args = parser.parse_args() + + # A fresh shared SQLite file for this run (unless a DSN is configured). + if 'A2A_CLUSTER_DSN' not in os.environ: + db_path = os.environ.setdefault( + 'A2A_CLUSTER_SQLITE', default_sqlite_path() + ) + with contextlib.suppress(FileNotFoundError): + Path(db_path).unlink() + + procs: list[subprocess.Popen] = [] + ports = [args.base_port + i for i in range(args.replicas)] + try: + for i, port in enumerate(ports): + cmd = [ + sys.executable, + '-m', + 'samples.clustermode.server', + '--host', + args.host, + '--port', + str(port), + '--replica-id', + f'replica-{i}', + '--ticks', + str(args.ticks), + ] + procs.append( + subprocess.Popen(cmd, env=os.environ.copy()) # noqa: S603 + ) + time.sleep(0.5) + + print('\nCluster running. Replica endpoints:') + for i, port in enumerate(ports): + print(f' replica-{i}: http://{args.host}:{port}/a2a/jsonrpc') + print('\nExercise it with:') + joined = ','.join(str(p) for p in ports) + print( + f' python -m samples.clustermode.exercise ' + f'--host {args.host} --ports {joined}' + ) + print('\nPress Ctrl+C to stop.\n') + + for p in procs: + p.wait() + except KeyboardInterrupt: + pass + finally: + for p in procs: + with contextlib.suppress(ProcessLookupError): + p.send_signal(signal.SIGINT) + for p in procs: + with contextlib.suppress(Exception): + p.wait(timeout=5) + + +if __name__ == '__main__': + main() diff --git a/samples/clustermode/server.py b/samples/clustermode/server.py new file mode 100644 index 000000000..8a2d2e0bc --- /dev/null +++ b/samples/clustermode/server.py @@ -0,0 +1,87 @@ +"""A single cluster-mode replica: one FastAPI app on one port. + +All replicas share the same database (via A2A_CLUSTER_DSN or the default SQLite +file), so they form one logical A2A service. Run several of these on different +ports and put any round-robin in front of them. +""" + +import argparse +import asyncio +import contextlib +import logging + +import uvicorn + +from fastapi import FastAPI + +from a2a.server.request_handlers import DefaultRequestHandler +from a2a.server.routes import ( + add_a2a_routes_to_fastapi, + create_agent_card_routes, + create_jsonrpc_routes, +) +from samples.clustermode.cluster_common import ( + SlowEchoAgent, + build_agent_card, + build_cluster_backends, + default_dsn, + init_schema, +) + + +logger = logging.getLogger(__name__) + + +async def serve(host: str, port: int, replica_id: str, ticks: int) -> None: + """Runs one replica bound to a shared store + event stream.""" + dsn = default_dsn() + await init_schema(dsn) + + base_url = f'http://{host}:{port}' + agent_card = build_agent_card(base_url) + store, stream = build_cluster_backends(dsn) + await store.initialize() + await stream.initialize() + + handler = DefaultRequestHandler( + agent_executor=SlowEchoAgent(replica_id, ticks=ticks), + task_store=store, + agent_card=agent_card, + event_stream=stream, + ) + + jsonrpc_routes = create_jsonrpc_routes( + request_handler=handler, + rpc_url='/a2a/jsonrpc', + ) + agent_card_routes = create_agent_card_routes(agent_card=agent_card) + + app = FastAPI() + add_a2a_routes_to_fastapi( + app, + agent_card_routes=agent_card_routes, + jsonrpc_routes=jsonrpc_routes, + ) + + logger.info( + 'Replica %s listening on %s (dsn=%s)', replica_id, base_url, dsn + ) + config = uvicorn.Config(app, host=host, port=port, log_level='warning') + await uvicorn.Server(config).serve() + + +def main() -> None: + """Parses args and runs one replica.""" + logging.basicConfig(level=logging.INFO) + parser = argparse.ArgumentParser(description='A2A cluster-mode replica') + parser.add_argument('--host', default='127.0.0.1') + parser.add_argument('--port', type=int, default=41241) + parser.add_argument('--replica-id', default='replica-0') + parser.add_argument('--ticks', type=int, default=10) + args = parser.parse_args() + with contextlib.suppress(KeyboardInterrupt): + asyncio.run(serve(args.host, args.port, args.replica_id, args.ticks)) + + +if __name__ == '__main__': + main()