Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/hooks/responses_id_security.py: 36%

1""" 

2Security hook to prevent user B from seeing response from user A. 

3 

4This hook uses the DBSpendUpdateWriter to batch-write response IDs to the database 

5instead of writing immediately on each request. 

6""" 

7 

8from collections.abc import AsyncGenerator, Callable, Mapping 

9from typing import TYPE_CHECKING, Any, Final, cast 

10 

11from fastapi import HTTPException 

12from pydantic import TypeAdapter, ValidationError 

13 

14from litellm._logging import verbose_proxy_logger 

15from litellm.integrations.custom_logger import CustomLogger 

16from litellm.proxy._types import LitellmUserRoles 

17from litellm.proxy.common_utils.encrypt_decrypt_utils import ( 

18 decrypt_value_helper, 

19 encrypt_value_helper, 

20) 

21from litellm.types.llms.openai import ( 

22 BaseLiteLLMOpenAIResponseObject, 

23 ResponsesAPIResponse, 

24) 

25from litellm.types.utils import ADDRESSED_RESPONSE_ID_FIELD, CallTypesLiteral, LLMResponseTypes, SpecialEnums 

26 

27if TYPE_CHECKING: 27 ↛ 28line 27 didn't jump to line 28 because the condition on line 27 was never true

28 from litellm.caching.caching import DualCache 

29 from litellm.proxy._types import UserAPIKeyAuth 

30 

31 

32_RESPONSES_API_PROVIDER_PREFIX: Final = "/openai" 

33_RESPONSES_API_CREATE_ROUTES: Final = frozenset({"/v1/responses", "/responses"}) 

34 

35_UNMANAGED_RESPONSE_ID_DETAIL: Final = ( 

36 "Forbidden. This response id was not issued by this proxy, so the proxy cannot tell who owns it. " 

37 "To let keys address responses this proxy did not issue, set " 

38 "general_settings::allow_unmanaged_response_ids to True in the config.yaml file." 

39) 

40_PROXY_ADMIN_ROLES: Final = frozenset({LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.PROXY_ADMIN.value}) 

41 

42 

43def _proxy_general_settings() -> Mapping[str, object]: 

44 from litellm.proxy.proxy_server import general_settings 

45 

46 return general_settings 

47 

48 

49def _proxy_signing_key() -> str | None: 

50 import os 

51 

52 from litellm.proxy.proxy_server import master_key 

53 

54 salt_key: Final = os.getenv("LITELLM_SALT_KEY", None) 

55 return master_key if salt_key is None else salt_key 

56 

57 

58_RESPONSE_PAYLOAD_ADAPTER: Final = TypeAdapter(Mapping[str, object]) 

59 

60 

61def _response_payload(response_obj: object) -> Mapping[str, object] | None: 

62 try: 

63 return _RESPONSE_PAYLOAD_ADAPTER.validate_python(response_obj) 

64 except ValidationError: 

65 return None 

66 

67 

68def _rewrite_advertised_id( 

69 event: BaseLiteLLMOpenAIResponseObject, 

70 rewrite: Callable[[str], str], 

71) -> BaseLiteLLMOpenAIResponseObject: 

72 event_id: Final = getattr(event, "id", None) 

73 if isinstance(event_id, str) and event_id.startswith("resp_"): 

74 setattr(event, "id", rewrite(event_id)) 

75 return event 

76 

77 nested: Final = getattr(event, "response", None) 

78 if isinstance(nested, ResponsesAPIResponse): 

79 setattr(nested, "id", rewrite(nested.id)) 

80 setattr(event, "response", nested) 

81 return event 

82 

83 payload: Final = _response_payload(nested) 

84 if payload is None: 

85 return event 

86 

87 payload_id: Final = payload.get("id") 

88 if not isinstance(payload_id, str): 

89 return event 

90 

91 rewritten: Final = {**payload, "id": rewrite(payload_id)} # mutable-ok: pydantic cannot serialize a frozen map 

92 setattr(event, "response", rewritten) 

93 return event 

94 

95 

96def _is_responses_api_create_route(request_route: str | None) -> bool: 

97 if request_route is None: 

98 return False 

99 canonical: Final = ( 

100 request_route[len(_RESPONSES_API_PROVIDER_PREFIX) :] 

101 if request_route.startswith(_RESPONSES_API_PROVIDER_PREFIX + "/") 

102 else request_route 

103 ) 

104 return canonical in _RESPONSES_API_CREATE_ROUTES 

105 

106 

107class ResponsesIDSecurity(CustomLogger): 

108 def __init__( 

109 self, 

110 general_settings_reader: Callable[[], Mapping[str, object]] = _proxy_general_settings, 

111 signing_key_reader: Callable[[], str | None] = _proxy_signing_key, 

112 ) -> None: 

113 self._general_settings_reader: Final = general_settings_reader 

114 self._signing_key_reader: Final = signing_key_reader 

115 

116 async def async_pre_call_hook( 

117 self, 

118 user_api_key_dict: "UserAPIKeyAuth", 

119 cache: "DualCache", 

120 data: dict, 

121 call_type: CallTypesLiteral, 

122 ) -> Exception | str | dict | None: 

