byokey_proxy/
error.rs

1//! API error type that maps [`ByokError`] variants to HTTP status codes.
2
3use axum::{
4    Json,
5    http::StatusCode,
6    response::{IntoResponse, Response},
7};
8use byokey_types::ByokError;
9use serde_json::json;
10
11/// Wrapper around [`ByokError`] that implements [`IntoResponse`].
12pub struct ApiError(pub ByokError);
13
14impl ApiError {
15    /// Returns `(status, error_type, error_code)` for the wrapped error.
16    fn classify(&self) -> (StatusCode, &'static str, &'static str) {
17        match &self.0 {
18            ByokError::Auth(_) => (
19                StatusCode::UNAUTHORIZED,
20                "authentication_error",
21                "invalid_api_key",
22            ),
23            ByokError::TokenNotFound(_) | ByokError::TokenExpired(_) => (
24                StatusCode::UNAUTHORIZED,
25                "authentication_error",
26                "token_not_found",
27            ),
28            ByokError::UnsupportedModel(_) => (
29                StatusCode::BAD_REQUEST,
30                "invalid_request_error",
31                "model_not_found",
32            ),
33            ByokError::UnsupportedProvider(_) => (
34                StatusCode::BAD_REQUEST,
35                "invalid_request_error",
36                "provider_not_found",
37            ),
38            ByokError::Translation(_) => (
39                StatusCode::BAD_REQUEST,
40                "invalid_request_error",
41                "translation_error",
42            ),
43            ByokError::Upstream { status, .. } => classify_upstream(*status),
44            ByokError::Http(_) => (StatusCode::BAD_GATEWAY, "server_error", "upstream_error"),
45            _ => (
46                StatusCode::INTERNAL_SERVER_ERROR,
47                "server_error",
48                "internal_error",
49            ),
50        }
51    }
52}
53
54fn classify_upstream(status: u16) -> (StatusCode, &'static str, &'static str) {
55    match status {
56        429 => (
57            StatusCode::TOO_MANY_REQUESTS,
58            "rate_limit_error",
59            "rate_limit_exceeded",
60        ),
61        401 => (
62            StatusCode::UNAUTHORIZED,
63            "authentication_error",
64            "invalid_api_key",
65        ),
66        403 => (
67            StatusCode::FORBIDDEN,
68            "permission_error",
69            "insufficient_quota",
70        ),
71        _ => (StatusCode::BAD_GATEWAY, "server_error", "upstream_error"),
72    }
73}
74
75impl IntoResponse for ApiError {
76    fn into_response(self) -> Response {
77        let (status, error_type, error_code) = self.classify();
78        let msg = self.0.to_string();
79        (
80            status,
81            Json(json!({
82                "error": {
83                    "message": msg,
84                    "type": error_type,
85                    "code": error_code,
86                }
87            })),
88        )
89            .into_response()
90    }
91}
92
93impl From<ByokError> for ApiError {
94    fn from(e: ByokError) -> Self {
95        Self(e)
96    }
97}
98
99#[cfg(test)]
100mod tests {
101    use super::*;
102    use byokey_types::ProviderId;
103    use http_body_util::BodyExt as _;
104
105    async fn extract_error_body(err: ApiError) -> (StatusCode, serde_json::Value) {
106        let resp = err.into_response();
107        let status = resp.status();
108        let bytes = resp.into_body().collect().await.unwrap().to_bytes();
109        let body: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
110        (status, body)
111    }
112
113    #[tokio::test]
114    async fn test_auth_error() {
115        let (status, body) =
116            extract_error_body(ApiError(ByokError::Auth("bad creds".into()))).await;
117        assert_eq!(status, StatusCode::UNAUTHORIZED);
118        assert_eq!(body["error"]["type"], "authentication_error");
119        assert_eq!(body["error"]["code"], "invalid_api_key");
120    }
121
122    #[tokio::test]
123    async fn test_token_not_found_error() {
124        let (status, body) =
125            extract_error_body(ApiError(ByokError::TokenNotFound(ProviderId::Claude))).await;
126        assert_eq!(status, StatusCode::UNAUTHORIZED);
127        assert_eq!(body["error"]["type"], "authentication_error");
128        assert_eq!(body["error"]["code"], "token_not_found");
129    }
130
131    #[tokio::test]
132    async fn test_unsupported_model_error() {
133        let (status, body) =
134            extract_error_body(ApiError(ByokError::UnsupportedModel("xyz".into()))).await;
135        assert_eq!(status, StatusCode::BAD_REQUEST);
136        assert_eq!(body["error"]["type"], "invalid_request_error");
137        assert_eq!(body["error"]["code"], "model_not_found");
138    }
139
140    #[tokio::test]
141    async fn test_translation_error() {
142        let (status, body) =
143            extract_error_body(ApiError(ByokError::Translation("bad format".into()))).await;
144        assert_eq!(status, StatusCode::BAD_REQUEST);
145        assert_eq!(body["error"]["type"], "invalid_request_error");
146        assert_eq!(body["error"]["code"], "translation_error");
147    }
148
149    #[tokio::test]
150    async fn test_upstream_429_error() {
151        let (status, body) = extract_error_body(ApiError(ByokError::Upstream {
152            retry_after: None,
153            status: 429,
154            body: "rate limited".into(),
155        }))
156        .await;
157        assert_eq!(status, StatusCode::TOO_MANY_REQUESTS);
158        assert_eq!(body["error"]["type"], "rate_limit_error");
159        assert_eq!(body["error"]["code"], "rate_limit_exceeded");
160    }
161
162    #[tokio::test]
163    async fn test_upstream_401_error() {
164        let (status, body) = extract_error_body(ApiError(ByokError::Upstream {
165            retry_after: None,
166            status: 401,
167            body: "unauthorized".into(),
168        }))
169        .await;
170        assert_eq!(status, StatusCode::UNAUTHORIZED);
171        assert_eq!(body["error"]["type"], "authentication_error");
172        assert_eq!(body["error"]["code"], "invalid_api_key");
173    }
174
175    #[tokio::test]
176    async fn test_upstream_403_error() {
177        let (status, body) = extract_error_body(ApiError(ByokError::Upstream {
178            retry_after: None,
179            status: 403,
180            body: "forbidden".into(),
181        }))
182        .await;
183        assert_eq!(status, StatusCode::FORBIDDEN);
184        assert_eq!(body["error"]["type"], "permission_error");
185        assert_eq!(body["error"]["code"], "insufficient_quota");
186    }
187
188    #[tokio::test]
189    async fn test_upstream_500_error() {
190        let (status, body) = extract_error_body(ApiError(ByokError::Upstream {
191            retry_after: None,
192            status: 500,
193            body: "server error".into(),
194        }))
195        .await;
196        assert_eq!(status, StatusCode::BAD_GATEWAY);
197        assert_eq!(body["error"]["type"], "server_error");
198        assert_eq!(body["error"]["code"], "upstream_error");
199    }
200
201    #[tokio::test]
202    async fn test_http_transport_error() {
203        let (status, body) =
204            extract_error_body(ApiError(ByokError::Http("connection refused".into()))).await;
205        assert_eq!(status, StatusCode::BAD_GATEWAY);
206        assert_eq!(body["error"]["type"], "server_error");
207        assert_eq!(body["error"]["code"], "upstream_error");
208    }
209
210    #[tokio::test]
211    async fn test_internal_error() {
212        let (status, body) =
213            extract_error_body(ApiError(ByokError::Config("bad config".into()))).await;
214        assert_eq!(status, StatusCode::INTERNAL_SERVER_ERROR);
215        assert_eq!(body["error"]["type"], "server_error");
216        assert_eq!(body["error"]["code"], "internal_error");
217    }
218}
byokey_proxy/error.rs

byokey_proxy/
error.rs
