Agent Bus：把多 Agent 协作写成可验证的任务契约

"""Offline contract exercise. Fixture sessions are NOT production authentication."""
from contextlib import contextmanager
from dataclasses import dataclass
import hashlib
import json
import sqlite3
import time
import uuid

FIXTURES = {
    "public-a": "Reliable tasks need explicit contracts.",
    "public-b": "Read the task before doing the work.",
}
WORKERS = {"worker-a", "worker-b"}
OPEN = {"queued", "running"}


class BusError(Exception):
    pass


def require(condition, code):
    if not condition:
        raise BusError(code)


def canonical(value):
    return json.dumps(value, sort_keys=True, separators=(",", ":"))


def digest(value):
    return hashlib.sha256(canonical(value).encode()).hexdigest()


def answer(fixture_id):
    text = FIXTURES[fixture_id]
    return {"fixture_id": fixture_id, "words": len(text.split()),
            "sha256": hashlib.sha256(text.encode()).hexdigest()}


@dataclass(frozen=True)
class Actor:
    name: str


class FixtureAuth:
    """Trusted test harness only; do not expose this object as a login API."""

    def __init__(self):
        self.sessions = {name: Actor(name) for name in
                         ["coordinator", "other", "clock", *sorted(WORKERS)]}

    def identify(self, session):
        for name, known in self.sessions.items():
            if session is known:
                return name
        raise BusError("unauthenticated")


