"""Tests for multi_agent_toolkit.py. Run: python test_multi_agent_toolkit.py -v"""
import io
import json
import os
import re
import sys
import threading
import unittest

import multi_agent_toolkit as m


def say(text):
    return lambda h, c: {"say": text}


class TestMocks(unittest.TestCase):
    def test_token_proxy_counts_words_and_punctuation(self):
        self.assertEqual(m.count_tokens("Hello, world!"), 4)
        self.assertEqual(m.count_tokens(""), 0)

    def test_scripted_llm_replays_and_records(self):
        llm = m.ScriptedLLM(["a", "b"])
        self.assertEqual([llm("p1"), llm("p2")], ["a", "b"])
        self.assertEqual(llm.calls, ["p1", "p2"])

    def test_scripted_llm_raises_when_script_ends(self):
        with self.assertRaises(RuntimeError):
            m.ScriptedLLM([])("x")

    def test_cost_meter_adds_up(self):
        c = m.CostMeter()
        c.add("one two", "three")
        c.add("a", "b c d")
        self.assertEqual((c.calls, c.tokens_in, c.tokens_out, c.tokens), (2, 3, 4, 7))


class TestAgents(unittest.TestCase):
    def setUp(self):
        self.billing = m.Agent("billing", say("refund done"))

    def test_handoff_tool_is_named_transfer_to_agent(self):
        a = m.Agent("triage", say("x"), handoffs=[self.billing])
        self.assertIn("transfer_to_billing", a.tools)

    def test_handoff_moves_control_and_final_speaker_is_specialist(self):
        pol = lambda h, c: {"call": "transfer_to_billing"} if not any(x["role"] == "tool" for x in h) else {"say": "no"}
        r = m.run_swarm(m.Agent("triage", pol, handoffs=[self.billing]), [{"role": "user", "content": "refund"}])
        self.assertEqual((r.agent, r.reply, r.handoffs), ("billing", "refund done", [("triage", "billing")]))

    def test_swarm_keeps_the_shared_history_across_handoffs(self):
        seen = []

        def spy(h, c):
            seen.append(len(h))
            return {"say": "ok"}

        pol = lambda h, c: {"call": "transfer_to_spy"} if not any(x["role"] == "tool" for x in h) else {"say": "no"}
        spy_agent = m.Agent("spy", spy)
        m.run_swarm(m.Agent("t", pol, handoffs=[spy_agent]), [{"role": "user", "content": "hi"}])
        self.assertGreaterEqual(seen[0], 3)  # user message + call + tool result are all visible

    def test_history_filter_limits_what_the_next_agent_sees(self):
        seen = []

        def spy(h, c):
            seen.append(list(h))
            return {"say": "ok"}

        pol = lambda h, c: {"call": "transfer_to_spy"} if not any(x["role"] == "tool" for x in h) else {"say": "no"}
        m.run_swarm(m.Agent("t", pol, handoffs=[m.Agent("spy", spy)]), [{"role": "user", "content": "secret"}],
                    history_filter=lambda h: [x for x in h if x["role"] == "user"])
        self.assertEqual([x["role"] for x in seen[0]], ["user"])

    def test_agent_as_tool_keeps_control_with_the_caller(self):
        def pol(h, c):
            if not any(x["role"] == "tool" for x in h):
                return {"call": "billing", "args": {"request": "refund 5"}}
            return {"say": "I got: " + h[-1]["content"]}

        caller = m.Agent("caller", pol, tools={"billing": self.billing.as_tool()})
        hist = [{"role": "user", "content": "go"}]
        kind, text = m.run_agent(caller, hist)
        self.assertEqual((kind, text), ("say", "I got: refund done"))
        self.assertEqual(hist[-1]["agent"], "caller")

    def test_agent_as_tool_sees_only_the_request_string(self):
        seen = []
        spy = m.Agent("spy", lambda h, c: (seen.append(list(h)), {"say": "ok"})[1])
        spy.as_tool()({}, "just this")
        self.assertEqual(seen[0], [{"role": "user", "content": "just this"}])

    def test_agent_as_tool_cannot_hand_off(self):
        pol = lambda h, c: {"call": "transfer_to_billing"}
        out = m.Agent("x", pol, handoffs=[self.billing]).as_tool()({}, "hi")
        self.assertTrue(out.startswith("error"))

    def test_unknown_tool_becomes_an_error_message_not_a_crash(self):
        calls = iter([{"call": "nope"}, {"say": "done"}])
        hist = [{"role": "user", "content": "x"}]
        m.run_agent(m.Agent("a", lambda h, c: next(calls)), hist)
        self.assertIn("unknown tool", hist[2]["content"])

    def test_policy_that_never_speaks_hits_the_step_limit(self):
        with self.assertRaises(m.MaxStepsExceeded):
            m.run_agent(m.Agent("a", lambda h, c: {"call": "noop"}, tools={"noop": lambda ctx: "x"}), [], max_steps=3)

    def test_bad_policy_output_is_rejected(self):
        with self.assertRaises(ValueError):
            m.run_agent(m.Agent("a", lambda h, c: {}), [])

    def test_ping_pong_handoff_loop_is_stopped(self):
        a, b = m.Agent("a", lambda h, c: {"call": "transfer_to_b"}), m.Agent("b", lambda h, c: {"call": "transfer_to_a"})
        a.tools["transfer_to_b"] = lambda ctx: m.Handoff(b)
        b.tools["transfer_to_a"] = lambda ctx: m.Handoff(a)
        with self.assertRaises(m.MaxStepsExceeded):
            m.run_swarm(a, [], max_turns=4)

    def test_tool_results_are_added_to_history_in_order(self):
        calls = iter([{"call": "t", "args": {"x": 2}}, {"say": "fin"}])
        hist = []
        m.run_agent(m.Agent("a", lambda h, c: next(calls), tools={"t": lambda ctx, x: x * 3}), hist)
        self.assertEqual([h["role"] for h in hist], ["assistant", "tool", "assistant"])
        self.assertEqual(hist[1]["content"], "6")


class TestOrchestrator(unittest.TestCase):
    plan = staticmethod(lambda t: [("a", "do 1"), ("b", "do 2"), ("c", "do 3")])
    synth = staticmethod(lambda t, rs: "+".join(r or "?" for r in rs))

    def test_results_keep_plan_order_even_in_parallel(self):
        ws = {"a": lambda i: "A", "b": lambda i: "B", "c": lambda i: "C"}
        self.assertEqual(m.run_orchestrator("t", self.plan, ws, self.synth).answer, "A+B+C")

    def test_failed_worker_is_retried_once(self):
        n = {"k": 0}

        def flaky(i):
            n["k"] += 1
            if n["k"] == 1:
                raise m.WorkerError("boom")
            return "B"

        r = m.run_orchestrator("t", self.plan, {"a": lambda i: "A", "b": flaky, "c": lambda i: "C"}, self.synth, parallel=False)
        self.assertEqual((r.answer, r.errors, n["k"]), ("A+B+C", [], 2))

    def test_worker_that_always_fails_is_reported_and_others_survive(self):
        def bad(i):
            raise m.WorkerError("down")

        r = m.run_orchestrator("t", self.plan, {"a": lambda i: "A", "b": bad, "c": lambda i: "C"}, self.synth)
        self.assertEqual(r.answer, "A+?+C")
        self.assertEqual(len(r.errors), 1)
        self.assertIn("b: down", r.errors[0])

    def test_max_retries_zero_means_no_second_try(self):
        n = {"k": 0}

        def bad(i):
            n["k"] += 1
            raise m.WorkerError("x")

        m.run_orchestrator("t", self.plan, {"a": bad, "b": bad, "c": bad}, self.synth, max_retries=0, parallel=False)
        self.assertEqual(n["k"], 3)

    def test_meter_counts_messages_both_ways(self):
        ws = {"a": lambda i: "A", "b": lambda i: "B", "c": lambda i: "C"}
        self.assertEqual(m.run_orchestrator("t", self.plan, ws, self.synth).meter.messages, 6)

    def test_workers_really_run_in_parallel_threads(self):
        barrier = threading.Barrier(3)

        def w(i):
            barrier.wait(timeout=3)
            return "x"

        r = m.run_orchestrator("t", self.plan, {"a": w, "b": w, "c": w}, self.synth, parallel=True)
        self.assertEqual(r.errors, [])

    def test_hybrid_workers_see_peer_notes_in_round_two(self):
        board = m.MessageBoard()
        seen = {}

        def worker(name):
            def f(ins, notes):
                seen.setdefault(name, []).append(list(notes))
                return f"{name}-note"
            return f

        m.run_hybrid("t", lambda t: [("a", "x"), ("b", "y")], {"a": worker("a"), "b": worker("b")}, board,
                     lambda t, rs: "|".join(rs), rounds=2)
        self.assertEqual(seen["a"][0], [])
        self.assertEqual(seen["a"][1], ["b-note"])


