from __future__ import annotations import unittest from src.agent_manager import AgentManager, ManagedAgentGroup, ManagedAgentRecord class TestManagedAgentRecordDefaults(unittest.TestCase): def test_defaults(self) -> None: rec = ManagedAgentRecord(agent_id="a1", prompt="do stuff") self.assertEqual(rec.agent_id, "a1") self.assertEqual(rec.prompt, "do stuff") self.assertIsNone(rec.parent_agent_id) self.assertIsNone(rec.group_id) self.assertIsNone(rec.child_index) self.assertIsNone(rec.label) self.assertIsNone(rec.resumed_from_session_id) self.assertIsNone(rec.session_id) self.assertIsNone(rec.session_path) self.assertEqual(rec.status, "running") self.assertEqual(rec.turns, 0) self.assertEqual(rec.tool_calls, 0) self.assertIsNone(rec.stop_reason) class TestManagedAgentGroupDefaults(unittest.TestCase): def test_defaults(self) -> None: grp = ManagedAgentGroup(group_id="g1") self.assertEqual(grp.group_id, "g1") self.assertIsNone(grp.label) self.assertIsNone(grp.parent_agent_id) self.assertEqual(grp.child_agent_ids, ()) self.assertEqual(grp.strategy, "serial") self.assertEqual(grp.status, "running") self.assertEqual(grp.completed_children, 0) self.assertEqual(grp.failed_children, 0) self.assertEqual(grp.batch_count, 0) self.assertEqual(grp.max_batch_size, 0) self.assertEqual(grp.dependency_skips, 0) class TestStartAgent(unittest.TestCase): def setUp(self) -> None: self.mgr = AgentManager() def test_increments_counter_and_returns_unique_ids(self) -> None: id1 = self.mgr.start_agent(prompt="task1") id2 = self.mgr.start_agent(prompt="task2") id3 = self.mgr.start_agent(prompt="task3") self.assertEqual(id1, "agent_1") self.assertEqual(id2, "agent_2") self.assertEqual(id3, "agent_3") self.assertEqual(len(self.mgr.records), 3) def test_record_stored_with_correct_fields(self) -> None: aid = self.mgr.start_agent(prompt="hello", label="my-label") rec = self.mgr.records[aid] self.assertEqual(rec.agent_id, aid) self.assertEqual(rec.prompt, "hello") self.assertEqual(rec.label, "my-label") self.assertEqual(rec.status, "running") def test_with_parent_agent_id_tracks_lineage(self) -> None: parent = self.mgr.start_agent(prompt="parent") child = self.mgr.start_agent(prompt="child", parent_agent_id=parent) rec = self.mgr.records[child] self.assertEqual(rec.parent_agent_id, parent) def test_with_group_id_registers_child(self) -> None: gid = self.mgr.start_group(label="grp") aid = self.mgr.start_agent(prompt="task", group_id=gid, child_index=0) grp = self.mgr.groups[gid] self.assertIn(aid, grp.child_agent_ids) self.assertEqual(self.mgr.records[aid].group_id, gid) self.assertEqual(self.mgr.records[aid].child_index, 0) def test_resumed_agents_tracked(self) -> None: aid = self.mgr.start_agent( prompt="resume", resumed_from_session_id="sess-old-123" ) rec = self.mgr.records[aid] self.assertEqual(rec.resumed_from_session_id, "sess-old-123") class TestStartGroup(unittest.TestCase): def setUp(self) -> None: self.mgr = AgentManager() def test_creates_group_with_strategy(self) -> None: gid = self.mgr.start_group(label="batch", strategy="parallel") self.assertEqual(gid, "group_1") grp = self.mgr.groups[gid] self.assertEqual(grp.label, "batch") self.assertEqual(grp.strategy, "parallel") self.assertEqual(grp.status, "running") def test_increments_group_counter(self) -> None: g1 = self.mgr.start_group(label="a") g2 = self.mgr.start_group(label="b") self.assertEqual(g1, "group_1") self.assertEqual(g2, "group_2") def test_parent_agent_id_stored(self) -> None: aid = self.mgr.start_agent(prompt="parent") gid = self.mgr.start_group(label="child-group", parent_agent_id=aid) self.assertEqual(self.mgr.groups[gid].parent_agent_id, aid) def test_default_strategy_is_serial(self) -> None: gid = self.mgr.start_group() self.assertEqual(self.mgr.groups[gid].strategy, "serial") class TestRegisterGroupChild(unittest.TestCase): def setUp(self) -> None: self.mgr = AgentManager() def test_adds_agent_to_group(self) -> None: gid = self.mgr.start_group(label="g") aid = self.mgr.start_agent(prompt="t") self.mgr.register_group_child(gid, aid, child_index=0) grp = self.mgr.groups[gid] self.assertIn(aid, grp.child_agent_ids) self.assertEqual(self.mgr.records[aid].group_id, gid) self.assertEqual(self.mgr.records[aid].child_index, 0) def test_duplicate_does_not_add_twice(self) -> None: gid = self.mgr.start_group(label="g") aid = self.mgr.start_agent(prompt="t", group_id=gid, child_index=0) # Already registered via start_agent; register again self.mgr.register_group_child(gid, aid, child_index=0) grp = self.mgr.groups[gid] self.assertEqual(grp.child_agent_ids.count(aid), 1) def test_unknown_group_is_noop(self) -> None: aid = self.mgr.start_agent(prompt="t") # Should not raise self.mgr.register_group_child("nonexistent", aid, child_index=0) def test_unknown_agent_does_not_crash(self) -> None: gid = self.mgr.start_group(label="g") # Agent does not exist; group gets the ID but record update is skipped self.mgr.register_group_child(gid, "fake_agent", child_index=0) self.assertIn("fake_agent", self.mgr.groups[gid].child_agent_ids) def test_updates_child_index_on_record(self) -> None: gid = self.mgr.start_group(label="g") aid = self.mgr.start_agent(prompt="t") self.mgr.register_group_child(gid, aid, child_index=5) self.assertEqual(self.mgr.records[aid].child_index, 5) self.assertEqual(self.mgr.records[aid].group_id, gid) class TestFinishAgent(unittest.TestCase): def setUp(self) -> None: self.mgr = AgentManager() def test_marks_completed_with_stats(self) -> None: aid = self.mgr.start_agent(prompt="work") self.mgr.finish_agent( aid, session_id="sess-1", session_path="/path/sess", turns=5, tool_calls=12, stop_reason="end_turn", ) rec = self.mgr.records[aid] self.assertEqual(rec.status, "completed") self.assertEqual(rec.session_id, "sess-1") self.assertEqual(rec.session_path, "/path/sess") self.assertEqual(rec.turns, 5) self.assertEqual(rec.tool_calls, 12) self.assertEqual(rec.stop_reason, "end_turn") def test_preserves_original_fields(self) -> None: aid = self.mgr.start_agent( prompt="p", parent_agent_id="parent_x", label="lbl" ) self.mgr.finish_agent( aid, session_id="s", session_path="/p", turns=1, tool_calls=2, stop_reason=None, ) rec = self.mgr.records[aid] self.assertEqual(rec.prompt, "p") self.assertEqual(rec.parent_agent_id, "parent_x") self.assertEqual(rec.label, "lbl") def test_unknown_agent_is_noop(self) -> None: # Should not raise self.mgr.finish_agent( "unknown_id", session_id=None, session_path=None, turns=0, tool_calls=0, stop_reason=None, ) self.assertEqual(len(self.mgr.records), 0) class TestFinishGroup(unittest.TestCase): def setUp(self) -> None: self.mgr = AgentManager() def test_updates_group_status_and_stats(self) -> None: gid = self.mgr.start_group(label="g", strategy="parallel") self.mgr.finish_group( gid, status="completed", completed_children=3, failed_children=1, batch_count=2, max_batch_size=4, dependency_skips=0, ) grp = self.mgr.groups[gid] self.assertEqual(grp.status, "completed") self.assertEqual(grp.completed_children, 3) self.assertEqual(grp.failed_children, 1) self.assertEqual(grp.batch_count, 2) self.assertEqual(grp.max_batch_size, 4) self.assertEqual(grp.dependency_skips, 0) # Preserved fields self.assertEqual(grp.label, "g") self.assertEqual(grp.strategy, "parallel") def test_unknown_group_is_noop(self) -> None: self.mgr.finish_group( "ghost", status="completed", completed_children=0, failed_children=0, ) self.assertEqual(len(self.mgr.groups), 0) def test_preserves_child_agent_ids(self) -> None: gid = self.mgr.start_group(label="g") aid = self.mgr.start_agent(prompt="t", group_id=gid, child_index=0) self.mgr.finish_group( gid, status="completed", completed_children=1, failed_children=0 ) self.assertIn(aid, self.mgr.groups[gid].child_agent_ids) class TestChildrenOf(unittest.TestCase): def setUp(self) -> None: self.mgr = AgentManager() def test_returns_only_children_of_specified_parent(self) -> None: p1 = self.mgr.start_agent(prompt="parent1") p2 = self.mgr.start_agent(prompt="parent2") c1 = self.mgr.start_agent(prompt="c1", parent_agent_id=p1) c2 = self.mgr.start_agent(prompt="c2", parent_agent_id=p1) c3 = self.mgr.start_agent(prompt="c3", parent_agent_id=p2) children_p1 = self.mgr.children_of(p1) children_p2 = self.mgr.children_of(p2) self.assertEqual(len(children_p1), 2) ids_p1 = {r.agent_id for r in children_p1} self.assertEqual(ids_p1, {c1, c2}) self.assertEqual(len(children_p2), 1) self.assertEqual(children_p2[0].agent_id, c3) def test_returns_empty_for_no_children(self) -> None: aid = self.mgr.start_agent(prompt="solo") self.assertEqual(self.mgr.children_of(aid), ()) def test_returns_empty_for_unknown_parent(self) -> None: self.assertEqual(self.mgr.children_of("nonexistent"), ()) class TestGroupChildren(unittest.TestCase): def setUp(self) -> None: self.mgr = AgentManager() def test_returns_sorted_members(self) -> None: gid = self.mgr.start_group(label="g") a2 = self.mgr.start_agent(prompt="b", group_id=gid, child_index=2) a0 = self.mgr.start_agent(prompt="a", group_id=gid, child_index=0) a1 = self.mgr.start_agent(prompt="c", group_id=gid, child_index=1) children = self.mgr.group_children(gid) self.assertEqual(len(children), 3) self.assertEqual(children[0].agent_id, a0) self.assertEqual(children[1].agent_id, a1) self.assertEqual(children[2].agent_id, a2) def test_none_child_index_sorted_last(self) -> None: gid = self.mgr.start_group(label="g") a_none = self.mgr.start_agent(prompt="x", group_id=gid) a0 = self.mgr.start_agent(prompt="y", group_id=gid, child_index=0) children = self.mgr.group_children(gid) self.assertEqual(children[0].agent_id, a0) self.assertEqual(children[1].agent_id, a_none) def test_empty_for_unknown_group(self) -> None: self.assertEqual(self.mgr.group_children("nope"), ()) class TestGroupSummary(unittest.TestCase): def setUp(self) -> None: self.mgr = AgentManager() def test_aggregates_statistics(self) -> None: gid = self.mgr.start_group(label="batch", strategy="parallel") a1 = self.mgr.start_agent( prompt="t1", group_id=gid, child_index=0, resumed_from_session_id="old-sess", ) a2 = self.mgr.start_agent(prompt="t2", group_id=gid, child_index=1) self.mgr.finish_agent( a1, session_id="s1", session_path="/p1", turns=3, tool_calls=5, stop_reason="end_turn", ) self.mgr.finish_agent( a2, session_id="s2", session_path="/p2", turns=2, tool_calls=4, stop_reason="max_turns", ) self.mgr.finish_group( gid, status="completed", completed_children=2, failed_children=0, batch_count=1, max_batch_size=2, ) summary = self.mgr.group_summary(gid) assert summary is not None self.assertEqual(summary["group_id"], gid) self.assertEqual(summary["label"], "batch") self.assertEqual(summary["strategy"], "parallel") self.assertEqual(summary["status"], "completed") self.assertEqual(summary["child_count"], 2) self.assertEqual(summary["completed_children"], 2) self.assertEqual(summary["failed_children"], 0) self.assertEqual(summary["resumed_children"], 1) self.assertEqual(summary["batch_count"], 1) self.assertEqual(summary["max_batch_size"], 2) self.assertEqual(summary["dependency_skips"], 0) self.assertEqual( summary["stop_reason_counts"], {"end_turn": 1, "max_turns": 1}, ) def test_running_agents_counted_as_na(self) -> None: gid = self.mgr.start_group(label="g") self.mgr.start_agent(prompt="t", group_id=gid, child_index=0) summary = self.mgr.group_summary(gid) assert summary is not None self.assertEqual(summary["stop_reason_counts"], {"n/a": 1}) def test_returns_none_for_unknown_group(self) -> None: self.assertIsNone(self.mgr.group_summary("unknown")) class TestCompletedRecords(unittest.TestCase): def setUp(self) -> None: self.mgr = AgentManager() def test_filters_only_completed(self) -> None: a1 = self.mgr.start_agent(prompt="t1") a2 = self.mgr.start_agent(prompt="t2") a3 = self.mgr.start_agent(prompt="t3") self.mgr.finish_agent( a1, session_id="s", session_path="/p", turns=1, tool_calls=1, stop_reason="done", ) self.mgr.finish_agent( a3, session_id="s2", session_path="/p2", turns=2, tool_calls=3, stop_reason="done", ) completed = self.mgr.completed_records() self.assertEqual(len(completed), 2) ids = {r.agent_id for r in completed} self.assertEqual(ids, {a1, a3}) def test_empty_when_none_completed(self) -> None: self.mgr.start_agent(prompt="running") self.assertEqual(self.mgr.completed_records(), ()) class TestSummaryLines(unittest.TestCase): def setUp(self) -> None: self.mgr = AgentManager() def test_empty_manager(self) -> None: lines = self.mgr.summary_lines() self.assertIn("- Managed agents: 0", lines) self.assertIn("- Completed agents: 0", lines) self.assertIn("- Child agents: 0", lines) self.assertIn("- Resumed agents: 0", lines) self.assertIn("- Agent groups: 0", lines) self.assertIn("- Completed groups: 0", lines) def test_basic_output_format(self) -> None: a1 = self.mgr.start_agent(prompt="task", label="worker-1") self.mgr.finish_agent( a1, session_id="s1", session_path="/p", turns=4, tool_calls=10, stop_reason="end_turn", ) lines = self.mgr.summary_lines() self.assertIn("- Managed agents: 1", lines) self.assertIn("- Completed agents: 1", lines) # Agent detail line detail = [l for l in lines if "worker-1" in l] self.assertEqual(len(detail), 1) self.assertIn("status=completed", detail[0]) self.assertIn("turns=4", detail[0]) self.assertIn("tool_calls=10", detail[0]) self.assertIn("stop=end_turn", detail[0]) def test_group_info_in_agent_line(self) -> None: gid = self.mgr.start_group(label="g") self.mgr.start_agent(prompt="t", group_id=gid, child_index=0, label="child-0") lines = self.mgr.summary_lines() detail = [l for l in lines if "child-0" in l] self.assertEqual(len(detail), 1) self.assertIn(f"group={gid}", detail[0]) self.assertIn("child_index=0", detail[0]) def test_resumed_from_in_agent_line(self) -> None: self.mgr.start_agent( prompt="t", label="res", resumed_from_session_id="old-sess-id", ) lines = self.mgr.summary_lines() detail = [l for l in lines if "res" in l] self.assertTrue(any("resumed_from=old-sess-id" in l for l in detail)) def test_agent_without_label_uses_id(self) -> None: aid = self.mgr.start_agent(prompt="no label") lines = self.mgr.summary_lines() detail = [l for l in lines if aid in l] self.assertEqual(len(detail), 1) def test_truncation_at_8_agents(self) -> None: for i in range(10): self.mgr.start_agent(prompt=f"task-{i}") lines = self.mgr.summary_lines() self.assertIn("- Managed agents: 10", lines) plus_line = [l for l in lines if "plus" in l and "managed agents" in l] self.assertEqual(len(plus_line), 1) self.assertIn("2 more managed agents", plus_line[0]) def test_truncation_at_6_groups(self) -> None: for i in range(8): self.mgr.start_group(label=f"grp-{i}") lines = self.mgr.summary_lines() plus_line = [l for l in lines if "plus" in l and "agent groups" in l] self.assertEqual(len(plus_line), 1) self.assertIn("2 more agent groups", plus_line[0]) def test_group_summary_line_format(self) -> None: gid = self.mgr.start_group(label="my-batch", strategy="parallel") a1 = self.mgr.start_agent(prompt="t1", group_id=gid, child_index=0) self.mgr.finish_agent( a1, session_id="s", session_path="/p", turns=1, tool_calls=2, stop_reason="end_turn", ) self.mgr.finish_group( gid, status="completed", completed_children=1, failed_children=0, batch_count=1, max_batch_size=1, ) lines = self.mgr.summary_lines() grp_line = [l for l in lines if "my-batch" in l and "group_status" in l] self.assertEqual(len(grp_line), 1) self.assertIn("group_status=completed", grp_line[0]) self.assertIn("children=1", grp_line[0]) self.assertIn("completed=1", grp_line[0]) self.assertIn("failed=0", grp_line[0]) self.assertIn("strategy=parallel", grp_line[0]) self.assertIn("stop_reasons=end_turn:1", grp_line[0]) def test_child_and_resumed_counts(self) -> None: p = self.mgr.start_agent(prompt="parent") self.mgr.start_agent(prompt="c1", parent_agent_id=p) self.mgr.start_agent( prompt="c2", parent_agent_id=p, resumed_from_session_id="old", ) lines = self.mgr.summary_lines() self.assertIn("- Child agents: 2", lines) self.assertIn("- Resumed agents: 1", lines) class TestMultipleAgentsAndGroupsInteraction(unittest.TestCase): """End-to-end scenario with multiple groups and cross-references.""" def test_full_lifecycle(self) -> None: mgr = AgentManager() # Parent agent spawns two groups parent = mgr.start_agent(prompt="orchestrate", label="orchestrator") g1 = mgr.start_group(label="build", parent_agent_id=parent, strategy="serial") g2 = mgr.start_group(label="test", parent_agent_id=parent, strategy="parallel") # Group 1 children b1 = mgr.start_agent(prompt="build-fe", group_id=g1, child_index=0, parent_agent_id=parent) b2 = mgr.start_agent(prompt="build-be", group_id=g1, child_index=1, parent_agent_id=parent) # Group 2 children, one resumed t1 = mgr.start_agent( prompt="test-unit", group_id=g2, child_index=0, parent_agent_id=parent, resumed_from_session_id="old-session", ) t2 = mgr.start_agent(prompt="test-e2e", group_id=g2, child_index=1, parent_agent_id=parent) # Finish agents for aid, turns, tc, sr in [ (b1, 3, 8, "end_turn"), (b2, 4, 10, "end_turn"), (t1, 2, 5, "end_turn"), (t2, 6, 15, "max_turns"), ]: mgr.finish_agent(aid, session_id=f"s-{aid}", session_path=f"/p/{aid}", turns=turns, tool_calls=tc, stop_reason=sr) mgr.finish_group(g1, status="completed", completed_children=2, failed_children=0, batch_count=2, max_batch_size=1) mgr.finish_group(g2, status="completed", completed_children=1, failed_children=1, batch_count=1, max_batch_size=2, dependency_skips=1) # Verify children_of children = mgr.children_of(parent) self.assertEqual(len(children), 4) # Verify group_children ordering g1_children = mgr.group_children(g1) self.assertEqual(g1_children[0].agent_id, b1) self.assertEqual(g1_children[1].agent_id, b2) g2_children = mgr.group_children(g2) self.assertEqual(g2_children[0].agent_id, t1) self.assertEqual(g2_children[1].agent_id, t2) # Verify completed records (parent is still running) completed = mgr.completed_records() self.assertEqual(len(completed), 4) # Verify group summaries s1 = mgr.group_summary(g1) assert s1 is not None self.assertEqual(s1["child_count"], 2) self.assertEqual(s1["resumed_children"], 0) self.assertEqual(s1["dependency_skips"], 0) s2 = mgr.group_summary(g2) assert s2 is not None self.assertEqual(s2["child_count"], 2) self.assertEqual(s2["resumed_children"], 1) self.assertEqual(s2["dependency_skips"], 1) self.assertEqual(s2["stop_reason_counts"], {"end_turn": 1, "max_turns": 1}) # Verify summary_lines produces output lines = mgr.summary_lines() self.assertIn("- Managed agents: 5", lines) self.assertIn("- Completed agents: 4", lines) self.assertIn("- Child agents: 4", lines) self.assertIn("- Resumed agents: 1", lines) self.assertIn("- Agent groups: 2", lines) self.assertIn("- Completed groups: 2", lines) class TestFrozenDataclasses(unittest.TestCase): def test_record_is_frozen(self) -> None: rec = ManagedAgentRecord(agent_id="a", prompt="p") with self.assertRaises(AttributeError): rec.status = "completed" # type: ignore[misc] def test_group_is_frozen(self) -> None: grp = ManagedAgentGroup(group_id="g") with self.assertRaises(AttributeError): grp.status = "completed" # type: ignore[misc] if __name__ == "__main__": unittest.main()