"""U4: TeamOrchestrator 无依赖子任务并行模式单元测试 (R1-R5). Covers: - TeamPlan.get_independent_subtasks() introspection - _rebalance_independent_subtasks() MAX_EXPERTS enforcement (R4) - topological_sort 同层并行 (existing behavior, validated here for U4 contract) - SharedWorkspace 路径唯一性 (phase_id is UUID → no collision) - 综合阶段等待所有并行子任务完成 (layer-sequential contract) """ from __future__ import annotations import asyncio import json from unittest.mock import AsyncMock, MagicMock import pytest from agentkit.core.handoff_transport import InProcessHandoffTransport from agentkit.core.protocol import TaskResult, TaskStatus from agentkit.core.shared_workspace import SharedWorkspace from agentkit.experts.config import ExpertConfig from agentkit.experts.expert import Expert from agentkit.experts.orchestrator import TeamOrchestrator from agentkit.experts.plan import PlanPhase, TeamPlan from agentkit.experts.team import ExpertTeam from tests.unit.experts._helpers import ( make_chat_stream_mock, make_execute_stream_mock, ) # ── Helpers ──────────────────────────────────────────────────────────── def _make_expert_config(name: str, is_lead: bool = False) -> ExpertConfig: return ExpertConfig( name=name, agent_type="expert", persona="测试专家", thinking_style="逻辑推理", bound_skills=["skill_a"], is_lead=is_lead, task_mode="llm_generate", prompt={"identity": "测试"}, llm={"model": "default"}, ) def _make_mock_expert(name: str, is_lead: bool = False) -> MagicMock: config = _make_expert_config(name=name, is_lead=is_lead) expert = MagicMock(spec=Expert) expert.config = config expert.is_active = True expert.team_id = None expert.get_capabilities_summary.return_value = { "name": name, "persona": config.persona, "thinking_style": config.thinking_style, "bound_skills": config.bound_skills, "is_lead": is_lead, } mock_agent = MagicMock() mock_agent.execute = AsyncMock( return_value=TaskResult( task_id="test", agent_name=name, status=TaskStatus.COMPLETED.value, output_data={"content": f"Result from {name}"}, error_message=None, started_at=None, completed_at=None, ) ) mock_agent.execute_stream = make_execute_stream_mock(f"Result from {name}") mock_agent._llm_gateway = None expert.agent = mock_agent return expert def _make_team_with_experts( expert_names: list[str] | None = None, lead_name: str = "lead", ) -> ExpertTeam: team = ExpertTeam() transport = AsyncMock(spec=InProcessHandoffTransport) team._handoff_transport = transport if expert_names is None: expert_names = [lead_name, "member1", "member2"] for name in expert_names: is_lead = name == lead_name expert = _make_mock_expert(name=name, is_lead=is_lead) team._experts[name] = expert if is_lead: team._lead_expert_name = name return team def _make_phase( id: str, name: str, assigned_expert: str, depends_on: list[str] | None = None ) -> PlanPhase: return PlanPhase( id=id, name=name, assigned_expert=assigned_expert, task_description=f"task for {name}", depends_on=depends_on or [], ) def _make_mock_llm_gateway_with_phases( phases_responses: list[list[dict]], synthesis_content: str = "综合结果", ) -> MagicMock: """Gateway that returns successive decompositions then synthesis. phases_responses: each element is a list of phase dicts to return on successive decomposition calls. After exhausted, returns synthesis. """ gateway = AsyncMock() decomp_responses = [MagicMock(content=json.dumps(phases)) for phases in phases_responses] synth_response = MagicMock(content=synthesis_content) call_count = [0] async def chat_side_effect(messages, model=None, **kwargs): call_count[0] += 1 if call_count[0] <= len(decomp_responses): return decomp_responses[call_count[0] - 1] return synth_response gateway.chat = AsyncMock(side_effect=chat_side_effect) gateway.chat_stream = make_chat_stream_mock(synthesis_content) return gateway # ── TeamPlan.get_independent_subtasks ───────────────────────────────── class TestGetIndependentSubtasks: """U4/R1: TeamPlan.get_independent_subtasks() returns phases with no depends_on.""" def test_returns_empty_for_no_phases(self): plan = TeamPlan(task="test", lead_expert="lead") assert plan.get_independent_subtasks() == [] def test_returns_only_independent_phases(self): plan = TeamPlan( task="test", lead_expert="lead", phases=[ _make_phase("p1", "A", "lead", depends_on=[]), _make_phase("p2", "B", "member1", depends_on=["p1"]), _make_phase("p3", "C", "member2", depends_on=[]), _make_phase("p4", "D", "member1", depends_on=["p1", "p3"]), ], ) independent = plan.get_independent_subtasks() assert len(independent) == 2 assert {ph.id for ph in independent} == {"p1", "p3"} def test_returns_all_when_no_dependencies(self): plan = TeamPlan( task="test", lead_expert="lead", phases=[ _make_phase("p1", "A", "lead", depends_on=[]), _make_phase("p2", "B", "member1", depends_on=[]), _make_phase("p3", "C", "member2", depends_on=[]), ], ) independent = plan.get_independent_subtasks() assert len(independent) == 3 # ── topological_sort 同层并行 contract (U4/R1, R3) ──────────────────── class TestParallelLayerContract: """U4/R1, R3: independent subtasks land in layer 0 (parallel); synthesis waits for all layers to complete (layer-sequential).""" def test_three_independent_subtasks_in_layer_0(self): plan = TeamPlan( task="test", lead_expert="lead", phases=[ _make_phase("p1", "A", "lead", depends_on=[]), _make_phase("p2", "B", "member1", depends_on=[]), _make_phase("p3", "C", "member2", depends_on=[]), ], ) layers = plan.topological_sort() assert len(layers) == 1 assert len(layers[0]) == 3 assert {ph.id for ph in layers[0]} == {"p1", "p2", "p3"} def test_shared_dependency_degrades_to_layered(self): """R3: subtasks sharing an upstream dependency are NOT parallel with that upstream — they land in later layers.""" plan = TeamPlan( task="test", lead_expert="lead", phases=[ _make_phase("p1", "Base", "lead", depends_on=[]), _make_phase("p2", "A", "member1", depends_on=["p1"]), _make_phase("p3", "B", "member2", depends_on=["p1"]), ], ) layers = plan.topological_sort() assert len(layers) == 2 assert len(layers[0]) == 1 # only Base assert layers[0][0].id == "p1" assert len(layers[1]) == 2 # A and B parallel (same layer) assert {ph.id for ph in layers[1]} == {"p2", "p3"} def test_synthesis_waits_for_all_layers(self): """R3: dependent phase in last layer — synthesis cannot run until its layer completes (validated via topological_sort layer count).""" plan = TeamPlan( task="test", lead_expert="lead", phases=[ _make_phase("p1", "A", "lead", depends_on=[]), _make_phase("p2", "B", "member1", depends_on=["p1"]), _make_phase("p3", "C", "member2", depends_on=["p2"]), ], ) layers = plan.topological_sort() assert len(layers) == 3 # strict sequential — no parallelism # ── SharedWorkspace 路径唯一性 (U4/R2) ───────────────────────────────── class TestSharedWorkspacePathUniqueness: """U4/R2: parallel subtask output paths must not collide. PlanPhase.id is UUID by default → unique per phase, ensuring ``{plan_id}/phase/{phase_id}/output`` never overlaps.""" def test_default_phase_ids_are_unique(self): phases = [ PlanPhase(name=f"phase_{i}", assigned_expert="lead", task_description="t") for i in range(5) ] ids = [ph.id for ph in phases] assert len(set(ids)) == len(ids), "phase_ids must be unique" def test_independent_subtasks_have_unique_output_paths(self): plan = TeamPlan( task="test", lead_expert="lead", phases=[ PlanPhase(name="A", assigned_expert="lead", task_description="t"), PlanPhase(name="B", assigned_expert="member1", task_description="t"), PlanPhase(name="C", assigned_expert="member2", task_description="t"), ], ) independent = plan.get_independent_subtasks() paths = {f"{plan.id}/phase/{ph.id}/output" for ph in independent} assert len(paths) == len(independent) # ── _rebalance_independent_subtasks (U4/R4 MAX_EXPERTS) ──────────────── class TestRebalanceIndependentSubtasks: """U4/R4: when Lead over-decomposes independent subtasks beyond MAX_INDEPENDENT_SUBTASKS, request re-decompose with merge hint.""" @pytest.mark.asyncio async def test_no_rebalance_when_within_limit(self): """5 independent subtasks ≤ 10 → no re-decompose.""" team = _make_team_with_experts() orchestrator = TeamOrchestrator(team) # 5 independent phases phases = [_make_phase(f"p{i}", f"phase_{i}", "lead", depends_on=[]) for i in range(5)] original_phases = list(phases) result = await orchestrator._rebalance_independent_subtasks( lead=team.lead_expert, task="test", phases=phases ) assert result is original_phases or result == original_phases @pytest.mark.asyncio async def test_rebalance_triggers_when_over_limit(self): """11 independent subtasks > 10 → re-decompose with merge hint.""" team = _make_team_with_experts() orchestrator = TeamOrchestrator(team) # Original: 11 independent (over limit) original_phases = [ _make_phase(f"p{i}", f"phase_{i}", "lead", depends_on=[]) for i in range(11) ] # Re-decomposed: 8 independent (within limit) — merge succeeded gateway = _make_mock_llm_gateway_with_phases( phases_responses=[ [ # first call (original _decompose_task already happened — # but _rebalance calls _decompose_task again, returning this) { "name": f"merged_{i}", "assigned_expert": "lead", "task_description": "t", "depends_on": [], } for i in range(8) ], ] ) team._experts["lead"].agent._llm_gateway = gateway result = await orchestrator._rebalance_independent_subtasks( lead=team.lead_expert, task="test", phases=original_phases ) # Should return the re-decomposed phases (8 independent) assert len(result) == 8 assert all(not ph.depends_on for ph in result) # Original phases should NOT be returned assert {ph.id for ph in result}.isdisjoint({ph.id for ph in original_phases}) @pytest.mark.asyncio async def test_rebalance_no_gateway_returns_original(self): """No LLM gateway → can't re-decompose, return original.""" team = _make_team_with_experts() orchestrator = TeamOrchestrator(team) # gateway is None by default in _make_mock_expert original_phases = [ _make_phase(f"p{i}", f"phase_{i}", "lead", depends_on=[]) for i in range(11) ] result = await orchestrator._rebalance_independent_subtasks( lead=team.lead_expert, task="test", phases=original_phases ) assert result is original_phases @pytest.mark.asyncio async def test_rebalance_still_over_limit_returns_new_anyway(self): """Re-decompose still produces >MAX independent → return new phases (MAX_PHASES truncation handles it in execute()).""" team = _make_team_with_experts() orchestrator = TeamOrchestrator(team) original_phases = [ _make_phase(f"p{i}", f"phase_{i}", "lead", depends_on=[]) for i in range(11) ] # Re-decompose returns 12 independent (still over) gateway = _make_mock_llm_gateway_with_phases( phases_responses=[ [ { "name": f"q{i}", "assigned_expert": "lead", "task_description": "t", "depends_on": [], } for i in range(12) ], ] ) team._experts["lead"].agent._llm_gateway = gateway result = await orchestrator._rebalance_independent_subtasks( lead=team.lead_expert, task="test", phases=original_phases ) # Should return the 12-phase re-decomposition (truncation happens later) assert len(result) == 12 @pytest.mark.asyncio async def test_rebalance_decompose_failure_returns_original(self): """Re-decompose raises LLMProviderError → return original phases.""" from agentkit.core.exceptions import LLMProviderError team = _make_team_with_experts() orchestrator = TeamOrchestrator(team) original_phases = [ _make_phase(f"p{i}", f"phase_{i}", "lead", depends_on=[]) for i in range(11) ] gateway = AsyncMock() gateway.chat = AsyncMock(side_effect=LLMProviderError("LLM down")) gateway.chat_stream = make_chat_stream_mock("err") team._experts["lead"].agent._llm_gateway = gateway result = await orchestrator._rebalance_independent_subtasks( lead=team.lead_expert, task="test", phases=original_phases ) assert result is original_phases @pytest.mark.asyncio async def test_rebalance_empty_redecompose_returns_original(self): """Re-decompose returns empty list → return original phases.""" team = _make_team_with_experts() orchestrator = TeamOrchestrator(team) original_phases = [ _make_phase(f"p{i}", f"phase_{i}", "lead", depends_on=[]) for i in range(11) ] # Gateway returns malformed JSON → _parse_phases returns [] gateway = AsyncMock() bad_response = MagicMock(content="not json") gateway.chat = AsyncMock(return_value=bad_response) gateway.chat_stream = make_chat_stream_mock("err") team._experts["lead"].agent._llm_gateway = gateway result = await orchestrator._rebalance_independent_subtasks( lead=team.lead_expert, task="test", phases=original_phases ) assert result is original_phases # ── Integration: full execute() with parallel subtasks ──────────────── class TestExecuteParallelSubtasks: """U4/R1, R3, R5: full execute() path with independent subtasks running in parallel + synthesis waits for all.""" @pytest.mark.asyncio async def test_three_parallel_subtasks_complete_then_synthesis(self): """3 independent subtasks → all run in layer 0 → synthesis runs after all complete.""" team = _make_team_with_experts() orchestrator = TeamOrchestrator(team) gateway = _make_mock_llm_gateway_with_phases( phases_responses=[ [ { "name": "A", "assigned_expert": "lead", "task_description": "a", "depends_on": [], }, { "name": "B", "assigned_expert": "member1", "task_description": "b", "depends_on": [], }, { "name": "C", "assigned_expert": "member2", "task_description": "c", "depends_on": [], }, ], ] ) team._experts["lead"].agent._llm_gateway = gateway result = await orchestrator.execute("并行任务") assert result["status"] == "completed" plan: TeamPlan = result["plan"] layers = plan.topological_sort() assert len(layers) == 1 # all 3 in layer 0 (parallel) assert len(layers[0]) == 3 # All phases completed completed = plan.completed_phases assert len(completed) == 3 # phase_results has all 3 phase_results: dict = result["phase_results"] assert len(phase_results) == 3 # ── U3: SharedWorkspace.lock_section TOCTOU protection ─────────────── class TestLockSectionToctou: """U3/KTD3: SharedWorkspace.lock_section guards parallel subtask read-then-write critical sections against TOCTOU races. lock_section is the async-context-manager form of lock/unlock: it polls lock() until acquired (or timeout), holds for the block, and releases on exit. Used by _finalize_phase to wrap workspace.write(). """ @pytest.mark.asyncio async def test_lock_section_acquires_and_releases(self): """Lock acquired on enter, released on exit — next caller can acquire.""" ws = SharedWorkspace() async with ws.lock_section("resource", "agent_a", timeout=1.0): # Held by agent_a — agent_b cannot acquire via raw lock() assert await ws.lock("resource", agent_id="agent_b") is False # Released after block — agent_b can now acquire assert await ws.lock("resource", agent_id="agent_b") is True await ws.unlock("resource", "agent_b") @pytest.mark.asyncio async def test_lock_section_calls_lock_and_unlock(self): """lock_section delegates to lock() on enter and unlock() on exit.""" ws = SharedWorkspace() lock_spy = AsyncMock(wraps=ws.lock) unlock_spy = AsyncMock(wraps=ws.unlock) ws.lock = lock_spy ws.unlock = unlock_spy async with ws.lock_section("k", "agent_a", timeout=1.0): pass lock_spy.assert_awaited_with("k", "agent_a", timeout=1.0) unlock_spy.assert_awaited_with("k", "agent_a") @pytest.mark.asyncio async def test_lock_section_raises_timeout_when_held(self): """Lock held by another agent → lock_section raises TimeoutError.""" ws = SharedWorkspace() await ws.lock("shared", "holder") # hold the lock with pytest.raises(TimeoutError): async with ws.lock_section("shared", "waiter", timeout=0.2): pass # pragma: no cover — never reached # Waiter never acquired, so holder still owns it. assert await ws.lock("shared", agent_id="other") is False @pytest.mark.asyncio async def test_parallel_writes_to_different_keys_do_not_conflict(self): """2 subtasks writing different keys concurrently both succeed.""" ws = SharedWorkspace() results: list[str] = [] async def writer(key: str, agent: str) -> None: async with ws.lock_section(key, agent, timeout=2.0): await ws.write(key, f"val-{agent}", agent) results.append(agent) await asyncio.gather(writer("k1", "a"), writer("k2", "b")) assert set(results) == {"a", "b"} assert (await ws.read("k1"))["value"] == "val-a" assert (await ws.read("k2"))["value"] == "val-b" @pytest.mark.asyncio async def test_same_key_writes_serialize(self): """2 subtasks writing the same key: the second waits for the first to release — enter/exit events never interleave.""" ws = SharedWorkspace() log: list[str] = [] async def writer(name: str) -> None: async with ws.lock_section("shared_key", name, timeout=2.0): log.append(f"{name}-enter") await asyncio.sleep(0.05) log.append(f"{name}-exit") await asyncio.gather(writer("a"), writer("b")) # One fully completes before the other begins (serialized). assert log in ( ["a-enter", "a-exit", "b-enter", "b-exit"], ["b-enter", "b-exit", "a-enter", "a-exit"], ) @pytest.mark.asyncio async def test_lock_section_releases_on_exception(self): """If the wrapped block raises, the lock is still released.""" ws = SharedWorkspace() with pytest.raises(RuntimeError, match="boom"): async with ws.lock_section("k", "agent_a", timeout=1.0): raise RuntimeError("boom") # Lock released despite exception — another agent can acquire. assert await ws.lock("k", agent_id="agent_b") is True await ws.unlock("k", "agent_b")