123 # MAP all the responses api response ids to the encrypted response ids 

124 responses_api_call_types: Final = { 

125 "aresponses", 

126 "aget_responses", 

127 "adelete_responses", 

128 "acancel_responses", 

129 "alist_input_items", 

130 } 

131 if call_type not in responses_api_call_types: 

132 return None 

133 addressed_id_field: Final = "previous_response_id" if call_type == "aresponses" else "response_id" 

134 retained_id: Final = data.get(ADDRESSED_RESPONSE_ID_FIELD) 

135 addressed_id: Final = ( 

136 retained_id if isinstance(retained_id, str) and retained_id else data.get(addressed_id_field) 

137 ) 

138 if not isinstance(addressed_id, str) or not addressed_id: 

139 return data 

140 authorized_id: Final = self._authorize_response_id(addressed_id, user_api_key_dict) 

141 data[addressed_id_field] = authorized_id 

142 data[ADDRESSED_RESPONSE_ID_FIELD] = addressed_id 

143 return data 

144 

145 def _authorize_response_id( 

146 self, 

147 response_id: str, 

148 user_api_key_dict: "UserAPIKeyAuth", 

149 ) -> str: 

150 if self._is_encrypted_response_id(response_id): 150 ↛ 151line 150 didn't jump to line 151 because the condition on line 150 was never true

151 original_response_id, user_id, team_id = self._decrypt_response_id(response_id) 

152 self.check_user_access_to_response_id(user_id, team_id, user_api_key_dict) 

153 return original_response_id 

154 

155 if self._unmanaged_response_ids_allowed(user_api_key_dict): 155 ↛ 158line 155 didn't jump to line 158 because the condition on line 155 was always true

156 return response_id 

157 

158 raise HTTPException(status_code=403, detail=_UNMANAGED_RESPONSE_ID_DETAIL) 

159 

160 def _unmanaged_response_ids_allowed(self, user_api_key_dict: "UserAPIKeyAuth") -> bool: 

161 general_settings: Final = self._general_settings_reader() 

162 

163 if general_settings.get("disable_responses_id_security", False): 163 ↛ 164line 163 didn't jump to line 164 because the condition on line 163 was never true

164 return True 

165 if general_settings.get("allow_unmanaged_response_ids", False): 165 ↛ 166line 165 didn't jump to line 166 because the condition on line 165 was never true

166 return True 

167 if self._get_signing_key() is None: 167 ↛ 168line 167 didn't jump to line 168 because the condition on line 167 was never true

168 return True 

169 return user_api_key_dict.user_role in _PROXY_ADMIN_ROLES 

170 

171 def check_user_access_to_response_id( 

172 self, 

173 response_id_user_id: str | None, 

174 response_id_team_id: str | None, 

175 user_api_key_dict: "UserAPIKeyAuth", 

176 ) -> bool: 

177 general_settings: Final = self._general_settings_reader() 

178 

179 if ( 

180 user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value 

181 or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN 

182 ): 

183 return True 

184 

185 if response_id_user_id and response_id_user_id != user_api_key_dict.user_id: 

186 if general_settings.get("disable_responses_id_security", False): 

187 verbose_proxy_logger.debug( 

188 "Responses ID Security is disabled. User %s is accessing response id %s which is not associated with them.", 

189 user_api_key_dict.user_id, 

190 response_id_user_id, 

191 ) 

192 return True 

193 raise HTTPException( 

194 status_code=403, 

195 detail="Forbidden. The response id is not associated with the user, who this key belongs to. To disable this security feature, set general_settings::disable_responses_id_security to True in the config.yaml file.", 

196 ) 

197 

198 if response_id_team_id and response_id_team_id != user_api_key_dict.team_id: 

199 if general_settings.get("disable_responses_id_security", False): 

200 verbose_proxy_logger.debug( 

201 "Responses ID Security is disabled. Response belongs to team %s but user %s is accessing it with team id %s.", 

202 response_id_team_id, 

203 user_api_key_dict.user_id, 

204 user_api_key_dict.team_id, 

205 ) 

206 return True 

207 raise HTTPException( 

208 status_code=403, 

209 detail="Forbidden. The response id is not associated with the team, who this key belongs to. To disable this security feature, set general_settings::disable_responses_id_security to True in the config.yaml file.", 

210 ) 

211 

212 return True 

213 

214 def _is_encrypted_response_id(self, response_id: str) -> bool: 

215 split_result: Final = response_id.split("resp_") 

216 if len(split_result) < 2: 216 ↛ 219line 216 didn't jump to line 219 because the condition on line 216 was always true

217 return False 

218 

219 remaining_string: Final = split_result[1] 

220 decrypted_value = decrypt_value_helper(value=remaining_string, key="response_id", return_original_value=True) 

221 

222 if decrypted_value is None: 

223 return False 

224 

225 if decrypted_value.startswith(SpecialEnums.LITELM_MANAGED_FILE_ID_PREFIX.value): 

226 return True 

227 return False 

228 

229 def _decrypt_response_id(self, response_id: str) -> tuple[str, str | None, str | None]: 

