404 lines
17 KiB
Python
404 lines
17 KiB
Python
"""Regression repro: native sub-agent harness must survive a cache refill.
|
|
|
|
Reproduces the production failure where polly (``claude-sdk``) spawns a
|
|
``claude_code`` sub-agent (``claude-native``) and the web UI then shows
|
|
*"Bridge closed: terminal resource not found or not running"*.
|
|
|
|
Root cause (see investigation notes): the child's harness is resolved by
|
|
swapping the parent spec to the named sub-agent's sub-spec
|
|
(``_find_spec_by_name``). That swap is gated on the in-memory dict
|
|
``_session_sub_agent_names``, populated ONLY in ``POST /v1/sessions`` from
|
|
``body["sub_agent_name"]``. The streaming dispatch path
|
|
(``_stream_message_to_harness`` -> ``_resolve_harness_config``) skips the
|
|
swap entirely: it derives the harness from the session's bound ``agent_id``
|
|
(the *parent* polly agent) and so resolves ``claude-sdk``. After a tunnel
|
|
reconnect / spec-cache eviction (when the in-memory map is empty), a turn
|
|
that takes this path asks ``HarnessProcessManager.get_client`` for
|
|
``claude-sdk``; the manager sees the harness change
|
|
``claude-native -> claude-sdk`` and respawns, tearing down the live
|
|
claude-native tmux terminal. The UI's bridge to ``terminal_claude_main``
|
|
then fails with the observed error.
|
|
|
|
The server already persists and returns ``sub_agent_name`` on
|
|
``GET /v1/sessions/{id}`` (``SessionResponse.sub_agent_name``); the runner
|
|
just drops it. The invariant under test: **a turn dispatched for a
|
|
sub-agent session asks the process manager to spawn the CHILD's harness
|
|
(``claude-native``) even when the in-memory sub-agent-name map is empty**
|
|
(the post-reconnect state), because the identity is recoverable from the
|
|
server snapshot.
|
|
|
|
This test FAILS on the buggy commit (the runner asks for ``claude-sdk``)
|
|
and PASSES once the streaming dispatch path is made sub-agent-aware.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
from typing import Any
|
|
|
|
import pytest
|
|
|
|
from omnigent.runner import create_runner_app
|
|
from omnigent.spec.types import AgentSpec, ExecutorSpec
|
|
|
|
# Reuse the proven harness/process-manager/client stubs from the sessions-native
|
|
# suite so this repro drives the exact same dispatch path the runner uses.
|
|
from tests.runner.test_app_sessions_native import (
|
|
NullServerClient,
|
|
_FakeProcessManager,
|
|
_runner_client,
|
|
_ScriptedHarnessClient,
|
|
_sse,
|
|
)
|
|
|
|
PARENT_AGENT_ID = "ag_polly"
|
|
CHILD_SESSION_ID = "conv_child_claude_code"
|
|
SUB_AGENT_NAME = "claude_code"
|
|
|
|
|
|
def _polly_spec_tree() -> AgentSpec:
|
|
"""Parent polly (claude-sdk) with a claude_code (claude-native) child.
|
|
|
|
Mirrors ``examples/polly/config.yaml`` + its ``claude_code`` sub-agent.
|
|
"""
|
|
child = AgentSpec(
|
|
spec_version=1,
|
|
name=SUB_AGENT_NAME,
|
|
executor=ExecutorSpec(type="omnigent", config={"harness": "claude-native"}),
|
|
)
|
|
return AgentSpec(
|
|
spec_version=1,
|
|
name="polly",
|
|
executor=ExecutorSpec(type="omnigent", config={"harness": "claude-sdk"}),
|
|
sub_agents=[child],
|
|
)
|
|
|
|
|
|
async def _parent_spec_resolver(agent_id: str, session_id: str | None = None) -> AgentSpec:
|
|
"""Resolve any agent_id to the parent polly tree (as the live server does).
|
|
|
|
A sub-agent session is bound to its *parent* agent_id server-side, so the
|
|
runner's spec_resolver returns the parent tree; only the sub-agent-name
|
|
swap turns it into the child spec.
|
|
"""
|
|
del agent_id, session_id
|
|
return _polly_spec_tree()
|
|
|
|
|
|
class _SubAgentSnapshotServer(NullServerClient):
|
|
"""Server client whose ``GET /v1/sessions/{child}`` carries sub_agent_name.
|
|
|
|
Mirrors ``SessionResponse.sub_agent_name`` (server routes/sessions.py): the
|
|
authoritative source the runner can use to recover the sub-agent identity
|
|
after the in-memory ``_session_sub_agent_names`` map is lost (e.g. after a
|
|
tunnel reconnect). All other endpoints fall through to the empty-200 base.
|
|
"""
|
|
|
|
class _Resp:
|
|
def __init__(self, payload: dict[str, Any]) -> None:
|
|
self.status_code = 200
|
|
self._payload = payload
|
|
self.text = json.dumps(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
|
|
# The bare session GET carries the snapshot the runner needs.
|
|
if url.rstrip("/").endswith(CHILD_SESSION_ID):
|
|
return self._Resp(
|
|
{
|
|
"agent_id": PARENT_AGENT_ID,
|
|
"sub_agent_name": SUB_AGENT_NAME,
|
|
"parent_session_id": "conv_parent_polly",
|
|
"created_at": 0,
|
|
"workspace": None,
|
|
}
|
|
)
|
|
if url.rstrip("/").endswith("/items"):
|
|
return self._Resp({"data": [], "has_more": False})
|
|
return super()._Response()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_subagent_turn_spawns_child_native_harness_without_prior_post() -> None:
|
|
"""A turn for a sub-agent session must spawn the CHILD (claude-native) harness.
|
|
|
|
Drives ``POST /v1/sessions/{child}/events?stream=true`` with NO prior
|
|
``POST /v1/sessions`` — exactly the post-reconnect state where the runner's
|
|
in-memory ``_session_sub_agent_names`` map is empty. The runner must still
|
|
ask the process manager for ``claude-native`` (the child's harness),
|
|
recovering ``sub_agent_name`` from the server snapshot.
|
|
|
|
Buggy commit: the runner derives the harness from the bound parent
|
|
``agent_id`` and asks for ``claude-sdk`` -> the manager respawns the
|
|
process, killing the native terminal -> "Bridge closed".
|
|
"""
|
|
hc = _ScriptedHarnessClient(
|
|
[
|
|
_sse({"type": "response.created", "response": {"id": "r1"}}),
|
|
_sse({"type": "response.completed", "response": {"id": "r1"}}),
|
|
]
|
|
)
|
|
pm = _FakeProcessManager(hc)
|
|
app = create_runner_app(
|
|
process_manager=pm, # type: ignore[arg-type]
|
|
spec_resolver=_parent_spec_resolver,
|
|
server_client=_SubAgentSnapshotServer(), # type: ignore[arg-type]
|
|
)
|
|
|
|
async with _runner_client(app) as client:
|
|
resp = await client.post(
|
|
f"/v1/sessions/{CHILD_SESSION_ID}/events",
|
|
params={"stream": "true"},
|
|
json={
|
|
"type": "message",
|
|
"role": "user",
|
|
"agent_id": PARENT_AGENT_ID,
|
|
"content": [{"type": "input_text", "text": "hi"}],
|
|
},
|
|
)
|
|
assert resp.status_code == 200, f"{resp.status_code} {resp.text}"
|
|
# Drain the SSE body so the streaming turn fully runs get_client.
|
|
_ = resp.text
|
|
|
|
harnesses = [h for (_conv, h, _env) in pm.get_client_calls]
|
|
assert harnesses, "the turn never asked the process manager for a harness"
|
|
assert all(h == "claude-native" for h in harnesses), (
|
|
f"runner asked the process manager to spawn {harnesses!r} for the "
|
|
"sub-agent session; expected only 'claude-native'. A 'claude-sdk' "
|
|
"spawn is the bug: it respawns the harness and tears down the live "
|
|
"claude-native terminal ('Bridge closed: terminal resource not found')."
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_subagent_background_turn_resolves_child_native_harness() -> None:
|
|
"""The fire-and-forget turn path must also resolve the CHILD harness.
|
|
|
|
``POST /v1/sessions/{child}/events`` with ``stream=false`` runs the
|
|
PRIMARY turn path (``_run_turn_bg`` -> ``_run_turn_bg_setup_and_stream``),
|
|
which derives the harness from the (swapped) cached spec and bakes it
|
|
into the ``TurnDispatch``. That path reads the sub-agent name to perform
|
|
the swap; with the in-memory map empty (post-reconnect) it must recover
|
|
the name from the server snapshot, otherwise it bakes the PARENT
|
|
``claude-sdk`` harness and the process manager respawns the child's
|
|
``claude-native`` terminal away ("Bridge closed").
|
|
|
|
This covers the gap the streaming test above does not: the background
|
|
path computes the harness itself rather than deferring to
|
|
``_resolve_harness_config``.
|
|
"""
|
|
hc = _ScriptedHarnessClient(
|
|
[
|
|
_sse({"type": "response.created", "response": {"id": "r1"}}),
|
|
_sse({"type": "response.completed", "response": {"id": "r1"}}),
|
|
]
|
|
)
|
|
pm = _FakeProcessManager(hc)
|
|
app = create_runner_app(
|
|
process_manager=pm, # type: ignore[arg-type]
|
|
spec_resolver=_parent_spec_resolver,
|
|
server_client=_SubAgentSnapshotServer(), # type: ignore[arg-type]
|
|
)
|
|
|
|
async with _runner_client(app) as client:
|
|
resp = await client.post(
|
|
f"/v1/sessions/{CHILD_SESSION_ID}/events",
|
|
json={
|
|
"type": "message",
|
|
"role": "user",
|
|
"agent_id": PARENT_AGENT_ID,
|
|
"content": [{"type": "input_text", "text": "hi"}],
|
|
},
|
|
)
|
|
# Fire-and-forget returns 202; the background turn runs get_client.
|
|
assert resp.status_code == 202, f"{resp.status_code} {resp.text}"
|
|
|
|
# Let the background turn task reach get_client.
|
|
for _ in range(200):
|
|
if pm.get_client_calls:
|
|
break
|
|
await asyncio.sleep(0.01)
|
|
|
|
harnesses = [h for (_conv, h, _env) in pm.get_client_calls]
|
|
assert harnesses, "the background turn never asked the process manager for a harness"
|
|
assert all(h == "claude-native" for h in harnesses), (
|
|
f"background turn asked the process manager to spawn {harnesses!r} for "
|
|
"the sub-agent session; expected only 'claude-native'. A 'claude-sdk' "
|
|
"spawn is the bug: it respawns the harness and tears down the live "
|
|
"claude-native terminal ('Bridge closed: terminal resource not found')."
|
|
)
|
|
|
|
|
|
class _CatchUpServer(_SubAgentSnapshotServer):
|
|
"""Like the snapshot server, but ``/items`` returns one fresh user message.
|
|
|
|
The catch-up scan only starts a turn when new user items arrived after the
|
|
last seen cursor — this provides exactly one so the scan dispatches.
|
|
"""
|
|
|
|
async def get(self, url: str, **kwargs: Any) -> Any:
|
|
del kwargs
|
|
u = url.rstrip("/")
|
|
if u.endswith(f"{CHILD_SESSION_ID}/items"):
|
|
return self._Resp(
|
|
{
|
|
"data": [
|
|
{
|
|
"id": "item_new1",
|
|
"type": "message",
|
|
"data": {
|
|
"role": "user",
|
|
"content": [{"type": "input_text", "text": "continue"}],
|
|
},
|
|
}
|
|
],
|
|
"has_more": False,
|
|
}
|
|
)
|
|
if u.endswith(CHILD_SESSION_ID):
|
|
return self._Resp(
|
|
{
|
|
"agent_id": PARENT_AGENT_ID,
|
|
"sub_agent_name": SUB_AGENT_NAME,
|
|
"parent_session_id": "conv_parent_polly",
|
|
"created_at": 0,
|
|
"workspace": None,
|
|
}
|
|
)
|
|
if u.endswith("/items"):
|
|
return self._Resp({"data": [], "has_more": False})
|
|
return super(_SubAgentSnapshotServer, self)._Response()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_reconnect_catch_up_scan_keeps_child_native_harness() -> None:
|
|
"""The reconnect catch-up scan must not flip a sub-agent off claude-native.
|
|
|
|
This drives the REAL ``app.state.catch_up_scan`` (the runner's
|
|
``on_reconnect`` callback) — the exact path that fired in production after
|
|
a Databricks Apps ingress WebSocket recycle. With the sub-agent session
|
|
known to the runner (it has history) but its spec cache holding the PARENT
|
|
spec, ``_is_native_harness`` returns False on the buggy code, so the scan
|
|
does NOT skip it, runs a catch-up turn, and resolves the parent
|
|
``claude-sdk`` harness — asking ``get_client`` for ``claude-sdk`` while the
|
|
live subprocess is ``claude-native`` (the respawn that kills the terminal).
|
|
|
|
The fix recovers ``sub_agent_name`` from the server snapshot so the turn
|
|
resolves ``claude-native`` instead.
|
|
"""
|
|
from omnigent.runner.app import _session_histories_ref
|
|
|
|
pm = _FakeProcessManager(
|
|
_ScriptedHarnessClient(
|
|
[
|
|
_sse({"type": "response.created", "response": {"id": "r1"}}),
|
|
_sse({"type": "response.completed", "response": {"id": "r1"}}),
|
|
]
|
|
)
|
|
)
|
|
app = create_runner_app(
|
|
process_manager=pm, # type: ignore[arg-type]
|
|
spec_resolver=_parent_spec_resolver,
|
|
server_client=_CatchUpServer(), # type: ignore[arg-type]
|
|
)
|
|
|
|
async with _runner_client(app):
|
|
# Field state after a reconnect: the runner knows the session (it has
|
|
# in-memory history) but never repopulated the sub-agent-name map for
|
|
# it. Seed history so the scan iterates this session.
|
|
_session_histories_ref[CHILD_SESSION_ID] = [
|
|
{"role": "user", "content": [{"type": "input_text", "text": "prev"}]}
|
|
]
|
|
try:
|
|
await app.state.catch_up_scan()
|
|
for _ in range(300):
|
|
if pm.get_client_calls:
|
|
break
|
|
await asyncio.sleep(0.01)
|
|
finally:
|
|
_session_histories_ref.pop(CHILD_SESSION_ID, None)
|
|
|
|
sub_calls = [h for (conv, h, _env) in pm.get_client_calls if conv == CHILD_SESSION_ID]
|
|
assert sub_calls, "catch-up scan did not dispatch a turn for the sub-agent session"
|
|
assert all(h == "claude-native" for h in sub_calls), (
|
|
f"reconnect catch-up scan asked the process manager to spawn {sub_calls!r} "
|
|
"for the sub-agent session; expected only 'claude-native'. A 'claude-sdk' "
|
|
"spawn is the production bug: the ingress WebSocket recycle triggers this "
|
|
"scan, which respawns the harness and tears down the live claude-native "
|
|
"terminal ('Bridge closed: terminal resource not found')."
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_resource_access_before_post_caches_child_not_parent_harness() -> None:
|
|
"""A resource request racing ahead of POST must not cache the PARENT harness.
|
|
|
|
Root-enabler race: ``_session_spec_cache`` is the source of truth for a
|
|
session's harness. Two paths populate it and race —
|
|
|
|
* ``POST /v1/sessions`` caches the swapped CHILD spec (claude-native);
|
|
* ``_resolve_session_spec_entry`` (hit by resource endpoints like
|
|
``GET /resources``, filesystem, terminal create) caches the spec the
|
|
bound ``agent_id`` resolves to — the PARENT (claude-sdk).
|
|
|
|
The web UI opens resource panels as it creates a sub-agent, so a resource
|
|
GET frequently lands BEFORE the runner's POST. ``_resolve_session_spec_entry``
|
|
early-returns once the key is cached, so whoever wins sticks. If the
|
|
resource path wins on buggy code it caches the parent, ``_is_native_harness``
|
|
goes False, and the next turn / reconnect flips the harness — tearing down
|
|
the native terminal.
|
|
|
|
This drives that exact ordering: a resource GET first (no prior POST), then
|
|
a turn. The turn must spawn ``claude-native``. On buggy code the cached
|
|
parent spec yields ``claude-sdk``.
|
|
"""
|
|
hc = _ScriptedHarnessClient(
|
|
[
|
|
_sse({"type": "response.created", "response": {"id": "r1"}}),
|
|
_sse({"type": "response.completed", "response": {"id": "r1"}}),
|
|
]
|
|
)
|
|
pm = _FakeProcessManager(hc)
|
|
app = create_runner_app(
|
|
process_manager=pm, # type: ignore[arg-type]
|
|
spec_resolver=_parent_spec_resolver,
|
|
server_client=_SubAgentSnapshotServer(), # type: ignore[arg-type]
|
|
)
|
|
|
|
async with _runner_client(app) as client:
|
|
# Resource access BEFORE any POST /v1/sessions — populates the spec
|
|
# cache via _resolve_session_spec_entry (the racing path).
|
|
rsrc = await client.get(f"/v1/sessions/{CHILD_SESSION_ID}/resources")
|
|
assert rsrc.status_code == 200, f"{rsrc.status_code} {rsrc.text}"
|
|
|
|
# Now a turn dispatches; it must resolve the CHILD harness despite the
|
|
# resource path having cached first.
|
|
resp = await client.post(
|
|
f"/v1/sessions/{CHILD_SESSION_ID}/events",
|
|
params={"stream": "true"},
|
|
json={
|
|
"type": "message",
|
|
"role": "user",
|
|
"agent_id": PARENT_AGENT_ID,
|
|
"content": [{"type": "input_text", "text": "hi"}],
|
|
},
|
|
)
|
|
assert resp.status_code == 200, f"{resp.status_code} {resp.text}"
|
|
_ = resp.text
|
|
|
|
harnesses = [h for (conv, h, _env) in pm.get_client_calls if conv == CHILD_SESSION_ID]
|
|
assert harnesses, "the turn never asked the process manager for a harness"
|
|
assert all(h == "claude-native" for h in harnesses), (
|
|
f"after a resource GET raced ahead of POST, the turn spawned {harnesses!r} "
|
|
"for the sub-agent session; expected only 'claude-native'. A 'claude-sdk' "
|
|
"spawn means the resource path cached the parent spec and it stuck — the "
|
|
"race that ultimately tears down the native terminal ('Bridge closed')."
|
|
)
|