wehub-resource-sync
247 lines
6.9 KiB
Python
247 lines
6.9 KiB
Python
"""TaskRuntime reserves capacity for non-subagent tasks under fan-out load."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from dataclasses import dataclass
|
|
|
|
import pytest
|
|
import structlog.testing
|
|
|
|
from opensquilla.gateway.routing import RouteEnvelope, SourceKind
|
|
from opensquilla.gateway.task_runtime import TaskRuntime
|
|
from opensquilla.session.models import AgentTaskRecord, AgentTaskStatus
|
|
|
|
|
|
@dataclass
|
|
class _StubStorage:
|
|
records: dict
|
|
|
|
@classmethod
|
|
def fresh(cls) -> _StubStorage:
|
|
return cls(records={})
|
|
|
|
async def create_agent_task(self, record: AgentTaskRecord) -> None:
|
|
self.records[record.task_id] = record
|
|
|
|
async def get_agent_task(self, task_id):
|
|
return self.records.get(task_id)
|
|
|
|
async def list_agent_tasks(self, **_):
|
|
return list(self.records.values())
|
|
|
|
async def update_agent_task(self, task_id, **fields):
|
|
rec = self.records.get(task_id)
|
|
if rec is None:
|
|
return
|
|
for k, v in fields.items():
|
|
setattr(rec, k, v)
|
|
|
|
|
|
def _envelope(session_key: str = "agent:s:main") -> RouteEnvelope:
|
|
return RouteEnvelope(
|
|
source_kind=SourceKind.SYSTEM,
|
|
source_name="test",
|
|
agent_id="s",
|
|
session_key=session_key,
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_subagent_reserved_slots_clamp_when_reserved_exceeds_max() -> None:
|
|
"""When config asks reserved >= max, clamp to max-1 (avoiding deadlock)."""
|
|
storage = _StubStorage.fresh()
|
|
|
|
async def handler(_):
|
|
return None
|
|
|
|
with structlog.testing.capture_logs() as logs:
|
|
rt = TaskRuntime(
|
|
storage=storage,
|
|
turn_handler=handler,
|
|
max_concurrency=1,
|
|
subagent_reserved_slots=2,
|
|
)
|
|
|
|
assert rt._subagent_reserved_slots == 0
|
|
assert {
|
|
"event": "task_runtime.subagent_reserved_slots_clamped",
|
|
"requested": 2,
|
|
"max_concurrency": 1,
|
|
"clamped_to": 0,
|
|
"log_level": "warning",
|
|
} in logs
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_subagent_reserved_slots_clamp_to_max_minus_one() -> None:
|
|
"""When max=4 and reserved=4, clamp to 3 so subagents can still acquire."""
|
|
storage = _StubStorage.fresh()
|
|
|
|
async def handler(_):
|
|
return None
|
|
|
|
rt = TaskRuntime(
|
|
storage=storage,
|
|
turn_handler=handler,
|
|
max_concurrency=4,
|
|
subagent_reserved_slots=4,
|
|
)
|
|
|
|
assert rt._subagent_reserved_slots == 3
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_default_run_kind_acquires_normally() -> None:
|
|
"""A non-subagent task acquires a slot and increments counters."""
|
|
storage = _StubStorage.fresh()
|
|
started = asyncio.Event()
|
|
finish = asyncio.Event()
|
|
|
|
async def handler(run):
|
|
started.set()
|
|
await finish.wait()
|
|
|
|
rt = TaskRuntime(
|
|
storage=storage,
|
|
turn_handler=handler,
|
|
max_concurrency=2,
|
|
subagent_reserved_slots=1,
|
|
)
|
|
|
|
handle = await rt.enqueue(_envelope("agent:s:main"), "hi", run_kind="default")
|
|
await asyncio.wait_for(started.wait(), timeout=1.0)
|
|
|
|
# While in-flight, counters reflect the running default-kind task.
|
|
assert rt._global_in_flight == 1
|
|
assert rt._subagent_in_flight == 0
|
|
|
|
finish.set()
|
|
await rt.wait(handle.task_id, timeout=1.0)
|
|
assert rt._global_in_flight == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_subagent_blocks_when_reserved_floor_would_be_violated() -> None:
|
|
"""With max=2 and reserved=1, only one subagent can be in flight at once."""
|
|
storage = _StubStorage.fresh()
|
|
sub_started = []
|
|
finish_first = asyncio.Event()
|
|
finish_second = asyncio.Event()
|
|
|
|
async def handler(run):
|
|
sub_started.append(run.task_id)
|
|
if len(sub_started) == 1:
|
|
await finish_first.wait()
|
|
else:
|
|
await finish_second.wait()
|
|
|
|
rt = TaskRuntime(
|
|
storage=storage,
|
|
turn_handler=handler,
|
|
max_concurrency=2,
|
|
subagent_reserved_slots=1,
|
|
)
|
|
|
|
h1 = await rt.enqueue(_envelope("agent:s:one"), "a", run_kind="subagent")
|
|
h2 = await rt.enqueue(_envelope("agent:s:two"), "b", run_kind="subagent")
|
|
|
|
# Give the loop a few cycles for h1 to start; h2 must remain queued
|
|
# because only 1 subagent slot is available (max=2 - reserved=1).
|
|
for _ in range(20):
|
|
await asyncio.sleep(0)
|
|
assert len(sub_started) == 1
|
|
assert rt._subagent_in_flight == 1
|
|
|
|
finish_first.set()
|
|
# After h1 finishes, h2 should be admitted.
|
|
for _ in range(50):
|
|
await asyncio.sleep(0)
|
|
if len(sub_started) == 2:
|
|
break
|
|
assert len(sub_started) == 2
|
|
|
|
finish_second.set()
|
|
await rt.wait(h1.task_id, timeout=1.0)
|
|
await rt.wait(h2.task_id, timeout=1.0)
|
|
assert rt._global_in_flight == 0
|
|
assert rt._subagent_in_flight == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_no_reservation_leaks_counters_after_handler_error() -> None:
|
|
"""An exception in the turn handler must still drop the in-flight count."""
|
|
storage = _StubStorage.fresh()
|
|
|
|
async def handler(_):
|
|
raise RuntimeError("boom")
|
|
|
|
rt = TaskRuntime(
|
|
storage=storage,
|
|
turn_handler=handler,
|
|
max_concurrency=2,
|
|
subagent_reserved_slots=1,
|
|
)
|
|
|
|
handle = await rt.enqueue(_envelope(), "x", run_kind="subagent")
|
|
record = await rt.wait(handle.task_id, timeout=1.0)
|
|
assert record.status == AgentTaskStatus.FAILED
|
|
assert rt._global_in_flight == 0
|
|
assert rt._subagent_in_flight == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_send_uses_default_run_kind_for_parent_wake() -> None:
|
|
"""TaskRuntime.send (parent wake path) defaults run_kind=default so wakes
|
|
don't consume reserved subagent capacity.
|
|
"""
|
|
storage = _StubStorage.fresh()
|
|
seen: list[str] = []
|
|
|
|
async def handler(run):
|
|
seen.append(run.run_kind)
|
|
|
|
rt = TaskRuntime(
|
|
storage=storage,
|
|
turn_handler=handler,
|
|
max_concurrency=2,
|
|
subagent_reserved_slots=1,
|
|
)
|
|
|
|
# Seed envelope cache by enqueueing once, then send.
|
|
h0 = await rt.enqueue(_envelope("agent:s:main"), "init", run_kind="default")
|
|
await rt.wait(h0.task_id, timeout=1.0)
|
|
|
|
h1 = await rt.send("agent:s:main", "wake")
|
|
await rt.wait(h1.task_id, timeout=1.0)
|
|
|
|
assert "subagent" not in seen
|
|
assert seen == ["default", "default"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_send_passes_stream_event_sink_to_parent_wake_task() -> None:
|
|
storage = _StubStorage.fresh()
|
|
seen_sinks = []
|
|
|
|
async def handler(run):
|
|
seen_sinks.append(run.stream_event_sink)
|
|
|
|
rt = TaskRuntime(
|
|
storage=storage,
|
|
turn_handler=handler,
|
|
max_concurrency=2,
|
|
subagent_reserved_slots=1,
|
|
)
|
|
|
|
h0 = await rt.enqueue(_envelope("agent:s:main"), "init", run_kind="default")
|
|
await rt.wait(h0.task_id, timeout=1.0)
|
|
|
|
async def sink(_event) -> None:
|
|
return None
|
|
|
|
h1 = await rt.send("agent:s:main", "wake", stream_event_sink=sink)
|
|
await rt.wait(h1.task_id, timeout=1.0)
|
|
|
|
assert seen_sinks == [None, sink]
|