From 5b48e7b3000426d5ed85b981779a9d7b41454175 Mon Sep 17 00:00:00 2001 From: Liang Jiaqing Date: Sun, 18 Oct 2026 15:44:36 +0800 Subject: [PATCH] Connect desktop GA instances to Hub with shared submit and cancel paths --- frontends/desktop_bridge.py | 24 +++++ frontends/hub.py | 14 ++- frontends/tests/test_bridge_hub.py | 130 ++++++++++++++++++++++++ frontends/tests/test_bridge_sessions.py | 2 + 4 files changed, 168 insertions(+), 2 deletions(-) create mode 100644 frontends/tests/test_bridge_hub.py diff --git a/frontends/desktop_bridge.py b/frontends/desktop_bridge.py index 5d0a7488..a4a8cd7c 100644 --- a/frontends/desktop_bridge.py +++ b/frontends/desktop_bridge.py @@ -764,11 +764,32 @@ class AgentManager: agent.inc_out = True agent.verbose = True threading.Thread(target=agent.run, daemon=True, name=f"GA-{sess.id}").start() + sess.agent = agent + from frontends import hub + hub.connect(agent, f"desktop-{sess.id}", + put_task=lambda text: self._hub_command(sess, "put", text), + abort=lambda: self._hub_command(sess, "abort")) return agent finally: with contextlib.suppress(Exception): os.chdir(old_cwd) + def _hub_command(self, sess: Session, op: str, text=None): + """Use the same entrance as Desktop; Hub reads state/output directly from GA.""" + try: + with self.mutation(): + if self.sessions.get(sess.id) is not sess: + return {"error": "session deleted", "code": "gone"} + if op == "put": + if not isinstance(text, str) or not text.strip(): + return {"error": "text must be non-empty", "code": "badop"} + if self._session_has_unfinished_work(sess): + return {"error": "session is busy", "code": "busy"} + return self.submit_prompt(sess.id, text) + return self.cancel(sess.id) + except MaintenanceConflict as exc: + return {"error": str(exc), "code": "busy"} + @staticmethod def _base_display_name(var: str, cfg: Optional[dict]) -> str: c = cfg or {} @@ -977,6 +998,9 @@ class AgentManager: if sess.agent and hasattr(sess.agent, "abort"): with contextlib.suppress(Exception): sess.agent.abort() + peer = getattr(sess.agent, "_hub", None) + if peer: + peer.close() if self._session_has_unfinished_work(sess): self._retired_sessions[sid] = sess emit_session_state(sess, "closed") diff --git a/frontends/hub.py b/frontends/hub.py index e2c41892..8d62a517 100644 --- a/frontends/hub.py +++ b/frontends/hub.py @@ -19,16 +19,26 @@ class HubClient: def __init__(self, name, put_task=None, get_outputs=None, abort=None, state=None, on_ev=None, sub=(), fixed=False, llm=None, inject=None): self.name, self.put_task, self.get_outputs, self.abort, self.llm, self.inject = name, put_task, get_outputs, abort, llm, inject self.state, self.on_ev, self.sub, self.fixed = state or dict, on_ev, list(sub), fixed # fixed -> stable name + self._closed, self._task = False, None self._tc, self._ws, self._lp = {}, None, None # (i,j) -> (fp, title): a finished step is never re-scanned def emit(self, topic, data=None, to=None): # thread-safe fire-and-forget from the host thread if self._ws and self._lp: asyncio.run_coroutine_threadsafe(self._ws.send(json.dumps( {'op': 'ev', 'topic': topic, 'to': to, 'data': data}, default=str)), self._lp) def start(self): - threading.Thread(target=lambda: asyncio.run(self._loop()), daemon=True).start(); return self + def run(): + try: asyncio.run(self._loop()) + except asyncio.CancelledError: pass # explicit close is normal shutdown + threading.Thread(target=run, daemon=True).start(); return self + def close(self): + """Retire a host instance without reconnecting it.""" + self._closed = True + if self._lp and self._lp.is_running() and self._task: + self._lp.call_soon_threadsafe(self._task.cancel) async def _loop(self): + self._lp, self._task = asyncio.get_running_loop(), asyncio.current_task() try: import websockets except ImportError: return - while True: + while not self._closed: try: async with websockets.connect(URL, open_timeout=3, max_size=None) as ws: self._ws, self._lp = ws, asyncio.get_running_loop() diff --git a/frontends/tests/test_bridge_hub.py b/frontends/tests/test_bridge_hub.py new file mode 100644 index 00000000..a7278d20 --- /dev/null +++ b/frontends/tests/test_bridge_hub.py @@ -0,0 +1,130 @@ +"""Real HubClient websocket transport; isolated GA with no model/API calls.""" +import asyncio +import json +import queue +import sys +from types import SimpleNamespace +from pathlib import Path + +from test_bridge_sessions import manager, tmp_ga_root, _mod, _make_stub +from frontends import hub + + +class FakeAgent: + def __init__(self): + self.llm_no = 0 + self.is_running = False + self.intervene = "" + self.all_outputs = [] + self.llmclient = SimpleNamespace(backend=SimpleNamespace(name="zero", history=[])) + + def run(self): + pass + + def list_llms(self): + return [(0, "zero", None), (1, "one", None)] + + def next_llm(self, no): + self.llm_no = no + self.llmclient.backend.name = ["zero", "one"][no] + + def put_task(self, text, **kwargs): + self.q = queue.Queue() + self.is_running = text == "hold" + self.all_outputs.append({"input": text, "outputs": ["answer:" + text]}) + if not self.is_running: + self.q.put({"done": "answer:" + text, "outputs": ["answer:" + text]}) + return self.q + + def abort(self): + self.is_running = False + if hasattr(self, "q"): + self.q.put({"done": "stopped"}) + + +def test_desktop_hub_roundtrip(manager, monkeypatch): + from websockets.asyncio.server import serve + monkeypatch.setattr(hub, "serve", lambda: None) + monkeypatch.setattr(manager, "ensure_ga_import_path", lambda: Path(manager.ga_root)) + monkeypatch.setitem(sys.modules, "agentmain", _make_stub("agentmain", {"GenericAgent": FakeAgent})) + monkeypatch.setattr(sys.modules["plan_state"], "sync_plan_path_from_text", lambda *a: None, raising=False) + + async def exercise(): + arrivals = asyncio.Queue() + peers = [] + + async def accepted(ws): + hello = json.loads(await ws.recv()) + replies = asyncio.Queue() + await arrivals.put((ws, hello, replies)) + async for message in ws: + await replies.put(json.loads(message)) + + async def rpc(peer, op, **kwargs): + ws, _, replies = peer + await ws.send(json.dumps({"op": op, "id": 1, **kwargs})) + return (await asyncio.wait_for(replies.get(), 3))["data"] + + async def wait_for(predicate): + for _ in range(150): + if predicate(): + return + await asyncio.sleep(.02) + raise AssertionError("timed out waiting for bridge/GA state") + + async with serve(accepted, "127.0.0.1", 0) as server: + port = server.sockets[0].getsockname()[1] + monkeypatch.setattr(hub, "URL", f"ws://127.0.0.1:{port}/ws") + sessions = [manager.create_session(), manager.create_session()] + try: + for sess in sessions: + manager.make_agent(sess) + peers.append(await asyncio.wait_for(arrivals.get(), 3)) + assert len({p[1]["name"] for p in peers}) == 2 + assert peers[0][1]["name"] == f"desktop-{sessions[0].id}" + assert set(peers[0][1]["caps"]) == {"get", "put", "abort", "llm", "inject"} + assert (await rpc(peers[0], "put_task", text="hello"))["accepted"] + await wait_for(lambda: not manager._session_has_unfinished_work(sessions[0])) + assert [m["content"] for m in sessions[0].messages] == ["hello", "answer:hello"] + assert not sessions[1].messages + assert (await rpc(peers[0], "get"))["tasks"][0]["input"] == "hello" + assert (await rpc(peers[0], "seg", i=0, j=0))["content"] == "answer:hello" + assert (await rpc(peers[0], "llm", no=1))["cur"] == 1 + assert manager._live_model(sessions[0])["llmNo"] == 1 + assert sessions[1].agent.llm_no == 0 + assert (await rpc(peers[0], "put_task", text="hold"))["accepted"] + await wait_for(lambda: sessions[0].agent.is_running) + assert (await rpc(peers[0], "put_task", text="extra"))["code"] == "busy" + assert (await rpc(peers[0], "inject", kind="intervene", text="steer"))["ok"] + assert sessions[0].agent.intervene == "steer" + await rpc(peers[0], "abort") + await wait_for(lambda: not manager._session_has_unfinished_work(sessions[0])) + assert sessions[0].status == "cancelled" + manager._maintenance_token = "test" + assert (await rpc(peers[1], "put_task", text="blocked"))["code"] == "busy" + manager._maintenance_token = None + assert manager.delete_session(sessions[0].id) + await asyncio.wait_for(peers[0][0].wait_closed(), 3) + assert sessions[0].agent._hub._closed + assert manager._hub_command(sessions[0], "put", "stale")["code"] == "gone" + assert (await rpc(peers[1], "get"))["nt"] == 0 + finally: + manager._maintenance_token = None + for sess in sessions: + if sess.agent: + sess.agent.abort() + sess.agent._hub.close() + for ws, _, _ in peers: + await asyncio.wait_for(ws.wait_closed(), 3) + + asyncio.run(exercise()) + + +def test_closed_hub_never_connects(monkeypatch): + import websockets + def unexpected(*args, **kwargs): + raise AssertionError("closed HubClient must not connect") + monkeypatch.setattr(websockets, "connect", unexpected) + client = hub.HubClient("closed") + client.close() + asyncio.run(client._loop()) diff --git a/frontends/tests/test_bridge_sessions.py b/frontends/tests/test_bridge_sessions.py index f05dbfc3..11be1ecd 100644 --- a/frontends/tests/test_bridge_sessions.py +++ b/frontends/tests/test_bridge_sessions.py @@ -1410,6 +1410,8 @@ def test_agent_constructor_owns_default_model(manager, monkeypatch, selected, ex def run(self): pass + from frontends import hub + monkeypatch.setattr(hub, "connect", lambda *a, **kw: None) monkeypatch.setattr(manager, "ensure_ga_import_path", lambda: Path(manager.ga_root)) monkeypatch.setitem(sys.modules, "agentmain", _make_stub("agentmain", {"GenericAgent": FakeAgent})) monkeypatch.setattr(_mod, "_settings_doc", lambda: {"ui": {"llmNo": 4}})