bge_m3_embedding_server/handler/
common.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//! Shared input validation, header utilities, and service-readiness helpers
16//! used by all handlers.
17
18use std::collections::BTreeMap;
19use std::fmt;
20use std::sync::atomic::Ordering;
21
22use axum::http::HeaderMap;
23use serde::Serialize;
24
25use crate::error::AppError;
26use crate::state::AppState;
27
28/// A sorted map of `X-*` HTTP request headers.
29///
30/// Keys are lowercase-normalized header names (e.g. `"x-request-id"`).
31/// Values are UTF-8-decoded header values; headers with non-UTF-8 values
32/// are silently skipped.
33///
34/// Serializes as a plain JSON object so it can be embedded as the
35/// `x_headers` field in structured log events.
36#[derive(Default, Serialize)]
37#[serde(transparent)]
38pub(super) struct XHeaders(pub(super) BTreeMap<String, String>);
39
40impl XHeaders {
41    /// Returns `true` when no `X-*` headers were present in the request.
42    pub(super) fn is_empty(&self) -> bool {
43        self.0.is_empty()
44    }
45}
46
47impl fmt::Display for XHeaders {
48    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
49        // Compact JSON — suitable for both text and JSON log formats.
50        match serde_json::to_string(&self.0) {
51            Ok(s) => f.write_str(&s),
52            Err(_) => f.write_str("{}"),
53        }
54    }
55}
56
57/// Collects all headers whose name starts with `x-` (case-insensitive) into
58/// an [`XHeaders`] map.
59///
60/// Header names are normalized: lowercase (axum guarantees this) and hyphens
61/// replaced with underscores so each key is a valid JSON identifier and can
62/// be referenced by log-processing tools that use JSON path notation
63/// (e.g. `x-example-project` → key `x_example_project`).
64/// Headers with non-UTF-8 values are silently skipped.
65pub(super) fn collect_x_headers(headers: &HeaderMap) -> XHeaders {
66    let mut map = BTreeMap::new();
67    for (name, value) in headers {
68        let name_str = name.as_str();
69        if name_str.starts_with("x-")
70            && let Ok(val) = value.to_str()
71        {
72            map.insert(name_str.replace('-', "_"), val.to_owned());
73        }
74    }
75    XHeaders(map)
76}
77
78/// Maximum characters allowed per individual input string (SEC-3).
79pub(super) const MAX_STRING_CHARS: usize = 32_768;
80
81/// Validates a batch of input texts against size and length constraints.
82///
83/// Returns [`AppError::InvalidRequest`] if:
84/// - `texts` is empty
85/// - `texts.len() > max_batch`
86/// - any individual text exceeds [`MAX_STRING_CHARS`] characters
87pub(super) fn validate_input(texts: &[String], max_batch: usize) -> Result<(), AppError> {
88    if texts.is_empty() {
89        return Err(AppError::InvalidRequest(
90            "input must not be empty".to_string(),
91        ));
92    }
93    if texts.len() > max_batch {
94        return Err(AppError::InvalidRequest(format!(
95            "batch size {} exceeds maximum {}",
96            texts.len(),
97            max_batch
98        )));
99    }
100    for (i, text) in texts.iter().enumerate() {
101        let char_count = text.chars().count();
102        if char_count > MAX_STRING_CHARS {
103            return Err(AppError::InvalidRequest(format!(
104                "input[{i}] length {char_count} exceeds maximum {MAX_STRING_CHARS} characters"
105            )));
106        }
107    }
108    Ok(())
109}
110
111/// Checks whether the service is ready to handle embedding requests.
112pub(super) fn check_ready(state: &AppState) -> Result<(), AppError> {
113    if !state.ready.load(Ordering::Acquire) {
114        return Err(AppError::ServiceUnavailable("model not ready".to_string()));
115    }
116    if state.pool.live_worker_count() == 0 {
117        return Err(AppError::ServiceUnavailable(
118            "no workers available".to_string(),
119        ));
120    }
121    Ok(())
122}
bge_m3_embedding_server/handler/common.rs

bge_m3_embedding_server/handler/
common.rs