230 """ 

231 Returns: 

232 - original_response_id: the original response id 

233 - user_id: the user id 

234 - team_id: the team id 

235 """ 

236 split_result: Final = response_id.split("resp_") 

237 if len(split_result) < 2: 

238 return response_id, None, None 

239 

240 remaining_string: Final = split_result[1] 

241 decrypted_value = decrypt_value_helper(value=remaining_string, key="response_id", return_original_value=True) 

242 

243 if decrypted_value is None: 

244 return response_id, None, None 

245 

246 if decrypted_value.startswith(SpecialEnums.LITELM_MANAGED_FILE_ID_PREFIX.value): 

247 # Expected format: "litellm_proxy:responses_api:response_id:{response_id};user_id:{user_id}" 

248 parts: Final = decrypted_value.split(";") 

249 

250 if len(parts) >= 2: 

251 # Extract response_id from "litellm_proxy:responses_api:response_id:{response_id}" 

252 response_id_part: Final = parts[0] 

253 original_response_id: Final = response_id_part.split("response_id:")[-1] 

254 

255 # Extract user_id from "user_id:{user_id}" 

256 user_id_part: Final = parts[1] 

257 user_id: Final = user_id_part.split("user_id:")[-1] 

258 

259 # Extract team_id from "team_id:{team_id}" 

260 team_id_part: Final = parts[2] 

261 team_id: Final = team_id_part.split("team_id:")[-1] 

262 

263 return original_response_id, user_id, team_id 

264 else: 

265 # Fallback if format is unexpected 

266 return response_id, None, None 

267 return response_id, None, None 

268 

269 def _get_signing_key(self) -> str | None: 

270 return self._signing_key_reader() 

271 

272 def _encrypt_response_id( 

273 self, 

274 response: BaseLiteLLMOpenAIResponseObject, 

275 user_api_key_dict: "UserAPIKeyAuth", 

276 request_cache: dict[str, str] | None = None, 

277 ) -> BaseLiteLLMOpenAIResponseObject: 

278 signing_key: Final = self._get_signing_key() 

279 if signing_key is None: 

280 verbose_proxy_logger.debug( 

281 "Response ID encryption is enabled but no signing key is configured. " 

282 "Please set LITELLM_SALT_KEY environment variable or configure a master_key. " 

283 "Skipping response ID encryption. " 

284 "See: https://docs.litellm.ai/docs/proxy/prod#5-set-litellm-salt-key" 

285 ) 

286 return response 

287 

288 def encrypt(original_id: str) -> str: 

289 cached: Final = request_cache.get(original_id) if request_cache is not None else None 

290 if cached is not None: 

291 return cached 

292 

293 managed_id: Final = SpecialEnums.LITELLM_MANAGED_RESPONSE_API_RESPONSE_ID_COMPLETE_STR.value.format( 

294 original_id, 

295 user_api_key_dict.user_id or "", 

296 user_api_key_dict.team_id or "", 

297 ) 

298 encrypted_id: Final = f"resp_{encrypt_value_helper(value=managed_id)}" 

299 if request_cache is not None: 

300 request_cache[original_id] = encrypted_id 

301 return encrypted_id 

302 

303 return _rewrite_advertised_id(response, encrypt) 

304 

305 async def async_post_call_success_hook( 

306 self, 

307 data: dict, 

308 user_api_key_dict: "UserAPIKeyAuth", 

309 response: LLMResponseTypes, 

310 ) -> LLMResponseTypes: 

311 """ 

312 Queue response IDs for batch processing instead of writing directly to DB. 

313 

314 This method adds response IDs to an in-memory queue, which are then 

315 batch-processed by the DBSpendUpdateWriter during regular database update cycles. 

316 """ 

317 general_settings: Final = self._general_settings_reader() 

318 

319 if general_settings.get("disable_responses_id_security", False): 319 ↛ 320line 319 didn't jump to line 320 because the condition on line 319 was never true

320 return response 

321 if isinstance(response, ResponsesAPIResponse): 321 ↛ 322line 321 didn't jump to line 322 because the condition on line 321 was never true

322 response = cast( 

323 ResponsesAPIResponse, 

324 self._encrypt_response_id(response, user_api_key_dict, request_cache=None), 

325 ) 

326 return response 

327 

328 async def async_post_call_streaming_iterator_hook( 

329 self, user_api_key_dict: "UserAPIKeyAuth", response: Any, request_data: dict 

330 ) -> AsyncGenerator[BaseLiteLLMOpenAIResponseObject, None]: 

331 general_settings: Final = self._general_settings_reader() 

332 

333 # Create a request-scoped cache for consistent encryption across streaming chunks. 

334 request_encryption_cache: Final[dict[str, str]] = {} 

335 

336 async for chunk in response: 

337 if ( 

338 isinstance(chunk, BaseLiteLLMOpenAIResponseObject) 

339 and _is_responses_api_create_route(user_api_key_dict.request_route) 

340 and not general_settings.get("disable_responses_id_security", False) 

341 ): 

342 chunk = self._encrypt_response_id(chunk, user_api_key_dict, request_encryption_cache) 

343 yield chunk