class TestBoards(unittest.TestCase):
    def test_subscriber_sees_only_its_topics_and_not_its_own_messages(self):
        b = m.MessageBoard()
        b.subscribe("x", ["t1"]); b.subscribe("y", ["t1", "t2"])
        b.publish("x", "t1", "hello"); b.publish("z", "t2", "other")
        self.assertEqual([x["content"] for x in b.read_new("x")], [])
        self.assertEqual([x["content"] for x in b.read_new("y")], ["hello", "other"])

    def test_read_new_advances_the_cursor(self):
        b = m.MessageBoard(); b.subscribe("x", ["t"])
        b.publish("y", "t", 1)
        self.assertEqual(len(b.read_new("x")), 1)
        self.assertEqual(b.read_new("x"), [])
        b.publish("y", "t", 2)
        self.assertEqual([x["content"] for x in b.read_new("x")], [2])

    def test_direct_message_reaches_only_the_addressee(self):
        b = m.MessageBoard(); b.subscribe("x", []); b.subscribe("y", ["t"])
        b.publish("z", "t", "private", to="x")
        self.assertEqual(len(b.read_new("x")), 1)
        self.assertEqual(len(b.read_new("y")), 0)

    def test_schema_rejects_free_text_and_missing_keys(self):
        b = m.MessageBoard(schemas={"prd": {"goals"}})
        with self.assertRaises(m.SchemaError):
            b.publish("pm", "prd", "text")
        with self.assertRaises(m.SchemaError):
            b.publish("pm", "prd", {"other": 1})
        self.assertEqual(b.publish("pm", "prd", {"goals": 1}), 1)

    def test_message_ids_are_sequential_and_cause_by_is_kept(self):
        b = m.MessageBoard()
        b.publish("a", "t", 1); b.publish("b", "t", 2, cause_by=1)
        self.assertEqual([x["id"] for x in b.history()], [1, 2])
        self.assertEqual(b.history("t")[1]["cause_by"], 1)

    def test_board_is_safe_under_many_threads(self):
        b = m.MessageBoard()
        ts = [threading.Thread(target=lambda: [b.publish("a", "t", i) for i in range(50)]) for _ in range(4)]
        [t.start() for t in ts]; [t.join() for t in ts]
        self.assertEqual(sorted(x["id"] for x in b.history()), list(range(1, 201)))


class TestTaskBoard(unittest.TestCase):
    def test_claim_is_exclusive(self):
        tb = m.TaskBoard(); tb.add("t1", "x")
        self.assertEqual(tb.claim("a"), "t1")
        self.assertIsNone(tb.claim("b"))

    def test_dependencies_block_claims_until_done(self):
        tb = m.TaskBoard(); tb.add("t1", "x"); tb.add("t2", "y", deps=["t1"])
        self.assertEqual(tb.claim("a"), "t1")
        self.assertIsNone(tb.claim("b"))
        tb.complete("t1", "a", "r")
        self.assertEqual(tb.claim("b"), "t2")

    def test_only_the_owner_can_complete(self):
        tb = m.TaskBoard(); tb.add("t1", "x"); tb.claim("a")
        with self.assertRaises(PermissionError):
            tb.complete("t1", "b", "r")

    def test_duplicate_task_ids_are_rejected(self):
        tb = m.TaskBoard(); tb.add("t1", "x")
        with self.assertRaises(KeyError):
            tb.add("t1", "y")

    def test_swarm_finishes_all_tasks_with_no_duplicate_claims(self):
        tb = m.TaskBoard()
        for i in range(12):
            tb.add(f"t{i:02d}", i)
        errs = m.run_board_swarm({n: (lambda p, r: p * 2) for n in "abcd"}, tb)
        self.assertEqual(errs, [])
        self.assertEqual(len(tb.claims), 12)
        self.assertEqual(len({t for t, _ in tb.claims}), 12)
        self.assertEqual(tb.results()["t05"], 10)

    def test_swarm_respects_a_dependency_chain(self):
        tb = m.TaskBoard()
        tb.add("t1", 1); tb.add("t2", 2, deps=["t1"]); tb.add("t3", 3, deps=["t2"])
        order = []
        m.run_board_swarm({n: (lambda p, r: order.append(p) or p) for n in "ab"}, tb)
        self.assertEqual(order, [1, 2, 3])

    def test_dependency_cycle_is_reported_not_hung(self):
        tb = m.TaskBoard(); tb.add("t1", 1, deps=["t2"]); tb.add("t2", 2, deps=["t1"])
        errs = m.run_board_swarm({"a": lambda p, r: p}, tb, timeout=2)
        self.assertIn("deadlock: dependency cycle", errs)

    def test_agent_failure_is_recorded(self):
        tb = m.TaskBoard(); tb.add("t1", 1)

        def bad(p, r):
            raise ValueError("oops")

        errs = m.run_board_swarm({"a": bad}, tb, timeout=2)
        self.assertTrue(any("oops" in e for e in errs))


def make_server(handler=None, **kw):
    ids = iter(range(1, 10000))
    card = m.make_agent_card("A", "d", "https://a.example/a2a", [{"id": "s", "name": "S", "description": "d", "tags": ["t"]}], **kw)
    h = handler or (lambda text, info: ("completed", "echo: " + text))
    return m.A2AServer(card, h, clock=lambda: "2026-01-01T00:00:00Z", id_factory=lambda: f"id-{next(ids)}")


def rpc(server, method, params, id=1, headers=None):
    return server.handle({"jsonrpc": "2.0", "id": id, "method": method, "params": params}, headers)


def msg(text="hi", **kw):
    return m.text_message(text, message_id=kw.pop("message_id", "m-" + text), **kw)


