Coverage for src/lilbee/server/anthropic_api/routes.py: 100%

1"""HTTP route for the Anthropic-compatible ``/v1/messages``.""" 

2 

3from __future__ import annotations 

4 

5import asyncio 

6import logging 

7import uuid 

8from collections.abc import AsyncGenerator 

9 

10from litestar import Request, Router, post 

11from litestar.background_tasks import BackgroundTask 

12from litestar.exceptions import NotAuthorizedException, ValidationException 

13from litestar.response import Response, Stream 

14 

15from lilbee.app.services import get_services 

16from lilbee.core.config import cfg 

17from lilbee.core.config.enums import ReasoningMode 

18from lilbee.retrieval.reasoning import effective_reasoning_cap 

19from lilbee.server.anthropic_api.errors import anthropic_error_body, anthropic_error_type 

20from lilbee.server.anthropic_api.models import ( 

21 _THINKING_DISABLED, 

22 AnthropicEventType, 

23 MessagesRequest, 

24 MessagesResponse, 

25) 

26from lilbee.server.anthropic_api.streaming import encode_anthropic_event, encode_anthropic_sse 

27from lilbee.server.anthropic_api.translate import ( 

28 canonical_stream_to_anthropic_events, 

29 canonical_to_messages_response, 

30 messages_to_canonical_request, 

31 resolve_reasoning_mode, 

32) 

33from lilbee.server.auth import auth_checked_in_handler, session_manager 

34from lilbee.server.chat_completions_api.errors import ( 

35 CompletionsErrorCode, 

36 classify_provider_error, 

37) 

38from lilbee.server.chat_dispatch.canonical import CanonicalChatRequest 

39from lilbee.server.chat_dispatch.concurrency import ( 

40 ChatBusyError, 

41 ChatSlotGuard, 

42 acquire_chat_slot_or_busy, 

43) 

44from lilbee.server.chat_dispatch.dispatch import ( 

45 dispatch_chat_stream, 

46 preflight_chat_request, 

47) 

48from lilbee.server.chat_dispatch.reasoning_cap import ( 

49 budget_capped_chars, 

50 cap_aware_chat, 

51 cap_aware_chat_stream, 

52) 

53from lilbee.server.handlers.sse import SSE_MEDIA_TYPE 

54from lilbee.server.validation_format import format_validation 

55 

56log = logging.getLogger(__name__) 

57 

58_INTERNAL_ERROR_MESSAGE = "Internal server error. Check the server logs for details." 

59 

60 

61async def _auth_before_request(request: Request) -> Response | None: 

62 """Reject an unauthenticated caller before Litestar parses the body. 

63 

64 Same rationale as the completions surface: the auth answer must ride this 

65 surface's own error envelope, and it must win over body validation. 

66 """ 

67 return _auth_failure(request) 

68 

69 

70@post("/v1/messages", status_code=200, before_request=_auth_before_request) 

71@auth_checked_in_handler 

72async def messages_endpoint(request: Request, data: MessagesRequest) -> Response | Stream: 

73 """``/v1/messages`` (stream + non-stream + tools), Anthropic wire format. 

74 

75 ``stream: true`` switches the 200 response from JSON to the Anthropic SSE 

76 event stream; the request body picks the arm, matching Anthropic's own 

77 contract. 

78 """ 

79 # Thinking is opt-in per request; the setting only presents it. 

80 mode = resolve_reasoning_mode(data.thinking, default=cfg.messages_reasoning) 

81 cap_chars = budget_capped_chars(effective_reasoning_cap(), _budget_tokens(data)) 

82 try: 

83 req = messages_to_canonical_request(data, mode=mode) 

84 except ValueError as exc: 

85 # Wire-valid but untranslatable (image content, bare tool choice). 

86 return _error_response(400, CompletionsErrorCode.INVALID_REQUEST, str(exc)) 

87 

88 preflight = await _preflight_resolved_model(req) 

89 if isinstance(preflight, Response): 

90 return preflight 

91 resolved_model = preflight 

92 

93 try: 

94 await acquire_chat_slot_or_busy(get_services().provider.max_concurrent_chats()) 

95 except ChatBusyError: 

96 return _error_response( 

97 429, 

98 CompletionsErrorCode.RATE_LIMIT_EXCEEDED, 

99 "Backend is busy. Retry in a moment.", 

100 headers={"Retry-After": "1"}, 

101 ) 

102 

103 guard = ChatSlotGuard() 

104 if req.stream: 

105 # The after-send hook frees the slot when a disconnect lands before the 

106 # generator's first iteration (its finally never runs in that case). 

107 return Stream( 

108 _gated_messages_stream( 

109 req, guard, model=resolved_model, mode=mode, cap_chars=cap_chars 

110 ), 

111 media_type=SSE_MEDIA_TYPE, 

112 background=BackgroundTask(guard.release), 

113 ) 

114 return await _run_non_stream( 

115 req, guard, canonical_model=resolved_model, mode=mode, cap_chars=cap_chars 

116 ) 

117 

118 

119def _budget_tokens(data: MessagesRequest) -> int | None: 

120 """The thinking budget this request asks for; ``disabled`` carries none.""" 

