Skip to main content

bge_m3_embedding_server/
error.rs

1// Copyright (c) 2026 J. Patrick Fulton
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! Application-level error types that map to HTTP status codes.
16
17use axum::{
18    Json,
19    http::StatusCode,
20    response::{IntoResponse, Response},
21};
22use serde_json::json;
23use tracing::error;
24
25/// Application-level errors that map to HTTP status codes.
26#[derive(Debug)]
27pub enum AppError {
28    /// The request was malformed or violates input constraints.
29    /// Maps to HTTP 400 Bad Request.
30    InvalidRequest(String),
31    /// The service is not yet ready (model loading) or has no live workers.
32    /// Maps to HTTP 503 Service Unavailable.
33    ServiceUnavailable(String),
34    /// An unexpected internal error occurred during embedding.
35    /// Maps to HTTP 500 Internal Server Error.
36    Internal(String),
37}
38
39impl IntoResponse for AppError {
40    fn into_response(self) -> Response {
41        let (status, error_type, code, message) = match self {
42            AppError::InvalidRequest(msg) => (
43                StatusCode::BAD_REQUEST,
44                "invalid_request_error",
45                400u16,
46                msg,
47            ),
48            AppError::ServiceUnavailable(msg) => (
49                StatusCode::SERVICE_UNAVAILABLE,
50                "service_unavailable",
51                503u16,
52                msg,
53            ),
54            AppError::Internal(msg) => (
55                StatusCode::INTERNAL_SERVER_ERROR,
56                "internal_error",
57                500u16,
58                msg,
59            ),
60        };
61
62        let body = json!({
63            "error": {
64                "message": message,
65                "type": error_type,
66                "code": code
67            }
68        });
69
70        (status, Json(body)).into_response()
71    }
72}
73
74impl From<anyhow::Error> for AppError {
75    fn from(err: anyhow::Error) -> Self {
76        // The in-band TRT JIT guard refuses dangerous, uncovered chunk shapes
77        // to avoid a process-killing pathological autotuner allocation. That
78        // refusal is a *retriable* condition (coverage may extend via adaptive
79        // warmup, or a peer task may already cover the shape), so it maps to
80        // 503 rather than a generic 500. Detection walks the error source
81        // chain because the worker wraps embed errors with `.context(...)`.
82        if crate::embedder::jit_guard::is_trt_shape_rejected(&err) {
83            error!(error = %err, "in-band TRT JIT guard refused request");
84            return AppError::ServiceUnavailable(
85                "embedding temporarily unavailable for this input size; please retry".to_string(),
86            );
87        }
88        error!(error = %err, "Internal error");
89        AppError::Internal("internal server error".to_string())
90    }
91}
92
93#[cfg(test)]
94mod tests {
95    use super::*;
96    use axum::body::to_bytes;
97    use axum::response::IntoResponse;
98
99    async fn response_parts(err: AppError) -> (StatusCode, serde_json::Value) {
100        let response = err.into_response();
101        let status = response.status();
102        let bytes = to_bytes(response.into_body(), usize::MAX)
103            .await
104            .expect("failed to read body");
105        let body: serde_json::Value =
106            serde_json::from_slice(&bytes).expect("body is not valid JSON");
107        (status, body)
108    }
109
110    #[tokio::test]
111    async fn invalid_request_serializes_as_400() {
112        let (status, body) =
113            response_parts(AppError::InvalidRequest("bad input".to_string())).await;
114        assert_eq!(status, StatusCode::BAD_REQUEST);
115        assert_eq!(body["error"]["code"], 400);
116        assert_eq!(body["error"]["type"], "invalid_request_error");
117        assert_eq!(body["error"]["message"], "bad input");
118    }
119
120    #[tokio::test]
121    async fn service_unavailable_serializes_as_503() {
122        let (status, body) =
123            response_parts(AppError::ServiceUnavailable("model not ready".to_string())).await;
124        assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
125        assert_eq!(body["error"]["code"], 503);
126        assert_eq!(body["error"]["type"], "service_unavailable");
127        assert_eq!(body["error"]["message"], "model not ready");
128    }
129
130    #[tokio::test]
131    async fn internal_error_serializes_as_500() {
132        let (status, body) =
133            response_parts(AppError::Internal("unexpected failure".to_string())).await;
134        assert_eq!(status, StatusCode::INTERNAL_SERVER_ERROR);
135        assert_eq!(body["error"]["code"], 500);
136        assert_eq!(body["error"]["type"], "internal_error");
137        assert_eq!(body["error"]["message"], "unexpected failure");
138    }
139
140    #[tokio::test]
141    async fn from_anyhow_error_produces_generic_message() {
142        let err = anyhow::anyhow!("secret path /var/models/onnx failed to load");
143        let app_err: AppError = err.into();
144        let (status, body) = response_parts(app_err).await;
145        assert_eq!(status, StatusCode::INTERNAL_SERVER_ERROR);
146        assert_eq!(
147            body["error"]["message"], "internal server error",
148            "internal details must not leak to client"
149        );
150    }
151
152    #[tokio::test]
153    async fn trt_jit_rejection_maps_to_503() {
154        // The in-band JIT guard refusal is a retriable condition and must map
155        // to 503, not the generic 500, even when wrapped in a context chain
156        // (as the worker does via `.context(...)`).
157        let base: anyhow::Error = crate::embedder::jit_guard::TrtJitRejection {
158            batch: 8,
159            seq: 8192,
160            guard_seq: 4096,
161            warmed_seq_ceiling: 2048,
162        }
163        .into();
164        let wrapped = base.context("Dual embed error");
165        let app_err: AppError = wrapped.into();
166        let (status, body) = response_parts(app_err).await;
167        assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
168        assert_eq!(body["error"]["code"], 503);
169        assert_eq!(body["error"]["type"], "service_unavailable");
170    }
171}