wehub-resource-sync
351 lines
13 KiB
Python
351 lines
13 KiB
Python
"""Native sub-agent completions must reach the parent inbox.
|
|
|
|
A native CLI sub-agent's completion is the forwarder's ``external_session_status:
|
|
idle`` (or ``failed``) event POSTed to the child's ``/events``. The runner turns
|
|
that into a ``sub_agent`` payload in the parent's async inbox (``sys_read_inbox``)
|
|
so the orchestrator wakes instead of busy-polling ``sys_session_get_history``.
|
|
|
|
Delivery is dropped whenever the runner's in-memory work entry for the child is
|
|
missing — a reconnect / restart wiped ``_subagent_work_by_child`` mid-turn, or a
|
|
``sys_session_create`` child never registered one (the server records a
|
|
``parent_session_id`` but no ``sub_agent_name``). The old code then returned
|
|
HTTP 204 and lost the completion. The fix rebuilds the entry from the server
|
|
snapshot before delivering, and returns 503 (so the forwarder retries) when
|
|
delivery still can't be confirmed on this runner.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from collections.abc import Iterator
|
|
from typing import Any
|
|
|
|
import pytest
|
|
|
|
from omnigent.runner import app as runner_app
|
|
from omnigent.runner import create_runner_app
|
|
from omnigent.spec.types import AgentSpec, ExecutorSpec
|
|
|
|
# Reuse the proven runner-turn stubs from the sessions-native suite.
|
|
from tests.runner.helpers import NullServerClient
|
|
from tests.runner.test_app_sessions_native import (
|
|
_FakeProcessManager,
|
|
_runner_client,
|
|
_ScriptedHarnessClient,
|
|
)
|
|
|
|
PARENT_SESSION_ID = "conv_parent_orchestrator"
|
|
CHILD_SESSION_ID = "conv_child_reviewer"
|
|
|
|
|
|
@pytest.fixture
|
|
def _clean_subagent_registry() -> Iterator[None]:
|
|
"""Snapshot and restore the process-wide sub-agent / inbox maps.
|
|
|
|
The sub-agent work registry and inbox queues live in module-level dicts on
|
|
``omnigent.runner.app`` that otherwise leak across tests. Clear them before
|
|
the test and restore the originals after.
|
|
"""
|
|
saved = (
|
|
dict(runner_app._subagent_work_by_child),
|
|
{k: set(v) for k, v in runner_app._subagent_work_by_parent.items()},
|
|
dict(runner_app._session_inboxes_ref),
|
|
set(runner_app._drained_delivered_subagent_children),
|
|
)
|
|
runner_app._subagent_work_by_child.clear()
|
|
runner_app._subagent_work_by_parent.clear()
|
|
runner_app._session_inboxes_ref.clear()
|
|
runner_app._drained_delivered_subagent_children.clear()
|
|
try:
|
|
yield
|
|
finally:
|
|
runner_app._subagent_work_by_child.clear()
|
|
runner_app._subagent_work_by_child.update(saved[0])
|
|
runner_app._subagent_work_by_parent.clear()
|
|
runner_app._subagent_work_by_parent.update(saved[1])
|
|
runner_app._session_inboxes_ref.clear()
|
|
runner_app._session_inboxes_ref.update(saved[2])
|
|
runner_app._drained_delivered_subagent_children.clear()
|
|
runner_app._drained_delivered_subagent_children.update(saved[3])
|
|
|
|
|
|
class _SnapshotServerClient(NullServerClient):
|
|
"""Server client whose ``GET /v1/sessions/{child}`` carries the sub-agent snapshot.
|
|
|
|
Mirrors ``SessionResponse`` (server routes/sessions.py): the authoritative
|
|
source the runner uses to rebuild a lost sub-agent work entry. The body is
|
|
configurable so a test can model a declared sub-agent (``sub_agent_name``
|
|
set), a ``sys_session_create`` child (``sub_agent_name`` null but
|
|
``parent_session_id`` set + ``agent_name``), or a top-level session (no
|
|
parent). All other endpoints fall through to the empty-200 base.
|
|
"""
|
|
|
|
def __init__(self, child_body: dict[str, Any]) -> None:
|
|
"""Configure the JSON body returned for the child session GET."""
|
|
self._child_body = child_body
|
|
|
|
class _Resp:
|
|
def __init__(self, payload: dict[str, Any]) -> None:
|
|
self.status_code = 200
|
|
self._payload = payload
|
|
|
|
def json(self) -> dict[str, Any]:
|
|
return self._payload
|
|
|
|
def raise_for_status(self) -> None:
|
|
return None
|
|
|
|
async def get(self, url: str, **kwargs: Any) -> Any:
|
|
del kwargs
|
|
if url.rstrip("/").endswith(CHILD_SESSION_ID):
|
|
return self._Resp(self._child_body)
|
|
if url.rstrip("/").endswith("/items"):
|
|
return self._Resp({"data": [], "has_more": False})
|
|
return self._Response()
|
|
|
|
|
|
def _child_snapshot(
|
|
*,
|
|
sub_agent_name: str | None,
|
|
parent_session_id: str | None,
|
|
agent_name: str | None = "cursor-native-ui",
|
|
) -> dict[str, Any]:
|
|
"""Build a child ``SessionResponse``-shaped body."""
|
|
return {
|
|
"id": CHILD_SESSION_ID,
|
|
"agent_id": "ag_reviewer",
|
|
"agent_name": agent_name,
|
|
"sub_agent_name": sub_agent_name,
|
|
"parent_session_id": parent_session_id,
|
|
"created_at": 0,
|
|
"workspace": None,
|
|
}
|
|
|
|
|
|
async def _post_native_idle(
|
|
*,
|
|
child_body: dict[str, Any],
|
|
seed_parent_inbox: bool,
|
|
register_work: bool,
|
|
output: str = "review complete: LGTM",
|
|
) -> tuple[int, list[dict[str, Any]]]:
|
|
"""POST a native ``external_session_status: idle`` and return (http, inbox items).
|
|
|
|
Models the forwarder reporting a finished native sub-agent turn.
|
|
``register_work`` seeds the in-memory work entry (the healthy case); leaving
|
|
it ``False`` models a reconnect-wiped map or a ``sys_session_create`` child
|
|
the dispatch never registered. ``seed_parent_inbox`` controls whether the
|
|
parent's inbox queue is present on this runner.
|
|
"""
|
|
if seed_parent_inbox:
|
|
runner_app._session_inboxes_ref[PARENT_SESSION_ID] = asyncio.Queue()
|
|
if register_work:
|
|
runner_app.register_subagent_work(
|
|
parent_session_id=PARENT_SESSION_ID,
|
|
child_session_id=CHILD_SESSION_ID,
|
|
agent="reviewer",
|
|
title="review",
|
|
)
|
|
|
|
pm = _FakeProcessManager(_ScriptedHarnessClient([]))
|
|
|
|
async def _resolver(agent_id: str, session_id: str | None = None) -> AgentSpec:
|
|
del agent_id, session_id
|
|
return AgentSpec(
|
|
spec_version=1,
|
|
name="reviewer",
|
|
executor=ExecutorSpec(type="omnigent", config={"harness": "claude-native"}),
|
|
)
|
|
|
|
app = create_runner_app(
|
|
process_manager=pm, # type: ignore[arg-type]
|
|
spec_resolver=_resolver,
|
|
server_client=_SnapshotServerClient(child_body), # type: ignore[arg-type]
|
|
)
|
|
|
|
async with _runner_client(app) as client:
|
|
resp = await client.post(
|
|
f"/v1/sessions/{CHILD_SESSION_ID}/events",
|
|
json={
|
|
"type": "external_session_status",
|
|
"data": {"status": "idle", "output": output},
|
|
},
|
|
)
|
|
|
|
inbox = runner_app._session_inboxes_ref.get(PARENT_SESSION_ID)
|
|
items: list[dict[str, Any]] = []
|
|
if inbox is not None:
|
|
while not inbox.empty():
|
|
items.append(inbox.get_nowait())
|
|
return resp.status_code, items
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_native_completion_recovers_reconnect_wiped_work_entry(
|
|
_clean_subagent_registry: None,
|
|
) -> None:
|
|
"""A declared native sub-agent still delivers after its work entry was lost.
|
|
|
|
With no in-memory work entry (a reconnect wiped it mid-turn), the idle edge
|
|
dropped silently on the old code. The fix rebuilds the entry from the
|
|
snapshot's ``parent_session_id`` + ``sub_agent_name`` and delivers.
|
|
"""
|
|
http, items = await _post_native_idle(
|
|
child_body=_child_snapshot(sub_agent_name="reviewer", parent_session_id=PARENT_SESSION_ID),
|
|
seed_parent_inbox=True,
|
|
register_work=False,
|
|
)
|
|
|
|
assert items, (
|
|
"native sub-agent reported idle but nothing was delivered to the parent "
|
|
"inbox: the work entry was missing (reconnect-wiped) and the idle edge "
|
|
f"was silently 204-acked. (http={http})"
|
|
)
|
|
payload = items[0]
|
|
assert payload["type"] == "sub_agent"
|
|
assert payload["conversation_id"] == CHILD_SESSION_ID
|
|
assert payload["status"] == "completed"
|
|
assert payload["output"] == "review complete: LGTM"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sys_session_create_child_without_sub_agent_name_delivers(
|
|
_clean_subagent_registry: None,
|
|
) -> None:
|
|
"""A ``sys_session_create`` child (no ``sub_agent_name``) still wakes the parent.
|
|
|
|
The child has ``agent_name: cursor-native-ui`` but ``sub_agent_name: null``,
|
|
and the dispatch never registered a work entry. The fix recovers the parent
|
|
link from the snapshot (keying on ``parent_session_id``) and labels the work
|
|
with the agent name.
|
|
"""
|
|
http, items = await _post_native_idle(
|
|
child_body=_child_snapshot(
|
|
sub_agent_name=None,
|
|
parent_session_id=PARENT_SESSION_ID,
|
|
agent_name="cursor-native-ui",
|
|
),
|
|
seed_parent_inbox=True,
|
|
register_work=False,
|
|
)
|
|
|
|
assert items, (
|
|
f"sys_session_create child reported idle but the parent inbox stayed empty. (http={http})"
|
|
)
|
|
assert items[0]["status"] == "completed"
|
|
assert items[0]["agent"] == "cursor-native-ui"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_healthy_registered_work_entry_still_delivers(
|
|
_clean_subagent_registry: None,
|
|
) -> None:
|
|
"""Control: the normal path (work entry present) keeps delivering.
|
|
|
|
Guards against the fix regressing the common case where dispatch already
|
|
registered the work entry on this runner.
|
|
"""
|
|
http, items = await _post_native_idle(
|
|
child_body=_child_snapshot(sub_agent_name="reviewer", parent_session_id=PARENT_SESSION_ID),
|
|
seed_parent_inbox=True,
|
|
register_work=True,
|
|
)
|
|
|
|
assert http == 204
|
|
assert items and items[0]["status"] == "completed"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_undeliverable_native_completion_returns_503_not_silent_204(
|
|
_clean_subagent_registry: None,
|
|
) -> None:
|
|
"""A recoverable sub-agent whose parent inbox is elsewhere must 503, not 204.
|
|
|
|
When the parent inbox is not on this runner (the parent lives on a different
|
|
runner, or the runner restarted and lost it), delivery cannot be confirmed.
|
|
The handler must return 503 so the forwarder retries and server-side recovery
|
|
re-routes to the parent's runner — instead of a silent 204 that drops it.
|
|
"""
|
|
http, items = await _post_native_idle(
|
|
child_body=_child_snapshot(sub_agent_name="reviewer", parent_session_id=PARENT_SESSION_ID),
|
|
seed_parent_inbox=False,
|
|
register_work=False,
|
|
)
|
|
|
|
assert http == 503, (
|
|
"an undeliverable native sub-agent completion was acked with "
|
|
f"http={http}; expected 503 so the forwarder retries. Items={items!r}"
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_replayed_idle_after_drain_does_not_redeliver(
|
|
_clean_subagent_registry: None,
|
|
) -> None:
|
|
"""The recovery must not re-deliver a child already delivered and drained.
|
|
|
|
Guards the snapshot-recovery arm against a duplicate: once a completion was
|
|
delivered and the parent drained it, the runner keeps a delivered tombstone.
|
|
A replayed idle whose snapshot *does* carry a ``parent_session_id`` (the
|
|
production shape) must NOT rebuild the work entry and re-enqueue — it stays a
|
|
benign already-delivered 204. (The existing suite's dedup test uses a stub
|
|
snapshot with no parent, so it would not catch a recovery-induced re-deliver.)
|
|
"""
|
|
child_body = _child_snapshot(sub_agent_name="reviewer", parent_session_id=PARENT_SESSION_ID)
|
|
# First completion delivers normally.
|
|
http1, items1 = await _post_native_idle(
|
|
child_body=child_body, seed_parent_inbox=True, register_work=True
|
|
)
|
|
assert http1 == 204
|
|
assert len(items1) == 1 # drained by the helper
|
|
|
|
# Mark the child delivered-and-drained, exactly as sys_read_inbox does.
|
|
runner_app.unregister_subagent_work(CHILD_SESSION_ID, remember_drained_delivery=True)
|
|
assert runner_app.get_subagent_work(CHILD_SESSION_ID) is None
|
|
|
|
# Replay the idle — snapshot carries a parent, so a naive recovery would
|
|
# rebuild the entry and re-deliver. The tombstone guard must prevent that.
|
|
pm = _FakeProcessManager(_ScriptedHarnessClient([]))
|
|
|
|
async def _resolver(agent_id: str, session_id: str | None = None) -> AgentSpec:
|
|
del agent_id, session_id
|
|
return AgentSpec(
|
|
spec_version=1,
|
|
name="reviewer",
|
|
executor=ExecutorSpec(type="omnigent", config={"harness": "claude-native"}),
|
|
)
|
|
|
|
app = create_runner_app(
|
|
process_manager=pm, # type: ignore[arg-type]
|
|
spec_resolver=_resolver,
|
|
server_client=_SnapshotServerClient(child_body), # type: ignore[arg-type]
|
|
)
|
|
inbox = runner_app._session_inboxes_ref[PARENT_SESSION_ID]
|
|
async with _runner_client(app) as client:
|
|
replay = await client.post(
|
|
f"/v1/sessions/{CHILD_SESSION_ID}/events",
|
|
json={"type": "external_session_status", "data": {"status": "idle", "output": "x"}},
|
|
)
|
|
|
|
assert replay.status_code == 204
|
|
assert inbox.qsize() == 0, "replayed idle re-delivered a duplicate to the parent inbox"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_top_level_session_idle_is_noop(
|
|
_clean_subagent_registry: None,
|
|
) -> None:
|
|
"""A top-level session (no parent) idle edge stays a quiet 204 no-op.
|
|
|
|
Ensures the recovery arm does not mis-classify a non-sub-agent sender as a
|
|
sub-agent and start 503-ing or fabricating inbox deliveries.
|
|
"""
|
|
http, items = await _post_native_idle(
|
|
child_body=_child_snapshot(sub_agent_name=None, parent_session_id=None),
|
|
seed_parent_inbox=True,
|
|
register_work=False,
|
|
)
|
|
|
|
assert http == 204
|
|
assert items == []
|