class TestAgentCard(unittest.TestCase):
    def test_made_card_is_valid_and_camel_case(self):
        card = make_server().card
        self.assertEqual(m.validate_agent_card(card), [])
        self.assertEqual(card["supportedInterfaces"][0]["protocolBinding"], "JSONRPC")
        self.assertEqual(card["supportedInterfaces"][0]["protocolVersion"], "1.0")

    def test_missing_required_fields_are_listed(self):
        errs = m.validate_agent_card({"name": "x"})
        for k in ("description", "supportedInterfaces", "version", "capabilities", "skills"):
            self.assertIn(f"missing {k}", errs)

    def test_snake_case_keys_are_flagged(self):
        card = dict(make_server().card, default_input_modes=["text/plain"])
        self.assertTrue(any("snake_case" in e for e in m.validate_agent_card(card)))

    def test_empty_interfaces_and_incomplete_skill_are_flagged(self):
        card = dict(make_server().card, supportedInterfaces=[], skills=[{"id": "x"}])
        errs = m.validate_agent_card(card)
        self.assertIn("supportedInterfaces is empty", errs)
        self.assertIn("skills[0] missing tags", errs)

    def test_card_is_served_at_the_well_known_path_only(self):
        s = make_server()
        self.assertEqual(s.get("/.well-known/agent-card.json")["name"], "A")
        self.assertIsNone(s.get("/agent.json"))

    def test_client_discovery_rejects_a_card_without_jsonrpc(self):
        s = make_server()
        s.card["supportedInterfaces"][0]["protocolBinding"] = "GRPC"
        with self.assertRaises(ValueError):
            m.A2AClient(s).discover()

    def test_task_states_and_terminal_sets_match_the_spec(self):
        self.assertEqual(len(m.TASK_STATES), 9)
        self.assertEqual(m.TERMINAL, {"TASK_STATE_COMPLETED", "TASK_STATE_FAILED", "TASK_STATE_CANCELED", "TASK_STATE_REJECTED"})
        self.assertEqual(m.INTERRUPTED, {"TASK_STATE_INPUT_REQUIRED", "TASK_STATE_AUTH_REQUIRED"})


class TestA2AErrorsAndShapes(unittest.TestCase):
    def test_error_codes_match_spec_table(self):
        want = {"TaskNotFoundError": -32001, "TaskNotCancelableError": -32002, "PushNotificationNotSupportedError": -32003,
                "UnsupportedOperationError": -32004, "ContentTypeNotSupportedError": -32005, "InvalidAgentResponseError": -32006,
                "ExtendedAgentCardNotConfiguredError": -32007, "ExtensionSupportRequiredError": -32008,
                "VersionNotSupportedError": -32009, "JSONParseError": -32700, "InvalidRequestError": -32600,
                "MethodNotFoundError": -32601, "InvalidParamsError": -32602, "InternalError": -32603}
        self.assertEqual({k: v[0] for k, v in m.RPC_ERRORS.items()}, want)

    def test_invalid_json_gives_parse_error_with_null_id(self):
        r = make_server().handle("{not json")
        self.assertEqual((r["error"]["code"], r["id"]), (-32700, None))

    def test_missing_method_or_version_is_invalid_request(self):
        s = make_server()
        self.assertEqual(s.handle({"jsonrpc": "2.0", "id": 1})["error"]["code"], -32600)
        self.assertEqual(s.handle({"jsonrpc": "1.0", "id": 1, "method": "GetTask"})["error"]["code"], -32600)
        self.assertEqual(s.handle([1, 2])["error"]["code"], -32600)

    def test_unknown_method_is_method_not_found_and_id_is_echoed(self):
        r = rpc(make_server(), "message/send", {}, id="abc")
        self.assertEqual((r["error"]["code"], r["id"]), (-32601, "abc"))

    def test_old_style_method_name_is_not_accepted(self):
        self.assertEqual(rpc(make_server(), "message/send", {"message": msg()})["error"]["code"], -32601)

    def test_response_has_result_xor_error(self):
        s = make_server()
        ok = rpc(s, "SendMessage", {"message": msg()})
        bad = rpc(s, "GetTask", {"id": "nope"})
        self.assertIn("result", ok); self.assertNotIn("error", ok)
        self.assertIn("error", bad); self.assertNotIn("result", bad)

    def test_validation_errors_carry_a_bad_request_detail(self):
        r = rpc(make_server(), "SendMessage", {"message": {"role": "ROLE_USER", "parts": []}})
        self.assertEqual(r["error"]["code"], -32602)
        d = r["error"]["data"][0]
        self.assertEqual(d["@type"], "type.googleapis.com/google.rpc.BadRequest")
        fields = {v["field"] for v in d["fieldViolations"]}
        self.assertEqual(fields, {"message.messageId", "message.parts"})

    def test_part_must_have_exactly_one_content_field(self):
        bad = {"messageId": "x", "role": "ROLE_USER", "parts": [{"text": "a", "url": "http://x"}]}
        self.assertEqual(rpc(make_server(), "SendMessage", {"message": bad})["error"]["code"], -32602)

    def test_client_message_must_use_role_user(self):
        bad = dict(msg(), role="ROLE_AGENT")
        self.assertEqual(rpc(make_server(), "SendMessage", {"message": bad})["error"]["code"], -32602)

    def test_params_must_be_an_object(self):
        self.assertEqual(rpc(make_server(), "SendMessage", "oops")["error"]["code"], -32602)

    def test_task_not_found_error_info_matches_the_spec_example(self):
        r = rpc(make_server(), "GetTask", {"id": "zzz"}, id=2)
        d = r["error"]["data"][0]
        self.assertEqual((r["error"]["code"], r["error"]["message"]), (-32001, "Task not found"))
        self.assertEqual((d["@type"], d["reason"], d["domain"]), ("type.googleapis.com/google.rpc.ErrorInfo", "TASK_NOT_FOUND", "a2a-protocol.org"))
        self.assertEqual(d["metadata"]["taskId"], "zzz")

    def test_unsupported_media_type_is_content_type_error(self):
        m_ = msg(); m_["parts"] = [{"text": "x", "mediaType": "image/png"}]
        self.assertEqual(rpc(make_server(), "SendMessage", {"message": m_})["error"]["code"], -32005)

    def test_unsupported_version_header_is_rejected(self):
        r = rpc(make_server(), "SendMessage", {"message": msg()}, headers={"A2A-Version": "9.9"})
        self.assertEqual(r["error"]["code"], -32009)

    def test_supported_version_header_passes(self):
        r = rpc(make_server(), "SendMessage", {"message": msg()}, headers={"A2A-Version": "1.0"})
        self.assertIn("result", r)

    def test_streaming_without_capability_is_unsupported_operation(self):
        self.assertEqual(rpc(make_server(), "SendStreamingMessage", {"message": msg()})["error"]["code"], -32004)

    def test_known_but_unimplemented_method_is_unsupported_not_unknown(self):
        self.assertEqual(rpc(make_server(), "SubscribeToTask", {"id": "x"})["error"]["code"], -32004)

    def test_handler_exception_marks_the_task_failed(self):
        def boom(t, i):
            raise RuntimeError("x")
        r = rpc(make_server(boom), "SendMessage", {"message": msg()})
        self.assertEqual(r["result"]["task"]["status"]["state"], "TASK_STATE_FAILED")

    def test_handler_returning_a_bad_kind_is_invalid_agent_response(self):
        r = rpc(make_server(lambda t, i: ("weird", "x")), "SendMessage", {"message": msg()})
        self.assertEqual(r["error"]["code"], -32006)

    def test_client_raises_a2a_error_with_the_right_name(self):
        c = m.A2AClient(make_server())
        with self.assertRaises(m.A2AError) as cm:
            c.get_task("nope")
        self.assertEqual((cm.exception.name, cm.exception.code), ("TaskNotFoundError", -32001))


