wehub-resource-sync
553 lines
18 KiB
Python
553 lines
18 KiB
Python
from __future__ import annotations
|
|
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
|
|
from opensquilla.gateway.subagent_announce import (
|
|
_build_subagent_group_outcome,
|
|
_build_terminal_group_payloads,
|
|
_format_parent_wake_message,
|
|
_tracker,
|
|
announce_subagent_completion,
|
|
close_subagent_spawn_group,
|
|
set_background_completion_manager,
|
|
)
|
|
from opensquilla.gateway.task_runtime import SubagentCompletionEvent
|
|
from opensquilla.session.models import AgentTaskStatus, SessionStatus
|
|
|
|
PARENT = "agent:main:webchat:parent"
|
|
PARENT_TASK = "task-parent"
|
|
|
|
|
|
class _SessionRow:
|
|
def __init__(
|
|
self,
|
|
session_key: str,
|
|
*,
|
|
status: str,
|
|
agent_id: str = "worker",
|
|
) -> None:
|
|
self.session_key = session_key
|
|
self.spawned_by = PARENT
|
|
self.parent_session_key = PARENT
|
|
self.agent_id = agent_id
|
|
self.status = status
|
|
self.origin = {"kind": "subagent", "parent_task_id": PARENT_TASK}
|
|
|
|
|
|
class _Storage:
|
|
def __init__(self, tasks_by_session: dict[str, list[SimpleNamespace]]) -> None:
|
|
self.tasks_by_session = tasks_by_session
|
|
self.batch_calls: list[tuple[str, ...]] = []
|
|
|
|
async def list_agent_tasks_for_sessions(
|
|
self,
|
|
session_keys: list[str],
|
|
limit_per_session: int = 100,
|
|
) -> dict[str, list[SimpleNamespace]]:
|
|
self.batch_calls.append(tuple(session_keys))
|
|
return {
|
|
key: list(self.tasks_by_session.get(key, []))[:limit_per_session]
|
|
for key in session_keys
|
|
}
|
|
|
|
|
|
class _SessionManager:
|
|
def __init__(
|
|
self,
|
|
rows: list[_SessionRow],
|
|
*,
|
|
tasks_by_session: dict[str, list[SimpleNamespace]],
|
|
transcripts: dict[str, str],
|
|
) -> None:
|
|
self.rows = rows
|
|
self._storage = _Storage(tasks_by_session)
|
|
self.transcripts = transcripts
|
|
self.messages: list[tuple[str, str, str, dict | None]] = []
|
|
self.finished: list[tuple[str, str]] = []
|
|
|
|
async def list_sessions(
|
|
self,
|
|
agent_id=None,
|
|
status=None,
|
|
limit=100,
|
|
offset=0,
|
|
spawned_by=None,
|
|
):
|
|
rows = self.rows
|
|
if spawned_by is not None:
|
|
rows = [row for row in rows if row.spawned_by == spawned_by]
|
|
return rows[offset : offset + limit]
|
|
|
|
async def read_transcript(self, session_key: str, limit: int = 50):
|
|
text = self.transcripts.get(session_key, "")
|
|
return [SimpleNamespace(role="assistant", content=text)] if text else []
|
|
|
|
async def append_message(
|
|
self,
|
|
key: str,
|
|
*,
|
|
role: str,
|
|
content: str,
|
|
provenance: dict | None = None,
|
|
) -> None:
|
|
self.messages.append((key, role, content, provenance))
|
|
|
|
async def get_session(self, key: str):
|
|
return SimpleNamespace(session_key=key, last_channel=None, last_to=None)
|
|
|
|
async def finish(self, session_key: str, *, status: SessionStatus) -> None:
|
|
self.finished.append((session_key, str(status)))
|
|
for row in self.rows:
|
|
if row.session_key == session_key:
|
|
row.status = str(status)
|
|
|
|
|
|
class _TaskRuntime:
|
|
def __init__(self) -> None:
|
|
self.sent: list[tuple[str, str, dict | None]] = []
|
|
|
|
async def send(self, session_key: str, message: str, provenance: dict | None = None):
|
|
self.sent.append((session_key, message, provenance))
|
|
return SimpleNamespace(task_id=f"wake-{len(self.sent)}")
|
|
|
|
|
|
class _BackgroundCompletion:
|
|
def __init__(self) -> None:
|
|
self.calls: list[str] = []
|
|
self.waiting: list[dict] = []
|
|
self.wakes: list[dict] = []
|
|
|
|
async def emit_waiting(self, **kwargs) -> None:
|
|
self.calls.append("waiting")
|
|
self.waiting.append(kwargs)
|
|
|
|
async def send_parent_wake(self, **kwargs) -> None:
|
|
self.calls.append("wake")
|
|
self.wakes.append(kwargs)
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _clean_tracker():
|
|
_tracker.evict(PARENT)
|
|
set_background_completion_manager(None)
|
|
yield
|
|
_tracker.evict(PARENT)
|
|
set_background_completion_manager(None)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_group_payloads_enrich_non_current_children_from_task_ledger() -> None:
|
|
child_done = "agent:worker:subagent:done"
|
|
child_failed = "agent:worker:subagent:failed"
|
|
manager = _SessionManager(
|
|
[
|
|
_SessionRow(child_done, status="done", agent_id="worker-a"),
|
|
_SessionRow(child_failed, status="failed", agent_id="worker-b"),
|
|
],
|
|
tasks_by_session={
|
|
child_done: [
|
|
SimpleNamespace(
|
|
task_id="task-default-newer",
|
|
agent_id="worker-a",
|
|
status=AgentTaskStatus.FAILED,
|
|
run_kind="default",
|
|
terminal_reason="followup_error",
|
|
created_at=100,
|
|
updated_at=200,
|
|
finished_at=300,
|
|
),
|
|
SimpleNamespace(
|
|
task_id="task-done",
|
|
agent_id="worker-a",
|
|
status=AgentTaskStatus.SUCCEEDED,
|
|
run_kind="subagent",
|
|
terminal_reason="done",
|
|
created_at=10,
|
|
updated_at=20,
|
|
finished_at=30,
|
|
)
|
|
],
|
|
child_failed: [
|
|
SimpleNamespace(
|
|
task_id="task-failed",
|
|
agent_id="worker-b",
|
|
status=AgentTaskStatus.FAILED,
|
|
run_kind="subagent",
|
|
terminal_reason="tool_error",
|
|
error_class="RuntimeError",
|
|
error_message="boom",
|
|
created_at=11,
|
|
updated_at=21,
|
|
finished_at=31,
|
|
)
|
|
],
|
|
},
|
|
transcripts={child_done: "done result", child_failed: "partial failure details"},
|
|
)
|
|
|
|
payloads = await _build_terminal_group_payloads(
|
|
parent_session_key=PARENT,
|
|
parent_task_id=PARENT_TASK,
|
|
session_manager=manager,
|
|
)
|
|
|
|
assert manager._storage.batch_calls == [(child_done, child_failed)]
|
|
assert payloads is not None
|
|
by_child = {payload["child_session_key"]: payload for payload in payloads}
|
|
assert by_child[child_done]["task_id"] == "task-done"
|
|
assert by_child[child_done]["agent_id"] == "worker-a"
|
|
assert by_child[child_done]["status"] == "succeeded"
|
|
assert by_child[child_failed]["task_id"] == "task-failed"
|
|
assert by_child[child_failed]["agent_id"] == "worker-b"
|
|
assert by_child[child_failed]["status"] == "failed"
|
|
assert by_child[child_failed]["terminal_reason"] == "tool_error"
|
|
assert by_child[child_failed]["error_class"] == "RuntimeError"
|
|
assert by_child[child_failed]["error_message"] == "boom"
|
|
|
|
wake_message = _format_parent_wake_message(PARENT_TASK, payloads)
|
|
assert "task_id=task-failed" in wake_message
|
|
assert "agent_id=worker-b" in wake_message
|
|
assert "error_class=RuntimeError" in wake_message
|
|
assert "error_message=boom" in wake_message
|
|
|
|
|
|
def test_parent_wake_bounds_aggregate_subagent_output() -> None:
|
|
payloads = []
|
|
for index in range(5):
|
|
payloads.append(
|
|
{
|
|
"child_session_key": f"agent:worker:subagent:{index}",
|
|
"task_id": f"task-{index}",
|
|
"agent_id": f"worker-{index}",
|
|
"status": "succeeded",
|
|
"terminal_reason": "completed",
|
|
"result": {"text": f"result-{index}-" + ("x" * 12_000)},
|
|
}
|
|
)
|
|
|
|
wake_message = _format_parent_wake_message(PARENT_TASK, payloads)
|
|
|
|
assert len(wake_message) < 20_000
|
|
assert wake_message.count("<untrusted_subagent_result>") == 5
|
|
assert "[subagent result truncated" in wake_message
|
|
assert "child_session_key=agent:worker:subagent:4" in wake_message
|
|
assert "task_id=task-4" in wake_message
|
|
|
|
|
|
def test_group_outcome_summary_counts_and_bounds_non_success_children() -> None:
|
|
long_error = "boom-" * 200
|
|
payloads = [
|
|
{
|
|
"child_session_key": "agent:worker:subagent:done",
|
|
"task_id": "task-done",
|
|
"agent_id": "worker-a",
|
|
"status": "succeeded",
|
|
"terminal_reason": "completed",
|
|
},
|
|
{
|
|
"child_session_key": "agent:worker:subagent:failed",
|
|
"task_id": "task-failed",
|
|
"agent_id": "worker-b",
|
|
"status": "failed",
|
|
"terminal_reason": "tool_error",
|
|
"error_class": "RuntimeError",
|
|
"error_message": long_error,
|
|
},
|
|
{
|
|
"child_session_key": "agent:worker:subagent:timeout",
|
|
"task_id": "task-timeout",
|
|
"agent_id": "worker-c",
|
|
"status": "timeout",
|
|
"terminal_reason": "timeout",
|
|
},
|
|
]
|
|
|
|
outcome = _build_subagent_group_outcome(payloads)
|
|
|
|
assert outcome["total"] == 3
|
|
assert outcome["succeeded"] == 1
|
|
assert outcome["failed"] == 1
|
|
assert outcome["timeout"] == 1
|
|
assert outcome["cancelled"] == 0
|
|
assert outcome["abandoned"] == 0
|
|
assert outcome["non_success"] == 2
|
|
assert outcome["runtime_partial_failure_disclosure_required"] is True
|
|
assert [child["child_session_key"] for child in outcome["failed_children"]] == [
|
|
"agent:worker:subagent:failed",
|
|
"agent:worker:subagent:timeout",
|
|
]
|
|
failed = outcome["failed_children"][0]
|
|
assert failed["task_id"] == "task-failed"
|
|
assert failed["agent_id"] == "worker-b"
|
|
assert failed["status"] == "failed"
|
|
assert failed["terminal_reason"] == "tool_error"
|
|
assert failed["error_class"] == "RuntimeError"
|
|
assert failed["error_message"].startswith("boom-")
|
|
assert len(failed["error_message"]) < len(long_error)
|
|
assert failed["error_message_truncated"] is True
|
|
|
|
|
|
def test_context_overflow_failure_is_sanitized_for_group_outcome_and_wake() -> None:
|
|
raw_error = (
|
|
"Context overflow is in the current turn's recent tool calls or "
|
|
"reasoning tail; history compaction cannot reduce it."
|
|
)
|
|
payloads = [
|
|
{
|
|
"child_session_key": "agent:worker:subagent:failed",
|
|
"task_id": "task-failed",
|
|
"agent_id": "worker-b",
|
|
"status": "failed",
|
|
"terminal_reason": "error",
|
|
"error_class": "current_turn_context_exhausted",
|
|
"error_message": raw_error,
|
|
}
|
|
]
|
|
|
|
outcome = _build_subagent_group_outcome(payloads)
|
|
failed = outcome["failed_children"][0]
|
|
wake_message = _format_parent_wake_message(PARENT_TASK, payloads, outcome=outcome)
|
|
|
|
assert failed["error_class"] == "provider_request_too_large"
|
|
assert "too large" in failed["error_message"].lower()
|
|
assert raw_error not in failed["error_message"]
|
|
assert "current_turn_context_exhausted" not in wake_message
|
|
assert "history compaction cannot reduce it" not in wake_message
|
|
assert "too large" in wake_message.lower()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_group_payloads_fall_back_when_task_ledger_is_unavailable() -> None:
|
|
child = "agent:worker:subagent:fallback"
|
|
manager = _SessionManager(
|
|
[_SessionRow(child, status="timeout", agent_id="worker")],
|
|
tasks_by_session={},
|
|
transcripts={child: "late output"},
|
|
)
|
|
manager._storage = SimpleNamespace()
|
|
|
|
payloads = await _build_terminal_group_payloads(
|
|
parent_session_key=PARENT,
|
|
parent_task_id=PARENT_TASK,
|
|
session_manager=manager,
|
|
)
|
|
|
|
assert payloads == [
|
|
{
|
|
"type": "subagent_completion",
|
|
"parent_session_key": PARENT,
|
|
"child_session_key": child,
|
|
"status": "timeout",
|
|
"terminal_reason": "timeout",
|
|
"parent_task_id": PARENT_TASK,
|
|
"result": {"text": "late output", "truncated": False, "source_role": "assistant"},
|
|
"agent_id": "worker",
|
|
}
|
|
]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_parent_wake_is_deferred_until_yield_and_sent_once() -> None:
|
|
child = "agent:worker:subagent:solo"
|
|
manager = _SessionManager(
|
|
[_SessionRow(child, status="done")],
|
|
tasks_by_session={
|
|
child: [
|
|
SimpleNamespace(
|
|
task_id="task-child",
|
|
agent_id="worker",
|
|
status=AgentTaskStatus.SUCCEEDED,
|
|
run_kind="subagent",
|
|
terminal_reason="done",
|
|
created_at=1,
|
|
updated_at=2,
|
|
finished_at=3,
|
|
)
|
|
]
|
|
},
|
|
transcripts={child: "child output"},
|
|
)
|
|
runtime = _TaskRuntime()
|
|
event = SubagentCompletionEvent(
|
|
parent_session_key=PARENT,
|
|
child_session_key=child,
|
|
task_id="task-child",
|
|
status=AgentTaskStatus.SUCCEEDED,
|
|
terminal_reason="done",
|
|
agent_id="worker",
|
|
parent_task_id=PARENT_TASK,
|
|
)
|
|
|
|
await announce_subagent_completion(
|
|
event,
|
|
session_manager=manager,
|
|
task_runtime=runtime,
|
|
)
|
|
assert runtime.sent == []
|
|
assert manager.messages
|
|
|
|
assert await close_subagent_spawn_group(
|
|
PARENT,
|
|
PARENT_TASK,
|
|
session_manager=manager,
|
|
task_runtime=runtime,
|
|
)
|
|
assert len(runtime.sent) == 1
|
|
assert "[SUBAGENT_COMPLETION_GROUP]" in runtime.sent[0][1]
|
|
assert "task_id=task-child" in runtime.sent[0][1]
|
|
|
|
await announce_subagent_completion(
|
|
event,
|
|
session_manager=manager,
|
|
task_runtime=runtime,
|
|
)
|
|
assert len(runtime.sent) == 1
|
|
|
|
assert await close_subagent_spawn_group(
|
|
PARENT,
|
|
PARENT_TASK,
|
|
session_manager=manager,
|
|
task_runtime=runtime,
|
|
)
|
|
assert len(runtime.sent) == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_parent_wake_provenance_carries_group_outcome_for_mixed_children() -> None:
|
|
child_done = "agent:worker:subagent:done"
|
|
child_failed = "agent:worker:subagent:failed"
|
|
manager = _SessionManager(
|
|
[
|
|
_SessionRow(child_done, status="done", agent_id="worker-a"),
|
|
_SessionRow(child_failed, status="failed", agent_id="worker-b"),
|
|
],
|
|
tasks_by_session={
|
|
child_done: [
|
|
SimpleNamespace(
|
|
task_id="task-done",
|
|
agent_id="worker-a",
|
|
status=AgentTaskStatus.SUCCEEDED,
|
|
run_kind="subagent",
|
|
terminal_reason="completed",
|
|
created_at=1,
|
|
updated_at=2,
|
|
finished_at=3,
|
|
)
|
|
],
|
|
child_failed: [
|
|
SimpleNamespace(
|
|
task_id="task-failed",
|
|
agent_id="worker-b",
|
|
status=AgentTaskStatus.FAILED,
|
|
run_kind="subagent",
|
|
terminal_reason="tool_error",
|
|
error_class="RuntimeError",
|
|
error_message="boom",
|
|
created_at=1,
|
|
updated_at=2,
|
|
finished_at=3,
|
|
)
|
|
],
|
|
},
|
|
transcripts={child_done: "done result", child_failed: "failed details"},
|
|
)
|
|
runtime = _TaskRuntime()
|
|
|
|
assert await close_subagent_spawn_group(
|
|
PARENT,
|
|
PARENT_TASK,
|
|
session_manager=manager,
|
|
task_runtime=runtime,
|
|
)
|
|
|
|
assert len(runtime.sent) == 1
|
|
_, wake_message, provenance = runtime.sent[0]
|
|
assert "Subagents: 1/2 succeeded" in wake_message
|
|
assert provenance is not None
|
|
assert provenance["runtime_partial_failure_disclosure_required"] is True
|
|
outcome = provenance["subagent_group_outcome"]
|
|
assert outcome["total"] == 2
|
|
assert outcome["succeeded"] == 1
|
|
assert outcome["failed"] == 1
|
|
assert outcome["non_success"] == 1
|
|
assert outcome["failed_children"] == [
|
|
{
|
|
"child_session_key": child_failed,
|
|
"task_id": "task-failed",
|
|
"agent_id": "worker-b",
|
|
"status": "failed",
|
|
"terminal_reason": "tool_error",
|
|
"error_class": "RuntimeError",
|
|
"error_message": "boom",
|
|
"error_message_truncated": False,
|
|
}
|
|
]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_close_emits_waiting_when_spawn_group_is_not_terminal() -> None:
|
|
child = "agent:worker:subagent:running"
|
|
manager = _SessionManager(
|
|
[_SessionRow(child, status="running")],
|
|
tasks_by_session={},
|
|
transcripts={},
|
|
)
|
|
runtime = _TaskRuntime()
|
|
background = _BackgroundCompletion()
|
|
set_background_completion_manager(background)
|
|
|
|
closed = await close_subagent_spawn_group(
|
|
PARENT,
|
|
PARENT_TASK,
|
|
session_manager=manager,
|
|
task_runtime=runtime,
|
|
)
|
|
|
|
assert closed is False
|
|
assert runtime.sent == []
|
|
assert background.waiting == [
|
|
{
|
|
"parent_session_key": PARENT,
|
|
"parent_task_id": PARENT_TASK,
|
|
"pending_count": 1,
|
|
}
|
|
]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_close_uses_background_completion_manager_for_parent_wake() -> None:
|
|
child = "agent:worker:subagent:solo"
|
|
manager = _SessionManager(
|
|
[_SessionRow(child, status="done")],
|
|
tasks_by_session={},
|
|
transcripts={child: "child output"},
|
|
)
|
|
runtime = _TaskRuntime()
|
|
background = _BackgroundCompletion()
|
|
set_background_completion_manager(background)
|
|
|
|
assert await close_subagent_spawn_group(
|
|
PARENT,
|
|
PARENT_TASK,
|
|
session_manager=manager,
|
|
task_runtime=runtime,
|
|
)
|
|
|
|
assert runtime.sent == []
|
|
assert background.calls == ["waiting", "wake"]
|
|
assert len(background.wakes) == 1
|
|
assert background.waiting == [
|
|
{
|
|
"parent_session_key": PARENT,
|
|
"parent_task_id": PARENT_TASK,
|
|
"pending_count": 0,
|
|
}
|
|
]
|
|
wake = background.wakes[0]
|
|
assert wake["parent_session_key"] == PARENT
|
|
assert wake["parent_task_id"] == PARENT_TASK
|
|
assert wake["task_runtime"] is runtime
|
|
assert "[SUBAGENT_COMPLETION_GROUP]" in wake["message"]
|