121 if data.thinking is None or data.thinking.type == _THINKING_DISABLED: 

122 return None 

123 return data.thinking.budget_tokens 

124 

125 

126async def _preflight_resolved_model(req: CanonicalChatRequest) -> str | Response: 

127 """Validate *req* before any streaming response starts (see completions).""" 

128 try: 

129 return await asyncio.to_thread(preflight_chat_request, req) 

130 except Exception as exc: 

131 classified = classify_provider_error(exc) 

132 if classified is None: 

133 return _internal_error_response() 

134 return _error_response(classified.http_status, classified.code, classified.message) 

135 

136 

137async def _run_non_stream( 

138 req: CanonicalChatRequest, 

139 guard: ChatSlotGuard, 

140 *, 

141 canonical_model: str, 

142 mode: ReasoningMode = ReasoningMode.SEPARATE, 

143 cap_chars: int = 0, 

144) -> Response: 

145 """Dispatch a non-streaming chat call, translating errors to the envelope.""" 

146 try: 

147 resp = await asyncio.to_thread( 

148 cap_aware_chat, req, canonical_model=canonical_model, cap_chars=cap_chars 

149 ) 

150 except Exception as exc: 

151 classified = classify_provider_error(exc) 

152 if classified is None: 

153 return _internal_error_response() 

154 return _error_response(classified.http_status, classified.code, classified.message) 

155 finally: 

156 await guard.release() 

157 body: MessagesResponse = canonical_to_messages_response( 

158 resp, response_id=_response_id(), mode=mode 

159 ) 

160 return Response(body.model_dump(), media_type="application/json") 

161 

162 

163async def _gated_messages_stream( 

164 req: CanonicalChatRequest, 

165 guard: ChatSlotGuard, 

166 *, 

167 model: str, 

168 mode: ReasoningMode = ReasoningMode.SEPARATE, 

169 cap_chars: int = 0, 

170) -> AsyncGenerator[bytes, None]: 

171 """Drive dispatch -> translate -> SSE-encode, freeing the slot on exit. 

172 

173 A mid-stream failure surfaces as Anthropic's ``event: error`` frame; the 

174 headers are already flushed at 200 by then, so the frame is the only 

175 channel left. 

176 """ 

177 response_id = _response_id() 

178 try: 

179 try: 

180 events = cap_aware_chat_stream( 

181 dispatch_chat_stream(req, canonical_model=model), 

182 req, 

183 canonical_model=model, 

184 cap_chars=cap_chars, 

185 ) 

186 pairs = canonical_stream_to_anthropic_events( 

187 events, model=model, response_id=response_id, mode=mode 

188 ) 

189 async for frame in encode_anthropic_sse(pairs): 

190 yield frame 

191 except Exception as exc: 

192 classified = classify_provider_error(exc) 

193 if classified is None: 

194 log.exception("anthropic messages stream failed") 

195 body = anthropic_error_body("api_error", _INTERNAL_ERROR_MESSAGE) 

196 else: 

197 body = anthropic_error_body( 

198 anthropic_error_type(classified.code), classified.message 

199 ) 

200 yield encode_anthropic_event(AnthropicEventType.ERROR, body) 

201 finally: 

202 await guard.release() 

203 

204 

205def _internal_error_response() -> Response: 

206 """Log and return the generic api_error 500 envelope.""" 

207 log.exception("messages_endpoint failed") 

208 return _error_response(500, CompletionsErrorCode.INTERNAL_ERROR, _INTERNAL_ERROR_MESSAGE) 

209 

210 

211def _error_response( 

212 status: int, 

213 code: CompletionsErrorCode, 

214 message: str, 

215 *, 

216 headers: dict[str, str] | None = None, 

217) -> Response: 

218 return Response( 

219 anthropic_error_body(anthropic_error_type(code), message), 

220 status_code=status, 

221 headers=headers or {}, 

222 media_type="application/json", 

223 ) 

224 

225 

226def _auth_failure(request: Request) -> Response | None: 

227 """Return a 401 envelope if the bearer token is missing/wrong, else None. 

228 

229 Claude Code sends ``ANTHROPIC_AUTH_TOKEN`` as a bearer Authorization 

230 header, which is exactly the session token check every /v1 route runs. 

231 """ 

232 auth_header = request.headers.get("authorization", "") 

233 try: 

234 authorized = session_manager.validate(auth_header) 

235 except NotAuthorizedException: 

236 authorized = False 

237 if authorized: 

238 return None 

239 return _error_response(401, CompletionsErrorCode.INVALID_API_KEY, "Missing or invalid API key.") 

240 

241 

242def _validation_exception_handler(_: Request, exc: ValidationException) -> Response: 

243 """Wrap Litestar's body-parse failures in the Anthropic error envelope.""" 

244 return _error_response(400, CompletionsErrorCode.INVALID_REQUEST, format_validation(exc)) 

245 

246 

247def _response_id() -> str: 

248 """Anthropic-style ``msg_*`` id.""" 

249 return f"msg_{uuid.uuid4().hex[:24]}" 

250 

251 

252anthropic_router = Router( 

253 path="/", 

254 route_handlers=[messages_endpoint], 

255 exception_handlers={ValidationException: _validation_exception_handler}, 

256)