class TestA2ATasks(unittest.TestCase):
    def test_completed_task_shape(self):
        r = rpc(make_server(), "SendMessage", {"message": msg("weather")})["result"]
        t = r["task"]
        self.assertEqual(list(r), ["task"])
        self.assertEqual(t["status"]["state"], "TASK_STATE_COMPLETED")
        self.assertEqual(t["artifacts"][0]["parts"], [{"text": "echo: weather"}])
        self.assertIn("artifactId", t["artifacts"][0])
        self.assertEqual(t["status"]["timestamp"], "2026-01-01T00:00:00Z")
        self.assertEqual([h["role"] for h in t["history"]], ["ROLE_USER", "ROLE_AGENT"])

    def test_server_generates_task_and_context_ids(self):
        t = rpc(make_server(), "SendMessage", {"message": msg()})["result"]["task"]
        self.assertTrue(t["id"].startswith("id-") and t["contextId"].startswith("id-"))
        self.assertNotEqual(t["id"], t["contextId"])

    def test_client_context_id_is_kept(self):
        t = rpc(make_server(), "SendMessage", {"message": msg(context_id="ctx-1")})["result"]["task"]
        self.assertEqual(t["contextId"], "ctx-1")

    def test_input_required_then_follow_up_completes_the_same_task(self):
        def h(text, info):
            return ("input_required", "from where?") if len(info["history"]) == 1 else ("completed", "booked")
        s = make_server(h)
        t1 = rpc(s, "SendMessage", {"message": msg("book")})["result"]["task"]
        self.assertEqual(t1["status"]["state"], "TASK_STATE_INPUT_REQUIRED")
        self.assertEqual(t1["status"]["message"]["role"], "ROLE_AGENT")
        t2 = rpc(s, "SendMessage", {"message": msg("SFO", task_id=t1["id"])})["result"]["task"]
        self.assertEqual((t2["id"], t2["status"]["state"]), (t1["id"], "TASK_STATE_COMPLETED"))
        self.assertEqual(t2["contextId"], t1["contextId"])

    def test_unknown_task_id_in_message_is_task_not_found(self):
        self.assertEqual(rpc(make_server(), "SendMessage", {"message": msg(task_id="ghost")})["error"]["code"], -32001)

    def test_mismatched_context_id_is_rejected(self):
        s = make_server(lambda t, i: ("input_required", "q"))
        t = rpc(s, "SendMessage", {"message": msg("a")})["result"]["task"]
        r = rpc(s, "SendMessage", {"message": msg("b", task_id=t["id"], context_id="other")})
        self.assertEqual(r["error"]["code"], -32602)

    def test_message_to_a_terminal_task_is_unsupported_operation(self):
        s = make_server()
        t = rpc(s, "SendMessage", {"message": msg("a")})["result"]["task"]
        r = rpc(s, "SendMessage", {"message": msg("b", task_id=t["id"])})
        self.assertEqual(r["error"]["code"], -32004)

    def test_duplicate_message_id_is_idempotent(self):
        calls = []
        s = make_server(lambda t, i: (calls.append(t), ("completed", "x"))[1])
        a = rpc(s, "SendMessage", {"message": msg("a", message_id="same")})
        b = rpc(s, "SendMessage", {"message": msg("a", message_id="same")})
        self.assertEqual((len(calls), a["result"]["task"]["id"]), (1, b["result"]["task"]["id"]))

    def test_failed_and_rejected_states_carry_a_message(self):
        for kind, state in (("failed", "TASK_STATE_FAILED"), ("rejected", "TASK_STATE_REJECTED")):
            t = rpc(make_server(lambda t_, i, k=kind: (k, "no")), "SendMessage", {"message": msg()})["result"]["task"]
            self.assertEqual((t["status"]["state"], t["status"]["message"]["parts"][0]["text"]), (state, "no"))

    def test_history_length_zero_omits_history_and_n_keeps_last_n(self):
        s = make_server(lambda t, i: ("input_required", "q"))
        t = rpc(s, "SendMessage", {"message": msg("a")})["result"]["task"]
        rpc(s, "SendMessage", {"message": msg("b", task_id=t["id"])})
        self.assertNotIn("history", rpc(s, "GetTask", {"id": t["id"], "historyLength": 0})["result"])
        self.assertEqual(len(rpc(s, "GetTask", {"id": t["id"], "historyLength": 1})["result"]["history"]), 1)
        self.assertEqual(len(rpc(s, "GetTask", {"id": t["id"]})["result"]["history"]), 4)

    def test_return_immediately_gives_submitted_and_run_pending_finishes(self):
        s = make_server()
        t = rpc(s, "SendMessage", {"message": msg("a"), "configuration": {"returnImmediately": True}})["result"]["task"]
        self.assertEqual(t["status"]["state"], "TASK_STATE_SUBMITTED")
        self.assertEqual(s.run_pending(), 1)
        self.assertEqual(rpc(s, "GetTask", {"id": t["id"]})["result"]["status"]["state"], "TASK_STATE_COMPLETED")

    def test_cancel_a_submitted_task(self):
        s = make_server()
        t = rpc(s, "SendMessage", {"message": msg("a"), "configuration": {"returnImmediately": True}})["result"]["task"]
        r = rpc(s, "CancelTask", {"id": t["id"]})["result"]
        self.assertEqual(r["status"]["state"], "TASK_STATE_CANCELED")
        s.run_pending()
        self.assertEqual(rpc(s, "GetTask", {"id": t["id"]})["result"]["status"]["state"], "TASK_STATE_CANCELED")

    def test_cancel_a_completed_task_is_not_cancelable(self):
        s = make_server()
        t = rpc(s, "SendMessage", {"message": msg("a")})["result"]["task"]
        self.assertEqual(rpc(s, "CancelTask", {"id": t["id"]})["error"]["code"], -32002)

    def test_cancel_unknown_task_and_missing_id(self):
        s = make_server()
        self.assertEqual(rpc(s, "CancelTask", {"id": "x"})["error"]["code"], -32001)
        self.assertEqual(rpc(s, "CancelTask", {})["error"]["code"], -32602)
        self.assertEqual(rpc(s, "GetTask", {})["error"]["code"], -32602)

    def test_list_tasks_filters_and_pages(self):
        s = make_server()
        for i in range(5):
            rpc(s, "SendMessage", {"message": msg(f"q{i}", context_id="c1" if i < 3 else "c2")})
        r = rpc(s, "ListTasks", {"contextId": "c1"})["result"]
        self.assertEqual((r["totalSize"], len(r["tasks"]), r["nextPageToken"]), (3, 3, ""))
        p1 = rpc(s, "ListTasks", {"pageSize": 2})["result"]
        p2 = rpc(s, "ListTasks", {"pageSize": 2, "pageToken": p1["nextPageToken"]})["result"]
        self.assertEqual((len(p1["tasks"]), p1["nextPageToken"], len(p2["tasks"])), (2, "2", 2))
        self.assertEqual(rpc(s, "ListTasks", {"status": "TASK_STATE_FAILED"})["result"]["totalSize"], 0)

    def test_list_tasks_hides_artifacts_unless_asked(self):
        s = make_server(); rpc(s, "SendMessage", {"message": msg()})
        self.assertNotIn("artifacts", rpc(s, "ListTasks", {})["result"]["tasks"][0])
        self.assertIn("artifacts", rpc(s, "ListTasks", {"includeArtifacts": True})["result"]["tasks"][0])

    def test_list_tasks_page_size_bounds(self):
        for bad in (0, 101, "x"):
            self.assertEqual(rpc(make_server(), "ListTasks", {"pageSize": bad})["error"]["code"], -32602)

    def test_client_round_trip_and_wire_log(self):
        s = make_server(); c = m.A2AClient(s); c.discover()
        t = c.send("hello")
        self.assertEqual(m.task_text(t), "echo: hello")
        self.assertEqual(c.log[0][0]["method"], "SendMessage")
        self.assertEqual(c.log[0][1]["id"], c.log[0][0]["id"])

    def test_a2a_agent_as_a_tool(self):
        c = m.A2AClient(make_server())
        self.assertEqual(m.a2a_tool(c)({}, "ping"), "echo: ping")
        c2 = m.A2AClient(make_server(lambda t, i: ("input_required", "which one?")))
        self.assertEqual(m.a2a_tool(c2)({}, "ping"), "[TASK_STATE_INPUT_REQUIRED] which one?")

    def test_every_response_is_json_serialisable(self):
        s = make_server()
        for r in (rpc(s, "SendMessage", {"message": msg()}), rpc(s, "GetTask", {"id": "x"}), rpc(s, "ListTasks", {})):
            json.dumps(r)

    def test_returned_task_is_a_copy(self):
        s = make_server()
        t = rpc(s, "SendMessage", {"message": msg()})["result"]["task"]
        t["status"]["state"] = "HACKED"
        self.assertEqual(rpc(s, "GetTask", {"id": t["id"]})["result"]["status"]["state"], "TASK_STATE_COMPLETED")


