sokuto_server/guardrails/
limits.rs

1//! # OOM 防壁・リソース制限モジュール
2//!
3//! リクエストの質問数、文字数、および推定トークン数を推論前に事前検証し、
4//! メモリ枯渇 (OOM) や GPU/CPU リソースの過負荷を未然に遮断する Fast-Fail 機構を提供する。
5
6use serde_json::Value;
7use sokuto_core::schema::SystemOneRequest;
8
9use crate::error::ServerError;
10
11/// デフォルトのリクエストあたり最大質問数。
12pub const DEFAULT_MAX_QUESTIONS: usize = 128;
13
14/// デフォルトの State 最大許容文字数 (約 64KB)。
15pub const DEFAULT_MAX_STATE_CHARS: usize = 65_536;
16
17/// デフォルトの質問あたり最大許容文字数 (約 8KB)。
18pub const DEFAULT_MAX_QUESTION_CHARS: usize = 8_192;
19
20/// デフォルトの推定総トークン数上限 (約 16,384 トークン)。
21pub const DEFAULT_MAX_ESTIMATED_TOKENS: usize = 16_384;
22
23/// デフォルトのリクエストボディ最大許容バイト数 (2MB)。
24pub const DEFAULT_MAX_BODY_BYTES: usize = 2 * 1024 * 1024;
25
26/// リソース制限ガードレールの設定構造体。
27#[derive(Debug, Clone, PartialEq, Eq)]
28pub struct LimitsConfig {
29    /// 1 リクエストで許容される最大質問数。
30    pub max_questions: usize,
31    /// State 文字列の最大許容文字数。
32    pub max_state_chars: usize,
33    /// 1 質問あたりの最大許容文字数 (Instructions + Criteria)。
34    pub max_question_chars: usize,
35    /// リクエスト全体の最大推定トークン数。
36    pub max_estimated_tokens: usize,
37    /// リクエストボディ全体の最大許容バイト数。
38    pub max_body_bytes: usize,
39}
40
41impl Default for LimitsConfig {
42    fn default() -> Self {
43        Self {
44            max_questions: DEFAULT_MAX_QUESTIONS,
45            max_state_chars: DEFAULT_MAX_STATE_CHARS,
46            max_question_chars: DEFAULT_MAX_QUESTION_CHARS,
47            max_estimated_tokens: DEFAULT_MAX_ESTIMATED_TOKENS,
48            max_body_bytes: DEFAULT_MAX_BODY_BYTES,
49        }
50    }
51}
52
53/// 文字列からトークン数を高速に安全側 (過大評価) で事前見積もりする。
54///
55/// UTF-8 バイト数および文字種に基づき、ゼロアロケーションで上限トークン数を算出する。
56/// 英語 (1 語 1〜2 トークン) および日本語 (1 文字 1〜2 トークン) の双対特性を吸収し、
57/// 最悪ケースとして `chars.len()` をベースとした安全側概算を行う。
58#[inline]
59pub fn estimate_tokens_fast(text: &str) -> usize {
60    // 空文字は 0 トークン
61    if text.is_empty() {
62        return 0;
63    }
64    // 日本語混じりでも 1 文字あたり最大 2 トークン程度に収まる安全側バウンド
65    let char_count = text.chars().count();
66    // 最低でも 1 トークン、文字数とバイト数から過大推定 (保守的バジェット)
67    char_count.max(text.len() / 3).max(1)
68}
69
70/// リクエスト全体の文字数および質問数を検証する。
71///
72/// # 引数
73/// - `req`: Jev 推論リクエスト。
74/// - `config`: リソース制限設定。
75///
76/// # 戻り値
77/// 検証成功時は `(推定総文字数, 推定総トークン数)` のタプルを返却し、上限超過時は `ServerError::PayloadTooLarge` を返却する。
78pub fn validate_limits(
79    req: &SystemOneRequest,
80    config: &LimitsConfig,
81) -> Result<(usize, usize), ServerError> {
82    // 1. 質問数上限チェック (L1 Fast-Fail)
83    let question_count = req.questions.len();
84    if question_count > config.max_questions {
85        return Err(ServerError::PayloadTooLarge(format!(
86            "リクエストに含まれる質問数 ({question_count}) が許容上限 ({}) を超過しています。",
87            config.max_questions
88        )));
89    }
90
91    // 2. State の文字数チェック (L1 Fast-Fail)
92    let state_chars = match &req.state {
93        Value::String(s) => s.chars().count(),
94        other => other.to_string().chars().count(),
95    };
96    if state_chars > config.max_state_chars {
97        return Err(ServerError::PayloadTooLarge(format!(
98            "State の文字数 ({state_chars}) が許容上限 ({}) を超過しています。",
99            config.max_state_chars
100        )));
101    }
102
103    // 3. 各質問の文字数および全体のトークン見積もり (L2 チェック)
104    let mut total_chars = state_chars;
105    let mut total_estimated_tokens = match &req.state {
106        Value::String(s) => estimate_tokens_fast(s),
107        other => estimate_tokens_fast(&other.to_string()),
108    };
109
110    for (q_key, question) in &req.questions {
111        let mut q_chars = question.instructions.chars().count();
112        let mut q_tokens = estimate_tokens_fast(&question.instructions);
113
114        if let Some(ref criteria) = question.criteria {
115            if let Some(map) = criteria.as_map() {
116                for (k, v) in map {
117                    q_chars += k.chars().count() + v.chars().count();
118                    q_tokens += estimate_tokens_fast(k) + estimate_tokens_fast(v);
119                }
120            } else if let Some(list) = criteria.as_list() {
121                for item in list {
122                    q_chars += item.chars().count();
123                    q_tokens += estimate_tokens_fast(item);
124                }
125            }
126        }
127
128        if q_chars > config.max_question_chars {
129            return Err(ServerError::PayloadTooLarge(format!(
130                "質問 '{q_key}' の文字数 ({q_chars}) が許容上限 ({}) を超過しています。",
131                config.max_question_chars
132            )));
133        }
134
135        total_chars += q_chars;
136        total_estimated_tokens += q_tokens.max(1);
137    }
138
139    if total_estimated_tokens > config.max_estimated_tokens {
140        return Err(ServerError::PayloadTooLarge(format!(
141            "リクエスト全体の推定トークン数 ({total_estimated_tokens}) が許容上限 ({}) を超過しています。",
142            config.max_estimated_tokens
143        )));
144    }
145
146    Ok((total_chars, total_estimated_tokens))
147}
148
149#[cfg(test)]
150mod tests {
151    use super::*;
152    use indexmap::IndexMap;
153    use sokuto_core::schema::Question;
154
155    #[test]
156    fn test_validate_limits_normal() {
157        let mut questions = IndexMap::new();
158        questions.insert(
159            "q1".to_string(),
160            Question::new_choice("test".to_string(), IndexMap::new()),
161        );
162        let req = SystemOneRequest::new(Value::String("hello world".to_string()), questions);
163        let config = LimitsConfig::default();
164
165        let res = validate_limits(&req, &config);
166        assert!(res.is_ok());
167    }
168
169    #[test]
170    fn test_validate_limits_exceed_questions() {
171        let mut questions = IndexMap::new();
172        for i in 0..10 {
173            questions.insert(
174                format!("q{i}"),
175                Question::new_choice("test".to_string(), IndexMap::new()),
176            );
177        }
178        let req = SystemOneRequest::new(Value::String("hello".to_string()), questions);
179        let config = LimitsConfig {
180            max_questions: 5,
181            ..Default::default()
182        };
183
184        let res = validate_limits(&req, &config);
185        assert!(matches!(res, Err(ServerError::PayloadTooLarge(_))));
186    }
187
188    #[test]
189    fn test_validate_limits_exceed_state_chars() {
190        let req = SystemOneRequest::new(
191            Value::String("A".repeat(100)),
192            IndexMap::from([(
193                "q1".to_string(),
194                Question::new_choice("test".to_string(), IndexMap::new()),
195            )]),
196        );
197        let config = LimitsConfig {
198            max_state_chars: 50,
199            ..Default::default()
200        };
201
202        let res = validate_limits(&req, &config);
203        assert!(matches!(res, Err(ServerError::PayloadTooLarge(_))));
204    }
205}
sokuto_server/guardrails/limits.rs

sokuto_server/guardrails/
limits.rs
