soma_codemode/execute/
budget.rs

1use serde_json::{json, Value};
2
3use crate::{CodeModeConfig, ToolError};
4
5#[derive(Debug, Clone)]
6pub(crate) struct RunBudget {
7    max_operations: u64,
8    operations: u64,
9    result_max_bytes: usize,
10    max_log_entries: usize,
11    max_log_bytes: usize,
12}
13
14impl RunBudget {
15    pub(crate) fn new(config: &CodeModeConfig) -> Self {
16        Self {
17            max_operations: crate::config::effective_max_calltool_per_run(config),
18            operations: 0,
19            result_max_bytes: crate::config::effective_calltool_result_max_bytes(config),
20            max_log_entries: config.max_log_entries,
21            max_log_bytes: config.max_log_bytes,
22        }
23    }
24
25    pub(crate) fn record_operation(&mut self, label: &str) -> Result<(), ToolError> {
26        self.operations = self.operations.saturating_add(1);
27        if self.operations > self.max_operations {
28            return Err(ToolError::Sdk {
29                sdk_kind: "budget_exceeded".to_string(),
30                message: format!(
31                    "Code Mode {label} exceeded the per-run operation budget of {}",
32                    self.max_operations
33                ),
34            });
35        }
36        Ok(())
37    }
38
39    pub(crate) fn cap_tool_result(&self, value: Value) -> Value {
40        let Ok(bytes) = serde_json::to_vec(&value) else {
41            return json!({"truncated": true, "reason": "tool result was not serializable"});
42        };
43        if bytes.len() <= self.result_max_bytes {
44            return value;
45        }
46        let serialized = serde_json::to_string(&value).unwrap_or_else(|_| "null".to_string());
47        let preview =
48            crate::util::utf8_prefix_by_bytes(&serialized, self.result_max_bytes.min(1024));
49        json!({
50            "truncated": true,
51            "original_size_bytes": bytes.len(),
52            "max_size_bytes": self.result_max_bytes,
53            "preview": preview,
54        })
55    }
56
57    pub(crate) fn cap_logs(&self, logs: Vec<String>) -> Vec<String> {
58        let mut capped = Vec::new();
59        let mut total = 0usize;
60        let max_entries = self.max_log_entries.max(1);
61        let max_bytes = self.max_log_bytes.max(1);
62        for log in logs {
63            let sanitized = crate::truncate::sanitize_log_text(&log, max_bytes.min(4096));
64            let next = sanitized.len();
65            if capped.len() >= max_entries || total.saturating_add(next) > max_bytes {
66                capped.push("[soma] Code Mode logs truncated".to_string());
67                break;
68            }
69            total = total.saturating_add(next);
70            capped.push(sanitized);
71        }
72        capped
73    }
74}
soma_codemode/execute/budget.rs

soma_codemode/execute/
budget.rs
