Coverage for src/lilbee/server/mcp_mount.py: 100%

1"""Mount the MCP tool server over streamable-http on the Litestar daemon.""" 

2 

3from __future__ import annotations 

4 

5import logging 

6from collections.abc import AsyncIterator, Callable 

7from contextlib import AbstractAsyncContextManager, asynccontextmanager 

8from typing import TYPE_CHECKING, cast 

9 

10from litestar.handlers import asgi 

11from litestar.types import ASGIApp, Receive, Scope, Send 

12from mcp.server.transport_security import TransportSecuritySettings 

13 

14from lilbee.app.endpoints import MCP_PATH 

15from lilbee.core.config import cfg 

16from lilbee.mcp_server import build_mcp_server, set_http_mounted 

17 

18if TYPE_CHECKING: 

19 from litestar import Litestar 

20 from litestar.handlers import ASGIRouteHandler 

21 

22log = logging.getLogger(__name__) 

23 

24_Lifespan = Callable[["Litestar"], AbstractAsyncContextManager[None]] 

25 

26_LOOPBACK_HOSTS = ("127.0.0.1", "::1", "localhost") 

27# Wildcard binds cannot be enumerated for a Host allowlist. Sentinels for 

28# comparison, not bind addresses. 

29_WILDCARD_BINDS = ("0.0.0.0", "::") # noqa: S104 

30 

31 

32def _fmt_host(host: str) -> str: 

33 """Bracket an IPv6 literal so it matches Host/Origin header syntax.""" 

34 return f"[{host}]" if ":" in host else host 

35 

36 

37def _transport_security() -> TransportSecuritySettings: 

38 """DNS-rebinding allowlist scoped to the configured bind host. 

39 

40 Defaults to loopback (the usual bind). When the daemon is bound to a 

41 specific non-loopback host, that host is added so the mount does not 

42 fail closed and reject every request. A wildcard bind (0.0.0.0 / ::) 

43 cannot be enumerated, so only loopback is allowed there. 

44 """ 

45 hosts = [f"{_fmt_host(h)}:*" for h in _LOOPBACK_HOSTS] 

46 origins = [ 

47 f"{scheme}://{_fmt_host(h)}:*" for h in _LOOPBACK_HOSTS for scheme in ("http", "https") 

48 ] 

49 bind = cfg.server_host 

50 if bind in _WILDCARD_BINDS: 

51 # The REST API on this port serves LAN clients fine, so without this 

52 # only /mcp fails, with an opaque transport-security rejection. 

53 log.warning( 

54 "Bound to %s, which cannot be enumerated for a Host allowlist, so %s " 

55 "accepts loopback Host headers only. Bind to a specific address to " 

56 "reach the MCP endpoint from other machines.", 

57 bind, 

58 MCP_PATH, 

59 ) 

60 elif bind and bind not in _LOOPBACK_HOSTS: 

61 hosts.append(f"{_fmt_host(bind)}:*") 

62 origins.extend(f"{scheme}://{_fmt_host(bind)}:*" for scheme in ("http", "https")) 

63 return TransportSecuritySettings( 

64 enable_dns_rebinding_protection=True, 

65 allowed_hosts=hosts, 

66 allowed_origins=origins, 

67 ) 

68 

69 

70def build_mcp_mount() -> tuple[ASGIRouteHandler, _Lifespan]: 

71 """Return the MCP route handler and the session-manager lifespan. 

72 

73 Each mount builds its own MCP server: the SDK caches a single session 

74 manager per server and ``run()`` is single-use, so a shared server's 

75 second lifespan would raise. 

76 """ 

77 # Mark MCP as served over the shared HTTP daemon so single-vault-only tools 

78 # (init, reset) refuse runtime vault-switch / teardown that would race 

79 # concurrent in-flight handlers on the process-global Services singleton. 

80 set_http_mounted(True) 

81 server = build_mcp_server() 

82 asgi_app = cast( 

83 "ASGIApp", 

84 server.streamable_http_app( 

85 streamable_http_path="/", 

86 transport_security=_transport_security(), 

87 ), 

88 ) 

89 manager = server.session_manager 

90 

91 async def _forward(scope: Scope, receive: Receive, send: Send) -> None: 

92 await asgi_app(scope, receive, send) 

93 

94 handler = asgi(MCP_PATH, is_mount=True, copy_scope=True)(_forward) 

95 

96 @asynccontextmanager 

97 async def _session_lifespan(app: Litestar) -> AsyncIterator[None]: 

98 async with manager.run(): 

99 yield 

100 

101 return handler, _session_lifespan