class TestLedgers(unittest.TestCase):
    def planner_factory(self, plans):
        def planner(task, ledger, reflection):
            v = 1 if ledger is None else ledger.version + 1
            return m.TaskLedger(task, plan=plans[min(v, len(plans)) - 1], version=v, guesses=["a guess"],
                                facts_given=([reflection] if reflection else []))
        return planner

    def test_happy_path_needs_no_replan(self):
        team = {"A": lambda i: "r1", "B": lambda i: "r2"}
        run = m.MagenticOrchestrator(team, self.planner_factory([[("A", "x"), ("B", "y")]])).run("t")
        self.assertEqual((run.status, run.answer, run.resets, len(run.ledgers)), ("done", "r2", 0, 1))

    def test_progress_entry_answers_the_five_questions(self):
        led = m.TaskLedger("t", plan=[("A", "x")])
        e = m.plan_follower_progress(led, [])
        self.assertEqual((e.satisfied, e.looping, e.progress, e.next_speaker, e.instruction), (False, False, True, "A", "x"))

    def test_loop_detection_needs_two_identical_entries(self):
        led = m.TaskLedger("t", plan=[("A", "x")])
        c = {"speaker": "A", "instruction": "x", "result": "ERROR", "ok": False}
        self.assertFalse(m.plan_follower_progress(led, [c]).looping)
        self.assertTrue(m.plan_follower_progress(led, [c, dict(c)]).looping)

    def test_stuck_agent_triggers_one_replan_then_succeeds(self):
        team = {"Bad": lambda i: "ERROR: nope", "Good": lambda i: "fine"}
        run = m.MagenticOrchestrator(team, self.planner_factory([[("Bad", "x")], [("Good", "x")]])).run("t")
        self.assertEqual((run.status, run.resets, len(run.ledgers)), ("done", 1, 2))
        self.assertEqual(run.ledgers[1].version, 2)
        self.assertIn("stalled", run.ledgers[1].facts_given[0])

    def test_stall_threshold_controls_how_many_failures_are_tolerated(self):
        team = {"Bad": lambda i: "ERROR", "Good": lambda i: "ok"}
        plans = [[("Bad", "x")], [("Good", "x")]]
        calls = []
        team["Bad"] = lambda i: (calls.append(1), "ERROR")[1]
        m.MagenticOrchestrator(team, self.planner_factory(plans), stall_threshold=0).run("t")
        n0 = len(calls); calls.clear()
        m.MagenticOrchestrator(team, self.planner_factory(plans), stall_threshold=4).run("t")
        self.assertLess(n0, len(calls))

    def test_a_failing_agent_is_called_threshold_plus_one_times_before_the_replan(self):
        for th in (0, 2, 4):
            calls = []
            team = {"Bad": lambda i: (calls.append(1), "ERROR")[1], "Good": lambda i: "ok"}
            run = m.MagenticOrchestrator(team, self.planner_factory([[("Bad", "x")], [("Good", "x")]]), stall_threshold=th).run("t")
            self.assertEqual((len(calls), run.resets, run.status), (th + 1, 1, "done"), th)

    def test_replan_leaves_a_marker_in_the_conversation_so_the_new_plan_starts_fresh(self):
        team = {"Bad": lambda i: "ERROR", "Good": lambda i: "ok"}
        run = m.MagenticOrchestrator(team, self.planner_factory([[("Bad", "x")], [("Good", "x")]])).run("t")
        self.assertEqual([c["speaker"] for c in run.conversation].count("orchestrator"), 1)
        self.assertEqual(run.conversation[-1]["speaker"], "Good")

    def test_gives_up_after_max_resets_with_the_best_guess(self):
        team = {"Bad": lambda i: "ERROR"}
        run = m.MagenticOrchestrator(team, self.planner_factory([[("Bad", "x")]]), max_resets=1).run("t")
        self.assertEqual((run.status, run.answer, run.resets, len(run.ledgers)), ("gave_up", "a guess", 2, 2))

    def test_max_steps_guard(self):
        team = {"A": lambda i: "ok"}
        orch = m.MagenticOrchestrator(team, self.planner_factory([[("A", "x")]]), progress_fn=lambda l, c: m.ProgressEntry(
            len(c) + 1, False, False, True, "A", "x"), max_steps=3)
        self.assertEqual(orch.run("t").status, "max_steps")

    def test_progress_ledger_records_every_inner_step(self):
        team = {"A": lambda i: "r1", "B": lambda i: "r2"}
        run = m.MagenticOrchestrator(team, self.planner_factory([[("A", "x"), ("B", "y")]])).run("t")
        self.assertEqual([p.next_speaker for p in run.progress], ["A", "B", None])
        self.assertTrue(run.progress[-1].satisfied)

    def test_ledger_render_lists_facts_and_plan(self):
        led = m.TaskLedger("find x", facts_given=["g"], facts_to_lookup=["l"], facts_to_derive=["d"], guesses=["e"], plan=[("A", "go")])
        txt = led.render()
        for s in ("TASK LEDGER v1", "given or verified facts", "facts to look up", "facts to derive", "educated guesses", "plan 1. A: go"):
            self.assertIn(s, txt)


class TestSharedMemory(unittest.TestCase):
    def test_versions_increase_and_get_returns_value_and_version(self):
        mem = m.SharedMemory()
        self.assertEqual(mem.get("k", "dflt"), ("dflt", 0))
        self.assertEqual(mem.put("k", 1, "a"), 1)
        self.assertEqual(mem.put("k", 2, "b"), 2)
        self.assertEqual(mem.get("k"), (2, 2))

    def test_stale_write_is_rejected(self):
        mem = m.SharedMemory(); mem.put("k", 1, "a")
        with self.assertRaises(m.VersionConflict):
            mem.put("k", 2, "b", expected_version=0)
        self.assertEqual(mem.put("k", 2, "b", expected_version=1), 2)

    def test_write_with_expected_zero_creates_exactly_once(self):
        mem = m.SharedMemory()
        mem.put("k", 1, "a", expected_version=0)
        with self.assertRaises(m.VersionConflict):
            mem.put("k", 1, "b", expected_version=0)

    def test_atomic_update_loses_nothing_under_threads(self):
        mem = m.SharedMemory()
        ts = [threading.Thread(target=lambda: [mem.update("n", lambda v: v + 1, "w", 0) for _ in range(100)]) for _ in range(4)]
        [t.start() for t in ts]; [t.join() for t in ts]
        self.assertEqual(mem.get("n")[0], 400)

    def test_log_records_authors_and_snapshot_is_a_copy(self):
        mem = m.SharedMemory(); mem.put("a", 1, "x"); mem.put("b", 2, "y")
        self.assertEqual(mem.log, [("a", 1, "x"), ("b", 1, "y")])
        snap = mem.snapshot(); snap["a"] = 99
        self.assertEqual(mem.get("a")[0], 1)

    def test_unlocked_race_loses_updates_and_the_lock_fixes_it(self):
        self.assertEqual(m.race_counter(False, workers=4), 1)
        self.assertEqual(m.race_counter(True, workers=4), 4)

    def test_team_memory_recalls_similar_tasks_with_insights(self):
        tm = m.TeamMemory()
        q1 = tm.record("migrate python backend to rust", ["plan", "port"], "ok")
        tm.record("write a poem about autumn", ["draft"], "ok")
        tm.distill(q1, "port tests first")
        got = tm.recall("migrate the backend to go")
        self.assertEqual([q["id"] for q in got["queries"]], [q1])
        self.assertEqual(got["insights"], ["port tests first"])
        self.assertEqual(got["interactions"], ["plan", "port"])

    def test_team_memory_recall_of_unrelated_task_is_empty(self):
        tm = m.TeamMemory(); tm.record("poem", ["x"], "ok")
        self.assertEqual(tm.recall("compiler")["queries"], [])

    def test_distilling_an_unknown_query_fails(self):
        with self.assertRaises(KeyError):
            m.TeamMemory().distill(9, "x")