class Bus:
    def __init__(self, path, auth, now=time.time):
        self.auth, self.now = auth, now
        self.db = sqlite3.connect(path, isolation_level=None, timeout=5)
        self.db.executescript("""
        CREATE TABLE IF NOT EXISTS tasks (
          id TEXT PRIMARY KEY, body TEXT NOT NULL);
        CREATE TABLE IF NOT EXISTS receipts (
          task_id TEXT, submission_id TEXT, hash TEXT, response TEXT,
          PRIMARY KEY (task_id, submission_id));
        CREATE TABLE IF NOT EXISTS audit (
          seq INTEGER PRIMARY KEY, task_id TEXT, actor TEXT,
          action TEXT, outcome TEXT, revision INTEGER, at REAL);
        """)

    def close(self):
        self.db.close()

    def _log(self, task_id, actor, action, outcome, revision=None):
        self.db.execute("INSERT INTO audit VALUES(NULL,?,?,?,?,?,?)",
                        (task_id, actor, action, outcome, revision, self.now()))

    @contextmanager
    def _write(self, session, action, task_id=""):
        actor = self.auth.identify(session)
        self.db.execute("BEGIN IMMEDIATE")
        try:
            yield actor
            self.db.commit()
        except Exception as exc:
            self.db.rollback()
            if isinstance(exc, BusError):
                self._log(task_id, actor, action, str(exc))
            raise

    def _load(self, task_id):
        row = self.db.execute("SELECT body FROM tasks WHERE id=?",
                              (task_id,)).fetchone()
        require(row is not None, "not_found")
        return json.loads(row[0])

    def _readable(self, task, actor):
        require(actor in {task["from"], task["to"]}, "forbidden")

    def _open(self, task, revision, states):
        require(type(revision) is int, "bad_revision")
        require(task["revision"] == revision, "version_conflict")
        require(task["status"] in states, "illegal_state")
        require(self.now() < task["expires_at"], "expired")

    def _save(self, task, actor, action):
        task["revision"] += 1
        self.db.execute("UPDATE tasks SET body=? WHERE id=?",
                        (canonical(task), task["task_id"]))
        self._log(task["task_id"], actor, action, "ok", task["revision"])
        return task

    def create(self, session, body):
        with self._write(session, "create") as actor:
            require(actor == "coordinator", "forbidden")
            keys = {"to", "type", "correlation_id", "data", "expires_at"}
            require(type(body) is dict and set(body) == keys, "bad_envelope")
            require(type(body["to"]) is str and body["to"] in WORKERS,
                    "bad_target")
            require(body["type"] == "count_words", "bad_type")
            correlation = body["correlation_id"]
            require(type(correlation) is str and 0 < len(correlation) <= 80,
                    "bad_correlation")
            data = body["data"]
            require(type(data) is dict and set(data) == {"fixture_id"},
                    "bad_data")
            require(type(data["fixture_id"]) is str, "bad_fixture")
            require(data["fixture_id"] in FIXTURES, "bad_fixture")
            expiry = body["expires_at"]
            require(type(expiry) in {int, float}, "bad_expiry")
            require(self.now() < expiry <= self.now() + 3600, "bad_expiry")
            task = json.loads(canonical(body))
            task.update(task_id=str(uuid.uuid4()), **{"from": actor},
                        status="queued", revision=1, result=None,
                        error=None, cancel=None)
            self.db.execute("INSERT INTO tasks VALUES(?,?)",
                            (task["task_id"], canonical(task)))
            self._log(task["task_id"], actor, "create", "ok", 1)
            return task

    def get(self, session, task_id):
        actor = self.auth.identify(session)
        task = self._load(task_id)
        self._readable(task, actor)
        return task

    def queue(self, session):
        actor = self.auth.identify(session)
        require(actor in WORKERS, "forbidden")
        tasks = [json.loads(row[0]) for row in
                 self.db.execute("SELECT body FROM tasks ORDER BY rowid")]
        return [t for t in tasks if t["to"] == actor
                and t["status"] == "queued" and self.now() < t["expires_at"]]

    def claim(self, session, task_id, revision):
        with self._write(session, "claim", task_id) as actor:
            task = self._load(task_id)
            require(actor == task["to"], "forbidden")
            self._open(task, revision, {"queued"})
            task["status"] = "running"
            return self._save(task, actor, "claim")

    def submit(self, session, task_id, revision, submission_id, result):
        with self._write(session, "submit", task_id) as actor:
            task = self._load(task_id)
            require(actor == task["to"], "forbidden")
            require(type(submission_id) is str
                    and 0 < len(submission_id) <= 80, "bad_submission_id")
            try:
                submission_id.encode("utf-8")
            except UnicodeEncodeError:
                raise BusError("bad_submission_id") from None
            require(type(result) is dict, "bad_result_schema")
            require(set(result) == {"fixture_id", "words", "sha256"},
                    "bad_result_schema")
            require(type(result["fixture_id"]) is str
                    and type(result["words"]) is int
                    and type(result["sha256"]) is str, "bad_result_schema")
            fingerprint = digest(result)
            receipt = self.db.execute(
                "SELECT hash,response FROM receipts "
                "WHERE task_id=? AND submission_id=?",
                (task_id, submission_id)).fetchone()
            if receipt:
                require(receipt[0] == fingerprint, "idempotency_conflict")
                self._log(task_id, actor, "submit", "replayed",
                          task["revision"])
                return json.loads(receipt[1])
            self._open(task, revision, {"running"})
            require(result == answer(task["data"]["fixture_id"]),
                    "result_verification_failed")
            task["status"], task["result"] = "succeeded", result
            self._save(task, actor, "submit")
            self.db.execute("INSERT INTO receipts VALUES(?,?,?,?)",
                            (task_id, submission_id, fingerprint,
                             canonical(task)))
            return task

    def fail(self, session, task_id, revision, error):
        with self._write(session, "fail", task_id) as actor:
            task = self._load(task_id)
            require(actor == task["to"], "forbidden")
            self._open(task, revision, {"running"})
            require(type(error) is dict
                    and error == {"code": "worker_failed", "retryable": False}
                    and type(error.get("retryable")) is bool, "bad_error")
            task["status"], task["error"] = "failed", error
            return self._save(task, actor, "fail")

    def cancel(self, session, task_id, revision):
        with self._write(session, "cancel", task_id) as actor:
            task = self._load(task_id)
            require(actor == task["from"], "forbidden")
            self._open(task, revision, OPEN)
            task["status"] = "cancelled"
            task["cancel"] = {"by": actor, "at": self.now()}
            return self._save(task, actor, "cancel")

    def expire(self, session):
        with self._write(session, "expire") as actor:
            require(actor == "clock", "forbidden")
            changed = []
            rows = self.db.execute("SELECT body FROM tasks").fetchall()
            for row in rows:
                task = json.loads(row[0])
                if task["status"] in OPEN and self.now() >= task["expires_at"]:
                    task["status"] = "expired"
                    task["error"] = {"code": "deadline_elapsed",
                                     "retryable": False}
                    changed.append(self._save(task, actor, "expire"))
            return changed

    def audit(self, session, task_id):
        self.get(session, task_id)
        return self.db.execute(
            "SELECT actor,action,outcome,revision FROM audit "
            "WHERE task_id=? ORDER BY seq", (task_id,)).fetchall()


