Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/_experimental/mcp_server/faults/list_outcomes.py: 59%

1"""Per-server outcomes for the aggregate MCP tools/list fan-out. 

2 

3The aggregate listing deliberately keeps serving the healthy subset when one server fails, but a 

4failed server must contribute a classified outcome instead of silently shrinking the list: an empty 

5contribution with no signal makes a broken upstream indistinguishable from a healthy server with no 

6tools. Outcomes carry only machine fields (category and status code) so nothing from an upstream 

7body crosses the trust boundary; classification is total, so any exception out of a server fetch 

8becomes an outcome, never a second failure. 

9""" 

10 

11from __future__ import annotations 

12 

13from collections.abc import Iterator 

14from typing import Final, Literal, NamedTuple, NoReturn, TypeAlias 

15 

16import httpx 

17import httpx2 

18from mcp.types import Tool as MCPTool 

19from pydantic import BaseModel, ConfigDict 

20from typing_extensions import assert_never 

21 

22from litellm.proxy._experimental.mcp_server.exceptions import ( 

23 MCPServerListError, 

24 MCPUpstreamAuthError, 

25) 

26from litellm.proxy._experimental.mcp_server.faults.traversal import iter_exception_tree 

27 

28ListFaultCategory: TypeAlias = Literal[ 

29 "auth_required", 

30 "forbidden", 

31 "timeout", 

32 "unreachable", 

33 "upstream_error", 

34 "internal", 

35] 

36 

37 

38class ServerListOk(BaseModel): 

39 model_config = ConfigDict(frozen=True) 

40 tag: Literal["ok"] = "ok" 

41 tool_count: int 

42 

43 

44class ServerListFault(BaseModel): 

45 """Why a server contributed nothing to a listing: the caller must authenticate upstream 

46 (``auth_required``/``forbidden``), the upstream did not answer (``timeout``/``unreachable``), 

47 the upstream answered outside its contract (``upstream_error``), or the gateway itself failed 

48 (``internal``). ``status_code`` is the upstream HTTP status when one exists.""" 

49 

50 model_config = ConfigDict(frozen=True) 

51 tag: ListFaultCategory 

52 status_code: int | None = None 

53 

54 

55ServerOutcome: TypeAlias = ServerListOk | ServerListFault 

56 

57SERVER_OUTCOMES_META_KEY: Final = "litellm.ai/server_outcomes" 

58"""The tools/list result ``_meta`` key carrying per-server outcomes. Prefixed with the litellm.ai 

59domain per the MCP spec's ``_meta`` key format so it cannot collide with spec-reserved names.""" 

60 

61 

62class AggregateToolListing(NamedTuple): 

63 tools: list[MCPTool] 

64 outcomes: dict[str, ServerOutcome] 

65 

66 

67def _iter_upstream_responses(exc: BaseException) -> Iterator[httpx.Response | httpx2.Response]: 

68 """Yield every upstream ``httpx``/``httpx2`` ``Response`` in the exception tree, in the shared traversal's deliberate 

69 order (explicit causes first, ExceptionGroup members in raise order, the incidental 

70 ``__context__`` chain last), so a response raised while handling the real failure can never 

71 shadow one on the explicit causal chain. Consumers apply their own predicate over the stream: 

72 selecting the first response and THEN testing it would miss a causal auth response sitting 

73 behind an unrelated earlier one.""" 

74 for current in iter_exception_tree(exc): 

75 response = getattr(current, "response", None) 

76 if isinstance(response, (httpx.Response, httpx2.Response)): 76 ↛ 77line 76 didn't jump to line 77 because the condition on line 76 was never true

77 yield response 

78 

79 

80def _find_upstream_response(exc: BaseException) -> httpx.Response | httpx2.Response | None: 

81 return next(_iter_upstream_responses(exc), None) 

82 

83 

84def upstream_auth_challenge(exc: BaseException) -> tuple[int, str | None] | None: 

85 """The first upstream 401/403 in deliberate order and its ``WWW-Authenticate`` challenge, both 

86 read from the SAME response, so the status that picks the carrier channel and the challenge that 

87 rides with it can never come from two different responses in the tree. Non-auth responses do not 

88 end the scan: a causal 401 behind an unrelated 5xx must still be found, or the client never 

89 receives the challenge it needs to re-authenticate.""" 

90 for response in _iter_upstream_responses(exc): 90 ↛ 91line 90 didn't jump to line 91 because the loop on line 90 never started

91 if response.status_code in (401, 403): 

92 return response.status_code, response.headers.get("www-authenticate") 

93 return None 

94 

95 

96def raise_classified_list_failure( 

97 exc: BaseException, 

98 server_name: str, 

99 suppress_challenge: bool = False, 

100) -> NoReturn: 

