diff --git a/tests/test_tui_gateway_ws.py b/tests/test_tui_gateway_ws.py index 3fba21777669..3c9792d70407 100644 --- a/tests/test_tui_gateway_ws.py +++ b/tests/test_tui_gateway_ws.py @@ -188,119 +188,30 @@ def test_ws_write_loop_stall_does_not_latch_transport(monkeypatch): loop.close() -def test_ws_transport_serializes_concurrent_sends(): - active_sends = 0 - max_active_sends = 0 - sent = [] +def test_ws_starts_mcp_discovery_before_ready(monkeypatch): + import tui_gateway.entry as entry + + calls = [] + events = [] + + monkeypatch.setattr(server, "_WS_ORPHAN_REAP_GRACE_S", 0) + monkeypatch.setattr(entry, "ensure_mcp_discovery_started", lambda: calls.append("mcp")) class FakeWS: + async def accept(self): + events.append("accept") + async def send_text(self, line): - nonlocal active_sends, max_active_sends - active_sends += 1 - max_active_sends = max(max_active_sends, active_sends) - try: - await asyncio.sleep(0.05) - sent.append(line) - finally: - active_sends -= 1 + if '"gateway.ready"' in line: + events.append(f"ready_after_{len(calls)}") - loop = asyncio.new_event_loop() - thread = threading.Thread(target=loop.run_forever, daemon=True) - thread.start() - try: - transport = ws_mod.WSTransport(FakeWS(), loop, peer="serialize-test") - with concurrent.futures.ThreadPoolExecutor(max_workers=2) as pool: - futures = [ - pool.submit(transport.write, {"idx": 1}), - pool.submit(transport.write, {"idx": 2}), - ] - assert [f.result(timeout=2) for f in futures] == [True, True] + async def receive_text(self): + raise ws_mod._WebSocketDisconnect() - assert len(sent) == 2 - assert max_active_sends == 1 - assert transport._closed is False - finally: - loop.call_soon_threadsafe(loop.stop) - thread.join(timeout=2) - loop.close() + async def close(self): + pass + asyncio.run(ws_mod.handle_ws(FakeWS())) -def test_ws_transport_preserves_cross_batch_order(): - async def scenario(): - entered = [] - first_entered = asyncio.Event() - release_first = asyncio.Event() - second_started = asyncio.Event() - - class FakeWS: - async def send_text(self, line): - entered.append(line) - if line == "A1": - first_entered.set() - await release_first.wait() - - transport = ws_mod.WSTransport( - FakeWS(), asyncio.get_running_loop(), peer="batch-order-test" - ) - first = asyncio.create_task(transport._safe_send_many(["A1", "A2"])) - await first_entered.wait() - - async def send_second(): - second_started.set() - await transport._safe_send_many(["B1", "B2"]) - - second = asyncio.create_task(send_second()) - await second_started.wait() - - # The second task has reached the transport. Without whole-batch - # serialization it runs B1/B2 before this task can resume. - assert entered == ["A1"] - - release_first.set() - await asyncio.gather(first, second) - assert entered == ["A1", "A2", "B1", "B2"] - - asyncio.run(scenario()) - - -def test_ws_write_async_keeps_drained_tokens_with_current_frame(): - async def scenario(): - entered = [] - first_entered = asyncio.Event() - release_first = asyncio.Event() - current_started = asyncio.Event() - - class FakeWS: - async def send_text(self, line): - entered.append(line) - if line == "A1": - first_entered.set() - await release_first.wait() - - transport = ws_mod.WSTransport( - FakeWS(), asyncio.get_running_loop(), peer="async-order-test" - ) - transport._pending_tokens.append("pending-token") - - first = asyncio.create_task(transport._safe_send_many(["A1", "A2"])) - await first_entered.wait() - - async def send_current(): - current_started.set() - await transport.write_async({"id": "current"}) - - current = asyncio.create_task(send_current()) - await current_started.wait() - later = asyncio.create_task(transport._safe_send_many(["later-batch"])) - - release_first.set() - await asyncio.gather(first, current, later) - assert entered == [ - "A1", - "A2", - "pending-token", - json.dumps({"id": "current"}), - "later-batch", - ] - - asyncio.run(scenario()) + assert calls == ["mcp"] + assert events == ["accept", "ready_after_1"] diff --git a/tui_gateway/entry.py b/tui_gateway/entry.py index e4c87be4a113..645fd482119d 100644 --- a/tui_gateway/entry.py +++ b/tui_gateway/entry.py @@ -338,57 +338,59 @@ def join_mcp_discovery(timeout: float | None = None) -> bool: _recovery_times: list[float] = [] + +def _has_configured_mcp_servers() -> bool: + """Return whether startup should attempt MCP discovery. + + Keep this cheap so non-MCP users do not pay the MCP SDK import cost. + """ + try: + from hermes_cli.config import read_raw_config + + mcp_servers = (read_raw_config() or {}).get("mcp_servers") + return isinstance(mcp_servers, dict) and len(mcp_servers) > 0 + except Exception: + # Be conservative: if we can't decide, fall back to attempting + # discovery. The caller starts it in the background. + return True + + +def ensure_mcp_discovery_started() -> None: + """Start background MCP discovery once for gateway entrypoints. + + ``main()`` covers the stdio/TUI path. WebSocket/Desktop entrypoints can + accept sessions without running ``main()``, so they must call this helper + before the first agent snapshots its tool list. + """ + global _mcp_discovery_thread + + if _mcp_discovery_thread is not None: + return + + if not _has_configured_mcp_servers(): + return + + def _discover_mcp_background() -> None: + try: + from tools.mcp_tool import discover_mcp_tools + + discover_mcp_tools() + except Exception: + logger.warning("Background MCP tool discovery failed", exc_info=True) + + import threading as _mcp_threading + + _mcp_discovery_thread = _mcp_threading.Thread( + target=_discover_mcp_background, + name="tui-mcp-discovery", + daemon=True, + ) + _mcp_discovery_thread.start() + + def main(): _install_sidecar_publisher() - # MCP tool discovery — runs in a background daemon thread so a slow or - # unreachable MCP server can't freeze TUI startup. Previously this ran - # inline before ``gateway.ready``, which meant any configured-but-down - # server stalled the whole shell on "summoning hermes…" for the full - # connect-retry backoff (e.g. a dead stdio/http server burns 1+2+4s of - # retries → ~7s of dead air before the composer appears). Discovery is - # idempotent and registers tools into the shared registry as servers - # connect. The agent isn't built until the first prompt, at which point - # ``_make_agent`` briefly joins this thread (``wait_for_mcp_discovery``, - # bounded) so already-spawning fast servers land in the tool snapshot — - # a dead server is simply not waited on past the bound. ``/reload-mcp`` - # rebuilds the snapshot for servers that connect later in the session. - # - # Cold-start guard: importing ``tools.mcp_tool`` transitively pulls the - # full MCP SDK (mcp, pydantic, httpx, jsonschema, starlette parsers — - # ~200ms on macOS). The overwhelming majority of users have no - # ``mcp_servers`` configured, in which case every byte of that import is - # wasted. Check the config first (cheap) and only spawn the discovery - # thread when there's actually MCP work to do, so the import cost stays - # off the path entirely for the common case. - try: - from hermes_cli.config import read_raw_config - _mcp_servers = (read_raw_config() or {}).get("mcp_servers") - _has_mcp_servers = isinstance(_mcp_servers, dict) and len(_mcp_servers) > 0 - except Exception: - # Be conservative: if we can't decide, fall back to attempting - # discovery (still backgrounded, so it can't block startup). - _has_mcp_servers = True - if _has_mcp_servers: - # Spawn via the shared owner in hermes_cli.mcp_startup instead of - # a hand-rolled thread, so the stdio TUI gets the same restart - # semantics as every other surface: a discovery run that completed - # with zero connected servers may be retried by a later spawn call - # instead of latching the process into a no-MCP-tools state. - # wait_for_mcp_discovery/mcp_discovery_in_flight/ - # join_mcp_discovery below already consult that owner. - global _mcp_discovery_enabled - _mcp_discovery_enabled = True - try: - from hermes_cli.mcp_startup import start_background_mcp_discovery - - start_background_mcp_discovery( - logger=logger, thread_name="tui-mcp-discovery" - ) - except Exception: - logger.warning( - "Background MCP tool discovery failed to start", exc_info=True - ) if not write_json({ "jsonrpc": "2.0",