class TestMerge3(unittest.TestCase):
    def lines(self, *xs):
        return [x + "\n" for x in xs]

    def test_no_changes(self):
        b = self.lines("a", "b")
        self.assertEqual(m.merge3(b, b, b), (b, 0))

    def test_one_sided_change_is_taken(self):
        b = self.lines("a", "b", "c")
        o = self.lines("a", "B", "c")
        self.assertEqual(m.merge3(b, o, b), (o, 0))
        self.assertEqual(m.merge3(b, b, o), (o, 0))

    def test_changes_far_apart_merge_cleanly(self):
        b = self.lines("a", "b", "c", "d", "e")
        merged, n = m.merge3(b, self.lines("A", "b", "c", "d", "e"), self.lines("a", "b", "c", "d", "E"))
        self.assertEqual((merged, n), (self.lines("A", "b", "c", "d", "E"), 0))

    def test_same_change_on_both_sides_is_not_a_conflict(self):
        b = self.lines("a", "b"); o = self.lines("a", "X")
        self.assertEqual(m.merge3(b, o, o), (o, 0))

    def test_different_changes_to_the_same_line_conflict_with_markers(self):
        b = self.lines("a", "b", "c")
        merged, n = m.merge3(b, self.lines("a", "X", "c"), self.lines("a", "Y", "c"))
        self.assertEqual(n, 1)
        self.assertEqual("".join(merged), "a\n<<<<<<< ours\nX\n=======\nY\n>>>>>>> theirs\nc\n")

    def test_touching_changes_conflict(self):
        b = self.lines("a", "b", "c")
        _, n = m.merge3(b, self.lines("X", "b", "c"), self.lines("a", "Y", "c"))
        self.assertEqual(n, 1)

    def test_one_line_apart_still_merges(self):
        b = self.lines("a", "b", "c")
        merged, n = m.merge3(b, self.lines("X", "b", "c"), self.lines("a", "b", "Z"))
        self.assertEqual((merged, n), (self.lines("X", "b", "Z"), 0))

    def test_two_different_insertions_at_the_same_point_conflict(self):
        b = self.lines("a", "c")
        _, n = m.merge3(b, self.lines("a", "b1", "c"), self.lines("a", "b2", "c"))
        self.assertEqual(n, 1)

    def test_delete_vs_edit_of_the_same_line_conflicts(self):
        b = self.lines("a", "b", "c")
        _, n = m.merge3(b, self.lines("a", "c"), self.lines("a", "B", "c"))
        self.assertEqual(n, 1)

    def test_two_clean_edits_and_one_conflict_in_one_file(self):
        b = self.lines("1", "2", "3", "4", "5", "6", "7")
        o = self.lines("one", "2", "3", "four", "5", "6", "7")
        t = self.lines("1", "2", "3", "FOUR", "5", "6", "seven")
        merged, n = m.merge3(b, o, t)
        self.assertEqual(n, 1)
        text = "".join(merged)
        self.assertTrue(text.startswith("one\n2\n3\n<<<<<<< ours\nfour\n=======\nFOUR\n"))
        self.assertTrue(text.endswith("seven\n"))

    def test_empty_base_add_add(self):
        merged, n = m.merge3([], self.lines("x"), self.lines("y"))
        self.assertEqual(n, 1)
        self.assertEqual(m.merge3([], self.lines("x"), self.lines("x")), (self.lines("x"), 0))


class TestWorktrees(unittest.TestCase):
    def setUp(self):
        self.repo = m.Repo({"app.py": "def add(a, b):\n    return a + b\n\n\ndef sub(a, b):\n    return a - b\n", "README": "hi\n"})
        self.addCleanup(self.repo.cleanup)

    def test_fork_copies_main_and_edits_do_not_leak_back(self):
        ws = self.repo.fork("a")
        ws.write("app.py", "changed\n")
        self.assertEqual(self.repo.snapshot()["app.py"].splitlines()[0], "def add(a, b):")
        self.assertEqual(ws.read("README"), "hi\n")

    def test_changes_report_added_modified_deleted(self):
        ws = self.repo.fork("a")
        ws.write("app.py", "x\n"); ws.write("new.txt", "n\n"); ws.delete("README")
        self.assertEqual(ws.changes(), {"app.py": "M", "new.txt": "A", "README": "D"})

    def test_duplicate_workspace_name_is_rejected(self):
        self.repo.fork("a")
        with self.assertRaises(KeyError):
            self.repo.fork("a")

    def test_path_escape_is_blocked(self):
        ws = self.repo.fork("a")
        with self.assertRaises(ValueError):
            ws.write("../evil.txt", "x")

    def test_edits_to_different_functions_merge_cleanly(self):
        a, b = self.repo.fork("a"), self.repo.fork("b")
        a.write("app.py", a.read("app.py").replace("return a + b", "return a + b  # a"))
        b.write("app.py", b.read("app.py").replace("return a - b", "return a - b  # b"))
        rep = m.merge_workspaces(self.repo, ["a", "b"])
        self.assertTrue(rep.clean)
        text = self.repo.snapshot()["app.py"]
        self.assertIn("# a", text); self.assertIn("# b", text)

    def test_same_line_edits_conflict_and_are_reported(self):
        a, b = self.repo.fork("a"), self.repo.fork("b")
        a.write("app.py", a.read("app.py").replace("return a - b", "return a - b - 1"))
        b.write("app.py", b.read("app.py").replace("return a - b", "return b - a"))
        rep = m.merge_workspaces(self.repo, ["a", "b"])
        self.assertFalse(rep.clean)
        self.assertIn("app.py", rep.conflicts)
        self.assertIn("<<<<<<< ours", self.repo.snapshot()["app.py"])

    def test_add_add_with_different_content_conflicts_and_same_content_does_not(self):
        a, b = self.repo.fork("a"), self.repo.fork("b")
        a.write("n.txt", "one\n"); b.write("n.txt", "two\n")
        self.assertFalse(m.merge_workspaces(self.repo, ["a", "b"]).clean)
        r2 = m.Repo({}); self.addCleanup(r2.cleanup)
        c, d = r2.fork("c"), r2.fork("d")
        c.write("n.txt", "same\n"); d.write("n.txt", "same\n")
        self.assertTrue(m.merge_workspaces(r2, ["c", "d"]).clean)

    def test_delete_is_applied_when_main_unchanged(self):
        a = self.repo.fork("a"); a.delete("README")
        rep = m.merge_workspaces(self.repo, ["a"])
        self.assertTrue(rep.clean)
        self.assertNotIn("README", self.repo.snapshot())

    def test_delete_vs_modify_conflicts(self):
        a, b = self.repo.fork("a"), self.repo.fork("b")
        a.delete("README"); b.write("README", "changed\n")
        rep = m.merge_workspaces(self.repo, ["b", "a"])
        self.assertIn("README", rep.conflicts)
        self.assertEqual(self.repo.snapshot()["README"], "changed\n")

    def test_modify_vs_delete_in_main_conflicts(self):
        a, b = self.repo.fork("a"), self.repo.fork("b")
        a.delete("README"); b.write("README", "changed\n")
        rep = m.merge_workspaces(self.repo, ["a", "b"])
        self.assertIn("README", rep.conflicts)

    def test_both_delete_is_fine(self):
        a, b = self.repo.fork("a"), self.repo.fork("b")
        a.delete("README"); b.delete("README")
        self.assertTrue(m.merge_workspaces(self.repo, ["a", "b"]).clean)

    def test_clean_text_merge_can_still_fail_the_combined_checks(self):
        a, b = self.repo.fork("a"), self.repo.fork("b")
        a.write("app.py", a.read("app.py").replace("def add(", "def plus("))          # rename
        b.write("app.py", b.read("app.py") + "\n\nprint(add(1, 2))\n")                # new call to the old name
        rep = m.merge_workspaces(self.repo, ["a", "b"],
                                 checks={"no call to a removed name": lambda f: "add(1, 2)" not in f["app.py"] or "def add(" in f["app.py"]})
        self.assertTrue(rep.clean)
        self.assertEqual(rep.checks, {"no call to a removed name": False})

    def test_merge_order_does_not_change_a_clean_result(self):
        a, b = self.repo.fork("a"), self.repo.fork("b")
        a.write("app.py", a.read("app.py").replace("return a + b", "return a + b  # a"))
        b.write("app.py", b.read("app.py").replace("return a - b", "return a - b  # b"))
        m.merge_workspaces(self.repo, ["a", "b"]); one = self.repo.snapshot()
        r2 = m.Repo({"app.py": a.base["app.py"], "README": "hi\n"}); self.addCleanup(r2.cleanup)
        a2, b2 = r2.fork("a"), r2.fork("b")
        a2.write("app.py", a.read("app.py")); b2.write("app.py", b.read("app.py"))
        m.merge_workspaces(r2, ["b", "a"])
        self.assertEqual(one, r2.snapshot())

    def test_without_isolation_the_last_write_wins(self):
        snap = m.shared_dir_demo([("a", "f.txt", "A\n"), ("b", "f.txt", "B\n")])
        self.assertEqual(snap["f.txt"], "B\n")

    def test_merged_files_are_listed_with_who_and_what(self):
        a = self.repo.fork("a"); a.write("new.txt", "n\n")
        rep = m.merge_workspaces(self.repo, ["a"])
        self.assertEqual(rep.merged, [("a", "new.txt", "A")])