class Worker:
    def __init__(self, bus, session):
        self.bus, self.session = bus, session

    def on_event(self, event):
        # Event content is only a hint. No instruction or result is executed.
        require(type(event) is dict and type(event.get("task_id")) is str,
                "bad_event")
        task = self.bus.get(self.session, event["task_id"])
        if task["status"] != "queued":
            return task
        task = self.bus.claim(self.session, task["task_id"], task["revision"])
        result = answer(task["data"]["fixture_id"])
        return self.bus.submit(self.session, task["task_id"], task["revision"],
                               "result-" + task["task_id"], result)

    def poll_once(self):
        return [self.on_event({"task_id": t["task_id"]})
                for t in self.bus.queue(self.session)]

"""No model, network, credentials, or external agent calls."""
from pathlib import Path
from tempfile import TemporaryDirectory
from task_bus import Bus, FixtureAuth, Worker


def main():
    with TemporaryDirectory() as tmp:
        auth = FixtureAuth()
        bus = Bus(Path(tmp) / "tasks.sqlite", auth, now=lambda: 1000)
        coordinator = auth.sessions["coordinator"]
        for target, fixture in [("worker-a", "public-a"),
                                ("worker-b", "public-b")]:
            bus.create(coordinator, {
                "to": target, "type": "count_words",
                "correlation_id": "public-demo", "data": {"fixture_id": fixture},
                "expires_at": 1060,
            })
        for name in ["worker-a", "worker-b"]:
            worker = Worker(bus, auth.sessions[name])
            for result in worker.poll_once():
                print(name, result["status"], "revision", result["revision"],
                      "words", result["result"]["words"])
                print("audit:", bus.audit(coordinator, result["task_id"]))
        bus.close()


if __name__ == "__main__":
    main()

from pathlib import Path
from tempfile import TemporaryDirectory
import unittest
import json
from task_bus import Actor, Bus, BusError, FixtureAuth, Worker, answer


