mirror of
https://github.com/NousResearch/hermes-agent.git
synced 2026-07-27 17:58:07 +00:00
fix(tui-gateway): start MCP discovery for websocket sessions
This commit is contained in:
LionGateOS
Teknium
parent
486c5ffc8d
commit
830ff5967a
2 changed files with 70 additions and 157 deletions
|
|
@ -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"]
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue