From b03db4e9f592c20623d413b97235b855559bf234 Mon Sep 17 00:00:00 2001 From: Daniel van Flymen Date: Thu, 2 Jul 2026 21:01:23 +0200 Subject: [PATCH] Implement conversation logging and compaction in Synapse, adding support for memory and SQLite logs; enhance task management with deferred handling and background processing. --- README.md | 73 +++++++- examples/README.md | 1 + examples/coding_team/README.md | 53 ++++++ examples/coding_team/__init__.py | 0 examples/coding_team/architect.py | 79 ++++++++ examples/coding_team/coder.py | 67 +++++++ examples/coding_team/common.py | 30 +++ synapse_p2p/__init__.py | 10 + synapse_p2p/client.py | 4 +- synapse_p2p/conversations.py | 215 ++++++++++++++++++++++ synapse_p2p/node.py | 189 +++++++++++++++++-- synapse_p2p/teams.py | 231 ++++++++++++++++++++++++ synapse_p2p/tests/test_conversations.py | 201 +++++++++++++++++++++ synapse_p2p/tests/test_deferred_ask.py | 98 ++++++++++ synapse_p2p/tests/test_teams.py | 145 +++++++++++++++ 15 files changed, 1376 insertions(+), 20 deletions(-) create mode 100644 examples/coding_team/README.md create mode 100644 examples/coding_team/__init__.py create mode 100644 examples/coding_team/architect.py create mode 100644 examples/coding_team/coder.py create mode 100644 examples/coding_team/common.py create mode 100644 synapse_p2p/conversations.py create mode 100644 synapse_p2p/teams.py create mode 100644 synapse_p2p/tests/test_conversations.py create mode 100644 synapse_p2p/tests/test_deferred_ask.py create mode 100644 synapse_p2p/tests/test_teams.py diff --git a/README.md b/README.md index 20f6540..baa3994 100644 --- a/README.md +++ b/README.md @@ -331,7 +331,7 @@ broadcast = await node.broadcast( ) ``` -A node with an `@node.ask` handler will ACK the conversation, run the handler, and reply with the result. Nodes without a handler fail quietly from the caller's point of view, just like any other broadcast recipient that cannot help. +A node with an `@node.ask` handler will ACK the conversation, then run the handler **in the background** and deliver the result later as a broadcast reply. The origin's RPC returns immediately with `{"accepted": True, "deferred": True}`, so a handler that wraps a slow LLM can take minutes without holding a socket open. If the handler raises, the node emits a `conversation.error` event instead of a reply. Nodes without a handler fail quietly from the caller's point of view, just like any other broadcast recipient that cannot help. The CLI wraps this flow: @@ -414,6 +414,70 @@ Why this is useful: --- +## 🗜️ Conversation memory: durability, sync, and compaction + +Conversation events live in a pluggable log. The default is in-memory; pass a SQLite log to survive restarts: + +```python +from synapse_p2p import Node, SqliteConversationLog + +node = Node( + name="architect", + conversation_log=SqliteConversationLog("architect-conversations.db"), +) +``` + +A late joiner (or a restarted node) can pull a conversation it missed from any peer: + +```python +added = await node.sync_conversation(peer, conversation_id) +``` + +Long conversations compact automatically. When a conversation grows past `conversation_max_events`, the node folds older events into a single `summary` event, keeping the opening message and the most recent `conversation_keep_recent` events verbatim: + +```python +node = Node( + name="architect", + conversation_max_events=100, + conversation_keep_recent=25, +) + + +@node.summarizer +async def summarize(events: list[ConversationEvent]) -> str: + transcript = "\n".join(f"{e.peer.name} [{e.kind}]: {e.payload}" for e in events) + return await my_llm_summarize(transcript) # bring your own model +``` + +Without a custom summarizer, Synapse uses a naive extractive digest. Compaction is **local**: each node compresses its own copy of the shared log, and gossip cannot resurrect compacted events. You can also compact on demand with `await node.compact_conversation(conversation_id)` and observe it with `@node.on("conversation.compacted")`. + +--- + +## 👷 Teams: an optional task layer + +`synapse_p2p.teams` layers a small task vocabulary on top of conversation events — nothing in the substrate is special-cased for it. A `Team` offers work; `Worker`s claim tasks matching their capabilities; the team grants each task to the first claimant, so exactly one worker runs it. + +```python +from synapse_p2p.teams import Assignment, Team, Worker + +# architect process +team = Team(node) +task = await team.offer("implement the parser", spec={"file": "parser.py"}, requires=["python"]) +result = await team.wait(task, timeout=600) + +# coder process +worker = Worker(coder_node) + +@worker.task +async def implement(assignment: Assignment) -> dict: + await assignment.progress("starting") + return {"diff": await my_agent.run(assignment.title, assignment.spec)} +``` + +Each task is one shared conversation (`task.offer` → `task.claim` → `task.grant` → `task.progress` → `task.done` / `task.failed`), so every peer can watch the work happen, late joiners can sync it, and long task threads compact like any other conversation. See [`examples/coding_team`](./examples/coding_team) for an architect on one model overseeing coders on another. + +--- + ## 🌅 Periodic tasks Nodes can wake up on a schedule: every few seconds, every weekday morning, or when the sun rises. @@ -602,6 +666,7 @@ See [`examples/`](./examples). Each example folder has its own README. | [`pydantic_ai_team`](./examples/pydantic_ai_team) | Pydantic AI agents behind Synapse nodes. | | [`periodic_tasks`](./examples/periodic_tasks) | Interval, cron, sunrise, and sunset jobs in a garden-caretaker node. | | [`stock_trading_team`](./examples/stock_trading_team) | Analyst/news/trader swarm with a dumb paper exchange API and market-hours periodic scans. | +| [`coding_team`](./examples/coding_team) | An architect and coder agents on different models, coordinated through the `synapse_p2p.teams` task layer. | The agent examples use `synapse.ask`, opt-in ACKs, shared conversation replies, and advertised `agent-card` artifacts. The stock example shows a periodic job that checks market hours before asking the swarm, so agent/model work only happens when the paper market is open. @@ -621,6 +686,8 @@ Built-in endpoints: | `_synapse.heartbeat` | update peer liveness | | `_synapse.broadcast.reply` | reply to a broadcast nonce | | `_synapse.conversation.event` | gossip a shared conversation event | +| `_synapse.conversation.sync` | serve a conversation's events to a late joiner | +| `_synapse.conversation.list` | list locally known conversation ids | | `_synapse.artifacts` | list advertised artifacts | | `_synapse.artifact.get` | fetch one advertised artifact | | `_node.info` | name, role, description, capabilities | @@ -633,6 +700,8 @@ Wire format: 1. 4-byte unsigned big-endian payload length 2. MsgPack payload bytes +Frames up to 4 MiB are accepted by default (`Node(max_upload_size=...)`, `Client(max_download_size=...)`), so results can carry real payloads like diffs and documents. + Request: ```python @@ -690,7 +759,7 @@ logger.enable("synapse_p2p") Synapse does **not** implement planning, memory, consensus, auth policy, NAT traversal, hosted registries, or UX. -Those belong above Synapse. +Those belong above Synapse. (Two exceptions are on the roadmap because they belong *in* the substrate: node identity + signed gossip for untrusted networks, and a relay mode so seeds can bridge peers that cannot dial each other.) Synapse is the substrate: diff --git a/examples/README.md b/examples/README.md index 60d59cb..92e1279 100644 --- a/examples/README.md +++ b/examples/README.md @@ -58,5 +58,6 @@ Most agent examples publish an `agent-card` artifact and use shared conversation | [`pydantic_ai_team`](./pydantic_ai_team) | Pydantic AI agents behind Synapse nodes. | | [`periodic_tasks`](./periodic_tasks) | Interval, cron, and solar jobs in a garden-caretaker node. | | [`stock_trading_team`](./stock_trading_team) | Analyst/news/trader swarm with a dumb paper exchange API and market-hours periodic scans. | +| [`coding_team`](./coding_team) | An architect and coder agents on different models: task offers, claims, grants, progress, and compacted task conversations via `synapse_p2p.teams`. | Each folder has its own README with exact run commands. diff --git a/examples/coding_team/README.md b/examples/coding_team/README.md new file mode 100644 index 0000000..608a308 --- /dev/null +++ b/examples/coding_team/README.md @@ -0,0 +1,53 @@ +# Coding team: an architect and coders on different models + +A minimal heterogeneous agent team built on `synapse_p2p.teams`: + +- **architect** — oversees the work. Offers tasks to the swarm, grants each to + the first capable claimant, collects results, and reviews them with its own + model (e.g. Claude). +- **coder-1 / coder-2** — claim tasks whose `requires` match their advertised + capabilities, implement them with their own model (e.g. GPT), and narrate + progress into the shared task conversation. + +Every task is one shared conversation (`conversation_id == task id`) that all +peers gossip: offer → claim → grant → progress → done. The architect compacts +long task threads automatically (`conversation_max_events=30`), folding old +progress chatter into a `summary` event. + +## Run it (offline) + +Without model env vars the agents use pydantic-ai's `TestModel`, so the whole +flow runs with no API keys. From the repo root, in three terminals: + +```bash +CODER_NAME=coder-1 uv run python -m examples.coding_team.coder +CODER_NAME=coder-2 uv run python -m examples.coding_team.coder +uv run python -m examples.coding_team.architect +``` + +Watch the swarm from a fourth terminal: + +```bash +uv run sn watch team.electron.network +``` + +## Run it with real models + +```bash +export ANTHROPIC_API_KEY=... OPENAI_API_KEY=... +ARCHITECT_MODEL=anthropic:claude-fable-5 uv run python -m examples.coding_team.architect +CODER_MODEL=openai:gpt-5.5 CODER_NAME=coder-1 uv run python -m examples.coding_team.coder +``` + +Any pydantic-ai model string works; the swarm doesn't care what model sits +behind a node. To push a coder's reasoning effort up, configure the agent's +`model_settings` in `common.py` (e.g. OpenAI's `reasoning_effort`). + +## What to look at + +- `synapse_p2p/teams.py` — the whole task vocabulary is ~200 lines of + conversation events. Nothing here is special-cased in the substrate. +- Deferred results: coders return immediately at the RPC layer and deliver + results as conversation events, so a task can take as long as the model needs. +- Late joiners can catch up on a task thread with + `node.sync_conversation(peer, task_id)`. diff --git a/examples/coding_team/__init__.py b/examples/coding_team/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/examples/coding_team/architect.py b/examples/coding_team/architect.py new file mode 100644 index 0000000..c37b420 --- /dev/null +++ b/examples/coding_team/architect.py @@ -0,0 +1,79 @@ +"""The architect node: offers tasks to the swarm, reviews what comes back. + +The architect and the coders can run different models entirely — the swarm +only sees offers, claims, progress, and results. Each task is one shared +conversation, and the architect compacts long task threads automatically. +""" + +import asyncio + +from examples.coding_team.common import ARCHITECT_MODEL, SWARM, make_agent, run_agent +from synapse_p2p import ConversationEvent, Node +from synapse_p2p.teams import Team + +brain, test_model = make_agent( + ARCHITECT_MODEL, + "You are the software architect overseeing a team of coder agents. " + "Review their submitted implementations and give a short verdict.", + "Architect: both implementations look solid; ship it.", +) + +node = Node( + name="architect", + role="architect", + swarm=SWARM, + capabilities=["architecture", "review", "coordination"], + mdns=True, + # Keep task conversations small: fold old progress chatter into summaries. + conversation_max_events=30, + conversation_keep_recent=10, +) + +team = Team(node) + +FEATURE = "Add a `sn swarm top`-style live dashboard to the CLI" +SUBTASKS = [ + ("Implement the dashboard rendering loop", {"module": "synapse_p2p/cli.py"}), + ("Add tests for the dashboard event stream", {"module": "synapse_p2p/tests/"}), +] + + +@node.on("conversation.task.progress") +async def on_progress(event: ConversationEvent) -> None: + print(f" progress from {event.peer.name}: {event.payload.get('message')}") + + +async def main() -> None: + await node.start() + await node.join(wait=1) + print(f"architect online, {len(node.peers)} peer(s) known") + + tasks = [] + for title, spec in SUBTASKS: + task = await team.offer(title, spec=spec, requires=["python", "implementation"]) + print(f"offered: {title} ({task.id})") + tasks.append(task) + + results = [] + for task in tasks: + result = await team.wait(task, timeout=600) + assignee = task.assignee.name if task.assignee else "?" + print(f"done: {task.title} by {assignee}") + results.append(result) + + verdict = await run_agent( + brain, + test_model, + f"Feature: {FEATURE}\nSubmissions: {results}\nGive your review verdict.", + ) + print(f"\narchitect verdict:\n{verdict}") + + for task in tasks: + kinds = [event.kind for event in node.conversation(task.id)] + print(f"conversation {task.id}: {kinds}") + + await node.stop() + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/examples/coding_team/coder.py b/examples/coding_team/coder.py new file mode 100644 index 0000000..28dd5a8 --- /dev/null +++ b/examples/coding_team/coder.py @@ -0,0 +1,67 @@ +"""A coder node: claims tasks it can do, implements them, narrates progress. + +Run several of these (they race to claim, the architect grants one): + + CODER_NAME=coder-1 uv run python -m examples.coding_team.coder + CODER_NAME=coder-2 uv run python -m examples.coding_team.coder +""" + +import asyncio +import os + +from examples.coding_team.common import CODER_MODEL, SWARM, make_agent, run_agent +from synapse_p2p import Node +from synapse_p2p.teams import Assignment, Worker + +brain, test_model = make_agent( + CODER_MODEL, + "You are a senior implementation engineer on a distributed agent team. " + "You receive one task with a spec. Reply with the implementation: code, " + "tests, and a one-paragraph note for the reviewing architect.", + "Coder: implemented the task with a small patch, added tests, all passing.", +) + +node = Node( + name=os.getenv("CODER_NAME", "coder-1"), + role="implementation", + swarm=SWARM, + capabilities=["python", "implementation", "tests"], + mdns=True, +) + +node.artifact( + "agent-card", + { + "name": node.name, + "role": node.role, + "model": CODER_MODEL or "test-model", + "capabilities": ["python", "implementation", "tests"], + "description": "Claims implementation tasks and returns patches with tests.", + }, + mime_type="application/vnd.synapse.agent-card+json", +) + +worker = Worker(node) + + +@worker.task +async def implement(assignment: Assignment) -> dict: + await assignment.progress(f"{node.name} starting", title=assignment.title) + prompt = f"Task: {assignment.title}\nSpec: {assignment.spec}" + output = await run_agent(brain, test_model, prompt) + await assignment.progress(f"{node.name} finished, submitting result") + return {"implementation": output, "by": node.name} + + +async def main() -> None: + await node.start() + try: + await node.join() + print(f"{node.name} online at {node.address}:{node.port}, waiting for offers") + await asyncio.Event().wait() + finally: + await node.stop() + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/examples/coding_team/common.py b/examples/coding_team/common.py new file mode 100644 index 0000000..dcf0cec --- /dev/null +++ b/examples/coding_team/common.py @@ -0,0 +1,30 @@ +import os +from contextlib import nullcontext +from typing import Any + +from pydantic_ai import Agent +from pydantic_ai.models.test import TestModel + +SWARM = "team.electron.network" + +# Set these to real models to run the demo live, e.g. +# ARCHITECT_MODEL=anthropic:claude-fable-5 CODER_MODEL=openai:gpt-5.5 +# Unset, the demo runs offline against pydantic-ai's TestModel. +ARCHITECT_MODEL = os.getenv("ARCHITECT_MODEL") +CODER_MODEL = os.getenv("CODER_MODEL") + + +def make_agent( + model: str | None, instructions: str, fallback: str +) -> tuple[Agent, TestModel | None]: + if model: + return Agent(model, instructions=instructions), None + test_model = TestModel(custom_output_text=fallback) + return Agent(test_model, instructions=instructions), test_model + + +async def run_agent(agent: Agent, test_model: TestModel | None, prompt: str) -> Any: + context = agent.override(model=test_model) if test_model is not None else nullcontext() + with context: + result = await agent.run(prompt) + return result.output diff --git a/synapse_p2p/__init__.py b/synapse_p2p/__init__.py index cc95196..ba00d23 100644 --- a/synapse_p2p/__init__.py +++ b/synapse_p2p/__init__.py @@ -20,6 +20,12 @@ """ from synapse_p2p.client import Client +from synapse_p2p.conversations import ( + BaseConversationLog, + MemoryConversationLog, + SqliteConversationLog, + default_summarizer, +) from synapse_p2p.messages import RemoteProcedureCall, RPCError, RPCRequest, RPCResponse from synapse_p2p.node import Capability, Node from synapse_p2p.schedules import CronSchedule, IntervalSchedule, SolarSchedule, cron, every, solar @@ -40,6 +46,7 @@ __all__ = [ "AdvertisedArtifact", "Capability", + "BaseConversationLog", "BaseRPCSerializer", "Broadcast", "BroadcastReply", @@ -47,6 +54,7 @@ "Connection", "ConversationEvent", "CronSchedule", + "MemoryConversationLog", "MessagePackRPCSerializer", "Node", "NodeKind", @@ -56,9 +64,11 @@ "RPCResponse", "RemoteProcedureCall", "ServedArtifact", + "SqliteConversationLog", "IntervalSchedule", "SolarSchedule", "cron", + "default_summarizer", "every", "solar", "__logo__", diff --git a/synapse_p2p/client.py b/synapse_p2p/client.py index 847d8b1..79dd5ee 100644 --- a/synapse_p2p/client.py +++ b/synapse_p2p/client.py @@ -14,7 +14,7 @@ def __init__( address: str = "127.0.0.1", port: int = 9999, serializer_class: type[BaseRPCSerializer] = MessagePackRPCSerializer, - max_download_size: int = 4096, + max_download_size: int = 4 * 1024 * 1024, timeout: float | None = 30, ) -> None: self.address = connect_address(address) @@ -29,7 +29,7 @@ def from_peer( peer: Peer, *, serializer_class: type[BaseRPCSerializer] = MessagePackRPCSerializer, - max_download_size: int = 4096, + max_download_size: int = 4 * 1024 * 1024, timeout: float | None = 30, ) -> "Client": return cls( diff --git a/synapse_p2p/conversations.py b/synapse_p2p/conversations.py new file mode 100644 index 0000000..0967636 --- /dev/null +++ b/synapse_p2p/conversations.py @@ -0,0 +1,215 @@ +"""Durable, compactable storage for shared conversation events. + +A conversation log stores every :class:`ConversationEvent` a node has seen and +supports *compaction*: folding old events into a single ``summary`` event so a +long-running conversation stays small enough to hand to an LLM. Synapse ships a +naive extractive summarizer; agents can plug in an LLM summarizer instead. +""" + +from __future__ import annotations + +import json +import sqlite3 +from collections.abc import Awaitable, Callable +from pathlib import Path + +from synapse_p2p.types import ConversationEvent + +Summarizer = Callable[[list[ConversationEvent]], Awaitable[str]] + +SUMMARY_KIND = "summary" + + +async def default_summarizer(events: list[ConversationEvent]) -> str: + """Summarize events without an LLM: keep a digest of who said what. + + Replace this with an LLM-backed callable via ``Node(summarizer=...)`` when a + conversation carries real prose worth compressing semantically. + """ + lines: list[str] = [] + counts: dict[str, int] = {} + for event in events: + counts[event.kind] = counts.get(event.kind, 0) + 1 + if event.kind == SUMMARY_KIND: + previous = str(event.payload.get(SUMMARY_KIND, "")) + if previous: + lines.append(previous) + continue + preview = json.dumps(event.payload, default=str) + if len(preview) > 200: + preview = preview[:200] + "…" + lines.append(f"{event.peer.name or event.peer.id} [{event.kind}]: {preview}") + header = ", ".join(f"{count} {kind}" for kind, count in sorted(counts.items())) + return f"({header})\n" + "\n".join(lines[-40:]) + + +class BaseConversationLog: + """Storage interface for conversation events. + + ``append`` must be idempotent by ``event_id`` and must keep remembering + compacted event ids so gossip cannot resurrect events a summary replaced. + """ + + def append(self, event: ConversationEvent) -> bool: + raise NotImplementedError + + def seen(self, event_id: str) -> bool: + raise NotImplementedError + + def events(self, conversation_id: str, *, since: float = 0.0) -> list[ConversationEvent]: + raise NotImplementedError + + def conversations(self) -> list[str]: + raise NotImplementedError + + def count(self, conversation_id: str) -> int: + raise NotImplementedError + + def compact( + self, + conversation_id: str, + removed_event_ids: list[str], + summary: ConversationEvent, + ) -> None: + raise NotImplementedError + + def close(self) -> None: + pass + + +class MemoryConversationLog(BaseConversationLog): + def __init__(self) -> None: + self._events: dict[str, list[ConversationEvent]] = {} + self._seen: set[str] = set() + + def append(self, event: ConversationEvent) -> bool: + if event.event_id in self._seen: + return False + self._seen.add(event.event_id) + self._events.setdefault(event.conversation_id, []).append(event) + return True + + def seen(self, event_id: str) -> bool: + return event_id in self._seen + + def events(self, conversation_id: str, *, since: float = 0.0) -> list[ConversationEvent]: + events = self._events.get(conversation_id, []) + if since: + events = [event for event in events if event.created_at > since] + return list(events) + + def conversations(self) -> list[str]: + return list(self._events) + + def count(self, conversation_id: str) -> int: + return len(self._events.get(conversation_id, [])) + + def compact( + self, + conversation_id: str, + removed_event_ids: list[str], + summary: ConversationEvent, + ) -> None: + removed = set(removed_event_ids) + kept = [ + event + for event in self._events.get(conversation_id, []) + if event.event_id not in removed + ] + self._seen.add(summary.event_id) + merged = kept + [summary] + merged.sort(key=lambda event: event.created_at) + self._events[conversation_id] = merged + + +class SqliteConversationLog(BaseConversationLog): + """Conversation log persisted to a single SQLite file. + + Survives restarts, so a node can serve ``_synapse.conversation.sync`` to + late joiners even after it has been rebooted. + """ + + def __init__(self, path: str | Path) -> None: + self.path = Path(path) + self._db = sqlite3.connect(self.path) + self._db.execute( + """ + CREATE TABLE IF NOT EXISTS events ( + event_id TEXT PRIMARY KEY, + conversation_id TEXT NOT NULL, + created_at REAL NOT NULL, + compacted INTEGER NOT NULL DEFAULT 0, + data TEXT NOT NULL + ) + """ + ) + self._db.execute( + "CREATE INDEX IF NOT EXISTS idx_events_conversation" + " ON events (conversation_id, created_at)" + ) + self._db.commit() + + def append(self, event: ConversationEvent) -> bool: + cursor = self._db.execute( + "INSERT OR IGNORE INTO events (event_id, conversation_id, created_at, data)" + " VALUES (?, ?, ?, ?)", + ( + event.event_id, + event.conversation_id, + event.created_at, + json.dumps(event.to_dict(), default=str), + ), + ) + self._db.commit() + return cursor.rowcount > 0 + + def seen(self, event_id: str) -> bool: + row = self._db.execute("SELECT 1 FROM events WHERE event_id = ?", (event_id,)).fetchone() + return row is not None + + def events(self, conversation_id: str, *, since: float = 0.0) -> list[ConversationEvent]: + rows = self._db.execute( + "SELECT data FROM events" + " WHERE conversation_id = ? AND compacted = 0 AND created_at > ?" + " ORDER BY created_at, rowid", + (conversation_id, since), + ).fetchall() + return [ConversationEvent.from_dict(json.loads(row[0])) for row in rows] + + def conversations(self) -> list[str]: + rows = self._db.execute( + "SELECT DISTINCT conversation_id FROM events WHERE compacted = 0" + ).fetchall() + return [row[0] for row in rows] + + def count(self, conversation_id: str) -> int: + row = self._db.execute( + "SELECT COUNT(*) FROM events WHERE conversation_id = ? AND compacted = 0", + (conversation_id,), + ).fetchone() + return int(row[0]) + + def compact( + self, + conversation_id: str, + removed_event_ids: list[str], + summary: ConversationEvent, + ) -> None: + self._db.executemany( + "UPDATE events SET compacted = 1 WHERE event_id = ?", + [(event_id,) for event_id in removed_event_ids], + ) + self._db.execute( + "INSERT OR REPLACE INTO events (event_id, conversation_id, created_at, data)" + " VALUES (?, ?, ?, ?)", + ( + summary.event_id, + summary.conversation_id, + summary.created_at, + json.dumps(summary.to_dict(), default=str), + ), + ) + self._db.commit() + + def close(self) -> None: + self._db.close() diff --git a/synapse_p2p/node.py b/synapse_p2p/node.py index 5bd7a44..619f219 100644 --- a/synapse_p2p/node.py +++ b/synapse_p2p/node.py @@ -14,6 +14,13 @@ from synapse_p2p import __logo__ from synapse_p2p.client import Client +from synapse_p2p.conversations import ( + SUMMARY_KIND, + BaseConversationLog, + MemoryConversationLog, + Summarizer, + default_summarizer, +) from synapse_p2p.exceptions import InvalidMessageError from synapse_p2p.framing import read_frame, write_frame from synapse_p2p.mdns import MdnsDiscovery @@ -70,7 +77,7 @@ def __init__( port: int = 0, advertise: str | None = "auto", serializer_class: type[BaseRPCSerializer] = MessagePackRPCSerializer, - max_upload_size: int = 4096, + max_upload_size: int = 4 * 1024 * 1024, node_id: str | None = None, name: str = "", role: str = "", @@ -83,6 +90,10 @@ def __init__( mdns: bool = False, heartbeat_interval: float | None = 5, peer_timeout: float = 20, + conversation_log: BaseConversationLog | None = None, + conversation_max_events: int | None = None, + conversation_keep_recent: int = 25, + summarizer: Summarizer | None = None, ) -> None: self.bind = bind self.advertise = advertise @@ -104,8 +115,12 @@ def __init__( self.peer_timeout = peer_timeout self.peers: dict[str, Peer] = {} self.broadcast_replies: dict[str, list[BroadcastReply]] = {} - self.conversation_events: dict[str, list[ConversationEvent]] = {} - self._seen_conversation_events: set[str] = set() + self.conversation_log = conversation_log or MemoryConversationLog() + self.conversation_max_events = conversation_max_events + self.conversation_keep_recent = conversation_keep_recent + self._summarizer: Summarizer = summarizer or default_summarizer + self._compacting: set[str] = set() + self._background: set[asyncio.Task] = set() self.artifact_directory: dict[str, ServedArtifact] = {} self.lifecycle_handlers: dict[str, list[Callable[[Any], Coroutine[Any, Any, None]]]] = {} self.endpoint_directory: dict[str, Callable] = {} @@ -353,10 +368,17 @@ async def start(self) -> asyncio.Server: return self._listener async def stop(self) -> None: - """Stop accepting connections and cancel periodic tasks.""" + """Stop accepting connections and cancel periodic and background tasks.""" await self.periodic_executor.stop() if self.mdns is not None: await self.mdns.stop() + + for task in list(self._background): + task.cancel() + if self._background: + await asyncio.gather(*self._background, return_exceptions=True) + self._background.clear() + if self._listener is None: return @@ -416,9 +438,16 @@ def decorator( return decorator + def _spawn(self, coroutine: Coroutine[Any, Any, Any], *, name: str) -> asyncio.Task: + """Run a coroutine in the background, keeping a strong reference until done.""" + task = asyncio.create_task(coroutine, name=name) + self._background.add(task) + task.add_done_callback(self._background.discard) + return task + def _emit_lifecycle(self, event: str, payload: Any) -> None: for handler in self.lifecycle_handlers.get(event, []): - asyncio.create_task(handler(payload), name=event) + self._spawn(handler(payload), name=event) def add_peer(self, peer: Peer, *, event: str = "peer.joined") -> None: if peer.id == self.node_id: @@ -545,19 +574,120 @@ def conversation(self, conversation: Broadcast | str) -> list[ConversationEvent] conversation_id = ( conversation.nonce if isinstance(conversation, Broadcast) else conversation ) - return list(self.conversation_events.get(conversation_id, [])) + return self.conversation_log.events(conversation_id) + + def conversations(self) -> list[str]: + """Return ids of all locally known conversations.""" + return self.conversation_log.conversations() def _remember_conversation_event(self, event: ConversationEvent) -> bool: - if event.event_id in self._seen_conversation_events: - return False self._validate_peer_membership(event.peer) - self._seen_conversation_events.add(event.event_id) - self.conversation_events.setdefault(event.conversation_id, []).append(event) + if not self.conversation_log.append(event): + return False self.add_peer(event.peer) self._emit_lifecycle("conversation.event", event) self._emit_lifecycle(f"conversation.{event.kind}", event) + self._maybe_compact(event.conversation_id) return True + def summarizer(self, wrapped: Summarizer) -> Summarizer: + """Register the coroutine used to summarize events during compaction.""" + self._summarizer = wrapped + return wrapped + + def _maybe_compact(self, conversation_id: str) -> None: + if self.conversation_max_events is None: + return + if conversation_id in self._compacting: + return + if self.conversation_log.count(conversation_id) <= self.conversation_max_events: + return + self._compacting.add(conversation_id) + self._spawn(self._compact_and_release(conversation_id), name="conversation.compact") + + async def _compact_and_release(self, conversation_id: str) -> None: + try: + await self.compact_conversation(conversation_id) + except Exception: + logger.exception("Could not compact conversation {}", conversation_id) + finally: + self._compacting.discard(conversation_id) + + async def compact_conversation( + self, + conversation: Broadcast | str, + *, + keep_recent: int | None = None, + ) -> ConversationEvent | None: + """Fold older events into one local ``summary`` event. + + Compaction is local: each node compresses its own copy of the shared + log. Gossip cannot resurrect compacted events because their ids stay + remembered by the conversation log. + """ + conversation_id = ( + conversation.nonce if isinstance(conversation, Broadcast) else conversation + ) + keep = self.conversation_keep_recent if keep_recent is None else keep_recent + events = self.conversation_log.events(conversation_id) + head = events[0] if events and events[0].kind == "message" else None + compactable = [event for event in events[: len(events) - keep] if event is not head] + if not compactable: + return None + + summary_text = await self._summarizer(compactable) + summary = ConversationEvent( + conversation_id=conversation_id, + event_id=random_hash(), + kind=SUMMARY_KIND, + peer=self.self_peer(), + payload={ + SUMMARY_KIND: summary_text, + "compacted_events": len(compactable), + "from": compactable[0].created_at, + "until": compactable[-1].created_at, + }, + parent_id=head.event_id if head is not None else None, + # Take the newest compacted timestamp so the summary sorts where + # the events it replaces used to sit. + created_at=compactable[-1].created_at, + ) + self.conversation_log.compact( + conversation_id, + [event.event_id for event in compactable], + summary, + ) + self._emit_lifecycle("conversation.compacted", summary) + return summary + + async def sync_conversation( + self, + peer: Peer, + conversation: Broadcast | str, + *, + since: float = 0.0, + ) -> int: + """Pull a conversation's events from a peer; return how many were new. + + Lets a late joiner (or a restarted node) catch up on a shared + conversation it missed. Synced events are stored and emitted locally + but not re-gossiped. + """ + conversation_id = ( + conversation.nonce if isinstance(conversation, Broadcast) else conversation + ) + response = await Client.from_peer(peer).call( + "_synapse.conversation.sync", conversation_id, since=since + ) + if not isinstance(response, dict): + return 0 + added = 0 + for data in response.get("events", []): + event = ConversationEvent.from_dict(data) + with contextlib.suppress(InvalidMessageError): + added += self._remember_conversation_event(event) + return added + async def emit_conversation_event( self, conversation: Broadcast | str, @@ -669,12 +799,30 @@ async def swarm_ask( ) -> Any: if self._ask_handler is None: raise RuntimeError("node has no ask handler") - if broadcast is not None: - await self.ack(broadcast) - result = await self._ask_handler(task, context or {}) - if broadcast is not None: - await self.reply(broadcast, result) - return result + if broadcast is None: + return await self._ask_handler(task, context or {}) + # Defer: ACK now, run the handler in the background, and deliver + # the result as a broadcast reply. Keeps the origin's RPC short no + # matter how long the agent behind the handler takes. + await self.ack(broadcast) + self._spawn(self._run_deferred_ask(task, context or {}, broadcast), name="synapse.ask") + return {"accepted": True, "deferred": True} + + async def _run_deferred_ask( + self, task: str, context: dict[str, Any], broadcast: Broadcast + ) -> None: + assert self._ask_handler is not None + try: + result = await self._ask_handler(task, context) + except asyncio.CancelledError: + raise + except Exception as e: + logger.exception("Ask handler raised for broadcast {}", broadcast.nonce) + await self.emit_conversation_event( + broadcast, "error", {"task": task, "error": str(e)} + ) + return + await self.reply(broadcast, result) def _register_system_endpoints(self) -> None: @self.endpoint("_synapse.ping", publish=False) @@ -719,6 +867,15 @@ async def conversation_event(event: dict[str, Any]) -> dict[str, Any]: await self._send_conversation_event(incoming) return {"ok": True, "stored": remembered} + @self.endpoint("_synapse.conversation.sync", publish=False) + async def conversation_sync(conversation_id: str, since: float = 0.0) -> dict[str, Any]: + events = self.conversation_log.events(conversation_id, since=since) + return {"events": [event.to_dict() for event in events]} + + @self.endpoint("_synapse.conversation.list", publish=False) + async def conversation_list() -> list[str]: + return self.conversation_log.conversations() + @self.endpoint("_synapse.join", publish=False) async def join(peer: dict) -> dict: incoming = Peer.from_dict(peer) diff --git a/synapse_p2p/teams.py b/synapse_p2p/teams.py new file mode 100644 index 0000000..c71f607 --- /dev/null +++ b/synapse_p2p/teams.py @@ -0,0 +1,231 @@ +"""A task layer for agent teams, built on shared conversation events. + +Synapse core stays neutral: it moves events and does not decide who works on +what. This module layers a small, explicit vocabulary on top: + +- ``task.offer`` — a :class:`Team` announces work and its requirements +- ``task.claim`` — a :class:`Worker` volunteers for an offered task +- ``task.grant`` — the offering team assigns the task to one claimant +- ``task.progress`` — the assignee narrates progress into the shared log +- ``task.done`` — the assignee delivers a result +- ``task.failed`` — the assignee reports an error + +Each task is its own conversation (``conversation_id == task id``), so the full +history of a task — offer, claims, progress, result — is one gossiped, durable, +compactable thread that any swarm member can watch or sync later. +""" + +from __future__ import annotations + +import asyncio +import time +from collections.abc import Awaitable, Callable +from dataclasses import dataclass, field +from typing import Any + +from loguru import logger + +from synapse_p2p.node import Node, new_nonce +from synapse_p2p.types import ConversationEvent, Peer + +TASK_OFFER = "task.offer" +TASK_CLAIM = "task.claim" +TASK_GRANT = "task.grant" +TASK_PROGRESS = "task.progress" +TASK_DONE = "task.done" +TASK_FAILED = "task.failed" + + +class TeamTaskError(RuntimeError): + """Raised when a task fails or times out while waiting for its result.""" + + +@dataclass(slots=True) +class TeamTask: + """The offering side's view of one unit of work.""" + + id: str + title: str + spec: dict[str, Any] = field(default_factory=dict) + requires: list[str] = field(default_factory=list) + status: str = "offered" + assignee: Peer | None = None + result: Any = None + error: str | None = None + progress: list[dict[str, Any]] = field(default_factory=list) + + +@dataclass(slots=True) +class Assignment: + """The worker side's view of a granted task, with a progress channel.""" + + id: str + title: str + spec: dict[str, Any] + node: Node + + async def progress(self, message: str, **data: Any) -> None: + await self.node.emit_conversation_event( + self.id, TASK_PROGRESS, {"task_id": self.id, "message": message, **data} + ) + + +class Team: + """Offer tasks to the swarm and collect results. + + The team grants each task to the first claimant, so exactly one worker + runs it even when many volunteer. + """ + + def __init__(self, node: Node) -> None: + self.node = node + self.tasks: dict[str, TeamTask] = {} + self._finished: dict[str, asyncio.Event] = {} + node.on(f"conversation.{TASK_CLAIM}")(self._on_claim) + node.on(f"conversation.{TASK_PROGRESS}")(self._on_progress) + node.on(f"conversation.{TASK_DONE}")(self._on_done) + node.on(f"conversation.{TASK_FAILED}")(self._on_failed) + + async def offer( + self, + title: str, + *, + spec: dict[str, Any] | None = None, + requires: list[str] | None = None, + ) -> TeamTask: + task = TeamTask(id=new_nonce(), title=title, spec=spec or {}, requires=requires or []) + self.tasks[task.id] = task + self._finished[task.id] = asyncio.Event() + await self.node.emit_conversation_event( + task.id, + TASK_OFFER, + {"task_id": task.id, "title": title, "spec": task.spec, "requires": task.requires}, + ) + return task + + async def wait(self, task: TeamTask, *, timeout: float | None = None) -> Any: + """Wait until a task finishes; return its result or raise TeamTaskError.""" + finished = self._finished[task.id] + try: + await asyncio.wait_for(finished.wait(), timeout) + except TimeoutError as e: + raise TeamTaskError(f"task {task.id} timed out: {task.title}") from e + if task.status == "failed": + raise TeamTaskError(task.error or f"task {task.id} failed") + return task.result + + async def _on_claim(self, event: ConversationEvent) -> None: + task = self.tasks.get(event.conversation_id) + if task is None or task.assignee is not None or task.status != "offered": + return + task.assignee = event.peer + task.status = "claimed" + await self.node.emit_conversation_event( + task.id, + TASK_GRANT, + {"task_id": task.id, "worker_id": event.peer.id, "worker_name": event.peer.name}, + parent_id=event.event_id, + ) + + async def _on_progress(self, event: ConversationEvent) -> None: + task = self.tasks.get(event.conversation_id) + if task is not None: + task.progress.append(dict(event.payload)) + + async def _on_done(self, event: ConversationEvent) -> None: + task = self.tasks.get(event.conversation_id) + if task is None or task.status in {"done", "failed"}: + return + task.result = event.payload.get("result") + task.status = "done" + self._finished[task.id].set() + + async def _on_failed(self, event: ConversationEvent) -> None: + task = self.tasks.get(event.conversation_id) + if task is None or task.status in {"done", "failed"}: + return + task.error = str(event.payload.get("error", "unknown error")) + task.status = "failed" + self._finished[task.id].set() + + +TaskHandler = Callable[[Assignment], Awaitable[Any]] + + +class Worker: + """Claim offered tasks the node is capable of, and run them when granted.""" + + def __init__(self, node: Node, *, claim_timeout: float = 60) -> None: + self.node = node + self.claim_timeout = claim_timeout + self._handler: TaskHandler | None = None + self._pending: dict[str, tuple[ConversationEvent, float]] = {} + node.on(f"conversation.{TASK_OFFER}")(self._on_offer) + node.on(f"conversation.{TASK_GRANT}")(self._on_grant) + + def task(self, wrapped: TaskHandler) -> TaskHandler: + """Decorator registering the coroutine that executes granted tasks.""" + self._handler = wrapped + return wrapped + + def _can_do(self, requires: list[str]) -> bool: + mine = {capability.name for capability in self.node.capabilities} + return set(requires) <= mine + + def _prune_pending(self) -> None: + deadline = time.time() - self.claim_timeout + stale = [task_id for task_id, (_, at) in self._pending.items() if at < deadline] + for task_id in stale: + self._pending.pop(task_id, None) + + async def _on_offer(self, event: ConversationEvent) -> None: + self._prune_pending() + if self._handler is None: + return + if event.peer.id == self.node.node_id: + return + if not self._can_do(list(event.payload.get("requires", []))): + return + task_id = str(event.payload["task_id"]) + self._pending[task_id] = (event, time.time()) + await self.node.emit_conversation_event( + task_id, + TASK_CLAIM, + { + "task_id": task_id, + "capabilities": [capability.name for capability in self.node.capabilities], + }, + parent_id=event.event_id, + ) + + async def _on_grant(self, event: ConversationEvent) -> None: + task_id = str(event.payload.get("task_id", "")) + pending = self._pending.pop(task_id, None) + if pending is None: + return + if event.payload.get("worker_id") != self.node.node_id: + return # another worker was granted the task + offer, _ = pending + assignment = Assignment( + id=task_id, + title=str(offer.payload.get("title", "")), + spec=dict(offer.payload.get("spec", {})), + node=self.node, + ) + self.node._spawn(self._run(assignment), name=f"task:{task_id}") + + async def _run(self, assignment: Assignment) -> None: + assert self._handler is not None + try: + result = await self._handler(assignment) + except asyncio.CancelledError: + raise + except Exception as e: + logger.exception("Task {} failed on {}", assignment.id, self.node.name) + await self.node.emit_conversation_event( + assignment.id, TASK_FAILED, {"task_id": assignment.id, "error": str(e)} + ) + return + await self.node.emit_conversation_event( + assignment.id, TASK_DONE, {"task_id": assignment.id, "result": result} + ) diff --git a/synapse_p2p/tests/test_conversations.py b/synapse_p2p/tests/test_conversations.py new file mode 100644 index 0000000..e18a4ea --- /dev/null +++ b/synapse_p2p/tests/test_conversations.py @@ -0,0 +1,201 @@ +import asyncio + +import pytest + +from synapse_p2p import ( + Client, + ConversationEvent, + MemoryConversationLog, + Node, + Peer, + SqliteConversationLog, + default_summarizer, +) + + +def make_event(event_id: str, *, conversation_id: str = "conv", created_at: float = 0.0): + return ConversationEvent( + conversation_id=conversation_id, + event_id=event_id, + kind="reply", + peer=Peer(id=f"peer-{event_id}", address="127.0.0.1", port=1, name=event_id), + payload={"result": event_id}, + created_at=created_at or float(len(event_id)), + ) + + +@pytest.mark.parametrize("backend", ["memory", "sqlite"]) +def test_log_appends_deduplicates_and_lists(backend, tmp_path): + log = ( + MemoryConversationLog() + if backend == "memory" + else SqliteConversationLog(tmp_path / "log.db") + ) + event = make_event("a") + + assert log.append(event) is True + assert log.append(event) is False + assert log.seen("a") is True + assert log.seen("missing") is False + assert log.count("conv") == 1 + assert log.conversations() == ["conv"] + assert log.events("conv")[0].payload == {"result": "a"} + log.close() + + +@pytest.mark.parametrize("backend", ["memory", "sqlite"]) +def test_log_compaction_replaces_events_and_blocks_resurrection(backend, tmp_path): + log = ( + MemoryConversationLog() + if backend == "memory" + else SqliteConversationLog(tmp_path / "log.db") + ) + old = make_event("old", created_at=1.0) + recent = make_event("recent", created_at=2.0) + log.append(old) + log.append(recent) + + summary = ConversationEvent( + conversation_id="conv", + event_id="summary-1", + kind="summary", + peer=old.peer, + payload={"summary": "old stuff"}, + created_at=1.0, + ) + log.compact("conv", ["old"], summary) + + remaining = {event.event_id for event in log.events("conv")} + assert remaining == {"summary-1", "recent"} + # Gossip re-delivering a compacted event must not resurrect it. + assert log.append(old) is False + assert log.seen("old") is True + log.close() + + +def test_sqlite_log_survives_restart(tmp_path): + path = tmp_path / "log.db" + log = SqliteConversationLog(path) + log.append(make_event("a")) + log.close() + + reopened = SqliteConversationLog(path) + assert reopened.count("conv") == 1 + assert reopened.seen("a") is True + reopened.close() + + +@pytest.mark.asyncio +async def test_default_summarizer_mentions_peers_and_folds_prior_summaries(): + events = [make_event("a"), make_event("b")] + events.append( + ConversationEvent( + conversation_id="conv", + event_id="s", + kind="summary", + peer=events[0].peer, + payload={"summary": "earlier: c replied"}, + ) + ) + text = await default_summarizer(events) + assert "a [reply]" in text + assert "earlier: c replied" in text + + +@pytest.mark.asyncio +async def test_node_auto_compacts_conversation_past_max_events(): + node = Node( + name="compactor", + swarm="foo.electron.network", + bind="127.0.0.1", + heartbeat_interval=None, + conversation_max_events=5, + conversation_keep_recent=2, + ) + compacted = asyncio.Event() + + @node.on("conversation.compacted") + async def on_compacted(event: ConversationEvent) -> None: + compacted.set() + + await node.start() + try: + for index in range(8): + await node.emit_conversation_event("conv", "note", {"index": index}) + await asyncio.wait_for(compacted.wait(), 1) + + events = node.conversation("conv") + assert len(events) <= 5 + summaries = [event for event in events if event.kind == "summary"] + assert summaries + assert summaries[0].payload["compacted_events"] >= 3 + # The most recent events are preserved verbatim. + assert events[-1].payload == {"index": 7} + finally: + await node.stop() + + +@pytest.mark.asyncio +async def test_manual_compaction_uses_custom_summarizer_and_keeps_head_message(): + node = Node( + name="compactor", + swarm="foo.electron.network", + bind="127.0.0.1", + heartbeat_interval=None, + ) + + @node.summarizer + async def summarize(events: list[ConversationEvent]) -> str: + return f"custom summary of {len(events)} events" + + await node.start() + try: + broadcast = await node.broadcast("team.question", "who can help?") + for index in range(6): + await node.emit_conversation_event(broadcast, "note", {"index": index}) + + summary = await node.compact_conversation(broadcast, keep_recent=2) + + assert summary is not None + assert summary.payload["summary"] == "custom summary of 4 events" + events = node.conversation(broadcast) + # Head message survives compaction so the conversation keeps its opening. + assert events[0].kind == "message" + assert [event.kind for event in events[1:]] == ["summary", "note", "note"] + finally: + await node.stop() + + +@pytest.mark.asyncio +async def test_late_joiner_syncs_conversation_from_peer(): + origin = Node( + name="origin", + swarm="foo.electron.network", + bind="127.0.0.1", + heartbeat_interval=None, + ) + late = Node( + name="late", + swarm="foo.electron.network", + bind="127.0.0.1", + heartbeat_interval=None, + ) + await origin.start() + + broadcast = await origin.broadcast("team.question", "who can help?") + await origin.emit_conversation_event(broadcast, "note", {"index": 1}) + + await late.start() + try: + added = await late.sync_conversation(origin.self_peer(), broadcast) + + assert added == 2 + events = late.conversation(broadcast) + assert [event.kind for event in events] == ["message", "note"] + # Syncing again is idempotent. + assert await late.sync_conversation(origin.self_peer(), broadcast) == 0 + listed = await Client.from_peer(origin.self_peer()).call("_synapse.conversation.list") + assert isinstance(listed, list) and broadcast.nonce in listed + finally: + await late.stop() + await origin.stop() diff --git a/synapse_p2p/tests/test_deferred_ask.py b/synapse_p2p/tests/test_deferred_ask.py new file mode 100644 index 0000000..0a6b780 --- /dev/null +++ b/synapse_p2p/tests/test_deferred_ask.py @@ -0,0 +1,98 @@ +import asyncio + +import pytest + +from synapse_p2p import Client, ConversationEvent, Node + + +def make_node(name: str) -> Node: + return Node( + name=name, + swarm="foo.electron.network", + bind="127.0.0.1", + heartbeat_interval=None, + ) + + +@pytest.mark.asyncio +async def test_swarm_ask_defers_slow_handlers_and_replies_later(): + origin = make_node("origin") + worker = make_node("worker") + started = asyncio.Event() + release = asyncio.Event() + + @worker.ask + async def handle(task: str, context: dict) -> dict: + started.set() + await release.wait() + return {"task": task} + + await origin.start() + await worker.start() + origin.add_peer(worker.self_peer()) + + try: + broadcast = await origin.broadcast("synapse.ask", "long job") + # The broadcast returns while the handler is still running. + await asyncio.wait_for(started.wait(), 1) + assert origin.replies(broadcast) == [] + + release.set() + for _ in range(50): + if origin.replies(broadcast): + break + await asyncio.sleep(0.02) + + replies = origin.replies(broadcast) + assert replies and replies[0].result == {"task": "long job"} + assert any(event.kind == "ack" for event in origin.conversation(broadcast)) + finally: + await worker.stop() + await origin.stop() + + +@pytest.mark.asyncio +async def test_swarm_ask_handler_error_becomes_conversation_error_event(): + origin = make_node("origin") + worker = make_node("worker") + errored = asyncio.Queue() + + @origin.on("conversation.error") + async def on_error(event: ConversationEvent) -> None: + await errored.put(event) + + @worker.ask + async def handle(task: str, context: dict) -> dict: + raise RuntimeError("no can do") + + await origin.start() + await worker.start() + origin.add_peer(worker.self_peer()) + + try: + await origin.broadcast("synapse.ask", "doomed job") + event = await asyncio.wait_for(errored.get(), 1) + + assert event.payload["error"] == "no can do" + assert event.peer.name == "worker" + finally: + await worker.stop() + await origin.stop() + + +@pytest.mark.asyncio +async def test_direct_synapse_ask_without_broadcast_stays_synchronous(): + worker = make_node("worker") + + @worker.ask + async def handle(task: str, context: dict) -> dict: + return {"task": task, "sync": True} + + server = await worker.start() + host, port = server.sockets[0].getsockname()[:2] + + try: + result = await Client(host, port).call("synapse.ask", "quick job") + assert result == {"task": "quick job", "sync": True} + finally: + await worker.stop() diff --git a/synapse_p2p/tests/test_teams.py b/synapse_p2p/tests/test_teams.py new file mode 100644 index 0000000..c39dae2 --- /dev/null +++ b/synapse_p2p/tests/test_teams.py @@ -0,0 +1,145 @@ +import asyncio + +import pytest + +from synapse_p2p import Capability, Node +from synapse_p2p.teams import Assignment, Team, TeamTaskError, Worker + + +def make_node(name: str, capabilities: list[str | Capability] | None = None) -> Node: + return Node( + name=name, + swarm="foo.electron.network", + capabilities=capabilities or [], + bind="127.0.0.1", + heartbeat_interval=None, + ) + + +async def connect(architect: Node, *workers: Node) -> None: + for worker in workers: + architect.add_peer(worker.self_peer()) + worker.add_peer(architect.self_peer()) + + +@pytest.mark.asyncio +async def test_offer_claim_grant_and_done_flow(): + architect_node = make_node("architect") + coder_node = make_node("coder", ["python"]) + team = Team(architect_node) + worker = Worker(coder_node) + + @worker.task + async def implement(assignment: Assignment) -> dict: + await assignment.progress("starting", step=1) + return {"diff": f"patch for {assignment.title}", "spec": assignment.spec} + + await architect_node.start() + await coder_node.start() + await connect(architect_node, coder_node) + + try: + task = await team.offer( + "implement the parser", spec={"file": "parser.py"}, requires=["python"] + ) + result = await team.wait(task, timeout=2) + + assert result == {"diff": "patch for implement the parser", "spec": {"file": "parser.py"}} + assert task.status == "done" + assert task.assignee is not None and task.assignee.name == "coder" + assert any(entry["message"] == "starting" for entry in task.progress) + + kinds = [event.kind for event in architect_node.conversation(task.id)] + assert kinds[0] == "task.offer" + assert {"task.claim", "task.grant", "task.progress", "task.done"} <= set(kinds) + finally: + await coder_node.stop() + await architect_node.stop() + + +@pytest.mark.asyncio +async def test_only_one_worker_is_granted_the_task(): + architect_node = make_node("architect") + coder_one = make_node("coder-1", ["python"]) + coder_two = make_node("coder-2", ["python"]) + team = Team(architect_node) + ran: list[str] = [] + + for node in (coder_one, coder_two): + worker = Worker(node) + + def handler_for(name: str): + async def implement(assignment: Assignment) -> str: + ran.append(name) + return f"done by {name}" + + return implement + + worker.task(handler_for(node.name)) + + await architect_node.start() + await coder_one.start() + await coder_two.start() + await connect(architect_node, coder_one, coder_two) + + try: + task = await team.offer("small fix", requires=["python"]) + result = await team.wait(task, timeout=2) + await asyncio.sleep(0.1) + + assert len(ran) == 1 + assert result == f"done by {ran[0]}" + finally: + await coder_two.stop() + await coder_one.stop() + await architect_node.stop() + + +@pytest.mark.asyncio +async def test_worker_without_required_capability_does_not_claim(): + architect_node = make_node("architect") + coder_node = make_node("docs-only", ["docs"]) + team = Team(architect_node) + worker = Worker(coder_node) + + @worker.task + async def implement(assignment: Assignment) -> str: + return "should never run" + + await architect_node.start() + await coder_node.start() + await connect(architect_node, coder_node) + + try: + task = await team.offer("rust rewrite", requires=["rust"]) + with pytest.raises(TeamTaskError, match="timed out"): + await team.wait(task, timeout=0.3) + assert task.assignee is None + finally: + await coder_node.stop() + await architect_node.stop() + + +@pytest.mark.asyncio +async def test_worker_failure_propagates_to_team_wait(): + architect_node = make_node("architect") + coder_node = make_node("coder", ["python"]) + team = Team(architect_node) + worker = Worker(coder_node) + + @worker.task + async def implement(assignment: Assignment) -> str: + raise ValueError("model refused") + + await architect_node.start() + await coder_node.start() + await connect(architect_node, coder_node) + + try: + task = await team.offer("impossible task", requires=["python"]) + with pytest.raises(TeamTaskError, match="model refused"): + await team.wait(task, timeout=2) + assert task.status == "failed" + finally: + await coder_node.stop() + await architect_node.stop()