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 ( 

12 ASGIApp, 

13 HTTPResponseBodyEvent, 

14 HTTPResponseStartEvent, 

15 HTTPScope, 

16 Receive, 

17 Scope, 

18 Send, 

19) 

20from mcp.server.transport_security import TransportSecuritySettings 

21 

22from lilbee.app.endpoints import MCP_PATH 

23from lilbee.core.config import cfg 

24from lilbee.mcp_server import build_mcp_server, set_http_mounted 

25 

26if TYPE_CHECKING: 

27 from litestar import Litestar 

28 from litestar.handlers import ASGIRouteHandler 

29 

30log = logging.getLogger(__name__) 

31 

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

33 

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

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

36# comparison, not bind addresses. 

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

38 

39 

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

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

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

43 

44 

45def _transport_security() -> TransportSecuritySettings: 

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

47 

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

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

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

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

52 """ 

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

54 origins = [ 

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

56 ] 

57 bind = cfg.server_host 

58 if bind in _WILDCARD_BINDS: 

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

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

61 log.warning( 

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

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

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

65 bind, 

66 MCP_PATH, 

67 ) 

68 elif bind and bind not in _LOOPBACK_HOSTS: 

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

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

71 return TransportSecuritySettings( 

72 enable_dns_rebinding_protection=True, 

73 allowed_hosts=hosts, 

74 allowed_origins=origins, 

75 ) 

76 

77 

78# Message-post path of the legacy HTTP+SSE transport, inside the mount. 

79_SSE_MESSAGE_PATH = "/messages/" 

80 

81 

82def _header(scope: HTTPScope, name: bytes) -> bytes | None: 

83 """Return the first value of *name* among the request headers, if any.""" 

84 return next((v for k, v in scope["headers"] if k == name), None) 

85 

86 

87def _is_legacy_sse_contact(scope: HTTPScope) -> bool: 

88 """True for a sessionless GET asking for SSE at the mount root. 

89 

90 That request is the legacy HTTP+SSE (2024-11-05) handshake, which clients 

91 fall back to when streamable-http first contact fails. A GET carrying a 

92 session ID is the streamable transport's own standalone stream instead. 

93 """ 

94 if scope["method"] != "GET" or scope["path"] not in ("", "/"): 

95 return False 

96 if _header(scope, b"mcp-session-id") is not None: 

97 return False 

98 accept = _header(scope, b"accept") or b"" 

99 return b"text/event-stream" in accept 

100 

101 

102async def _method_not_allowed(send: Send) -> None: 

103 """Answer 405: a sessionless GET with no SSE accept has no stream to serve.""" 

104 start: HTTPResponseStartEvent = { 

105 "type": "http.response.start", 

106 "status": 405, 

107 "headers": [(b"allow", b"GET, POST, DELETE")], 

108 } 

109 body: HTTPResponseBodyEvent = { 

110 "type": "http.response.body", 

111 "body": b"", 

112 "more_body": False, 

113 } 

114 await send(start) 

115 await send(body) 

116 

117 

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

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

120 

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

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

123 second lifespan would raise. 

124 """ 

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

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

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

128 set_http_mounted(True) 

129 server = build_mcp_server() 

130 security = _transport_security() 

131 streamable_app = cast( 

132 "ASGIApp", 

133 server.streamable_http_app( 

134 streamable_http_path="/", 

135 transport_security=security, 

136 ), 

137 ) 

138 # The SDK ships the two transports as separate apps and no combined mount, 

139 # so this dispatcher serves both at one endpoint per the spec's 

140 # backwards-compatibility flow: without the legacy app a falling-back 

141 # client meets a 400 and reports the server as disconnected. 

142 sse_app = cast( 

143 "ASGIApp", 

144 server.sse_app( 

145 sse_path="/", 

146 message_path=_SSE_MESSAGE_PATH, 

147 transport_security=security, 

148 ), 

149 ) 

150 manager = server.session_manager 

151 

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

153 if scope["type"] != "http": 

154 await streamable_app(scope, receive, send) 

155 return 

156 # The union does not narrow through the "type" comparison above. 

157 http_scope = cast("HTTPScope", scope) 

158 # Starlette route matching needs the bare mount root spelled "/". 

159 if http_scope["path"] == "": 

160 http_scope["path"] = "/" 

161 if http_scope["path"].startswith(_SSE_MESSAGE_PATH) or _is_legacy_sse_contact(http_scope): 

162 # The endpoint event advertises root_path + message path, and 

163 # the mount strips MCP_PATH from the path it forwards. 

164 http_scope["root_path"] = MCP_PATH 

165 await sse_app(scope, receive, send) 

166 return 

167 if ( 

168 http_scope["method"] == "GET" 

169 and http_scope["path"] == "/" 

170 and _header(http_scope, b"mcp-session-id") is None 

171 ): 

172 await _method_not_allowed(send) 

173 return 

174 await streamable_app(scope, receive, send) 

175 

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

177 

178 @asynccontextmanager 

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

180 async with manager.run(): 

181 yield 

182 

183 return handler, _session_lifespan