570 lines
21 KiB
Python
570 lines
21 KiB
Python
"""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")
|