class TestFailures(unittest.TestCase):
    spec = m.TeamSpec()

    def test_clean_run_has_no_detections(self):
        self.assertEqual(m.detected_modes(m.run_team(self.spec), self.spec), [])

    def test_there_are_exactly_14_modes_in_3_categories(self):
        self.assertEqual(len(m.MAST_MODES), 14)
        self.assertEqual({v[1] for v in m.MAST_MODES.values()}, {"System design", "Inter-agent misalignment", "Task verification"})
        counts = {}
        for v in m.MAST_MODES.values():
            counts[v[1]] = counts.get(v[1], 0) + 1
        self.assertEqual(counts, {"System design": 5, "Inter-agent misalignment": 6, "Task verification": 3})

    def test_category_follows_the_id_prefix(self):
        for k in m.MAST_MODES:
            self.assertEqual(m.mode_category(k), {"1": "System design", "2": "Inter-agent misalignment", "3": "Task verification"}[k[0]])

    def test_each_injected_fault_is_detected_alone(self):
        for mode in sorted(m.MAST_MODES):
            with self.subTest(mode=mode):
                self.assertIn(mode, m.detected_modes(m.run_team(self.spec, {mode}), self.spec))

    def test_each_fault_is_detected_without_any_unrelated_mode(self):
        for mode in sorted(m.MAST_MODES):
            with self.subTest(mode=mode):
                self.assertEqual(m.detected_modes(m.run_team(self.spec, {mode}), self.spec), [mode])

    def test_two_faults_together_are_both_found(self):
        got = m.detected_modes(m.run_team(self.spec, {"1.3", "2.3"}), self.spec)
        self.assertEqual(got, ["1.3", "2.3"])

    def test_unknown_fault_id_is_an_error(self):
        with self.assertRaises(ValueError):
            m.run_team(self.spec, {"9.9"})

    def test_detections_carry_evidence_and_time(self):
        d = m.classify(m.run_team(self.spec, {"1.2"}), self.spec)[0]
        self.assertEqual(d.mode, "1.2")
        self.assertIn("reviewer", d.evidence)
        self.assertIsInstance(d.t, int)

    def test_trace_events_have_increasing_time_and_known_types(self):
        tr = m.run_team(self.spec, {"1.5"})
        self.assertEqual([e["t"] for e in tr], list(range(len(tr))))
        self.assertTrue({e["type"] for e in tr} <= {"message", "action", "finish"})

    def test_premature_termination_is_not_the_same_as_missing_verification(self):
        self.assertEqual(m.detected_modes(m.run_team(self.spec, {"3.1"}), self.spec), ["3.1"])
        self.assertEqual(m.detected_modes(m.run_team(self.spec, {"3.2"}), self.spec), ["3.2"])

    def test_a_trace_with_no_verification_at_all_is_flagged(self):
        tr = [e for e in m.run_team(self.spec) if e["data"].get("does") != "verify"]
        self.assertIn("3.2", m.detected_modes(tr, self.spec))

    def test_classifier_uses_the_spec_constraints(self):
        loose = m.TeamSpec(constraints={})
        self.assertNotIn("1.1", m.detected_modes(m.run_team(self.spec, {"1.1"}), loose))

    def test_role_table_is_what_decides_a_role_violation(self):
        tr = m.run_team(self.spec)
        wide = m.TeamSpec(roles={"planner": {"plan", "ask", "say"}, "coder": {"implement", "run_tests", "ask", "say"},
                                 "reviewer": {"verify", "implement", "say"}})
        self.assertNotIn("1.2", m.detected_modes(m.run_team(wide, {"1.2"}), wide))
        self.assertIn("1.2", m.detected_modes(m.run_team(self.spec, {"1.2"}), self.spec))
        self.assertEqual(m.detected_modes(tr, self.spec), [])

    def test_failure_profile_counts_modes_and_category_shares(self):
        prof = m.failure_profile([["1.3", "1.5"], ["2.4"], ["1.3"]])
        self.assertEqual(prof["modes"], {"1.3": 2, "1.5": 1, "2.4": 1})
        self.assertEqual(prof["categories"], {"Inter-agent misalignment": 25.0, "System design": 75.0})

    def test_failure_profile_of_nothing_is_empty(self):
        self.assertEqual(m.failure_profile([]), {"modes": {}, "categories": {}})

    def test_injection_study_covers_all_14(self):
        st = m.injection_study()
        self.assertEqual(len(st["rows"]), 14)
        self.assertTrue(all(r["hit"] and not r["extras"] for r in st["rows"]))
        self.assertEqual(st["clean"], [])


