diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
new file mode 100644
index 0000000..1b6ed83
--- /dev/null
+++ b/.github/workflows/test.yml
@@ -0,0 +1,31 @@
+name: Tests
+
+on:
+ push:
+ branches: [main, develop]
+ pull_request:
+ branches: [main]
+
+jobs:
+ test:
+ runs-on: ubuntu-latest
+ strategy:
+ matrix:
+ python-version: ["3.11", "3.12", "3.13"]
+
+ steps:
+ - uses: actions/checkout@v4
+
+ - name: Set up Python ${{ matrix.python-version }}
+ uses: actions/setup-python@v5
+ with:
+ python-version: ${{ matrix.python-version }}
+
+ - name: Install dependencies
+ run: |
+ python -m pip install --upgrade pip
+ pip install -e ".[dev]"
+
+ - name: Run tests
+ run: |
+ python -m pytest tests/ -v --tb=short
diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..87cd8c9
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,18 @@
+__pycache__/
+*.py[cod]
+*$py.class
+*.egg-info/
+dist/
+build/
+.eggs/
+*.egg
+.env
+.env.local
+*.pem
+*.key
+.agent_memory/
+.pytest_cache/
+.mypy_cache/
+.ruff_cache/
+.pytest_cache/
+*.egg-info/
diff --git a/README.md b/README.md
new file mode 100644
index 0000000..5efddf2
--- /dev/null
+++ b/README.md
@@ -0,0 +1,173 @@
+# Agent Framework Starter
+
+[](https://github.com/genbrain-ai/agent-framework-starter/actions/workflows/test.yml)
+[](https://www.python.org/downloads/)
+[](LICENSE)
+
+**A minimal Python framework for building multi-agent AI teams using NATS JetStream messaging.**
+
+Built by [GenBrain AI](https://genbrain.ai) — the company behind [agent.ceo](https://agent.ceo).
+
+---
+
+## Architecture
+
+```mermaid
+graph LR
+ subgraph Control Plane
+ O[Orchestrator]
+ end
+
+ subgraph Message Bus
+ N[NATS JetStream]
+ end
+
+ subgraph Agent Team
+ A1[Agent 1]
+ A2[Agent 2]
+ AN[Agent N]
+ end
+
+ O -->|delegate & broadcast| N
+ N -->|task assignments| A1
+ N -->|task assignments| A2
+ N -->|task assignments| AN
+ A1 -->|results| N
+ A2 -->|results| N
+ AN -->|results| N
+ N -->|collect results| O
+```
+
+## Quick Start
+
+```python
+import asyncio
+from src.agent import Agent
+from src.messaging import Message
+from src.orchestrator import Orchestrator
+
+class Worker(Agent):
+ async def handle_message(self, message: Message):
+ print(f"Working on: {message.payload['title']}")
+ await self.send_message(message.from_agent, {"status": "done"})
+
+async def main():
+ orch = Orchestrator(org_id="myteam")
+ orch.add_agent(Worker(name="worker", role="analyst"))
+ await orch.start_all()
+ orch.delegate_task("Analyze Q4 data", agent_name="worker")
+
+asyncio.run(main())
+```
+
+## Features
+
+- **Async-first** — built on `asyncio` with `async/await` throughout
+- **NATS JetStream messaging** — durable, at-least-once delivery between agents
+- **Task delegation** — assign, track, and verify work across agents
+- **Orchestrator** — manage agent lifecycles, health monitoring, and broadcasting
+- **Agent memory** — persistent key-value storage with file locking
+- **Status validation** — task lifecycle with enforced state transitions
+- **Extensible** — subclass `Agent` and override `handle_message()` to build anything
+
+## Installation
+
+### Prerequisites
+
+- Python 3.11+
+- Docker (for NATS)
+
+### Install
+
+```bash
+git clone https://github.com/genbrain-ai/agent-framework-starter.git
+cd agent-framework-starter
+pip install -e ".[dev]"
+```
+
+### Start NATS
+
+```bash
+docker run -d --name nats -p 4222:4222 -p 8222:8222 nats:latest -js
+```
+
+### Run the Examples
+
+```bash
+# Simplest agent — connects and responds to messages
+python -m examples.hello_agent
+
+# Two-agent team — CEO delegates to a Worker
+python -m examples.two_agent_team
+
+# Security scanner agent
+python -m examples.security_agent
+
+# DevOps deployment pipeline agent
+python -m examples.devops_agent
+```
+
+### Run the Tests
+
+```bash
+python -m pytest tests/ -v
+```
+
+## Project Structure
+
+```
+agent-framework-starter/
+ src/
+ agent.py # Base Agent class
+ messaging.py # NATS JetStream helpers + Message model
+ task_manager.py # Task lifecycle management
+ memory.py # Persistent key-value store
+ orchestrator.py # Multi-agent coordinator
+ examples/
+ hello_agent.py # Minimal agent example
+ two_agent_team.py # Delegation pattern demo
+ security_agent.py # Security review agent
+ devops_agent.py # Deployment pipeline agent
+ tests/
+ test_agent.py
+ test_messaging.py
+ test_orchestrator.py
+ test_task_manager.py
+ test_memory.py
+ docs/
+ architecture.md # System design & diagrams
+ getting-started.md# Step-by-step setup guide
+ patterns.md # Common multi-agent patterns
+```
+
+## Documentation
+
+- [Architecture](docs/architecture.md) — how multi-agent systems work, component overview, Mermaid diagrams
+- [Getting Started](docs/getting-started.md) — install, run NATS, launch your first agent
+- [Patterns](docs/patterns.md) — delegation, verification, fan-out, pipelines, heartbeats
+
+## Production Use
+
+This starter kit demonstrates the foundational patterns. For production-grade AI agent orchestration with Kubernetes-native deployment, automatic scaling, built-in security review, LLM integration, and a web dashboard:
+
+**[Try agent.ceo](https://agent.ceo) — 1 agent-week free.**
+
+Full documentation at [docs.agent.ceo](https://docs.agent.ceo).
+
+## Contributing
+
+Contributions are welcome! Please:
+
+1. Fork the repository
+2. Create a feature branch (`git checkout -b feature/my-feature`)
+3. Write tests for your changes
+4. Ensure all tests pass (`pytest tests/ -v`)
+5. Submit a pull request
+
+## License
+
+MIT License. See [LICENSE](LICENSE) for details.
+
+---
+
+Built with care by [GenBrain AI](https://genbrain.ai).
diff --git a/docs/architecture.md b/docs/architecture.md
new file mode 100644
index 0000000..fa8db9b
--- /dev/null
+++ b/docs/architecture.md
@@ -0,0 +1,143 @@
+# Architecture
+
+## How Multi-Agent Systems Work
+
+A multi-agent system is a group of autonomous agents that collaborate to solve problems no single agent could handle alone. Each agent has a specific **role** (e.g., CEO, DevOps, Security Reviewer) and communicates with other agents through a shared **message bus**.
+
+This framework uses [NATS JetStream](https://nats.io/) as the message bus — a lightweight, high-performance messaging system that provides:
+
+- **Pub/Sub messaging** — agents publish and subscribe to subjects
+- **Durable delivery** — messages are persisted so agents can reconnect without missing work
+- **At-least-once semantics** — critical for task delegation where you cannot afford to lose a message
+
+## System Diagram
+
+```mermaid
+graph TB
+ subgraph Orchestrator
+ O[Orchestrator]
+ TM[Task Manager]
+ HB[Heartbeat Monitor]
+ end
+
+ subgraph Message Bus
+ NATS[NATS JetStream]
+ end
+
+ subgraph Agents
+ A1[Agent 1
CEO]
+ A2[Agent 2
Worker]
+ A3[Agent 3
DevOps]
+ A4[Agent N
...]
+ end
+
+ O --> TM
+ O --> HB
+ O -->|delegate task| NATS
+ NATS -->|task_assignment| A1
+ NATS -->|task_assignment| A2
+ NATS -->|task_assignment| A3
+ NATS -->|task_assignment| A4
+ A1 -->|task_result| NATS
+ A2 -->|task_result| NATS
+ A3 -->|task_result| NATS
+ A4 -->|task_result| NATS
+ NATS -->|results| O
+ A1 <-->|peer messages| NATS
+ A2 <-->|peer messages| NATS
+```
+
+## Core Components
+
+### Agent (`src/agent.py`)
+
+The base building block. Every agent:
+
+1. **Connects** to NATS on startup
+2. **Subscribes** to its inbox subject (`{org_id}.agents.{name}.inbox`)
+3. **Processes** incoming messages via `handle_message()`
+4. **Sends** messages to other agents via `send_message()`
+
+Subclass `Agent` and override `handle_message()` to build any kind of agent.
+
+### Messaging (`src/messaging.py`)
+
+Handles NATS connection management and message serialization. Every message is a structured `Message` object with:
+
+- `id` — unique identifier
+- `from_agent` / `to_agent` — routing information
+- `type` — message category (e.g., `task_assignment`, `task_result`, `broadcast`)
+- `payload` — arbitrary JSON data
+- `timestamp` — when the message was created
+
+### Task Manager (`src/task_manager.py`)
+
+Tracks units of work flowing through the system. Tasks have a defined lifecycle:
+
+```
+ASSIGNED → IN_PROGRESS → COMPLETED
+ ↓ ↓
+ FAILED ← FAILED
+ ↓
+ ASSIGNED (retry)
+```
+
+Status transitions are validated — you cannot skip from `ASSIGNED` directly to `COMPLETED` without going through `IN_PROGRESS`.
+
+### Orchestrator (`src/orchestrator.py`)
+
+The control plane that manages multiple agents:
+
+- **Registry** — add/remove agents
+- **Lifecycle** — start/stop all agents as a group
+- **Delegation** — assign tasks to specific agents
+- **Broadcasting** — send a message to every agent
+- **Health monitoring** — track agent liveness via heartbeats
+
+### Memory Store (`src/memory.py`)
+
+Persistent key-value storage using JSON files. Agents can save and restore state across restarts. Thread-safe via file locking.
+
+## Message Flow
+
+A typical task delegation flow:
+
+1. **Orchestrator** creates a `Task` and assigns it to an agent
+2. **Orchestrator** publishes a `task_assignment` message to the agent's inbox
+3. **Agent** receives the message, processes the task
+4. **Agent** publishes a `task_result` message back to the assigner
+5. **Assigner** verifies the result and marks the task complete
+
+```mermaid
+sequenceDiagram
+ participant O as Orchestrator
+ participant NATS as NATS JetStream
+ participant W as Worker Agent
+
+ O->>NATS: publish task_assignment
+ NATS->>W: deliver to worker.inbox
+ W->>W: process task
+ W->>NATS: publish task_result
+ NATS->>O: deliver result
+ O->>O: verify & complete task
+```
+
+## Design Principles
+
+1. **Agents are autonomous** — each agent runs independently and makes its own decisions
+2. **Communication is asynchronous** — agents don't block waiting for responses
+3. **Tasks have clear ownership** — every task has an assigner and an assignee
+4. **Verification is separate from execution** — the assigner verifies the assignee's work
+5. **Failure is expected** — agents can fail tasks, and tasks can be retried
+
+## Production Use
+
+This starter kit demonstrates the core patterns. For production-grade agent orchestration with:
+
+- Kubernetes-native deployment
+- Automatic agent scaling
+- Built-in security review
+- LLM integration
+- Web dashboard
+
+Try [agent.ceo](https://agent.ceo) by GenBrain AI.
diff --git a/docs/getting-started.md b/docs/getting-started.md
new file mode 100644
index 0000000..9a4876f
--- /dev/null
+++ b/docs/getting-started.md
@@ -0,0 +1,131 @@
+# Getting Started
+
+Get a multi-agent system running in under 5 minutes.
+
+## Prerequisites
+
+- Python 3.11+
+- Docker (for running NATS)
+
+## Step 1: Clone and Install
+
+```bash
+git clone https://github.com/genbrain-ai/agent-framework-starter.git
+cd agent-framework-starter
+pip install -e ".[dev]"
+```
+
+## Step 2: Start NATS
+
+Start a NATS server with JetStream enabled:
+
+```bash
+docker run -d --name nats \
+ -p 4222:4222 \
+ -p 8222:8222 \
+ nats:latest -js
+```
+
+Verify NATS is running:
+
+```bash
+curl http://localhost:8222/healthz
+# Should return: {"status":"ok"}
+```
+
+## Step 3: Run the Hello Agent
+
+The simplest possible agent — it connects, listens for messages, and responds with a greeting.
+
+```bash
+python -m examples.hello_agent
+```
+
+You should see:
+
+```
+Starting HelloAgent... (press Ctrl+C to stop)
+[hello] Ready and listening on demo.agents.hello.inbox
+```
+
+The agent is now connected and waiting for messages. Press `Ctrl+C` to stop it.
+
+## Step 4: Run the Two-Agent Team
+
+This example shows delegation: a CEO agent creates a task, sends it to a Worker, and the Worker reports back.
+
+```bash
+python -m examples.two_agent_team
+```
+
+You should see output like:
+
+```
+============================================================
+Two-Agent Team Demo: CEO delegates to Worker
+============================================================
+[ceo] CEO agent ready. Delegating initial task...
+[worker] Worker agent ready. Waiting for tasks...
+[ceo] Delegated task 'Analyze quarterly revenue data...' to worker
+[worker] Received task: 'Analyze quarterly revenue data...'
+[worker] Processing...
+[worker] Task complete. Reporting result to 'ceo'...
+[ceo] Received task result from 'worker':
+ Task ID: ...
+ Status: completed
+ Result: { ... }
+[ceo] Task verified and marked complete!
+[ceo] All tasks complete. Shutting down team...
+Demo complete!
+```
+
+## Step 5: Run the Tests
+
+```bash
+python -m pytest tests/ -v
+```
+
+## Step 6: Build Your Own Agent
+
+Create a new file `my_agent.py`:
+
+```python
+import asyncio
+from src.agent import Agent
+from src.messaging import Message
+
+
+class MyAgent(Agent):
+ async def handle_message(self, message: Message) -> None:
+ print(f"Received: {message.payload}")
+
+ # Do your custom logic here
+ result = {"processed": True, "input": message.payload}
+
+ await self.send_message(
+ to_agent=message.from_agent,
+ payload=result,
+ msg_type="result",
+ )
+
+
+async def main():
+ agent = MyAgent(name="my-agent", role="processor", org_id="myorg")
+ await agent.run()
+
+
+if __name__ == "__main__":
+ asyncio.run(main())
+```
+
+## More Examples
+
+- **Security Agent** (`examples/security_agent.py`) — scans repos for security issues
+- **DevOps Agent** (`examples/devops_agent.py`) — simulates a deployment pipeline
+
+## Next Steps
+
+- Read the [Architecture Guide](architecture.md) to understand how the system works
+- Explore [Common Patterns](patterns.md) for delegation, fan-out, and pipelines
+- For production use, check out [agent.ceo](https://agent.ceo)
+- Full documentation at [docs.agent.ceo](https://docs.agent.ceo)
diff --git a/docs/patterns.md b/docs/patterns.md
new file mode 100644
index 0000000..bc058fb
--- /dev/null
+++ b/docs/patterns.md
@@ -0,0 +1,278 @@
+# Common Patterns
+
+Proven patterns for building multi-agent systems. Each pattern includes a description and runnable code snippet.
+
+## 1. Delegation
+
+The most fundamental pattern: one agent assigns work to another and waits for a result.
+
+**When to use:** Any time you need to break work into specialized subtasks.
+
+```python
+from src.agent import Agent
+from src.messaging import Message
+from src.task_manager import TaskManager
+
+
+class ManagerAgent(Agent):
+ def __init__(self, **kwargs):
+ super().__init__(**kwargs)
+ self.task_manager = TaskManager()
+
+ async def on_start(self):
+ task = self.task_manager.assign(
+ title="Review PR #42",
+ assignee="reviewer",
+ assigner=self.name,
+ )
+ await self.send_message(
+ to_agent="reviewer",
+ payload=task.model_dump(),
+ msg_type="task_assignment",
+ )
+
+ async def handle_message(self, message: Message):
+ if message.type == "task_result":
+ task_id = message.payload["task_id"]
+ print(f"Task {task_id} completed: {message.payload['result']}")
+```
+
+## 2. Verification
+
+The assigner verifies the assignee's work before marking it complete. This prevents agents from self-certifying their output.
+
+**When to use:** Any time quality matters (code review, data analysis, deployment approval).
+
+```python
+class VerifyingManager(Agent):
+ async def handle_message(self, message: Message):
+ if message.type == "task_result":
+ result = message.payload.get("result", {})
+
+ # Verify the work meets requirements
+ if self._verify(result):
+ print(f"Verified: {result}")
+ await self.send_message(
+ to_agent=message.from_agent,
+ payload={"status": "approved"},
+ msg_type="verification_result",
+ )
+ else:
+ # Send back for rework
+ await self.send_message(
+ to_agent=message.from_agent,
+ payload={
+ "status": "rejected",
+ "feedback": "Missing test coverage",
+ },
+ msg_type="verification_result",
+ )
+
+ def _verify(self, result: dict) -> bool:
+ return result.get("tests_passing", False)
+```
+
+## 3. Fan-Out (Broadcast)
+
+Send the same task or query to multiple agents and collect their responses. Useful for parallel processing or gathering diverse perspectives.
+
+**When to use:** Parallel data processing, multi-reviewer code review, distributed search.
+
+```python
+from src.orchestrator import Orchestrator
+
+
+async def fan_out_example():
+ orch = Orchestrator(org_id="myorg")
+ orch.add_agent(worker_1)
+ orch.add_agent(worker_2)
+ orch.add_agent(worker_3)
+
+ await orch.start_all()
+
+ # Send the same work to all agents
+ await orch.broadcast(
+ payload={
+ "task": "Analyze dataset chunk",
+ "dataset_url": "s3://data/chunk-*.parquet",
+ },
+ msg_type="parallel_task",
+ )
+```
+
+### Collecting Fan-Out Results
+
+```python
+class CollectorAgent(Agent):
+ def __init__(self, expected_count: int, **kwargs):
+ super().__init__(**kwargs)
+ self.results = []
+ self.expected = expected_count
+
+ async def handle_message(self, message: Message):
+ if message.type == "partial_result":
+ self.results.append(message.payload)
+ print(f"Collected {len(self.results)}/{self.expected}")
+
+ if len(self.results) >= self.expected:
+ merged = self._merge_results(self.results)
+ print(f"All results collected: {merged}")
+
+ def _merge_results(self, results: list[dict]) -> dict:
+ # Combine partial results into a final answer
+ return {"total_records": sum(r.get("count", 0) for r in results)}
+```
+
+## 4. Sequential Pipeline
+
+Pass work through a chain of agents, each adding their contribution. Each agent processes the output of the previous one.
+
+**When to use:** Multi-stage processing (e.g., build -> test -> deploy), content pipelines.
+
+```python
+PIPELINE = ["builder", "tester", "deployer"]
+
+
+class PipelineAgent(Agent):
+ """Agent that processes work and forwards to the next stage."""
+
+ def __init__(self, next_agent: str | None = None, **kwargs):
+ super().__init__(**kwargs)
+ self.next_agent = next_agent
+
+ async def handle_message(self, message: Message):
+ if message.type == "pipeline_work":
+ # Process this stage
+ result = await self._process(message.payload)
+
+ if self.next_agent:
+ # Forward to next stage
+ await self.send_message(
+ to_agent=self.next_agent,
+ payload={
+ **message.payload,
+ f"{self.role}_result": result,
+ },
+ msg_type="pipeline_work",
+ )
+ else:
+ # Final stage — report back to originator
+ await self.send_message(
+ to_agent=message.payload.get("originator", message.from_agent),
+ payload={
+ **message.payload,
+ f"{self.role}_result": result,
+ "pipeline_complete": True,
+ },
+ msg_type="pipeline_result",
+ )
+
+ async def _process(self, data: dict) -> dict:
+ # Override in each stage
+ return {"status": "processed"}
+
+
+# Wire up the pipeline
+builder = PipelineAgent(name="builder", role="build", next_agent="tester", org_id="myorg")
+tester = PipelineAgent(name="tester", role="test", next_agent="deployer", org_id="myorg")
+deployer = PipelineAgent(name="deployer", role="deploy", next_agent=None, org_id="myorg")
+```
+
+## 5. Heartbeat Monitoring
+
+Track agent liveness by sending periodic heartbeats. The orchestrator detects agents that stop responding.
+
+**When to use:** Any production deployment where you need to detect agent failures.
+
+```python
+import asyncio
+
+
+class HeartbeatAgent(Agent):
+ """Agent that sends periodic heartbeats to the orchestrator."""
+
+ async def on_start(self):
+ asyncio.create_task(self._heartbeat_loop())
+
+ async def _heartbeat_loop(self):
+ while self._running:
+ await self.send_message(
+ to_agent="orchestrator",
+ payload={
+ "agent": self.name,
+ "status": "alive",
+ "uptime_seconds": self._uptime(),
+ },
+ msg_type="heartbeat",
+ )
+ await asyncio.sleep(10)
+
+ def _uptime(self) -> float:
+ # Track uptime since start
+ return 0.0 # implement with time.monotonic()
+```
+
+The orchestrator side:
+
+```python
+from src.orchestrator import Orchestrator
+
+orch = Orchestrator(org_id="myorg", heartbeat_interval=10.0)
+
+# Check health at any time
+health = orch.get_agent_health()
+# {'worker-1': True, 'worker-2': True, 'worker-3': False}
+
+for name, healthy in health.items():
+ if not healthy:
+ print(f"ALERT: Agent '{name}' is not responding!")
+```
+
+## 6. Agent Memory
+
+Persist agent state across restarts using the MemoryStore.
+
+**When to use:** Long-running agents that need to remember context, configuration, or intermediate results.
+
+```python
+from src.memory import MemoryStore
+
+
+class StatefulAgent(Agent):
+ def __init__(self, **kwargs):
+ super().__init__(**kwargs)
+ self.memory = MemoryStore(f".memory/{self.name}")
+
+ async def on_start(self):
+ # Restore previous state
+ state = self.memory.load("state")
+ if state:
+ print(f"Restored state: processed {state['tasks_completed']} tasks")
+ else:
+ self.memory.save("state", {"tasks_completed": 0})
+
+ async def handle_message(self, message: Message):
+ # Process message...
+ state = self.memory.load("state")
+ state["tasks_completed"] += 1
+ self.memory.save("state", state)
+```
+
+## Combining Patterns
+
+Real systems combine these patterns. For example, a CI/CD pipeline might use:
+
+- **Delegation** — CEO assigns a deploy task
+- **Sequential Pipeline** — build -> test -> deploy stages
+- **Verification** — CEO verifies the deployment succeeded
+- **Heartbeat** — monitor all agents throughout
+- **Memory** — track deployment history
+
+See `examples/devops_agent.py` for a concrete implementation.
+
+## Learn More
+
+- [Architecture Guide](architecture.md) — system design and component overview
+- [Getting Started](getting-started.md) — run your first agent
+- [agent.ceo](https://agent.ceo) — production-grade agent orchestration
+- [docs.agent.ceo](https://docs.agent.ceo) — full documentation
diff --git a/examples/__init__.py b/examples/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/examples/devops_agent.py b/examples/devops_agent.py
new file mode 100644
index 0000000..6b2bc7e
--- /dev/null
+++ b/examples/devops_agent.py
@@ -0,0 +1,252 @@
+#!/usr/bin/env python3
+"""DevOps agent that simulates a deployment pipeline.
+
+Receives deployment requests and walks through build, test, and deploy
+stages — sending status updates at each step.
+
+Usage:
+ # Start NATS first: docker run -p 4222:4222 nats:latest
+ python -m examples.devops_agent
+"""
+
+import asyncio
+import logging
+import signal
+import time
+import uuid
+
+from src.agent import Agent
+from src.messaging import Message
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s %(name)s %(message)s")
+
+
+class DeploymentPipeline:
+ """Simulates a multi-stage deployment pipeline."""
+
+ STAGES = ["build", "test", "deploy"]
+
+ def __init__(self, deploy_id: str, service: str, version: str):
+ self.deploy_id = deploy_id
+ self.service = service
+ self.version = version
+ self.current_stage: str | None = None
+ self.status = "pending"
+ self.started_at = time.time()
+ self.logs: list[str] = []
+
+ async def run_stage(self, stage: str) -> dict:
+ """Simulate running a pipeline stage.
+
+ Returns:
+ Stage result with status, duration, and logs.
+ """
+ self.current_stage = stage
+ self.status = "running"
+ stage_start = time.time()
+
+ self.logs.append(f"[{stage}] Starting {stage} stage for {self.service}:{self.version}")
+
+ if stage == "build":
+ result = await self._run_build()
+ elif stage == "test":
+ result = await self._run_tests()
+ elif stage == "deploy":
+ result = await self._run_deploy()
+ else:
+ result = {"success": False, "error": f"Unknown stage: {stage}"}
+
+ duration = round(time.time() - stage_start, 2)
+ result["duration_seconds"] = duration
+ result["stage"] = stage
+
+ self.logs.append(
+ f"[{stage}] Completed in {duration}s — "
+ f"{'SUCCESS' if result.get('success') else 'FAILED'}"
+ )
+ return result
+
+ async def _run_build(self) -> dict:
+ """Simulate building the service."""
+ self.logs.append("[build] Pulling base image python:3.12-slim")
+ await asyncio.sleep(0.5)
+ self.logs.append("[build] Installing dependencies from requirements.txt")
+ await asyncio.sleep(0.3)
+ self.logs.append("[build] Compiling application")
+ await asyncio.sleep(0.2)
+ image_tag = f"{self.service}:{self.version}"
+ self.logs.append(f"[build] Built image: {image_tag}")
+ return {
+ "success": True,
+ "image": image_tag,
+ "size_mb": 142,
+ }
+
+ async def _run_tests(self) -> dict:
+ """Simulate running the test suite."""
+ self.logs.append("[test] Running unit tests...")
+ await asyncio.sleep(0.4)
+ self.logs.append("[test] 47 tests passed, 0 failed")
+ self.logs.append("[test] Running integration tests...")
+ await asyncio.sleep(0.3)
+ self.logs.append("[test] 12 integration tests passed")
+ self.logs.append("[test] Code coverage: 87%")
+ return {
+ "success": True,
+ "unit_tests": {"passed": 47, "failed": 0},
+ "integration_tests": {"passed": 12, "failed": 0},
+ "coverage_pct": 87,
+ }
+
+ async def _run_deploy(self) -> dict:
+ """Simulate deploying to the target environment."""
+ self.logs.append("[deploy] Applying Kubernetes manifests")
+ await asyncio.sleep(0.3)
+ self.logs.append("[deploy] Rolling update started")
+ await asyncio.sleep(0.5)
+ self.logs.append("[deploy] Waiting for pods to become ready...")
+ await asyncio.sleep(0.3)
+ self.logs.append("[deploy] 3/3 pods ready")
+ self.logs.append("[deploy] Health check passed")
+ url = f"https://{self.service}.example.com"
+ return {
+ "success": True,
+ "url": url,
+ "replicas": 3,
+ "health": "passing",
+ }
+
+ def summary(self) -> dict:
+ total_duration = round(time.time() - self.started_at, 2)
+ return {
+ "deploy_id": self.deploy_id,
+ "service": self.service,
+ "version": self.version,
+ "status": self.status,
+ "total_duration_seconds": total_duration,
+ "logs": self.logs,
+ }
+
+
+class DevOpsAgent(Agent):
+ """Agent that processes deployment requests through a simulated pipeline."""
+
+ def __init__(self, **kwargs):
+ super().__init__(**kwargs)
+ self.pipelines: dict[str, DeploymentPipeline] = {}
+
+ async def on_start(self) -> None:
+ print(f"[{self.name}] DevOps agent ready. Waiting for deployment requests...")
+ print(f"[{self.name}] Send type='deploy_request' with payload={{'service': '...', 'version': '...'}}")
+
+ async def handle_message(self, message: Message) -> None:
+ if message.type == "deploy_request":
+ await self._handle_deploy(message)
+ elif message.type == "status_request":
+ await self._handle_status(message)
+ else:
+ print(f"[{self.name}] Unknown message type: {message.type}")
+
+ async def _handle_deploy(self, message: Message) -> None:
+ service = message.payload.get("service", "unknown-service")
+ version = message.payload.get("version", "latest")
+ deploy_id = str(uuid.uuid4())[:8]
+
+ print(f"\n[{self.name}] === Deployment {deploy_id} ===")
+ print(f"[{self.name}] Service: {service}, Version: {version}")
+
+ pipeline = DeploymentPipeline(deploy_id, service, version)
+ self.pipelines[deploy_id] = pipeline
+
+ # Notify requester that pipeline started
+ await self.send_message(
+ to_agent=message.from_agent,
+ payload={
+ "deploy_id": deploy_id,
+ "status": "started",
+ "service": service,
+ "version": version,
+ },
+ msg_type="deploy_status",
+ )
+
+ # Run each stage, sending progress updates
+ all_passed = True
+ for stage in DeploymentPipeline.STAGES:
+ print(f"\n[{self.name}] >> Stage: {stage}")
+
+ # Send stage-start notification
+ await self.send_message(
+ to_agent=message.from_agent,
+ payload={
+ "deploy_id": deploy_id,
+ "stage": stage,
+ "status": "running",
+ },
+ msg_type="deploy_stage_update",
+ )
+
+ result = await pipeline.run_stage(stage)
+
+ # Print logs
+ for log_line in pipeline.logs[-3:]:
+ print(f"[{self.name}] {log_line}")
+
+ # Send stage-complete notification
+ await self.send_message(
+ to_agent=message.from_agent,
+ payload={
+ "deploy_id": deploy_id,
+ "stage": stage,
+ "status": "completed" if result["success"] else "failed",
+ "result": result,
+ },
+ msg_type="deploy_stage_update",
+ )
+
+ if not result["success"]:
+ all_passed = False
+ break
+
+ # Send final result
+ pipeline.status = "completed" if all_passed else "failed"
+ summary = pipeline.summary()
+
+ print(f"\n[{self.name}] === Deployment {deploy_id} {'SUCCEEDED' if all_passed else 'FAILED'} ===")
+
+ await self.send_message(
+ to_agent=message.from_agent,
+ payload=summary,
+ msg_type="deploy_result",
+ )
+
+ async def _handle_status(self, message: Message) -> None:
+ deploy_id = message.payload.get("deploy_id")
+ if deploy_id and deploy_id in self.pipelines:
+ summary = self.pipelines[deploy_id].summary()
+ else:
+ summary = {
+ "active_pipelines": len(self.pipelines),
+ "pipeline_ids": list(self.pipelines.keys()),
+ }
+
+ await self.send_message(
+ to_agent=message.from_agent,
+ payload=summary,
+ msg_type="deploy_status",
+ )
+
+
+async def main() -> None:
+ agent = DevOpsAgent(name="devops", role="deployment-engineer", org_id="demo")
+
+ loop = asyncio.get_running_loop()
+ for sig in (signal.SIGINT, signal.SIGTERM):
+ loop.add_signal_handler(sig, lambda: asyncio.create_task(agent.stop()))
+
+ print("Starting DevOps Agent... (press Ctrl+C to stop)")
+ await agent.run()
+
+
+if __name__ == "__main__":
+ asyncio.run(main())
diff --git a/examples/hello_agent.py b/examples/hello_agent.py
new file mode 100644
index 0000000..684adf8
--- /dev/null
+++ b/examples/hello_agent.py
@@ -0,0 +1,51 @@
+#!/usr/bin/env python3
+"""Simplest possible agent — connects, logs messages, responds with 'Hello!'
+
+Usage:
+ # Start NATS first: docker run -p 4222:4222 nats:latest
+ python -m examples.hello_agent
+"""
+
+import asyncio
+import logging
+import signal
+
+from src.agent import Agent
+from src.messaging import Message
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s %(name)s %(message)s")
+
+
+class HelloAgent(Agent):
+ """An agent that greets anyone who sends it a message."""
+
+ async def on_start(self) -> None:
+ print(f"[{self.name}] Ready and listening on {self.inbox_subject}")
+
+ async def handle_message(self, message: Message) -> None:
+ print(f"[{self.name}] Got message from '{message.from_agent}': {message.payload}")
+
+ # Reply with a greeting
+ await self.send_message(
+ to_agent=message.from_agent,
+ payload={"text": f"Hello from {self.name}! I received your message."},
+ msg_type="greeting",
+ )
+ print(f"[{self.name}] Sent greeting back to '{message.from_agent}'")
+
+
+async def main() -> None:
+ agent = HelloAgent(name="hello", role="greeter", org_id="demo")
+
+ # Handle Ctrl+C gracefully
+ loop = asyncio.get_running_loop()
+ for sig in (signal.SIGINT, signal.SIGTERM):
+ loop.add_signal_handler(sig, lambda: asyncio.create_task(agent.stop()))
+
+ print("Starting HelloAgent... (press Ctrl+C to stop)")
+ await agent.run()
+ print("HelloAgent stopped.")
+
+
+if __name__ == "__main__":
+ asyncio.run(main())
diff --git a/examples/security_agent.py b/examples/security_agent.py
new file mode 100644
index 0000000..132b2d9
--- /dev/null
+++ b/examples/security_agent.py
@@ -0,0 +1,258 @@
+#!/usr/bin/env python3
+"""Security review agent that scans repositories for common issues.
+
+Receives a repo path, performs basic security checks (file listing,
+dependency analysis, pattern scanning), and reports findings.
+
+Usage:
+ # Start NATS first: docker run -p 4222:4222 nats:latest
+ python -m examples.security_agent
+"""
+
+import asyncio
+import logging
+import os
+import re
+import signal
+from pathlib import Path
+
+from src.agent import Agent
+from src.messaging import Message
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s %(name)s %(message)s")
+
+# Patterns that suggest potential security issues
+SENSITIVE_FILE_PATTERNS = [
+ r"\.env$",
+ r"\.env\.local$",
+ r"credentials\.json$",
+ r"secrets\.ya?ml$",
+ r"\.pem$",
+ r"\.key$",
+ r"id_rsa",
+ r"\.p12$",
+]
+
+RISKY_CODE_PATTERNS = [
+ (r"eval\s*\(", "Use of eval() — potential code injection"),
+ (r"exec\s*\(", "Use of exec() — potential code injection"),
+ (r"subprocess\.call\(.*shell\s*=\s*True", "Shell=True in subprocess — command injection risk"),
+ (r"pickle\.loads?\(", "Pickle deserialization — potential arbitrary code execution"),
+ (r"password\s*=\s*['\"][^'\"]+['\"]", "Hardcoded password detected"),
+ (r"api_key\s*=\s*['\"][^'\"]+['\"]", "Hardcoded API key detected"),
+ (r"token\s*=\s*['\"][A-Za-z0-9_\-]{20,}['\"]", "Possible hardcoded token"),
+]
+
+KNOWN_VULNERABLE_PACKAGES = {
+ "requests": {"below": "2.31.0", "reason": "CVE-2023-32681 — CRLF injection"},
+ "urllib3": {"below": "2.0.7", "reason": "CVE-2023-45803 — request body leak"},
+ "cryptography": {"below": "41.0.6", "reason": "CVE-2023-49083 — NULL dereference"},
+ "pillow": {"below": "10.2.0", "reason": "CVE-2023-50447 — arbitrary code execution"},
+}
+
+
+class SecurityAgent(Agent):
+ """Agent that performs basic security reviews on code repositories."""
+
+ async def on_start(self) -> None:
+ print(f"[{self.name}] Security agent ready. Send a scan request to get started.")
+
+ async def handle_message(self, message: Message) -> None:
+ if message.type == "scan_request":
+ repo_path = message.payload.get("repo_path", ".")
+ print(f"\n[{self.name}] Starting security scan of: {repo_path}")
+
+ findings = await self._run_scan(repo_path)
+
+ await self.send_message(
+ to_agent=message.from_agent,
+ payload={
+ "repo_path": repo_path,
+ "findings": findings,
+ "summary": self._summarize(findings),
+ },
+ msg_type="scan_result",
+ )
+ print(f"[{self.name}] Scan complete. Sent results to '{message.from_agent}'.")
+ else:
+ print(f"[{self.name}] Unknown message type: {message.type}")
+
+ async def _run_scan(self, repo_path: str) -> list[dict]:
+ """Execute all security checks and collect findings."""
+ findings: list[dict] = []
+ root = Path(repo_path)
+
+ if not root.exists():
+ findings.append({
+ "severity": "error",
+ "check": "path_validation",
+ "message": f"Path does not exist: {repo_path}",
+ })
+ return findings
+
+ # Run checks
+ findings.extend(self._check_sensitive_files(root))
+ findings.extend(self._check_code_patterns(root))
+ findings.extend(self._check_dependencies(root))
+ findings.extend(self._check_gitignore(root))
+
+ return findings
+
+ def _check_sensitive_files(self, root: Path) -> list[dict]:
+ """Scan for files that should not be committed."""
+ findings = []
+ for dirpath, _, filenames in os.walk(root):
+ # Skip hidden dirs and common non-source dirs
+ rel_dir = os.path.relpath(dirpath, root)
+ if any(part.startswith(".") for part in Path(rel_dir).parts if part != "."):
+ continue
+ if "node_modules" in rel_dir or "__pycache__" in rel_dir:
+ continue
+
+ for fname in filenames:
+ for pattern in SENSITIVE_FILE_PATTERNS:
+ if re.search(pattern, fname, re.IGNORECASE):
+ findings.append({
+ "severity": "high",
+ "check": "sensitive_file",
+ "message": f"Sensitive file found: {os.path.join(rel_dir, fname)}",
+ "file": os.path.join(rel_dir, fname),
+ })
+ return findings
+
+ def _check_code_patterns(self, root: Path) -> list[dict]:
+ """Scan Python files for risky code patterns."""
+ findings = []
+ for py_file in root.rglob("*.py"):
+ rel_path = py_file.relative_to(root)
+ if any(part.startswith(".") for part in rel_path.parts):
+ continue
+
+ try:
+ content = py_file.read_text(encoding="utf-8", errors="ignore")
+ except (OSError, PermissionError):
+ continue
+
+ for line_num, line in enumerate(content.splitlines(), 1):
+ for pattern, description in RISKY_CODE_PATTERNS:
+ if re.search(pattern, line):
+ findings.append({
+ "severity": "medium",
+ "check": "code_pattern",
+ "message": description,
+ "file": str(rel_path),
+ "line": line_num,
+ })
+ return findings
+
+ def _check_dependencies(self, root: Path) -> list[dict]:
+ """Check requirements files for known vulnerable packages."""
+ findings = []
+ req_files = list(root.glob("requirements*.txt")) + list(root.glob("**/requirements*.txt"))
+
+ for req_file in req_files:
+ try:
+ content = req_file.read_text(encoding="utf-8")
+ except (OSError, PermissionError):
+ continue
+
+ for line in content.splitlines():
+ line = line.strip()
+ if not line or line.startswith("#") or line.startswith("-"):
+ continue
+
+ # Parse package==version or package>=version
+ match = re.match(r"^([a-zA-Z0-9_-]+)\s*[=<>!]+\s*([0-9.]+)", line)
+ if match:
+ pkg_name = match.group(1).lower()
+ if pkg_name in KNOWN_VULNERABLE_PACKAGES:
+ vuln = KNOWN_VULNERABLE_PACKAGES[pkg_name]
+ findings.append({
+ "severity": "high",
+ "check": "dependency_vulnerability",
+ "message": (
+ f"Package '{pkg_name}' may be vulnerable "
+ f"(versions below {vuln['below']}): {vuln['reason']}"
+ ),
+ "file": str(req_file.relative_to(root)),
+ })
+
+ # Check pyproject.toml too
+ pyproject = root / "pyproject.toml"
+ if pyproject.exists():
+ try:
+ content = pyproject.read_text(encoding="utf-8")
+ for pkg_name in KNOWN_VULNERABLE_PACKAGES:
+ if pkg_name in content:
+ findings.append({
+ "severity": "info",
+ "check": "dependency_check",
+ "message": (
+ f"Package '{pkg_name}' found in pyproject.toml — "
+ "verify version is not vulnerable"
+ ),
+ "file": "pyproject.toml",
+ })
+ except (OSError, PermissionError):
+ pass
+
+ return findings
+
+ def _check_gitignore(self, root: Path) -> list[dict]:
+ """Verify .gitignore includes common sensitive patterns."""
+ findings = []
+ gitignore = root / ".gitignore"
+
+ if not gitignore.exists():
+ findings.append({
+ "severity": "medium",
+ "check": "missing_gitignore",
+ "message": "No .gitignore file found — sensitive files may be committed",
+ })
+ return findings
+
+ content = gitignore.read_text(encoding="utf-8")
+ recommended = [".env", "*.pem", "*.key", "__pycache__"]
+ for pattern in recommended:
+ if pattern not in content:
+ findings.append({
+ "severity": "low",
+ "check": "gitignore_missing_pattern",
+ "message": f".gitignore is missing recommended pattern: {pattern}",
+ })
+
+ return findings
+
+ def _summarize(self, findings: list[dict]) -> dict:
+ """Produce a severity-bucketed summary."""
+ counts = {"high": 0, "medium": 0, "low": 0, "info": 0, "error": 0}
+ for f in findings:
+ sev = f.get("severity", "info")
+ counts[sev] = counts.get(sev, 0) + 1
+ return {
+ "total_findings": len(findings),
+ "by_severity": counts,
+ "verdict": (
+ "FAIL — critical issues found"
+ if counts["high"] > 0
+ else "WARN — review recommended"
+ if counts["medium"] > 0
+ else "PASS"
+ ),
+ }
+
+
+async def main() -> None:
+ agent = SecurityAgent(name="security", role="security-reviewer", org_id="demo")
+
+ loop = asyncio.get_running_loop()
+ for sig in (signal.SIGINT, signal.SIGTERM):
+ loop.add_signal_handler(sig, lambda: asyncio.create_task(agent.stop()))
+
+ print("Starting Security Agent... (press Ctrl+C to stop)")
+ print("Send a message with type='scan_request' and payload={'repo_path': '/path/to/repo'}")
+ await agent.run()
+
+
+if __name__ == "__main__":
+ asyncio.run(main())
diff --git a/examples/two_agent_team.py b/examples/two_agent_team.py
new file mode 100644
index 0000000..c36beaf
--- /dev/null
+++ b/examples/two_agent_team.py
@@ -0,0 +1,162 @@
+#!/usr/bin/env python3
+"""CEO + Worker two-agent team demonstrating task delegation via NATS.
+
+The CEO agent creates a task, sends it to the Worker agent, the Worker
+processes the task and reports back with a result.
+
+Usage:
+ # Start NATS first: docker run -p 4222:4222 nats:latest
+ python -m examples.two_agent_team
+"""
+
+import asyncio
+import json
+import logging
+import signal
+
+from src.agent import Agent
+from src.messaging import Message
+from src.task_manager import Task, TaskManager, TaskStatus
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s %(name)s %(message)s")
+
+
+class CEOAgent(Agent):
+ """Agent that creates and delegates tasks, then waits for results."""
+
+ def __init__(self, **kwargs):
+ super().__init__(**kwargs)
+ self.task_manager = TaskManager()
+ self.pending_tasks: dict[str, str] = {} # task_id -> assignee
+
+ async def on_start(self) -> None:
+ print(f"[{self.name}] CEO agent ready. Delegating initial task...")
+
+ # Create and send a task to the worker
+ task = self.task_manager.assign(
+ title="Analyze quarterly revenue data and produce a summary report",
+ assignee="worker",
+ assigner=self.name,
+ )
+ self.pending_tasks[task.id] = task.assignee
+
+ await self.send_message(
+ to_agent="worker",
+ payload=task.model_dump(),
+ msg_type="task_assignment",
+ )
+ print(f"[{self.name}] Delegated task '{task.title}' to worker (id={task.id})")
+
+ async def handle_message(self, message: Message) -> None:
+ if message.type == "task_result":
+ task_id = message.payload.get("task_id")
+ status = message.payload.get("status")
+ result = message.payload.get("result")
+
+ print(f"\n[{self.name}] Received task result from '{message.from_agent}':")
+ print(f" Task ID: {task_id}")
+ print(f" Status: {status}")
+ print(f" Result: {json.dumps(result, indent=2)}")
+
+ # Update task in our tracker
+ if task_id in self.pending_tasks:
+ if status == "completed":
+ self.task_manager.update_status(task_id, TaskStatus.IN_PROGRESS)
+ self.task_manager.complete(task_id, result=result)
+ print(f"[{self.name}] Task verified and marked complete!")
+ else:
+ self.task_manager.update_status(task_id, TaskStatus.IN_PROGRESS)
+ self.task_manager.fail(task_id, reason=str(result))
+ print(f"[{self.name}] Task failed: {result}")
+ del self.pending_tasks[task_id]
+
+ # If all tasks are done, we can stop
+ if not self.pending_tasks:
+ print(f"\n[{self.name}] All tasks complete. Shutting down team...")
+ # Give worker a moment to see the stop signal
+ await asyncio.sleep(0.5)
+ await self.stop()
+ else:
+ print(f"[{self.name}] Got unexpected message type: {message.type}")
+
+
+class WorkerAgent(Agent):
+ """Agent that receives tasks, processes them, and reports results."""
+
+ async def on_start(self) -> None:
+ print(f"[{self.name}] Worker agent ready. Waiting for tasks...")
+
+ async def handle_message(self, message: Message) -> None:
+ if message.type == "task_assignment":
+ task_data = message.payload
+ task_title = task_data.get("title", "Unknown task")
+ task_id = task_data.get("id", "unknown")
+
+ print(f"\n[{self.name}] Received task: '{task_title}' (id={task_id})")
+ print(f"[{self.name}] Processing...")
+
+ # Simulate work
+ await asyncio.sleep(1)
+
+ # Produce a result
+ result = {
+ "summary": "Revenue increased 23% QoQ driven by enterprise segment.",
+ "key_metrics": {
+ "total_revenue": "$4.2M",
+ "growth_rate": "23%",
+ "top_segment": "Enterprise",
+ },
+ "recommendations": [
+ "Increase enterprise sales team by 2 headcount",
+ "Invest in self-serve onboarding to reduce CAC",
+ ],
+ }
+
+ print(f"[{self.name}] Task complete. Reporting result to '{message.from_agent}'...")
+
+ await self.send_message(
+ to_agent=message.from_agent,
+ payload={
+ "task_id": task_id,
+ "status": "completed",
+ "result": result,
+ },
+ msg_type="task_result",
+ )
+ else:
+ print(f"[{self.name}] Ignoring message type: {message.type}")
+
+
+async def main() -> None:
+ ceo = CEOAgent(name="ceo", role="executive", org_id="demo")
+ worker = WorkerAgent(name="worker", role="analyst", org_id="demo")
+
+ # Handle Ctrl+C
+ loop = asyncio.get_running_loop()
+ for sig in (signal.SIGINT, signal.SIGTERM):
+ loop.add_signal_handler(
+ sig,
+ lambda: asyncio.create_task(_shutdown(ceo, worker)),
+ )
+
+ print("=" * 60)
+ print("Two-Agent Team Demo: CEO delegates to Worker")
+ print("=" * 60)
+
+ # Start both agents concurrently
+ await asyncio.gather(
+ ceo.run(),
+ worker.run(),
+ return_exceptions=True,
+ )
+
+ print("\nDemo complete!")
+
+
+async def _shutdown(*agents: Agent) -> None:
+ for agent in agents:
+ await agent.stop()
+
+
+if __name__ == "__main__":
+ asyncio.run(main())
diff --git a/pyproject.toml b/pyproject.toml
new file mode 100644
index 0000000..f2fc9eb
--- /dev/null
+++ b/pyproject.toml
@@ -0,0 +1,17 @@
+[project]
+name = "agent-framework-starter"
+version = "0.1.0"
+description = "Minimal starter kit for building AI agent teams"
+requires-python = ">=3.11"
+license = {text = "MIT"}
+dependencies = [
+ "nats-py>=2.7.0",
+ "pydantic>=2.0",
+]
+
+[project.optional-dependencies]
+dev = ["pytest>=8.0", "pytest-asyncio>=0.23"]
+
+[build-system]
+requires = ["setuptools>=68.0"]
+build-backend = "setuptools.build_meta"
diff --git a/pytest.ini b/pytest.ini
new file mode 100644
index 0000000..78c5011
--- /dev/null
+++ b/pytest.ini
@@ -0,0 +1,3 @@
+[pytest]
+asyncio_mode = auto
+testpaths = tests
diff --git a/src/__init__.py b/src/__init__.py
new file mode 100644
index 0000000..6ad4cc1
--- /dev/null
+++ b/src/__init__.py
@@ -0,0 +1,21 @@
+"""Agent Framework Starter — a minimal Python framework for multi-agent AI teams."""
+
+from src.agent import Agent
+from src.messaging import Message, connect, create_stream, publish, subscribe
+from src.task_manager import Task, TaskManager, TaskStatus
+from src.memory import MemoryStore
+from src.orchestrator import Orchestrator
+
+__all__ = [
+ "Agent",
+ "Message",
+ "connect",
+ "create_stream",
+ "publish",
+ "subscribe",
+ "Task",
+ "TaskManager",
+ "TaskStatus",
+ "MemoryStore",
+ "Orchestrator",
+]
diff --git a/src/agent.py b/src/agent.py
new file mode 100644
index 0000000..9598c39
--- /dev/null
+++ b/src/agent.py
@@ -0,0 +1,179 @@
+"""Base Agent class — the fundamental building block for AI agent teams."""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+from typing import Any
+
+import nats
+from nats.aio.client import Client as NATSClient
+
+from src.messaging import Message
+
+logger = logging.getLogger(__name__)
+
+
+class Agent:
+ """Asynchronous agent that connects to NATS and processes inbox messages.
+
+ Subclass this and override ``handle_message`` to build custom agents.
+
+ Example::
+
+ class GreeterAgent(Agent):
+ async def handle_message(self, message: Message) -> None:
+ print(f"Got: {message.payload}")
+ await self.send_message(message.from_agent, {
+ "text": "Hello from Greeter!"
+ })
+
+ agent = GreeterAgent(name="greeter", role="greeter", org_id="myorg")
+ asyncio.run(agent.run())
+ """
+
+ def __init__(
+ self,
+ name: str,
+ role: str,
+ org_id: str = "default",
+ nats_url: str = "nats://localhost:4222",
+ nats_token: str | None = None,
+ ) -> None:
+ self.name = name
+ self.role = role
+ self.org_id = org_id
+ self.nats_url = nats_url
+ self.nats_token = nats_token
+
+ self._nc: NATSClient | None = None
+ self._sub: Any = None
+ self._running = False
+
+ # ---- properties -------------------------------------------------------
+
+ @property
+ def inbox_subject(self) -> str:
+ """NATS subject this agent listens on."""
+ return f"{self.org_id}.agents.{self.name}.inbox"
+
+ @property
+ def is_connected(self) -> bool:
+ return self._nc is not None and self._nc.is_connected
+
+ # ---- lifecycle --------------------------------------------------------
+
+ async def connect(self) -> None:
+ """Establish a connection to the NATS server."""
+ opts: dict[str, Any] = {"servers": [self.nats_url]}
+ if self.nats_token:
+ opts["token"] = self.nats_token
+
+ self._nc = await nats.connect(**opts)
+ logger.info("Agent '%s' connected to NATS at %s", self.name, self.nats_url)
+
+ async def disconnect(self) -> None:
+ """Gracefully close the NATS connection."""
+ if self._sub:
+ await self._sub.unsubscribe()
+ self._sub = None
+ if self._nc and not self._nc.is_closed:
+ await self._nc.drain()
+ self._nc = None
+ self._running = False
+ logger.info("Agent '%s' disconnected", self.name)
+
+ async def run(self) -> None:
+ """Main run loop — connect, subscribe, and process messages until stopped.
+
+ Call ``stop()`` from another coroutine (or signal handler) to exit
+ the loop cleanly.
+ """
+ await self.connect()
+
+ assert self._nc is not None
+ self._sub = await self._nc.subscribe(self.inbox_subject, cb=self._on_message)
+ self._running = True
+
+ logger.info(
+ "Agent '%s' (role=%s) listening on %s",
+ self.name,
+ self.role,
+ self.inbox_subject,
+ )
+ await self.on_start()
+
+ try:
+ while self._running:
+ await asyncio.sleep(0.1)
+ except asyncio.CancelledError:
+ pass
+ finally:
+ await self.disconnect()
+
+ async def stop(self) -> None:
+ """Signal the agent to stop its run loop."""
+ self._running = False
+
+ # ---- messaging --------------------------------------------------------
+
+ async def send_message(
+ self,
+ to_agent: str,
+ payload: dict[str, Any],
+ msg_type: str = "message",
+ ) -> None:
+ """Send a message to another agent's inbox.
+
+ Args:
+ to_agent: Name of the target agent.
+ payload: Arbitrary JSON-serializable data.
+ msg_type: Message type identifier.
+ """
+ if not self._nc or self._nc.is_closed:
+ raise RuntimeError(f"Agent '{self.name}' is not connected to NATS")
+
+ msg = Message(
+ from_agent=self.name,
+ to_agent=to_agent,
+ type=msg_type,
+ payload=payload,
+ )
+ subject = f"{self.org_id}.agents.{to_agent}.inbox"
+ await self._nc.publish(subject, msg.to_bytes())
+ logger.debug("Agent '%s' → '%s': %s", self.name, to_agent, msg_type)
+
+ # ---- hooks (override in subclasses) -----------------------------------
+
+ async def on_start(self) -> None:
+ """Called once after the agent subscribes to its inbox.
+
+ Override to run initialization logic (e.g. announce presence).
+ """
+
+ async def handle_message(self, message: Message) -> None:
+ """Process an incoming message. Override in subclasses.
+
+ Args:
+ message: The decoded Message received on this agent's inbox.
+ """
+ logger.info(
+ "Agent '%s' received message from '%s': %s",
+ self.name,
+ message.from_agent,
+ message.type,
+ )
+
+ # ---- internal ---------------------------------------------------------
+
+ async def _on_message(self, raw_msg: Any) -> None:
+ """Internal NATS callback — deserialize and dispatch."""
+ try:
+ msg = Message.from_bytes(raw_msg.data)
+ await self.handle_message(msg)
+ except Exception:
+ logger.exception(
+ "Agent '%s' failed to process message on %s",
+ self.name,
+ self.inbox_subject,
+ )
diff --git a/src/memory.py b/src/memory.py
new file mode 100644
index 0000000..c4f4511
--- /dev/null
+++ b/src/memory.py
@@ -0,0 +1,118 @@
+"""Agent memory/persistence using JSON files with thread-safe access."""
+
+from __future__ import annotations
+
+import fcntl
+import json
+import os
+from pathlib import Path
+from typing import Any
+
+
+class MemoryStore:
+ """Persistent key-value store backed by JSON files.
+
+ Each key maps to a separate JSON file in the configured directory,
+ ensuring that concurrent reads/writes to different keys don't
+ contend. File-level locking (``fcntl.flock``) is used so multiple
+ processes (or threads) can safely access the same key.
+
+ Example::
+
+ mem = MemoryStore("/tmp/agent-memory")
+ mem.save("config", {"model": "claude-4", "temperature": 0.7})
+ config = mem.load("config")
+ print(config) # {'model': 'claude-4', 'temperature': 0.7}
+ mem.delete("config")
+ """
+
+ def __init__(self, directory: str | Path = ".agent_memory") -> None:
+ self._dir = Path(directory)
+ self._dir.mkdir(parents=True, exist_ok=True)
+
+ # ---- public API -------------------------------------------------------
+
+ def save(self, key: str, value: Any) -> None:
+ """Persist a value under the given key.
+
+ Args:
+ key: Identifier for the stored value. Used as the filename
+ (with ``.json`` extension).
+ value: Any JSON-serializable Python object.
+ """
+ path = self._key_path(key)
+ with open(path, "w", encoding="utf-8") as fh:
+ fcntl.flock(fh, fcntl.LOCK_EX)
+ try:
+ json.dump(value, fh, indent=2, default=str)
+ finally:
+ fcntl.flock(fh, fcntl.LOCK_UN)
+
+ def load(self, key: str) -> Any:
+ """Load a previously saved value.
+
+ Args:
+ key: Identifier to look up.
+
+ Returns:
+ The deserialized Python object, or ``None`` if the key
+ does not exist.
+ """
+ path = self._key_path(key)
+ if not path.exists():
+ return None
+ with open(path, "r", encoding="utf-8") as fh:
+ fcntl.flock(fh, fcntl.LOCK_SH)
+ try:
+ return json.load(fh)
+ finally:
+ fcntl.flock(fh, fcntl.LOCK_UN)
+
+ def delete(self, key: str) -> bool:
+ """Remove a stored key.
+
+ Args:
+ key: Identifier to delete.
+
+ Returns:
+ ``True`` if the key existed and was removed, ``False``
+ otherwise.
+ """
+ path = self._key_path(key)
+ try:
+ os.remove(path)
+ return True
+ except FileNotFoundError:
+ return False
+
+ def list_keys(self) -> list[str]:
+ """Return all stored keys (alphabetically sorted).
+
+ Returns:
+ List of key names (without the ``.json`` extension).
+ """
+ keys = [
+ p.stem
+ for p in self._dir.iterdir()
+ if p.is_file() and p.suffix == ".json"
+ ]
+ return sorted(keys)
+
+ def clear(self) -> int:
+ """Delete all stored keys.
+
+ Returns:
+ Number of keys removed.
+ """
+ keys = self.list_keys()
+ for key in keys:
+ self.delete(key)
+ return len(keys)
+
+ # ---- internal ---------------------------------------------------------
+
+ def _key_path(self, key: str) -> Path:
+ """Map a key name to its on-disk path."""
+ # Sanitize key to prevent path traversal
+ safe_key = key.replace("/", "_").replace("\\", "_").replace("..", "_")
+ return self._dir / f"{safe_key}.json"
diff --git a/src/messaging.py b/src/messaging.py
new file mode 100644
index 0000000..9f90755
--- /dev/null
+++ b/src/messaging.py
@@ -0,0 +1,144 @@
+"""NATS JetStream messaging helpers for agent communication."""
+
+from __future__ import annotations
+
+import json
+import uuid
+from datetime import datetime, timezone
+from typing import Any, Callable, Awaitable
+
+import nats
+from nats.aio.client import Client as NATSClient
+from nats.js.api import StreamConfig
+from nats.js.client import JetStreamContext
+from pydantic import BaseModel, Field
+
+
+class Message(BaseModel):
+ """Structured message exchanged between agents."""
+
+ id: str = Field(default_factory=lambda: str(uuid.uuid4()))
+ from_agent: str
+ to_agent: str
+ type: str = "message"
+ payload: dict[str, Any] = Field(default_factory=dict)
+ timestamp: str = Field(
+ default_factory=lambda: datetime.now(timezone.utc).isoformat()
+ )
+
+ def to_bytes(self) -> bytes:
+ """Serialize to bytes for NATS transmission."""
+ return self.model_dump_json().encode("utf-8")
+
+ @classmethod
+ def from_bytes(cls, data: bytes) -> "Message":
+ """Deserialize from NATS bytes."""
+ return cls.model_validate_json(data)
+
+
+async def connect(url: str = "nats://localhost:4222", token: str | None = None) -> NATSClient:
+ """Connect to a NATS server with optional token authentication.
+
+ Args:
+ url: NATS server URL. Defaults to localhost.
+ token: Optional authentication token.
+
+ Returns:
+ Connected NATS client.
+ """
+ options: dict[str, Any] = {"servers": [url]}
+ if token:
+ options["token"] = token
+
+ nc = await nats.connect(**options)
+ return nc
+
+
+async def create_stream(
+ js: JetStreamContext,
+ stream_name: str,
+ subjects: list[str],
+) -> None:
+ """Create or update a JetStream stream.
+
+ If the stream already exists, this is a no-op. If it exists with
+ different subjects the stream config is updated.
+
+ Args:
+ js: JetStream context from a connected NATS client.
+ stream_name: Name of the stream to create.
+ subjects: List of subject patterns the stream should capture.
+ """
+ try:
+ info = await js.find_stream_info_by_subject(subjects[0])
+ # Stream already exists — check if we need to update subjects
+ if set(info.config.subjects or []) != set(subjects):
+ info.config.subjects = subjects
+ await js.update_stream(info.config)
+ except Exception:
+ # Stream does not exist yet — create it
+ await js.add_stream(
+ StreamConfig(
+ name=stream_name,
+ subjects=subjects,
+ retention="limits",
+ max_msgs=10_000,
+ max_age=86_400_000_000_000, # 24 hours in nanoseconds
+ )
+ )
+
+
+async def publish(
+ js: JetStreamContext,
+ subject: str,
+ data: Message | dict[str, Any],
+) -> None:
+ """Publish a message to a JetStream subject.
+
+ Args:
+ js: JetStream context.
+ subject: NATS subject to publish to.
+ data: Message object or dict to serialize and publish.
+ """
+ if isinstance(data, Message):
+ raw = data.to_bytes()
+ else:
+ raw = json.dumps(data).encode("utf-8")
+ await js.publish(subject, raw)
+
+
+async def subscribe(
+ js: JetStreamContext,
+ subject: str,
+ handler: Callable[[Message], Awaitable[None]],
+ durable_name: str | None = None,
+) -> Any:
+ """Create a durable pull subscription and start consuming messages.
+
+ Args:
+ js: JetStream context.
+ subject: Subject pattern to subscribe to.
+ handler: Async callback invoked for each received Message.
+ durable_name: Optional durable consumer name for resumable delivery.
+
+ Returns:
+ The NATS subscription object (can be used to unsubscribe).
+ """
+ sub = await js.subscribe(subject, durable=durable_name)
+
+ async def _dispatch() -> None:
+ async for raw_msg in sub.messages:
+ try:
+ msg = Message.from_bytes(raw_msg.data)
+ await handler(msg)
+ await raw_msg.ack()
+ except Exception as exc:
+ # Log but don't crash — let the consumer continue
+ print(f"[messaging] Error handling message on {subject}: {exc}")
+ await raw_msg.nak()
+
+ # The caller is responsible for running _dispatch in a task
+ import asyncio
+
+ asyncio.create_task(_dispatch())
+ return sub
diff --git a/src/orchestrator.py b/src/orchestrator.py
new file mode 100644
index 0000000..91de7e1
--- /dev/null
+++ b/src/orchestrator.py
@@ -0,0 +1,284 @@
+"""Multi-agent orchestrator — manages agent lifecycles and coordination."""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import time
+from typing import Any
+
+from src.agent import Agent
+from src.messaging import Message
+from src.task_manager import Task, TaskManager, TaskStatus
+
+logger = logging.getLogger(__name__)
+
+
+class AgentNotFoundError(KeyError):
+ """Raised when referencing an agent name that is not registered."""
+
+
+class Orchestrator:
+ """Manages multiple Agent instances, delegates tasks, and monitors health.
+
+ The orchestrator maintains a registry of agents, starts/stops them as a
+ group, and provides helpers for task delegation and broadcasting.
+
+ Example::
+
+ orch = Orchestrator(org_id="myorg")
+ orch.add_agent(worker_agent)
+ orch.add_agent(reviewer_agent)
+
+ await orch.start_all()
+
+ task = orch.delegate_task(
+ title="Analyze PR #42",
+ agent_name="worker",
+ assigner="orchestrator",
+ )
+
+ await orch.broadcast({"type": "status_request"})
+ await orch.stop_all()
+ """
+
+ def __init__(
+ self,
+ org_id: str = "default",
+ nats_url: str = "nats://localhost:4222",
+ heartbeat_interval: float = 10.0,
+ ) -> None:
+ self.org_id = org_id
+ self.nats_url = nats_url
+ self.heartbeat_interval = heartbeat_interval
+
+ self._agents: dict[str, Agent] = {}
+ self._agent_tasks: dict[str, asyncio.Task[None]] = {}
+ self._heartbeats: dict[str, float] = {}
+ self._heartbeat_task: asyncio.Task[None] | None = None
+ self._task_manager = TaskManager()
+ self._running = False
+
+ # ---- agent registry ---------------------------------------------------
+
+ def add_agent(self, agent: Agent) -> None:
+ """Register an agent with the orchestrator.
+
+ Args:
+ agent: Agent instance to manage. Its ``org_id`` is set to
+ match the orchestrator's.
+ """
+ agent.org_id = self.org_id
+ self._agents[agent.name] = agent
+ self._heartbeats[agent.name] = time.monotonic()
+ logger.info("Orchestrator registered agent '%s' (role=%s)", agent.name, agent.role)
+
+ def remove_agent(self, name: str) -> Agent:
+ """Unregister an agent.
+
+ If the agent is currently running, it will be stopped first.
+
+ Args:
+ name: Name of the agent to remove.
+
+ Returns:
+ The removed Agent instance.
+
+ Raises:
+ AgentNotFoundError: If the name is not registered.
+ """
+ agent = self._get_agent(name)
+ # Cancel its run task if active
+ task = self._agent_tasks.pop(name, None)
+ if task and not task.done():
+ task.cancel()
+ self._agents.pop(name)
+ self._heartbeats.pop(name, None)
+ logger.info("Orchestrator removed agent '%s'", name)
+ return agent
+
+ def list_agents(self) -> list[dict[str, Any]]:
+ """Return a summary of all registered agents."""
+ now = time.monotonic()
+ return [
+ {
+ "name": a.name,
+ "role": a.role,
+ "connected": a.is_connected,
+ "last_heartbeat_age_s": round(now - self._heartbeats.get(a.name, 0), 1),
+ }
+ for a in self._agents.values()
+ ]
+
+ # ---- lifecycle --------------------------------------------------------
+
+ async def start_all(self) -> None:
+ """Start all registered agents and the heartbeat monitor."""
+ self._running = True
+ for name, agent in self._agents.items():
+ if name not in self._agent_tasks or self._agent_tasks[name].done():
+ self._agent_tasks[name] = asyncio.create_task(
+ self._run_agent(agent), name=f"agent-{name}"
+ )
+ self._heartbeat_task = asyncio.create_task(
+ self._heartbeat_loop(), name="heartbeat-monitor"
+ )
+ logger.info("Orchestrator started %d agents", len(self._agents))
+
+ async def stop_all(self) -> None:
+ """Stop all agents and the heartbeat monitor gracefully."""
+ self._running = False
+
+ # Stop heartbeat
+ if self._heartbeat_task and not self._heartbeat_task.done():
+ self._heartbeat_task.cancel()
+ try:
+ await self._heartbeat_task
+ except asyncio.CancelledError:
+ pass
+
+ # Stop each agent
+ for name, agent in self._agents.items():
+ await agent.stop()
+
+ # Await all agent tasks
+ tasks = list(self._agent_tasks.values())
+ if tasks:
+ await asyncio.gather(*tasks, return_exceptions=True)
+ self._agent_tasks.clear()
+
+ logger.info("Orchestrator stopped all agents")
+
+ # ---- task delegation --------------------------------------------------
+
+ def delegate_task(
+ self,
+ title: str,
+ agent_name: str,
+ assigner: str = "orchestrator",
+ ) -> Task:
+ """Create and assign a task to a specific agent.
+
+ Args:
+ title: Description of the work to do.
+ agent_name: Name of the agent to assign the task to.
+ assigner: Name of the assigning entity.
+
+ Returns:
+ The created Task object.
+
+ Raises:
+ AgentNotFoundError: If agent_name is not registered.
+ """
+ self._get_agent(agent_name) # validate agent exists
+ task = self._task_manager.assign(
+ title=title, assignee=agent_name, assigner=assigner
+ )
+ logger.info(
+ "Orchestrator delegated task '%s' to '%s' (id=%s)",
+ title,
+ agent_name,
+ task.id,
+ )
+ return task
+
+ async def send_task_to_agent(self, task: Task) -> None:
+ """Send a task to the assigned agent via NATS messaging.
+
+ The agent must be connected. The task is serialized as a message
+ with type ``task_assignment``.
+
+ Args:
+ task: Task to send (the ``assignee`` field determines the target).
+ """
+ agent = self._get_agent(task.assignee)
+ await agent.send_message(
+ to_agent=task.assignee,
+ payload=task.model_dump(),
+ msg_type="task_assignment",
+ )
+
+ def get_task(self, task_id: str) -> Task:
+ """Retrieve a task by its ID."""
+ return self._task_manager.get(task_id)
+
+ def list_tasks(
+ self,
+ assignee: str | None = None,
+ status: TaskStatus | None = None,
+ ) -> list[Task]:
+ """List tasks, optionally filtered by assignee or status."""
+ return self._task_manager.list_tasks(assignee=assignee, status=status)
+
+ # ---- broadcast --------------------------------------------------------
+
+ async def broadcast(self, payload: dict[str, Any], msg_type: str = "broadcast") -> None:
+ """Send a message to every registered agent.
+
+ Uses the first connected agent as the sender. If no agent is
+ connected, a RuntimeError is raised.
+
+ Args:
+ payload: Data to broadcast.
+ msg_type: Message type for all outgoing messages.
+ """
+ sender = self._find_connected_agent()
+ for name in self._agents:
+ await sender.send_message(to_agent=name, payload=payload, msg_type=msg_type)
+ logger.info("Orchestrator broadcast '%s' to %d agents", msg_type, len(self._agents))
+
+ # ---- heartbeat --------------------------------------------------------
+
+ def record_heartbeat(self, agent_name: str) -> None:
+ """Record a heartbeat from an agent (call from message handlers)."""
+ self._heartbeats[agent_name] = time.monotonic()
+
+ def get_agent_health(self) -> dict[str, bool]:
+ """Check which agents have sent a heartbeat recently.
+
+ Returns:
+ Dict mapping agent name to healthy (True) or stale (False).
+ """
+ now = time.monotonic()
+ threshold = self.heartbeat_interval * 3 # miss 3 beats = unhealthy
+ return {
+ name: (now - self._heartbeats.get(name, 0)) < threshold
+ for name in self._agents
+ }
+
+ # ---- internal ---------------------------------------------------------
+
+ def _get_agent(self, name: str) -> Agent:
+ try:
+ return self._agents[name]
+ except KeyError:
+ raise AgentNotFoundError(
+ f"Agent '{name}' is not registered with the orchestrator"
+ ) from None
+
+ def _find_connected_agent(self) -> Agent:
+ for agent in self._agents.values():
+ if agent.is_connected:
+ return agent
+ raise RuntimeError("No connected agents available for sending messages")
+
+ async def _run_agent(self, agent: Agent) -> None:
+ """Wrapper that runs an agent and logs errors."""
+ try:
+ await agent.run()
+ except asyncio.CancelledError:
+ pass
+ except Exception:
+ logger.exception("Agent '%s' crashed", agent.name)
+
+ async def _heartbeat_loop(self) -> None:
+ """Periodically log agent health status."""
+ try:
+ while self._running:
+ await asyncio.sleep(self.heartbeat_interval)
+ health = self.get_agent_health()
+ unhealthy = [n for n, ok in health.items() if not ok]
+ if unhealthy:
+ logger.warning("Unhealthy agents: %s", ", ".join(unhealthy))
+ except asyncio.CancelledError:
+ pass
diff --git a/src/task_manager.py b/src/task_manager.py
new file mode 100644
index 0000000..b9a670d
--- /dev/null
+++ b/src/task_manager.py
@@ -0,0 +1,180 @@
+"""Simple in-memory task management for agent workflows."""
+
+from __future__ import annotations
+
+import uuid
+from datetime import datetime, timezone
+from enum import Enum
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+
+class TaskStatus(str, Enum):
+ """Valid lifecycle states for a task."""
+
+ ASSIGNED = "assigned"
+ IN_PROGRESS = "in_progress"
+ COMPLETED = "completed"
+ FAILED = "failed"
+
+
+# Legal transitions: from_status -> set of allowed target statuses
+_TRANSITIONS: dict[TaskStatus, set[TaskStatus]] = {
+ TaskStatus.ASSIGNED: {TaskStatus.IN_PROGRESS, TaskStatus.FAILED},
+ TaskStatus.IN_PROGRESS: {TaskStatus.COMPLETED, TaskStatus.FAILED},
+ TaskStatus.COMPLETED: set(), # terminal
+ TaskStatus.FAILED: {TaskStatus.ASSIGNED}, # allow retry
+}
+
+
+class Task(BaseModel):
+ """A unit of work assigned from one agent to another."""
+
+ id: str = Field(default_factory=lambda: str(uuid.uuid4()))
+ title: str
+ assignee: str
+ assigner: str
+ status: TaskStatus = TaskStatus.ASSIGNED
+ result: dict[str, Any] | None = None
+ created_at: str = Field(
+ default_factory=lambda: datetime.now(timezone.utc).isoformat()
+ )
+ updated_at: str = Field(
+ default_factory=lambda: datetime.now(timezone.utc).isoformat()
+ )
+
+
+class InvalidTransitionError(Exception):
+ """Raised when a task status transition is not allowed."""
+
+
+class TaskNotFoundError(KeyError):
+ """Raised when a task ID is not found in the store."""
+
+
+class TaskManager:
+ """In-memory task store with status-transition validation.
+
+ Example::
+
+ tm = TaskManager()
+ task = tm.assign("Deploy v2", assignee="devops", assigner="ceo")
+ tm.update_status(task.id, TaskStatus.IN_PROGRESS)
+ tm.complete(task.id, result={"url": "https://example.com"})
+ """
+
+ def __init__(self) -> None:
+ self._tasks: dict[str, Task] = {}
+
+ # ---- mutations --------------------------------------------------------
+
+ def assign(self, title: str, assignee: str, assigner: str) -> Task:
+ """Create a new task in ASSIGNED status.
+
+ Args:
+ title: Human-readable description of the work.
+ assignee: Agent name responsible for execution.
+ assigner: Agent name that created the task.
+
+ Returns:
+ The newly created Task.
+ """
+ task = Task(title=title, assignee=assignee, assigner=assigner)
+ self._tasks[task.id] = task
+ return task
+
+ def update_status(self, task_id: str, new_status: TaskStatus) -> Task:
+ """Transition a task to a new status with validation.
+
+ Args:
+ task_id: ID of the task to update.
+ new_status: Target status.
+
+ Returns:
+ The updated Task.
+
+ Raises:
+ TaskNotFoundError: If the task ID does not exist.
+ InvalidTransitionError: If the transition is not allowed.
+ """
+ task = self._get(task_id)
+ allowed = _TRANSITIONS.get(task.status, set())
+ if new_status not in allowed:
+ raise InvalidTransitionError(
+ f"Cannot transition task '{task.title}' from "
+ f"{task.status.value} to {new_status.value}. "
+ f"Allowed: {', '.join(s.value for s in allowed) or 'none (terminal)'}"
+ )
+ task.status = new_status
+ task.updated_at = datetime.now(timezone.utc).isoformat()
+ return task
+
+ def complete(
+ self, task_id: str, result: dict[str, Any] | None = None
+ ) -> Task:
+ """Mark a task as completed with optional result data.
+
+ Args:
+ task_id: ID of the task to complete.
+ result: Optional result payload.
+
+ Returns:
+ The completed Task.
+ """
+ task = self.update_status(task_id, TaskStatus.COMPLETED)
+ task.result = result
+ return task
+
+ def fail(self, task_id: str, reason: str | None = None) -> Task:
+ """Mark a task as failed.
+
+ Args:
+ task_id: ID of the task that failed.
+ reason: Optional failure description.
+
+ Returns:
+ The failed Task.
+ """
+ task = self.update_status(task_id, TaskStatus.FAILED)
+ task.result = {"error": reason} if reason else None
+ return task
+
+ # ---- queries ----------------------------------------------------------
+
+ def get(self, task_id: str) -> Task:
+ """Retrieve a task by ID.
+
+ Raises:
+ TaskNotFoundError: If the task does not exist.
+ """
+ return self._get(task_id)
+
+ def list_tasks(
+ self,
+ assignee: str | None = None,
+ status: TaskStatus | None = None,
+ ) -> list[Task]:
+ """List tasks with optional filters.
+
+ Args:
+ assignee: Filter by assignee name.
+ status: Filter by task status.
+
+ Returns:
+ List of matching tasks.
+ """
+ tasks = list(self._tasks.values())
+ if assignee:
+ tasks = [t for t in tasks if t.assignee == assignee]
+ if status:
+ tasks = [t for t in tasks if t.status == status]
+ return tasks
+
+ # ---- internal ---------------------------------------------------------
+
+ def _get(self, task_id: str) -> Task:
+ try:
+ return self._tasks[task_id]
+ except KeyError:
+ raise TaskNotFoundError(f"Task '{task_id}' not found") from None
diff --git a/tests/__init__.py b/tests/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/conftest.py b/tests/conftest.py
new file mode 100644
index 0000000..2df2e78
--- /dev/null
+++ b/tests/conftest.py
@@ -0,0 +1,13 @@
+"""Shared fixtures for the test suite."""
+
+import pytest
+
+
+@pytest.fixture
+def org_id():
+ return "test-org"
+
+
+@pytest.fixture
+def nats_url():
+ return "nats://localhost:4222"
diff --git a/tests/test_agent.py b/tests/test_agent.py
new file mode 100644
index 0000000..049e9d7
--- /dev/null
+++ b/tests/test_agent.py
@@ -0,0 +1,144 @@
+"""Tests for the base Agent class."""
+
+import asyncio
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+
+from src.agent import Agent
+from src.messaging import Message
+
+
+class EchoAgent(Agent):
+ """Test agent that records received messages."""
+
+ def __init__(self, **kwargs):
+ super().__init__(**kwargs)
+ self.received: list[Message] = []
+
+ async def handle_message(self, message: Message) -> None:
+ self.received.append(message)
+
+
+class TestAgentCreation:
+ def test_agent_attributes(self):
+ agent = Agent(name="test", role="worker", org_id="myorg")
+ assert agent.name == "test"
+ assert agent.role == "worker"
+ assert agent.org_id == "myorg"
+
+ def test_default_org_id(self):
+ agent = Agent(name="a", role="r")
+ assert agent.org_id == "default"
+
+ def test_inbox_subject(self):
+ agent = Agent(name="bot", role="worker", org_id="acme")
+ assert agent.inbox_subject == "acme.agents.bot.inbox"
+
+ def test_is_connected_false_initially(self):
+ agent = Agent(name="a", role="r")
+ assert agent.is_connected is False
+
+ def test_custom_nats_url(self):
+ agent = Agent(name="a", role="r", nats_url="nats://custom:4222")
+ assert agent.nats_url == "nats://custom:4222"
+
+ def test_nats_token_stored(self):
+ agent = Agent(name="a", role="r", nats_token="secret")
+ assert agent.nats_token == "secret"
+
+
+class TestAgentMessageHandling:
+ @pytest.mark.asyncio
+ async def test_handle_message_default_does_not_crash(self):
+ agent = Agent(name="a", role="r")
+ msg = Message(from_agent="b", to_agent="a", payload={"x": 1})
+ # Default handler should log, not raise
+ await agent.handle_message(msg)
+
+ @pytest.mark.asyncio
+ async def test_subclass_receives_messages(self):
+ agent = EchoAgent(name="echo", role="worker")
+ msg = Message(from_agent="sender", to_agent="echo", payload={"data": "hi"})
+ await agent.handle_message(msg)
+ assert len(agent.received) == 1
+ assert agent.received[0].payload == {"data": "hi"}
+
+ @pytest.mark.asyncio
+ async def test_on_message_deserializes_and_dispatches(self):
+ agent = EchoAgent(name="echo", role="worker")
+ msg = Message(from_agent="sender", to_agent="echo", type="test")
+ raw = MagicMock()
+ raw.data = msg.to_bytes()
+ await agent._on_message(raw)
+ assert len(agent.received) == 1
+ assert agent.received[0].type == "test"
+
+ @pytest.mark.asyncio
+ async def test_on_message_handles_bad_data(self):
+ agent = EchoAgent(name="echo", role="worker")
+ raw = MagicMock()
+ raw.data = b"not valid json"
+ # Should not raise — just log
+ await agent._on_message(raw)
+ assert len(agent.received) == 0
+
+
+class TestAgentSendMessage:
+ @pytest.mark.asyncio
+ async def test_send_message_raises_when_disconnected(self):
+ agent = Agent(name="a", role="r", org_id="org")
+ with pytest.raises(RuntimeError, match="not connected"):
+ await agent.send_message("other", {"hello": "world"})
+
+ @pytest.mark.asyncio
+ async def test_send_message_publishes_to_nats(self):
+ agent = Agent(name="sender", role="r", org_id="org")
+ mock_nc = AsyncMock()
+ mock_nc.is_closed = False
+ mock_nc.is_connected = True
+ agent._nc = mock_nc
+
+ await agent.send_message("target", {"key": "val"}, msg_type="custom")
+
+ mock_nc.publish.assert_called_once()
+ call_args = mock_nc.publish.call_args
+ assert call_args[0][0] == "org.agents.target.inbox"
+
+ # Verify the published data is a valid Message
+ published_data = call_args[0][1]
+ decoded = Message.from_bytes(published_data)
+ assert decoded.from_agent == "sender"
+ assert decoded.to_agent == "target"
+ assert decoded.type == "custom"
+ assert decoded.payload == {"key": "val"}
+
+
+class TestAgentLifecycle:
+ @pytest.mark.asyncio
+ async def test_stop_sets_running_false(self):
+ agent = Agent(name="a", role="r")
+ agent._running = True
+ await agent.stop()
+ assert agent._running is False
+
+ @pytest.mark.asyncio
+ async def test_on_start_default_is_noop(self):
+ agent = Agent(name="a", role="r")
+ # Should not raise
+ await agent.on_start()
+
+ @pytest.mark.asyncio
+ async def test_disconnect_drains_connection(self):
+ agent = Agent(name="a", role="r")
+ mock_nc = AsyncMock()
+ mock_nc.is_closed = False
+ agent._nc = mock_nc
+ agent._sub = AsyncMock()
+
+ await agent.disconnect()
+
+ agent._sub is None
+ mock_nc.drain.assert_called_once()
+ assert agent._nc is None
+ assert agent._running is False
diff --git a/tests/test_memory.py b/tests/test_memory.py
new file mode 100644
index 0000000..1c2f8e5
--- /dev/null
+++ b/tests/test_memory.py
@@ -0,0 +1,93 @@
+"""Tests for the MemoryStore."""
+
+import json
+import os
+import tempfile
+
+import pytest
+
+from src.memory import MemoryStore
+
+
+@pytest.fixture
+def store(tmp_path):
+ return MemoryStore(directory=tmp_path / "test_memory")
+
+
+class TestMemoryStore:
+ def test_save_and_load(self, store):
+ store.save("key1", {"name": "test", "value": 42})
+ result = store.load("key1")
+ assert result == {"name": "test", "value": 42}
+
+ def test_load_nonexistent_returns_none(self, store):
+ assert store.load("nonexistent") is None
+
+ def test_save_overwrites(self, store):
+ store.save("key", "first")
+ store.save("key", "second")
+ assert store.load("key") == "second"
+
+ def test_delete_existing(self, store):
+ store.save("key", "value")
+ assert store.delete("key") is True
+ assert store.load("key") is None
+
+ def test_delete_nonexistent(self, store):
+ assert store.delete("ghost") is False
+
+ def test_list_keys_empty(self, store):
+ assert store.list_keys() == []
+
+ def test_list_keys(self, store):
+ store.save("beta", 1)
+ store.save("alpha", 2)
+ store.save("gamma", 3)
+ keys = store.list_keys()
+ assert keys == ["alpha", "beta", "gamma"] # sorted
+
+ def test_clear(self, store):
+ store.save("a", 1)
+ store.save("b", 2)
+ count = store.clear()
+ assert count == 2
+ assert store.list_keys() == []
+
+ def test_various_value_types(self, store):
+ store.save("string", "hello")
+ store.save("number", 3.14)
+ store.save("bool", True)
+ store.save("null", None)
+ store.save("list", [1, 2, 3])
+ store.save("nested", {"a": {"b": [1, 2]}})
+
+ assert store.load("string") == "hello"
+ assert store.load("number") == 3.14
+ assert store.load("bool") is True
+ assert store.load("null") is None
+ assert store.load("list") == [1, 2, 3]
+ assert store.load("nested") == {"a": {"b": [1, 2]}}
+
+ def test_path_traversal_prevention(self, store):
+ # Keys with path-traversal characters should be sanitized
+ store.save("../evil", "data")
+ # Should not create files outside the memory dir
+ result = store.load("../evil")
+ # The key gets sanitized to "__evil"
+ assert result is not None or store.load("__evil") is not None
+
+ def test_creates_directory(self, tmp_path):
+ new_dir = tmp_path / "deep" / "nested" / "dir"
+ store = MemoryStore(directory=new_dir)
+ store.save("key", "value")
+ assert store.load("key") == "value"
+ assert new_dir.exists()
+
+ def test_files_are_json(self, store):
+ store.save("testkey", {"x": 1})
+ # Find the file and verify it's valid JSON
+ files = list(store._dir.glob("*.json"))
+ assert len(files) == 1
+ with open(files[0]) as f:
+ data = json.load(f)
+ assert data == {"x": 1}
diff --git a/tests/test_messaging.py b/tests/test_messaging.py
new file mode 100644
index 0000000..8d0aad0
--- /dev/null
+++ b/tests/test_messaging.py
@@ -0,0 +1,147 @@
+"""Tests for NATS messaging helpers."""
+
+import json
+from datetime import datetime, timezone
+
+import pytest
+
+from src.messaging import Message
+
+
+class TestMessage:
+ def test_create_with_defaults(self):
+ msg = Message(from_agent="a", to_agent="b")
+ assert msg.from_agent == "a"
+ assert msg.to_agent == "b"
+ assert msg.type == "message"
+ assert msg.payload == {}
+ assert msg.id # auto-generated UUID
+ assert msg.timestamp # auto-generated timestamp
+
+ def test_create_with_all_fields(self):
+ msg = Message(
+ id="custom-id",
+ from_agent="sender",
+ to_agent="receiver",
+ type="task",
+ payload={"key": "value"},
+ timestamp="2025-01-01T00:00:00+00:00",
+ )
+ assert msg.id == "custom-id"
+ assert msg.type == "task"
+ assert msg.payload == {"key": "value"}
+ assert msg.timestamp == "2025-01-01T00:00:00+00:00"
+
+ def test_unique_ids(self):
+ msg1 = Message(from_agent="a", to_agent="b")
+ msg2 = Message(from_agent="a", to_agent="b")
+ assert msg1.id != msg2.id
+
+ def test_to_bytes(self):
+ msg = Message(from_agent="a", to_agent="b", payload={"x": 1})
+ data = msg.to_bytes()
+ assert isinstance(data, bytes)
+ parsed = json.loads(data)
+ assert parsed["from_agent"] == "a"
+ assert parsed["to_agent"] == "b"
+ assert parsed["payload"] == {"x": 1}
+
+ def test_from_bytes(self):
+ original = Message(
+ from_agent="sender",
+ to_agent="receiver",
+ type="test",
+ payload={"numbers": [1, 2, 3]},
+ )
+ data = original.to_bytes()
+ restored = Message.from_bytes(data)
+
+ assert restored.id == original.id
+ assert restored.from_agent == original.from_agent
+ assert restored.to_agent == original.to_agent
+ assert restored.type == original.type
+ assert restored.payload == original.payload
+ assert restored.timestamp == original.timestamp
+
+ def test_roundtrip_preserves_data(self):
+ payload = {
+ "nested": {"a": 1, "b": [True, False, None]},
+ "text": "hello world",
+ }
+ msg = Message(from_agent="x", to_agent="y", payload=payload)
+ restored = Message.from_bytes(msg.to_bytes())
+ assert restored.payload == payload
+
+ def test_from_bytes_invalid_json_raises(self):
+ with pytest.raises(Exception):
+ Message.from_bytes(b"not json")
+
+ def test_from_bytes_missing_required_fields_raises(self):
+ with pytest.raises(Exception):
+ Message.from_bytes(json.dumps({"id": "1"}).encode())
+
+ def test_payload_defaults_to_empty_dict(self):
+ raw = json.dumps({
+ "id": "1",
+ "from_agent": "a",
+ "to_agent": "b",
+ "type": "msg",
+ "timestamp": "2025-01-01T00:00:00",
+ }).encode()
+ msg = Message.from_bytes(raw)
+ assert msg.payload == {}
+
+ def test_timestamp_format(self):
+ msg = Message(from_agent="a", to_agent="b")
+ # Should be a valid ISO format timestamp
+ dt = datetime.fromisoformat(msg.timestamp)
+ assert dt.tzinfo is not None # timezone-aware
+
+
+class TestConnectOptions:
+ """Test that connect() builds the right options (without actually connecting)."""
+
+ @pytest.mark.asyncio
+ async def test_connect_default_url(self):
+ """Verify default URL is localhost:4222."""
+ from unittest.mock import patch, AsyncMock
+
+ with patch("src.messaging.nats.connect", new_callable=AsyncMock) as mock_connect:
+ mock_connect.return_value = AsyncMock()
+ from src.messaging import connect
+ await connect()
+ mock_connect.assert_called_once_with(servers=["nats://localhost:4222"])
+
+ @pytest.mark.asyncio
+ async def test_connect_custom_url(self):
+ from unittest.mock import patch, AsyncMock
+
+ with patch("src.messaging.nats.connect", new_callable=AsyncMock) as mock_connect:
+ mock_connect.return_value = AsyncMock()
+ from src.messaging import connect
+ await connect(url="nats://remote:4223")
+ mock_connect.assert_called_once_with(servers=["nats://remote:4223"])
+
+ @pytest.mark.asyncio
+ async def test_connect_with_token(self):
+ from unittest.mock import patch, AsyncMock
+
+ with patch("src.messaging.nats.connect", new_callable=AsyncMock) as mock_connect:
+ mock_connect.return_value = AsyncMock()
+ from src.messaging import connect
+ await connect(url="nats://localhost:4222", token="mytoken")
+ mock_connect.assert_called_once_with(
+ servers=["nats://localhost:4222"], token="mytoken"
+ )
+
+ @pytest.mark.asyncio
+ async def test_connect_without_token_omits_key(self):
+ from unittest.mock import patch, AsyncMock
+
+ with patch("src.messaging.nats.connect", new_callable=AsyncMock) as mock_connect:
+ mock_connect.return_value = AsyncMock()
+ from src.messaging import connect
+ await connect()
+ call_kwargs = mock_connect.call_args
+ # token should NOT be in the call
+ assert "token" not in (call_kwargs.kwargs if call_kwargs.kwargs else {})
diff --git a/tests/test_orchestrator.py b/tests/test_orchestrator.py
new file mode 100644
index 0000000..849bdec
--- /dev/null
+++ b/tests/test_orchestrator.py
@@ -0,0 +1,179 @@
+"""Tests for the Orchestrator class."""
+
+import asyncio
+import time
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+
+from src.agent import Agent
+from src.orchestrator import AgentNotFoundError, Orchestrator
+from src.task_manager import TaskStatus
+
+
+@pytest.fixture
+def orch():
+ return Orchestrator(org_id="test-org")
+
+
+@pytest.fixture
+def make_agent():
+ """Factory for creating test agents."""
+
+ def _make(name: str, role: str = "worker") -> Agent:
+ return Agent(name=name, role=role, org_id="test-org")
+
+ return _make
+
+
+class TestAddRemoveAgents:
+ def test_add_agent(self, orch, make_agent):
+ agent = make_agent("alpha")
+ orch.add_agent(agent)
+ assert "alpha" in [a["name"] for a in orch.list_agents()]
+
+ def test_add_sets_org_id(self, orch, make_agent):
+ agent = make_agent("alpha")
+ agent.org_id = "other"
+ orch.add_agent(agent)
+ assert agent.org_id == "test-org"
+
+ def test_add_multiple_agents(self, orch, make_agent):
+ orch.add_agent(make_agent("a"))
+ orch.add_agent(make_agent("b"))
+ orch.add_agent(make_agent("c"))
+ names = [a["name"] for a in orch.list_agents()]
+ assert set(names) == {"a", "b", "c"}
+
+ def test_remove_agent(self, orch, make_agent):
+ orch.add_agent(make_agent("alpha"))
+ removed = orch.remove_agent("alpha")
+ assert removed.name == "alpha"
+ assert "alpha" not in [a["name"] for a in orch.list_agents()]
+
+ def test_remove_nonexistent_raises(self, orch):
+ with pytest.raises(AgentNotFoundError):
+ orch.remove_agent("ghost")
+
+ def test_list_agents_shows_health(self, orch, make_agent):
+ orch.add_agent(make_agent("alpha"))
+ agents = orch.list_agents()
+ assert len(agents) == 1
+ assert agents[0]["name"] == "alpha"
+ assert agents[0]["connected"] is False
+ assert "last_heartbeat_age_s" in agents[0]
+
+
+class TestDelegateTask:
+ def test_delegate_creates_task(self, orch, make_agent):
+ orch.add_agent(make_agent("worker"))
+ task = orch.delegate_task("Do something", "worker")
+ assert task.title == "Do something"
+ assert task.assignee == "worker"
+ assert task.assigner == "orchestrator"
+ assert task.status == TaskStatus.ASSIGNED
+
+ def test_delegate_custom_assigner(self, orch, make_agent):
+ orch.add_agent(make_agent("worker"))
+ task = orch.delegate_task("Fix bug", "worker", assigner="ceo")
+ assert task.assigner == "ceo"
+
+ def test_delegate_to_nonexistent_raises(self, orch):
+ with pytest.raises(AgentNotFoundError):
+ orch.delegate_task("Task", "nobody")
+
+ def test_delegate_returns_retrievable_task(self, orch, make_agent):
+ orch.add_agent(make_agent("worker"))
+ task = orch.delegate_task("Analyze data", "worker")
+ retrieved = orch.get_task(task.id)
+ assert retrieved.id == task.id
+ assert retrieved.title == "Analyze data"
+
+ def test_list_tasks_by_assignee(self, orch, make_agent):
+ orch.add_agent(make_agent("a"))
+ orch.add_agent(make_agent("b"))
+ orch.delegate_task("Task 1", "a")
+ orch.delegate_task("Task 2", "b")
+ orch.delegate_task("Task 3", "a")
+
+ a_tasks = orch.list_tasks(assignee="a")
+ assert len(a_tasks) == 2
+ assert all(t.assignee == "a" for t in a_tasks)
+
+ def test_list_tasks_by_status(self, orch, make_agent):
+ orch.add_agent(make_agent("w"))
+ t1 = orch.delegate_task("T1", "w")
+ orch.delegate_task("T2", "w")
+
+ # Move t1 to in_progress
+ orch._task_manager.update_status(t1.id, TaskStatus.IN_PROGRESS)
+
+ assigned = orch.list_tasks(status=TaskStatus.ASSIGNED)
+ in_progress = orch.list_tasks(status=TaskStatus.IN_PROGRESS)
+ assert len(assigned) == 1
+ assert len(in_progress) == 1
+
+
+class TestBroadcast:
+ @pytest.mark.asyncio
+ async def test_broadcast_sends_to_all(self, orch, make_agent):
+ a1 = make_agent("a")
+ a2 = make_agent("b")
+ orch.add_agent(a1)
+ orch.add_agent(a2)
+
+ # Mock a connected agent as sender
+ a1._nc = AsyncMock()
+ a1._nc.is_closed = False
+ a1._nc.is_connected = True
+
+ await orch.broadcast({"msg": "hello"}, msg_type="ping")
+
+ # send_message should be called for each agent
+ assert a1._nc.publish.call_count == 2
+
+ @pytest.mark.asyncio
+ async def test_broadcast_no_connected_agents_raises(self, orch, make_agent):
+ orch.add_agent(make_agent("a"))
+ with pytest.raises(RuntimeError, match="No connected agents"):
+ await orch.broadcast({"msg": "hello"})
+
+
+class TestHeartbeat:
+ def test_record_heartbeat(self, orch, make_agent):
+ orch.add_agent(make_agent("a"))
+ orch.record_heartbeat("a")
+ health = orch.get_agent_health()
+ assert health["a"] is True
+
+ def test_stale_heartbeat_shows_unhealthy(self, orch, make_agent):
+ orch.add_agent(make_agent("a"))
+ # Set heartbeat far in the past
+ orch._heartbeats["a"] = time.monotonic() - 999
+ health = orch.get_agent_health()
+ assert health["a"] is False
+
+ def test_fresh_heartbeat_shows_healthy(self, orch, make_agent):
+ orch.add_agent(make_agent("a"))
+ orch.record_heartbeat("a")
+ health = orch.get_agent_health()
+ assert health["a"] is True
+
+
+class TestSendTaskToAgent:
+ @pytest.mark.asyncio
+ async def test_send_task_publishes_message(self, orch, make_agent):
+ agent = make_agent("worker")
+ orch.add_agent(agent)
+
+ # Mock NATS connection
+ agent._nc = AsyncMock()
+ agent._nc.is_closed = False
+ agent._nc.is_connected = True
+
+ task = orch.delegate_task("Test task", "worker")
+ await orch.send_task_to_agent(task)
+
+ agent._nc.publish.assert_called_once()
+ call_args = agent._nc.publish.call_args
+ assert call_args[0][0] == "test-org.agents.worker.inbox"
diff --git a/tests/test_task_manager.py b/tests/test_task_manager.py
new file mode 100644
index 0000000..0fdbe7b
--- /dev/null
+++ b/tests/test_task_manager.py
@@ -0,0 +1,121 @@
+"""Tests for the TaskManager."""
+
+import pytest
+
+from src.task_manager import (
+ InvalidTransitionError,
+ Task,
+ TaskManager,
+ TaskNotFoundError,
+ TaskStatus,
+)
+
+
+@pytest.fixture
+def tm():
+ return TaskManager()
+
+
+class TestTaskCreation:
+ def test_assign_creates_task(self, tm):
+ task = tm.assign("Do work", assignee="worker", assigner="boss")
+ assert task.title == "Do work"
+ assert task.assignee == "worker"
+ assert task.assigner == "boss"
+ assert task.status == TaskStatus.ASSIGNED
+
+ def test_assign_generates_unique_ids(self, tm):
+ t1 = tm.assign("A", assignee="w", assigner="b")
+ t2 = tm.assign("B", assignee="w", assigner="b")
+ assert t1.id != t2.id
+
+ def test_get_task(self, tm):
+ task = tm.assign("X", assignee="w", assigner="b")
+ retrieved = tm.get(task.id)
+ assert retrieved.id == task.id
+
+ def test_get_nonexistent_raises(self, tm):
+ with pytest.raises(TaskNotFoundError):
+ tm.get("nonexistent-id")
+
+
+class TestStatusTransitions:
+ def test_assigned_to_in_progress(self, tm):
+ task = tm.assign("T", assignee="w", assigner="b")
+ updated = tm.update_status(task.id, TaskStatus.IN_PROGRESS)
+ assert updated.status == TaskStatus.IN_PROGRESS
+
+ def test_in_progress_to_completed(self, tm):
+ task = tm.assign("T", assignee="w", assigner="b")
+ tm.update_status(task.id, TaskStatus.IN_PROGRESS)
+ updated = tm.update_status(task.id, TaskStatus.COMPLETED)
+ assert updated.status == TaskStatus.COMPLETED
+
+ def test_assigned_to_failed(self, tm):
+ task = tm.assign("T", assignee="w", assigner="b")
+ updated = tm.update_status(task.id, TaskStatus.FAILED)
+ assert updated.status == TaskStatus.FAILED
+
+ def test_failed_to_assigned_retry(self, tm):
+ task = tm.assign("T", assignee="w", assigner="b")
+ tm.update_status(task.id, TaskStatus.FAILED)
+ updated = tm.update_status(task.id, TaskStatus.ASSIGNED)
+ assert updated.status == TaskStatus.ASSIGNED
+
+ def test_cannot_go_from_assigned_to_completed(self, tm):
+ task = tm.assign("T", assignee="w", assigner="b")
+ with pytest.raises(InvalidTransitionError):
+ tm.update_status(task.id, TaskStatus.COMPLETED)
+
+ def test_cannot_transition_from_completed(self, tm):
+ task = tm.assign("T", assignee="w", assigner="b")
+ tm.update_status(task.id, TaskStatus.IN_PROGRESS)
+ tm.update_status(task.id, TaskStatus.COMPLETED)
+ with pytest.raises(InvalidTransitionError):
+ tm.update_status(task.id, TaskStatus.ASSIGNED)
+
+ def test_complete_shortcut(self, tm):
+ task = tm.assign("T", assignee="w", assigner="b")
+ tm.update_status(task.id, TaskStatus.IN_PROGRESS)
+ completed = tm.complete(task.id, result={"output": "done"})
+ assert completed.status == TaskStatus.COMPLETED
+ assert completed.result == {"output": "done"}
+
+ def test_fail_shortcut(self, tm):
+ task = tm.assign("T", assignee="w", assigner="b")
+ failed = tm.fail(task.id, reason="timeout")
+ assert failed.status == TaskStatus.FAILED
+ assert failed.result == {"error": "timeout"}
+
+
+class TestListTasks:
+ def test_list_all(self, tm):
+ tm.assign("A", assignee="w1", assigner="b")
+ tm.assign("B", assignee="w2", assigner="b")
+ assert len(tm.list_tasks()) == 2
+
+ def test_filter_by_assignee(self, tm):
+ tm.assign("A", assignee="w1", assigner="b")
+ tm.assign("B", assignee="w2", assigner="b")
+ tm.assign("C", assignee="w1", assigner="b")
+ result = tm.list_tasks(assignee="w1")
+ assert len(result) == 2
+ assert all(t.assignee == "w1" for t in result)
+
+ def test_filter_by_status(self, tm):
+ t1 = tm.assign("A", assignee="w", assigner="b")
+ tm.assign("B", assignee="w", assigner="b")
+ tm.update_status(t1.id, TaskStatus.IN_PROGRESS)
+ result = tm.list_tasks(status=TaskStatus.IN_PROGRESS)
+ assert len(result) == 1
+ assert result[0].id == t1.id
+
+ def test_filter_by_both(self, tm):
+ t1 = tm.assign("A", assignee="w1", assigner="b")
+ tm.assign("B", assignee="w2", assigner="b")
+ tm.update_status(t1.id, TaskStatus.FAILED)
+ result = tm.list_tasks(assignee="w1", status=TaskStatus.FAILED)
+ assert len(result) == 1
+
+ def test_empty_list(self, tm):
+ assert tm.list_tasks() == []