101 """The one place a failed server fetch chooses its carrier: an upstream 401/403 travels as 

102 ``MCPUpstreamAuthError`` with the upstream's own challenge preserved (a challenge is only ever 

103 fabricated at the HTTP edge, and only for a 401), everything else as ``MCPServerListError`` with 

104 a classified fault. Every fetch site delegates here so the two channels cannot drift apart per 

105 call site. ``suppress_challenge`` is for dcr_bridge servers, whose upstream challenge points 

106 clients at the wrong protected-resource metadata and must never relay.""" 

107 auth: Final = upstream_auth_challenge(exc) 

108 if auth is not None: 108 ↛ 109line 108 didn't jump to line 109 because the condition on line 108 was never true

109 status_code, challenge = auth 

110 raise MCPUpstreamAuthError( 

111 status_code=status_code, 

112 www_authenticate=None if suppress_challenge else challenge, 

113 server_name=server_name, 

114 ) from exc 

115 raise MCPServerListError(classify_list_exception(exc), server_name) from exc 

116 

117 

118def classify_list_exception(exc: BaseException) -> ServerListFault: 

119 """Classify a per-server listing failure into exactly one outcome. Total: an exception this 

120 function cannot recognize is the gateway's own fault (``internal``), never a re-raise.""" 

121 if isinstance(exc, MCPServerListError) and isinstance(exc.fault, ServerListFault): 

122 return exc.fault 

123 if isinstance(exc, MCPUpstreamAuthError): 123 ↛ 124line 123 didn't jump to line 124 because the condition on line 123 was never true

124 tag: Final = "forbidden" if exc.status_code == 403 else "auth_required" 

125 return ServerListFault(tag=tag, status_code=exc.status_code) 

126 if isinstance(exc, TimeoutError): 126 ↛ 127line 126 didn't jump to line 127 because the condition on line 126 was never true

127 return ServerListFault(tag="timeout") 

128 if isinstance(exc, ConnectionError): 128 ↛ 129line 128 didn't jump to line 129 because the condition on line 128 was never true

129 return ServerListFault(tag="unreachable") 

130 auth: Final = upstream_auth_challenge(exc) 

131 if auth is not None: 131 ↛ 132line 131 didn't jump to line 132 because the condition on line 131 was never true

132 status_code, _ = auth 

133 return ServerListFault( 

134 tag="forbidden" if status_code == 403 else "auth_required", 

135 status_code=status_code, 

136 ) 

137 response: Final = _find_upstream_response(exc) 

138 if response is not None: 138 ↛ 139line 138 didn't jump to line 139 because the condition on line 138 was never true

139 return ServerListFault(tag="upstream_error", status_code=response.status_code) 

140 if isinstance(exc, (httpx.TimeoutException, httpx2.TimeoutException)): 140 ↛ 141line 140 didn't jump to line 141 because the condition on line 140 was never true

141 return ServerListFault(tag="timeout") 

142 if isinstance(exc, (httpx.TransportError, httpx2.TransportError)): 

143 return ServerListFault(tag="unreachable") 

144 return ServerListFault(tag="internal") 

145 

146 

147def outcome_wire_value(outcome: ServerOutcome) -> dict[str, object]: 

148 """The client-visible form of one outcome, for the tools/list result ``_meta`` and the REST 

149 response: category plus status code only, never upstream prose or URLs.""" 

150 match outcome.tag: 

151 case "ok": 151 ↛ 152line 151 didn't jump to line 152 because the pattern on line 151 never matched

152 return {"status": "ok", "tool_count": outcome.tool_count} 

153 case "auth_required" | "forbidden" | "timeout" | "unreachable" | "upstream_error" | "internal": 153 ↛ 158line 153 didn't jump to line 158 because the pattern on line 153 always matched

154 return { 

155 "status": outcome.tag, 

156 **({"http_status": outcome.status_code} if outcome.status_code is not None else {}), 

157 } 

158 case _: 

159 assert_never(outcome.tag) 

160 

161 

162def list_fault_http_status(fault: ServerListFault) -> int: 

163 """The truthful HTTP status for a single-upstream listing fault per RFC 9110: the upstream's own 

164 401/403 for auth, 504 for a timeout, 502 for an unreachable or misbehaving upstream, and 500 only 

165 for the gateway's own failure.""" 

166 match fault.tag: 

167 case "auth_required": 

168 return fault.status_code or 401 

169 case "forbidden": 

170 return 403 

171 case "timeout": 

172 return 504 

173 case "unreachable" | "upstream_error": 

174 return 502 

175 case "internal": 

176 return 500 

177 case _: 

178 assert_never(fault.tag)