class TestSimulation(unittest.TestCase):
    def test_same_seed_gives_same_numbers_and_other_seed_differs(self):
        t = m.SimTask("parallel", 8)
        a, b = m.simulate("single", t, seed=1), m.simulate("single", t, seed=1)
        self.assertEqual(a, b)
        self.assertNotEqual(a["success"], m.simulate("single", t, seed=2)["success"])

    def test_single_agent_success_matches_the_hand_computed_product(self):
        P = m.SimParams()
        exact = 1.0
        for i in range(8):
            exact *= max(0.02, min(1, P.skill - P.dilution * i))
        got = m.simulate("single", m.SimTask("parallel", 8), P, trials=20000, seed=3)["success"]
        self.assertAlmostEqual(got, exact, delta=0.015)

    def test_perfect_skill_gives_perfect_success_for_every_architecture(self):
        P = m.SimParams(skill=1.0, dilution=0.0, shared_hard=0.0, handoff_loss=0.0, swarm_context_loss=0.0)
        for kind in ("parallel", "sequential"):
            for arch in m.ARCHS:
                self.assertEqual(m.simulate(arch, m.SimTask(kind, 6), P, trials=50)["success"], 1.0, (kind, arch))

    def test_zero_skill_floor_keeps_success_near_zero(self):
        P = m.SimParams(skill=0.0)
        self.assertLess(m.simulate("single", m.SimTask("parallel", 4), P, trials=500)["success"], 0.01)

    def test_more_skill_means_more_success(self):
        lo = m.simulate("single", m.SimTask("sequential", 8), m.SimParams(skill=0.85), trials=3000)["success"]
        hi = m.simulate("single", m.SimTask("sequential", 8), m.SimParams(skill=0.97), trials=3000)["success"]
        self.assertGreater(hi, lo)

    def test_parallel_work_finishes_in_fewer_rounds_with_workers(self):
        t = m.SimTask("parallel", 8)
        self.assertLess(m.simulate("orchestrator", t, trials=10)["rounds"], m.simulate("single", t, trials=10)["rounds"])
        self.assertLess(m.simulate("swarm", t, trials=10)["rounds"], m.simulate("single", t, trials=10)["rounds"])

    def test_sequential_work_is_not_faster_with_an_orchestrator(self):
        t = m.SimTask("sequential", 8)
        self.assertGreater(m.simulate("orchestrator", t, trials=10)["rounds"], m.simulate("single", t, trials=10)["rounds"])

    def test_handoff_loss_hurts_the_sequential_orchestrator(self):
        t = m.SimTask("sequential", 8)
        lo = m.simulate("orchestrator", t, m.SimParams(handoff_loss=0.0), trials=4000)["success"]
        hi = m.simulate("orchestrator", t, m.SimParams(handoff_loss=0.3), trials=4000)["success"]
        self.assertGreater(lo, hi + 0.1)

    def test_default_orchestrator_is_worse_than_single_on_sequential_and_better_on_parallel(self):
        seq, par = m.SimTask("sequential", 8), m.SimTask("parallel", 8)
        self.assertLess(m.simulate("orchestrator", seq, trials=4000)["success"], m.simulate("single", seq, trials=4000)["success"])
        self.assertGreater(m.simulate("orchestrator", par, trials=4000)["success"], m.simulate("single", par, trials=4000)["success"])

    def test_verification_helps_and_costs_tokens(self):
        t = m.SimTask("parallel", 8)
        a, b = m.simulate("orchestrator", t, trials=3000), m.simulate("orchestrator+verify", t, trials=3000)
        self.assertGreater(b["success"], a["success"] + 0.2)
        self.assertGreater(b["tokens"], a["tokens"])

    def test_a_single_agent_with_a_check_beats_an_unchecked_team_on_sequential_work(self):
        t = m.SimTask("sequential", 8)
        self.assertGreater(m.simulate("single+verify", t, trials=3000)["success"], m.simulate("orchestrator", t, trials=3000)["success"] + 0.3)

    def test_independent_voting_costs_about_voters_times_a_single_run(self):
        t = m.SimTask("parallel", 8)
        ratio = m.simulate("independent", t, trials=5)["tokens"] / m.simulate("single", t, trials=5)["tokens"]
        self.assertGreater(ratio, 4.5)

    def test_shared_hard_items_cap_the_voting_gain(self):
        t = m.SimTask("parallel", 8)
        free = m.simulate("independent", t, m.SimParams(shared_hard=0.0), trials=3000)["success"]
        capped = m.simulate("independent", t, m.SimParams(shared_hard=0.2), trials=3000)["success"]
        self.assertGreater(free, capped + 0.2)

    def test_multi_agent_systems_use_more_tokens_than_one_agent_by_default(self):
        for kind in ("parallel", "sequential"):
            t = m.SimTask(kind, 8)
            base = m.simulate("single", t, trials=5)["tokens"]
            for arch in ("independent", "orchestrator", "orchestrator+verify", "hybrid"):
                self.assertGreater(m.simulate(arch, t, trials=5)["tokens"], base, (kind, arch))

    def test_rounds_follow_the_simple_formulas(self):
        par, seq = m.SimTask("parallel", 8), m.SimTask("sequential", 8)
        r = lambda a, t: m.simulate(a, t, trials=3)["rounds"]
        self.assertEqual([r("single", par), r("single+verify", par), r("independent", par), r("orchestrator", par), r("swarm", par)],
                         [8, 16, 9, 4, 2])
        self.assertEqual([r("orchestrator", seq), r("orchestrator+verify", seq), r("hybrid", seq), r("swarm", seq)], [17, 17, 17, 8])

    def test_voting_needs_at_least_two_agreeing_voters(self):
        P = m.SimParams(voters=1, shared_hard=0.0, skill=1.0, dilution=0.0)
        self.assertEqual(m.simulate("independent", m.SimTask("parallel", 3), P, trials=20)["success"], 0.0)
        self.assertEqual(m.simulate("independent", m.SimTask("parallel", 3), m.SimParams(voters=2, shared_hard=0.0, skill=1.0, dilution=0.0), trials=20)["success"], 1.0)

    def test_a_broken_chain_hurts_sequential_voting_more_than_parallel_voting(self):
        P = m.SimParams(shared_hard=0.0)
        par = m.simulate("independent", m.SimTask("parallel", 8), P, trials=4000)["success"]
        seq = m.simulate("independent", m.SimTask("sequential", 8), P, trials=4000)["success"]
        self.assertGreater(par, seq + 0.1)

    def test_checking_every_item_costs_at_least_one_message_and_one_step_per_item(self):
        P = m.SimParams(); t = m.SimTask("parallel", 8)
        extra = m.simulate("orchestrator+verify", t, P, trials=3000)["tokens"] - m.simulate("orchestrator", t, P, trials=3000)["tokens"]
        self.assertGreaterEqual(extra, 8 * (P.msg_tokens + P.step_tokens))

    def test_task_validation(self):
        with self.assertRaises(ValueError):
            m.SimTask("weird")
        with self.assertRaises(ValueError):
            m.SimTask("parallel", 0)
        with self.assertRaises(ValueError):
            m.simulate("nope", m.SimTask("parallel", 2), trials=1)

    def test_one_item_task_works_for_all_architectures(self):
        for arch in m.ARCHS:
            r = m.simulate(arch, m.SimTask("parallel", 1), trials=20)
            self.assertTrue(0 <= r["success"] <= 1 and r["rounds"] >= 1)

    def test_comparison_table_has_every_row_and_formats(self):
        rows = m.compare_architectures(trials=50)
        self.assertEqual(len(rows), 2 * len(m.ARCHS))
        self.assertEqual(len(m.format_table(rows).splitlines()), 1 + len(rows))


class TestDemo(unittest.TestCase):
    def test_demo_runs_and_has_all_nine_sections(self):
        buf = []
        m.demo(out=buf.append)
        text = "\n".join(buf)
        for i in range(1, 10):
            self.assertIn(f"== {i}.", text)
        self.assertIn("counter = 1", text)
        self.assertIn("clean run: found nothing", text)

    def test_demo_is_deterministic(self):
        a, b = [], []
        m.demo(out=a.append); m.demo(out=b.append)
        self.assertEqual(a, b)

    def test_command_line_runs(self):
        import subprocess
        r = subprocess.run([sys.executable, "multi_agent_toolkit.py"], capture_output=True, text=True, cwd=os.path.dirname(os.path.abspath(__file__)))
        self.assertEqual(r.returncode, 0)
        self.assertIn("== 9.", r.stdout)


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