Skip to main content

soma_codemode/
config.rs

1#![allow(dead_code)]
2
3use serde::{Deserialize, Serialize};
4
5pub const SERVICE: &str = "code_mode";
6pub const MAX_SOURCE_BYTES: usize = 20_000;
7pub(crate) const MAX_SNIPPET_RESOLVES_PER_RUN: usize = 32;
8pub(crate) const MAX_INTERNAL_CALLS_PER_RUN: usize = 32;
9pub(crate) const MAX_SNIPPET_RESOLVED_BYTES_PER_RUN: usize = 256 * 1024;
10
11const DEFAULT_MAX_CALLTOOL_PER_RUN: u64 = 512;
12const MAX_CALLTOOL_PER_RUN_CEILING: u64 = 2048;
13const DEFAULT_CALLTOOL_RESULT_MAX_MIB: usize = 8;
14
15static MAX_CALLS_PER_RUN_CONFIG_DEFAULT: std::sync::OnceLock<Option<u64>> =
16    std::sync::OnceLock::new();
17static CALLTOOL_RESULT_MAX_MIB_CONFIG_DEFAULT: std::sync::OnceLock<Option<usize>> =
18    std::sync::OnceLock::new();
19
20#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
21#[serde(rename_all = "snake_case")]
22pub enum CodeModeResultShapePolicy {
23    #[default]
24    Off,
25    Truncate,
26}
27
28#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
29pub struct SemanticSearchConfig {
30    #[serde(default, skip_serializing_if = "Option::is_none")]
31    pub tei_url: Option<String>,
32    #[serde(default = "default_semantic_search_blend_weight")]
33    pub blend_weight: f32,
34}
35
36impl Default for SemanticSearchConfig {
37    fn default() -> Self {
38        Self {
39            tei_url: None,
40            blend_weight: default_semantic_search_blend_weight(),
41        }
42    }
43}
44
45impl SemanticSearchConfig {
46    #[must_use]
47    pub fn is_configured(&self) -> bool {
48        self.tei_url
49            .as_deref()
50            .is_some_and(|url| !url.trim().is_empty())
51    }
52}
53
54#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
55pub struct CodeModeConfig {
56    #[serde(default)]
57    pub enabled: bool,
58    #[serde(default = "default_true")]
59    pub trace_params: bool,
60    #[serde(default)]
61    pub result_shape_policy: CodeModeResultShapePolicy,
62    #[serde(default = "default_code_mode_timeout_ms")]
63    pub timeout_ms: u64,
64    #[serde(default = "default_code_mode_max_response_bytes")]
65    pub max_response_bytes: usize,
66    #[serde(default = "default_code_mode_max_response_tokens")]
67    pub max_response_tokens: usize,
68    #[serde(default = "default_token_estimate_divisor")]
69    pub token_estimate_divisor: u32,
70    #[serde(default = "default_max_log_entries")]
71    pub max_log_entries: usize,
72    #[serde(default = "default_max_log_bytes")]
73    pub max_log_bytes: usize,
74    #[serde(default)]
75    pub semantic_search: SemanticSearchConfig,
76    #[serde(default)]
77    pub max_calls_per_run: Option<u64>,
78    #[serde(default)]
79    pub calltool_result_max_mib: Option<usize>,
80}
81
82impl Default for CodeModeConfig {
83    fn default() -> Self {
84        Self {
85            enabled: false,
86            trace_params: true,
87            result_shape_policy: CodeModeResultShapePolicy::Off,
88            timeout_ms: default_code_mode_timeout_ms(),
89            max_response_bytes: default_code_mode_max_response_bytes(),
90            max_response_tokens: default_code_mode_max_response_tokens(),
91            token_estimate_divisor: default_token_estimate_divisor(),
92            max_log_entries: default_max_log_entries(),
93            max_log_bytes: default_max_log_bytes(),
94            semantic_search: SemanticSearchConfig::default(),
95            max_calls_per_run: None,
96            calltool_result_max_mib: None,
97        }
98    }
99}
100
101impl CodeModeConfig {
102    pub fn validate(&self) -> Result<(), String> {
103        if !(1..=60_000).contains(&self.timeout_ms) {
104            return Err("timeout_ms must be between 1 and 60000".into());
105        }
106        if !(1024..=1024 * 1024).contains(&self.max_response_bytes) {
107            return Err("max_response_bytes must be between 1024 and 1048576".into());
108        }
109        if !(256..=256_000).contains(&self.max_response_tokens) {
110            return Err("max_response_tokens must be between 256 and 256000".into());
111        }
112        if !(1..=64).contains(&self.token_estimate_divisor) {
113            return Err("token_estimate_divisor must be between 1 and 64".into());
114        }
115        if !(0.0..=1.0).contains(&self.semantic_search.blend_weight) {
116            return Err("semantic_search.blend_weight must be between 0 and 1".into());
117        }
118        Ok(())
119    }
120}
121
122pub fn install_call_budget_config_defaults(
123    max_calls_per_run: Option<u64>,
124    calltool_result_max_mib: Option<usize>,
125) {
126    let _ = MAX_CALLS_PER_RUN_CONFIG_DEFAULT.set(max_calls_per_run);
127    let _ = CALLTOOL_RESULT_MAX_MIB_CONFIG_DEFAULT.set(calltool_result_max_mib);
128}
129
130pub(crate) fn max_calltool_per_run() -> u64 {
131    let Some(raw) = crate::home::env_non_empty("SOMA_CODE_MODE_MAX_CALLS_PER_RUN") else {
132        return MAX_CALLS_PER_RUN_CONFIG_DEFAULT
133            .get()
134            .copied()
135            .flatten()
136            .filter(|value| *value > 0)
137            .map_or(DEFAULT_MAX_CALLTOOL_PER_RUN, |value| {
138                value.min(MAX_CALLTOOL_PER_RUN_CEILING)
139            });
140    };
141    raw.trim()
142        .parse::<u64>()
143        .ok()
144        .filter(|value| *value > 0)
145        .map_or(DEFAULT_MAX_CALLTOOL_PER_RUN, |value| {
146            value.min(MAX_CALLTOOL_PER_RUN_CEILING)
147        })
148}
149
150pub(crate) fn effective_max_calltool_per_run(config: &CodeModeConfig) -> u64 {
151    crate::home::env_non_empty("SOMA_CODE_MODE_MAX_CALLS_PER_RUN")
152        .and_then(|raw| raw.trim().parse::<u64>().ok())
153        .filter(|value| *value > 0)
154        .or(config.max_calls_per_run.filter(|value| *value > 0))
155        .map_or_else(max_calltool_per_run, |value| {
156            value.min(MAX_CALLTOOL_PER_RUN_CEILING)
157        })
158}
159
160pub(crate) fn calltool_result_max_bytes() -> usize {
161    let default_bytes = CALLTOOL_RESULT_MAX_MIB_CONFIG_DEFAULT
162        .get()
163        .copied()
164        .flatten()
165        .filter(|mib| *mib > 0)
166        .map(|mib| mib.saturating_mul(1024 * 1024))
167        .unwrap_or(DEFAULT_CALLTOOL_RESULT_MAX_MIB * 1024 * 1024);
168    crate::home::env_non_empty("SOMA_CODE_MODE_CALLTOOL_RESULT_MAX_MIB")
169        .and_then(|raw| raw.trim().parse::<usize>().ok())
170        .filter(|mib| *mib > 0)
171        .map(|mib| mib.saturating_mul(1024 * 1024))
172        .unwrap_or(default_bytes)
173}
174
175pub(crate) fn effective_calltool_result_max_bytes(config: &CodeModeConfig) -> usize {
176    crate::home::env_non_empty("SOMA_CODE_MODE_CALLTOOL_RESULT_MAX_MIB")
177        .and_then(|raw| raw.trim().parse::<usize>().ok())
178        .filter(|mib| *mib > 0)
179        .or(config.calltool_result_max_mib.filter(|mib| *mib > 0))
180        .map(|mib| mib.saturating_mul(1024 * 1024))
181        .unwrap_or_else(calltool_result_max_bytes)
182}
183
184fn default_true() -> bool {
185    true
186}
187
188fn default_code_mode_timeout_ms() -> u64 {
189    30_000
190}
191
192fn default_code_mode_max_response_bytes() -> usize {
193    24 * 1024
194}
195
196fn default_code_mode_max_response_tokens() -> usize {
197    6_000
198}
199
200fn default_token_estimate_divisor() -> u32 {
201    4
202}
203
204fn default_max_log_entries() -> usize {
205    1000
206}
207
208fn default_max_log_bytes() -> usize {
209    65_536
210}
211
212fn default_semantic_search_blend_weight() -> f32 {
213    0.5
214}