class ContractTests(unittest.TestCase):
    def setUp(self):
        self.tmp = TemporaryDirectory()
        self.auth = FixtureAuth()
        self.now = 1000
        self.path = Path(self.tmp.name) / "tasks.sqlite"
        self.bus = Bus(self.path, self.auth, lambda: self.now)
        self.c = self.auth.sessions["coordinator"]
        self.a = self.auth.sessions["worker-a"]
        self.b = self.auth.sessions["worker-b"]

    def tearDown(self):
        self.bus.close()
        self.tmp.cleanup()

    def body(self, target="worker-a"):
        return {"to": target, "type": "count_words", "correlation_id": "test",
                "data": {"fixture_id": "public-a"}, "expires_at": 1060}

    def create(self):
        return self.bus.create(self.c, self.body())

    def running(self):
        task = self.create()
        return self.bus.claim(self.a, task["task_id"], 1)

    def rejected(self, code, fn, *args):
        with self.assertRaisesRegex(BusError, "^" + code + "$"):
            fn(*args)

    def test_forged_actor(self):
        self.rejected("unauthenticated", self.bus.create,
                      Actor("coordinator"), self.body())

    def test_body_from_rejected(self):
        body = self.body()
        body["from"] = "coordinator"
        self.rejected("bad_envelope", self.bus.create, self.c, body)

    def test_bad_target_types(self):
        for target in [[], {}, None, 42]:
            self.rejected("bad_target", self.bus.create,
                          self.c, self.body(target))

    def test_worker_cannot_delegate(self):
        self.rejected("forbidden", self.bus.create, self.a, self.body())

    def test_wrong_agent_cannot_read_or_audit(self):
        task = self.create()
        for operation in [self.bus.get, self.bus.audit]:
            self.rejected("forbidden", operation, self.b, task["task_id"])

    def test_queue_is_scoped(self):
        self.create()
        self.assertEqual(self.bus.queue(self.b), [])
        self.assertEqual(len(self.bus.queue(self.a)), 1)

    def test_wrong_agent_cannot_claim(self):
        task = self.create()
        self.rejected("forbidden", self.bus.claim, self.b, task["task_id"], 1)

    def test_success_without_claim_rejected(self):
        task = self.create()
        self.rejected("illegal_state", self.bus.submit, self.a,
                      task["task_id"], 1, "r1", answer("public-a"))

    def test_version_conflict(self):
        task = self.running()
        self.rejected("version_conflict", self.bus.submit, self.a,
                      task["task_id"], 1, "r1", answer("public-a"))

    def test_only_one_claim_across_connections(self):
        task = self.create()
        other = Bus(self.path, self.auth, lambda: self.now)
        try:
            self.bus.claim(self.a, task["task_id"], 1)
            self.rejected("version_conflict", other.claim,
                          self.a, task["task_id"], 1)
        finally:
            other.close()

    def test_result_replay(self):
        task = self.running()
        args = (self.a, task["task_id"], 2, "r1", answer("public-a"))
        first = self.bus.submit(*args)
        self.assertEqual(self.bus.submit(*args), first)
        self.assertEqual(first["revision"], 3)

    def test_submission_id_rejects_surrogates(self):
        task = self.running()
        keys = ["\ud800", "\udfff", "receipt-\ud800-x", "\ud83d\ude00",
                json.loads('"\\ud800"')]
        for key in keys:
            with self.subTest(key=ascii(key)):
                self.rejected("bad_submission_id", self.bus.submit,
                              self.a, task["task_id"], 2, key,
                              answer("public-a"))
                self.assertEqual(self.bus.get(self.a, task["task_id"]), task)
                count = self.bus.db.execute(
                    "SELECT count(*) FROM receipts").fetchone()[0]
                self.assertEqual(count, 0)
                self.assertFalse(self.bus.db.in_transaction)
                self.assertEqual(self.bus.audit(self.c, task["task_id"])[-1],
                                 ("worker-a", "submit", "bad_submission_id",
                                  None))

    def test_unicode_submission_id_replay(self):
        task = self.running()
        key = "交付-" + json.loads('"\\ud83d\\ude00"')
        args = (self.a, task["task_id"], 2, key, answer("public-a"))
        first = self.bus.submit(*args)
        self.assertEqual(first["status"], "succeeded")
        self.assertEqual(self.bus.submit(*args), first)
        self.assertEqual(self.bus.db.execute(
            "SELECT submission_id FROM receipts").fetchone()[0], key)

    def test_invalid_submission_id_preserves_existing_receipt(self):
        task = self.running()
        args = (self.a, task["task_id"], 2, "valid-key", answer("public-a"))
        first = self.bus.submit(*args)
        self.rejected("bad_submission_id", self.bus.submit,
                      self.a, task["task_id"], 2, "\ud800", answer("public-a"))
        self.assertEqual(self.bus.get(self.a, task["task_id"]), first)
        self.assertEqual(self.bus.submit(*args), first)
        self.assertEqual(self.bus.db.execute(
            "SELECT count(*) FROM receipts").fetchone()[0], 1)

    def test_submission_id_type_and_length_boundaries(self):
        task = self.running()
        for key in [None, True, 7, [], "", "x" * 81]:
            with self.subTest(key=key):
                self.rejected("bad_submission_id", self.bus.submit,
                              self.a, task["task_id"], 2, key,
                              answer("public-a"))
        args = (self.a, task["task_id"], 2, "字" * 80, answer("public-a"))
        self.assertEqual(self.bus.submit(*args)["status"], "succeeded")

    def test_same_key_different_result(self):
        task = self.running()
        self.bus.submit(self.a, task["task_id"], 2, "r1", answer("public-a"))
        changed = dict(answer("public-a"), words=99)
        self.rejected("idempotency_conflict", self.bus.submit,
                      self.a, task["task_id"], 2, "r1", changed)

    def test_second_submission_key_rejected(self):
        task = self.running()
        self.bus.submit(self.a, task["task_id"], 2, "r1", answer("public-a"))
        self.rejected("illegal_state", self.bus.submit, self.a,
                      task["task_id"], 3, "r2", answer("public-a"))

    def test_wrong_worker_cannot_replay_receipt(self):
        task = self.running()
        self.bus.submit(self.a, task["task_id"], 2, "r1", answer("public-a"))
        self.rejected("forbidden", self.bus.submit, self.b,
                      task["task_id"], 2, "r1", answer("public-a"))

    def test_string_is_not_success(self):
        task = self.running()
        self.rejected("bad_result_schema", self.bus.submit,
                      self.a, task["task_id"], 2, "r1", "done")

    def test_boolean_is_not_word_count(self):
        task = self.running()
        result = dict(answer("public-a"), words=True)
        self.rejected("bad_result_schema", self.bus.submit,
                      self.a, task["task_id"], 2, "r1", result)

    def test_schema_valid_but_wrong_answer(self):
        task = self.running()
        result = dict(answer("public-a"), words=99)
        self.rejected("result_verification_failed", self.bus.submit,
                      self.a, task["task_id"], 2, "r1", result)
        self.assertEqual(self.bus.get(self.a, task["task_id"])["revision"], 2)

    def test_cancel_running_rejects_late_result(self):
        task = self.running()
        self.bus.cancel(self.c, task["task_id"], 2)
        self.rejected("version_conflict", self.bus.submit, self.a,
                      task["task_id"], 2, "r1", answer("public-a"))
        self.rejected("illegal_state", self.bus.submit, self.a,
                      task["task_id"], 3, "r1", answer("public-a"))

    def test_worker_cannot_cancel(self):
        task = self.running()
        self.rejected("forbidden", self.bus.cancel, self.a, task["task_id"], 2)

    def test_expired_before_sweep_is_not_claimable(self):
        task = self.create()
        self.now = 1060
        self.assertEqual(self.bus.queue(self.a), [])
        self.rejected("expired", self.bus.claim, self.a, task["task_id"], 1)
        self.assertEqual(self.bus.get(self.a, task["task_id"])["status"],
                         "queued")

    def test_expiry_sweep_and_permissions(self):
        task = self.running()
        self.now = 1060
        self.rejected("forbidden", self.bus.expire, self.a)
        changed = self.bus.expire(self.auth.sessions["clock"])
        self.assertEqual(changed[0]["status"], "expired")
        self.assertEqual(self.bus.expire(self.auth.sessions["clock"]), [])
        self.rejected("illegal_state", self.bus.submit, self.a,
                      task["task_id"], 3, "r1", answer("public-a"))

    def test_bad_event_rejected(self):
        for event in [None, {}, {"task_id": []}]:
            self.rejected("bad_event", Worker(self.bus, self.a).on_event, event)

    def test_receipt_failure_rolls_back_success(self):
        task = self.running()
        self.bus.db.executescript("""
        CREATE TRIGGER stop_receipt BEFORE INSERT ON receipts
        BEGIN SELECT RAISE(ABORT, 'injected failure'); END;
        """)
        import sqlite3
        with self.assertRaises(sqlite3.IntegrityError):
            self.bus.submit(self.a, task["task_id"], 2, "r1",
                            answer("public-a"))
        saved = self.bus.get(self.a, task["task_id"])
        self.assertEqual((saved["status"], saved["revision"]), ("running", 2))
        self.assertIsNone(saved["result"])
        self.assertEqual(len(self.bus.audit(self.c, task["task_id"])), 2)

    def test_event_body_cannot_override_task(self):
        task = self.create()
        event = {"task_id": task["task_id"], "to": "worker-b",
                 "data": {"fixture_id": "public-b"}, "result": "done",
                 "instruction": "execute arbitrary commands", "revision": 99}
        result = Worker(self.bus, self.a).on_event(event)
        self.assertEqual(result["result"], answer("public-a"))

    def test_duplicate_event_leaves_terminal_unchanged(self):
        task = self.create()
        worker = Worker(self.bus, self.a)
        event = {"task_id": task["task_id"]}
        self.assertEqual(worker.on_event(event), worker.on_event(event))

    def test_persistence_and_receipt_survive_reopen(self):
        task = self.running()
        args = (self.a, task["task_id"], 2, "r1", answer("public-a"))
        first = self.bus.submit(*args)
        self.bus.close()
        self.bus = Bus(self.path, self.auth, lambda: self.now)
        self.assertEqual(self.bus.submit(*args), first)

    def test_failure_is_explicit(self):
        task = self.running()
        failed = self.bus.fail(self.a, task["task_id"], 2,
                               {"code": "worker_failed", "retryable": False})
        self.assertEqual(failed["status"], "failed")
        self.assertIsNone(failed["result"])

    def test_audit_has_mutation_and_rejection(self):
        task = self.running()
        self.rejected("forbidden", self.bus.claim, self.b, task["task_id"], 2)
        trail = self.bus.audit(self.c, task["task_id"])
        self.assertEqual(trail[-1], ("worker-b", "claim", "forbidden", None))
        self.assertEqual([r[1] for r in trail], ["create", "claim", "claim"])

    def test_get_is_read_only_and_returns_copy(self):
        task = self.create()
        before = self.bus.db.total_changes
        view = self.bus.get(self.a, task["task_id"])
        view["status"] = "succeeded"
        self.assertEqual(self.bus.db.total_changes, before)
        self.assertEqual(self.bus.get(self.a, task["task_id"])["status"],
                         "queued")


if __name__ == "__main__":
    unittest.main(verbosity=2)