wehub-resource-sync
596 lines
20 KiB
Python
596 lines
20 KiB
Python
"""DAG-parallelism contract tests for MetaOrchestrator (M7).
|
|
|
|
Three contracts:
|
|
1. Two steps with disjoint depends_on run concurrently — their
|
|
ToolUseStartEvents both arrive before either's ToolResultEvent.
|
|
2. When one step in a parallel batch fails, the sibling task is
|
|
cancelled — its ToolResultEvent never arrives.
|
|
3. A purely-linear DAG keeps the same event order as the old
|
|
linear-topo scheduler (backwards compat).
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from collections.abc import AsyncIterator
|
|
from typing import Any
|
|
|
|
import pytest
|
|
|
|
from opensquilla.engine.types import (
|
|
AgentEvent,
|
|
DoneEvent,
|
|
TextDeltaEvent,
|
|
ToolResultEvent,
|
|
ToolUseStartEvent,
|
|
)
|
|
from opensquilla.skills.meta.orchestrator import MetaOrchestrator
|
|
from opensquilla.skills.meta.parser import parse_meta_plan
|
|
from opensquilla.skills.meta.types import MetaMatch, MetaResult
|
|
from opensquilla.skills.types import SkillLayer, SkillSpec
|
|
|
|
|
|
def _meta_spec(steps: list[dict[str, Any]]) -> SkillSpec:
|
|
return SkillSpec(
|
|
name="meta-test",
|
|
description="t",
|
|
layer=SkillLayer.BUNDLED,
|
|
always=False,
|
|
triggers=["t"],
|
|
content="fallback",
|
|
kind="meta",
|
|
composition_raw={"steps": steps},
|
|
)
|
|
|
|
|
|
def _skill(name: str) -> SkillSpec:
|
|
return SkillSpec(
|
|
name=name,
|
|
description=f"{name} d",
|
|
layer=SkillLayer.BUNDLED,
|
|
always=False,
|
|
triggers=[],
|
|
content=name.upper(),
|
|
kind="skill",
|
|
)
|
|
|
|
|
|
class _FakeLoader:
|
|
def __init__(self, specs: list[SkillSpec]) -> None:
|
|
self._by_name = {s.name: s for s in specs}
|
|
|
|
def get_by_name(self, name: str) -> SkillSpec | None:
|
|
return self._by_name.get(name)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_independent_steps_run_concurrently() -> None:
|
|
"""Two steps with no deps both start before either finishes."""
|
|
spec = _meta_spec(
|
|
[
|
|
{"id": "a", "skill": "skill_a", "with": {}},
|
|
{"id": "b", "skill": "skill_b", "with": {}}, # no depends_on
|
|
],
|
|
)
|
|
plan = parse_meta_plan(spec)
|
|
assert plan is not None
|
|
|
|
started = asyncio.Event()
|
|
second_started_first = asyncio.Event()
|
|
|
|
async def runner(system_prompt: str, _u: str) -> AsyncIterator[AgentEvent]:
|
|
# First runner blocks until the second has also started — proves
|
|
# they overlap. With a serial scheduler this hangs forever.
|
|
if "SKILL_A" in system_prompt:
|
|
started.set()
|
|
await asyncio.wait_for(second_started_first.wait(), timeout=2.0)
|
|
yield TextDeltaEvent(text="A done")
|
|
else:
|
|
await asyncio.wait_for(started.wait(), timeout=2.0)
|
|
second_started_first.set()
|
|
yield TextDeltaEvent(text="B done")
|
|
yield DoneEvent(text="")
|
|
|
|
orch = MetaOrchestrator(
|
|
agent_runner=runner,
|
|
skill_loader=_FakeLoader([_skill("skill_a"), _skill("skill_b")]),
|
|
)
|
|
|
|
starts: list[str] = []
|
|
results: list[str] = []
|
|
final: MetaResult | None = None
|
|
async for ev in orch.iter_events(MetaMatch(plan=plan, inputs={})):
|
|
if isinstance(ev, MetaResult):
|
|
final = ev
|
|
elif isinstance(ev, ToolUseStartEvent) and ev.tool_name.startswith("meta-step:"):
|
|
starts.append(ev.tool_name)
|
|
elif isinstance(ev, ToolResultEvent) and ev.tool_name.startswith("meta-step:"):
|
|
results.append(ev.tool_name)
|
|
|
|
assert final is not None and final.ok, final.error if final else "no final"
|
|
# Both starts must arrive before either result.
|
|
assert set(starts) == {"meta-step:a", "meta-step:b"}
|
|
assert set(results) == {"meta-step:a", "meta-step:b"}
|
|
# In the linear scheduler, starts[0] would always produce results[0]
|
|
# before starts[1]. The parallel scheduler interleaves them, which we
|
|
# already proved by the asyncio.Event coupling above — if we got here
|
|
# without TimeoutError, the steps ran concurrently.
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_failing_step_cancels_sibling() -> None:
|
|
"""When one parallel step raises, the sibling task is cancelled."""
|
|
spec = _meta_spec(
|
|
[
|
|
{"id": "a", "skill": "skill_a", "with": {}},
|
|
{"id": "b", "skill": "skill_b", "with": {}},
|
|
],
|
|
)
|
|
plan = parse_meta_plan(spec)
|
|
assert plan is not None
|
|
|
|
sibling_completed = False
|
|
|
|
async def runner(system_prompt: str, _u: str) -> AsyncIterator[AgentEvent]:
|
|
nonlocal sibling_completed
|
|
if "SKILL_A" in system_prompt:
|
|
# Fail fast.
|
|
raise RuntimeError("a failed")
|
|
# b is slow — must be cancelled before it finishes.
|
|
try:
|
|
await asyncio.sleep(5.0)
|
|
except asyncio.CancelledError:
|
|
raise
|
|
sibling_completed = True
|
|
yield TextDeltaEvent(text="should not arrive")
|
|
yield DoneEvent(text="")
|
|
|
|
orch = MetaOrchestrator(
|
|
agent_runner=runner,
|
|
skill_loader=_FakeLoader([_skill("skill_a"), _skill("skill_b")]),
|
|
)
|
|
|
|
final: MetaResult | None = None
|
|
error_results: list[ToolResultEvent] = []
|
|
async for ev in orch.iter_events(MetaMatch(plan=plan, inputs={})):
|
|
if isinstance(ev, MetaResult):
|
|
final = ev
|
|
elif isinstance(ev, ToolResultEvent) and ev.is_error:
|
|
error_results.append(ev)
|
|
|
|
assert final is not None and final.ok is False
|
|
assert sibling_completed is False
|
|
assert len(error_results) >= 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_failure_closes_sibling_brackets() -> None:
|
|
"""Every yielded ToolUseStart for a meta-step has a matching ToolResult.
|
|
|
|
Concretely: when two siblings A/B both emit ToolUseStartEvents and A
|
|
fails while B is still sleeping, the orchestrator must emit a
|
|
ToolResultEvent for B carrying ``is_error=True`` and a "cancelled"
|
|
message so the WebUI does not leak an open in-progress card.
|
|
"""
|
|
spec = _meta_spec(
|
|
[
|
|
{"id": "a", "skill": "skill_a", "with": {}},
|
|
{"id": "b", "skill": "skill_b", "with": {}},
|
|
],
|
|
)
|
|
plan = parse_meta_plan(spec)
|
|
assert plan is not None
|
|
|
|
b_running = asyncio.Event()
|
|
|
|
async def runner(system_prompt: str, _u: str) -> AsyncIterator[AgentEvent]:
|
|
if "SKILL_A" in system_prompt:
|
|
# Make sure B has emitted its ToolUseStart before A fails — we
|
|
# want to exercise the path where the outer loop yielded both
|
|
# starts already.
|
|
await asyncio.wait_for(b_running.wait(), timeout=2.0)
|
|
raise RuntimeError("a blew up mid-flight")
|
|
# B emits a TextDelta (forces _run_one to enter the streaming
|
|
# body so its ToolUseStart has been pushed) then signals A and
|
|
# sleeps. The sleep is interruptible so cancellation lands.
|
|
yield TextDeltaEvent(text="b: starting")
|
|
b_running.set()
|
|
try:
|
|
await asyncio.sleep(5.0)
|
|
except asyncio.CancelledError:
|
|
raise
|
|
yield TextDeltaEvent(text="b: should not arrive")
|
|
yield DoneEvent(text="")
|
|
|
|
orch = MetaOrchestrator(
|
|
agent_runner=runner,
|
|
skill_loader=_FakeLoader([_skill("skill_a"), _skill("skill_b")]),
|
|
)
|
|
|
|
starts: dict[str, ToolUseStartEvent] = {}
|
|
results: dict[str, ToolResultEvent] = {}
|
|
final: MetaResult | None = None
|
|
async for ev in orch.iter_events(MetaMatch(plan=plan, inputs={})):
|
|
if isinstance(ev, MetaResult):
|
|
final = ev
|
|
elif isinstance(ev, ToolUseStartEvent) and ev.tool_name.startswith(
|
|
"meta-step:",
|
|
):
|
|
starts[ev.tool_use_id] = ev
|
|
elif isinstance(ev, ToolResultEvent) and ev.tool_name.startswith(
|
|
"meta-step:",
|
|
):
|
|
# Keep the FIRST result we see per id — a real ToolResult
|
|
# from the failing task should not be overwritten by a
|
|
# synthetic one (and the orchestrator should not emit two
|
|
# results for the same start anyway).
|
|
results.setdefault(ev.tool_use_id, ev)
|
|
|
|
assert final is not None and final.ok is False
|
|
# Close-bracket invariant: every start has a matching result with
|
|
# the same tool_use_id.
|
|
assert set(starts.keys()) == set(results.keys()), (
|
|
f"unbalanced brackets: starts={set(starts.keys())} "
|
|
f"results={set(results.keys())}"
|
|
)
|
|
# Both meta-step:a and meta-step:b must be present.
|
|
assert "meta_step_a" in starts
|
|
assert "meta_step_b" in starts
|
|
|
|
# Sibling B (cancelled) must carry is_error=True and a "cancelled" hint.
|
|
b_result = results["meta_step_b"]
|
|
assert b_result.is_error is True
|
|
assert "cancel" in str(b_result.result).lower()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_linear_dag_event_order_preserved() -> None:
|
|
"""A→B→C chain keeps deterministic event ordering."""
|
|
spec = _meta_spec(
|
|
[
|
|
{"id": "a", "skill": "skill_a", "with": {}},
|
|
{"id": "b", "skill": "skill_b", "depends_on": ["a"], "with": {}},
|
|
{"id": "c", "skill": "skill_c", "depends_on": ["b"], "with": {}},
|
|
],
|
|
)
|
|
plan = parse_meta_plan(spec)
|
|
assert plan is not None
|
|
|
|
async def runner(system_prompt: str, _u: str) -> AsyncIterator[AgentEvent]:
|
|
for letter in ("A", "B", "C"):
|
|
if f"SKILL_{letter}" in system_prompt:
|
|
yield TextDeltaEvent(text=f"{letter}-done")
|
|
break
|
|
yield DoneEvent(text="")
|
|
|
|
orch = MetaOrchestrator(
|
|
agent_runner=runner,
|
|
skill_loader=_FakeLoader(
|
|
[_skill("skill_a"), _skill("skill_b"), _skill("skill_c")],
|
|
),
|
|
)
|
|
|
|
ordering: list[tuple[str, str]] = []
|
|
final: MetaResult | None = None
|
|
async for ev in orch.iter_events(MetaMatch(plan=plan, inputs={})):
|
|
if isinstance(ev, MetaResult):
|
|
final = ev
|
|
elif isinstance(ev, ToolUseStartEvent) and ev.tool_name.startswith("meta-step:"):
|
|
ordering.append(("start", ev.tool_name))
|
|
elif isinstance(ev, ToolResultEvent) and ev.tool_name.startswith("meta-step:"):
|
|
ordering.append(("end", ev.tool_name))
|
|
|
|
assert final is not None and final.ok
|
|
# Strict interleaving: start(a) → end(a) → start(b) → end(b) → start(c) → end(c)
|
|
assert ordering == [
|
|
("start", "meta-step:a"),
|
|
("end", "meta-step:a"),
|
|
("start", "meta-step:b"),
|
|
("end", "meta-step:b"),
|
|
("start", "meta-step:c"),
|
|
("end", "meta-step:c"),
|
|
]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_max_parallelism_caps_concurrent_steps() -> None:
|
|
"""With max_parallelism=2 and 4 independent steps, no more than 2 are
|
|
in-flight at any moment."""
|
|
|
|
spec = _meta_spec(
|
|
[
|
|
{"id": f"s{i}", "skill": f"skill_{i}", "with": {}}
|
|
for i in range(4)
|
|
],
|
|
)
|
|
plan = parse_meta_plan(spec)
|
|
assert plan is not None
|
|
|
|
in_flight = 0
|
|
high_water_mark = 0
|
|
lock = asyncio.Lock()
|
|
|
|
async def runner(_s: str, _u: str) -> AsyncIterator[AgentEvent]:
|
|
nonlocal in_flight, high_water_mark
|
|
async with lock:
|
|
in_flight += 1
|
|
high_water_mark = max(high_water_mark, in_flight)
|
|
try:
|
|
await asyncio.sleep(0.05)
|
|
yield TextDeltaEvent(text="done")
|
|
yield DoneEvent(text="")
|
|
finally:
|
|
async with lock:
|
|
in_flight -= 1
|
|
|
|
loader = _FakeLoader([_skill(f"skill_{i}") for i in range(4)])
|
|
orch = MetaOrchestrator(
|
|
agent_runner=runner,
|
|
skill_loader=loader,
|
|
max_parallelism=2,
|
|
)
|
|
|
|
async for _ in orch.iter_events(MetaMatch(plan=plan, inputs={})):
|
|
pass
|
|
|
|
assert high_water_mark <= 2, f"high_water_mark={high_water_mark}"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_max_parallelism_none_unbounded() -> None:
|
|
"""max_parallelism=None preserves the current 5-way fan-out."""
|
|
|
|
spec = _meta_spec(
|
|
[
|
|
{"id": f"s{i}", "skill": f"skill_{i}", "with": {}}
|
|
for i in range(5)
|
|
],
|
|
)
|
|
plan = parse_meta_plan(spec)
|
|
assert plan is not None
|
|
|
|
started_count = 0
|
|
target_lock = asyncio.Lock()
|
|
all_started = asyncio.Event()
|
|
|
|
async def runner(_s: str, _u: str) -> AsyncIterator[AgentEvent]:
|
|
nonlocal started_count
|
|
async with target_lock:
|
|
started_count += 1
|
|
if started_count == 5:
|
|
all_started.set()
|
|
await asyncio.wait_for(all_started.wait(), timeout=1.0)
|
|
yield TextDeltaEvent(text="done")
|
|
yield DoneEvent(text="")
|
|
|
|
loader = _FakeLoader([_skill(f"skill_{i}") for i in range(5)])
|
|
orch = MetaOrchestrator(
|
|
agent_runner=runner,
|
|
skill_loader=loader,
|
|
max_parallelism=None,
|
|
)
|
|
|
|
final: MetaResult | None = None
|
|
async for ev in orch.iter_events(MetaMatch(plan=plan, inputs={})):
|
|
if isinstance(ev, MetaResult):
|
|
final = ev
|
|
|
|
assert final is not None and final.ok
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_orchestrator_default_parallelism_caps_at_four() -> None:
|
|
spec = _meta_spec(
|
|
[
|
|
{"id": f"s{i}", "skill": f"skill_{i}", "with": {}}
|
|
for i in range(5)
|
|
],
|
|
)
|
|
plan = parse_meta_plan(spec)
|
|
assert plan is not None
|
|
|
|
in_flight = 0
|
|
high_water_mark = 0
|
|
lock = asyncio.Lock()
|
|
|
|
async def runner(_s: str, _u: str) -> AsyncIterator[AgentEvent]:
|
|
nonlocal in_flight, high_water_mark
|
|
async with lock:
|
|
in_flight += 1
|
|
high_water_mark = max(high_water_mark, in_flight)
|
|
try:
|
|
await asyncio.sleep(0.05)
|
|
yield TextDeltaEvent(text="done")
|
|
yield DoneEvent(text="")
|
|
finally:
|
|
async with lock:
|
|
in_flight -= 1
|
|
|
|
orch = MetaOrchestrator(
|
|
agent_runner=runner,
|
|
skill_loader=_FakeLoader([_skill(f"skill_{i}") for i in range(5)]),
|
|
)
|
|
|
|
async for _ in orch.iter_events(MetaMatch(plan=plan, inputs={})):
|
|
pass
|
|
|
|
assert high_water_mark <= 4, f"high_water_mark={high_water_mark}"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_run_dag_emits_three_callback_types() -> None:
|
|
"""C3: scheduler must externalise step begin / done / failover events."""
|
|
from opensquilla.skills.meta.scheduler import run_dag
|
|
from opensquilla.skills.meta.types import MetaMatch, MetaPlan, MetaStep
|
|
|
|
plan = MetaPlan(
|
|
name="cb-test",
|
|
triggers=("x",),
|
|
priority=10,
|
|
steps=(
|
|
MetaStep(id="a", skill="alpha", kind="agent", on_failure="b"),
|
|
MetaStep(id="b", skill="beta", kind="agent"),
|
|
),
|
|
)
|
|
match = MetaMatch(plan=plan, inputs={})
|
|
begin_calls: list[tuple[str, str]] = []
|
|
finish_calls: list[tuple[str, str]] = []
|
|
failover_calls: list[tuple[str, str, str]] = []
|
|
|
|
async def begin_cb(step_id: str, effective_skill: str, rendered_inputs: dict) -> None:
|
|
begin_calls.append((step_id, effective_skill))
|
|
|
|
async def finish_cb(
|
|
step_id: str,
|
|
status: str,
|
|
output_text: str | None,
|
|
error: str | None,
|
|
) -> None:
|
|
finish_calls.append((step_id, status))
|
|
|
|
async def failover_cb(failed_step_id: str, substitute_step_id: str, error: str) -> None:
|
|
failover_calls.append((failed_step_id, substitute_step_id, error))
|
|
|
|
async def dispatch_stub(step, effective_skill, inputs, outputs):
|
|
if step.id == "a":
|
|
raise RuntimeError("alpha exploded")
|
|
from opensquilla.skills.meta.events import _StepDone
|
|
yield _StepDone(text="beta-output")
|
|
|
|
async def preface_stub(step_id: str, effective_skill: str):
|
|
if False:
|
|
yield None
|
|
return
|
|
|
|
async for _ in run_dag(
|
|
match,
|
|
dispatch_step_stream=dispatch_stub,
|
|
yield_skill_view_preface=preface_stub,
|
|
max_parallelism=4,
|
|
on_step_begin=begin_cb,
|
|
on_step_finish=finish_cb,
|
|
on_step_failover=failover_cb,
|
|
):
|
|
pass
|
|
|
|
assert ("a", "alpha") in begin_calls
|
|
assert ("b", "beta") in begin_calls
|
|
assert ("b", "ok") in finish_calls
|
|
assert len(failover_calls) == 1
|
|
assert failover_calls[0][:2] == ("a", "b")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_run_dag_on_step_begin_gets_rendered_inputs() -> None:
|
|
from opensquilla.skills.meta.scheduler import run_dag
|
|
from opensquilla.skills.meta.types import MetaMatch, MetaPlan, MetaStep
|
|
|
|
plan = MetaPlan(
|
|
name="rendered-inputs-test",
|
|
triggers=("x",),
|
|
priority=10,
|
|
steps=(
|
|
MetaStep(
|
|
id="a",
|
|
skill="alpha",
|
|
kind="agent",
|
|
with_args={"q": "hello {{ inputs.user_message }}"},
|
|
),
|
|
),
|
|
)
|
|
match = MetaMatch(plan=plan, inputs={"user_message": "world"})
|
|
begin_inputs: list[dict[str, Any]] = []
|
|
|
|
async def begin_cb(
|
|
step_id: str,
|
|
effective_skill: str,
|
|
rendered_inputs: dict[str, Any],
|
|
) -> None:
|
|
begin_inputs.append(rendered_inputs)
|
|
|
|
async def dispatch_stub(step, effective_skill, inputs, outputs):
|
|
from opensquilla.skills.meta.events import _StepDone
|
|
yield _StepDone(text="ok")
|
|
|
|
async def preface_stub(step_id: str, effective_skill: str):
|
|
if False:
|
|
yield None
|
|
return
|
|
|
|
async for _ in run_dag(
|
|
match,
|
|
dispatch_step_stream=dispatch_stub,
|
|
yield_skill_view_preface=preface_stub,
|
|
on_step_begin=begin_cb,
|
|
):
|
|
pass
|
|
|
|
assert begin_inputs == [{"q": "hello world"}]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_run_dag_tool_result_includes_scoped_step_usage() -> None:
|
|
from opensquilla.engine.types import ToolResultEvent
|
|
from opensquilla.engine.usage import UsageTracker
|
|
from opensquilla.skills.meta.events import _StepDone
|
|
from opensquilla.skills.meta.scheduler import run_dag
|
|
from opensquilla.skills.meta.types import MetaMatch, MetaPlan, MetaStep
|
|
|
|
plan = MetaPlan(
|
|
name="step-usage-test",
|
|
triggers=("x",),
|
|
priority=10,
|
|
steps=(
|
|
MetaStep(id="a", skill="alpha", kind="llm_chat"),
|
|
MetaStep(id="b", skill="beta", kind="llm_chat"),
|
|
),
|
|
)
|
|
match = MetaMatch(plan=plan, inputs={})
|
|
tracker = UsageTracker()
|
|
results: dict[str, ToolResultEvent] = {}
|
|
|
|
async def dispatch_stub(step, effective_skill, inputs, outputs):
|
|
tracker.add(
|
|
"session-a",
|
|
input_tokens=10 if step.id == "a" else 20,
|
|
output_tokens=1 if step.id == "a" else 2,
|
|
model_id=f"model-{step.id}",
|
|
billed_cost=0.123 if step.id == "a" else 0.0,
|
|
)
|
|
yield _StepDone(text=f"{step.id}-done")
|
|
|
|
async def preface_stub(step_id: str, effective_skill: str):
|
|
if False:
|
|
yield None
|
|
return
|
|
|
|
async for event in run_dag(
|
|
match,
|
|
dispatch_step_stream=dispatch_stub,
|
|
yield_skill_view_preface=preface_stub,
|
|
max_parallelism=2,
|
|
usage_tracker=tracker,
|
|
session_key="session-a",
|
|
usage_scope_prefix="run-1",
|
|
):
|
|
if isinstance(event, ToolResultEvent):
|
|
results[event.tool_name] = event
|
|
|
|
usage_a = (results["meta-step:a"].arguments or {})["usage"]
|
|
usage_b = (results["meta-step:b"].arguments or {})["usage"]
|
|
assert usage_a["input_tokens"] == 10
|
|
assert usage_a["output_tokens"] == 1
|
|
assert usage_a["total_tokens"] == 11
|
|
assert usage_a["model"] == "model-a"
|
|
assert usage_a["cost_usd"] == pytest.approx(0.123)
|
|
assert usage_a["billed_cost"] == pytest.approx(0.123)
|
|
assert usage_a["billed_cost_usd"] == pytest.approx(0.123)
|
|
assert usage_a["cost_source"] == "provider_billed"
|
|
assert usage_a["is_provider_billed"] is True
|
|
assert usage_b["input_tokens"] == 20
|
|
assert usage_b["output_tokens"] == 2
|
|
assert usage_b["total_tokens"] == 22
|
|
assert usage_b["model"] == "model-b"
|
|
assert usage_b["billed_cost"] == 0.0
|
|
assert usage_b["is_provider_billed"] is False
|