soma_codemode/
truncate.rs

1#![allow(dead_code)]
2
3use std::sync::LazyLock;
4
5use regex::Regex;
6use serde_json::{json, Value};
7
8use crate::types::CodeModeExecutionResponse;
9
10static SECRET_REGEX: LazyLock<Regex> = LazyLock::new(|| {
11    Regex::new(
12        r"(?:redaction-canary-[A-Za-z0-9_-]{20,}|sk-[A-Za-z0-9_-]{20,}|ghp_[A-Za-z0-9]{20,}|github_pat_[A-Za-z0-9_]{20,}|glpat-[A-Za-z0-9_-]{20,}|xox[bp]-[A-Za-z0-9-]+|eyJ[A-Za-z0-9_-]+\.[A-Za-z0-9_-]+)",
13    )
14    .expect("secret regex is valid")
15});
16
17pub fn redact_secret_like_segments(input: &str) -> String {
18    let split_redacted = input
19        .split_whitespace()
20        .map(|segment| {
21            if segment.starts_with("sk-")
22                || segment.starts_with("redaction-canary-")
23                || segment.starts_with("ghp_")
24                || segment.starts_with("github_pat_")
25                || segment.starts_with("glpat-")
26                || segment.starts_with("xoxb-")
27                || segment.starts_with("xoxp-")
28                || segment.starts_with("eyJ")
29            {
30                "[REDACTED]".to_string()
31            } else {
32                segment.to_string()
33            }
34        })
35        .collect::<Vec<_>>()
36        .join(" ");
37    SECRET_REGEX
38        .replace_all(&split_redacted, "[REDACTED]")
39        .into_owned()
40}
41
42pub(crate) fn sanitize_log_text(input: &str, max_len: usize) -> String {
43    let mut value = input.to_string();
44    value.retain(|ch| {
45        !matches!(
46            ch,
47            '\u{0000}'..='\u{001F}'
48                | '\u{007F}'..='\u{009F}'
49                | '\u{202A}'..='\u{202E}'
50                | '\u{2066}'..='\u{2069}'
51        )
52    });
53    for marker in ["<system>", "[INST]", "###", "<<"] {
54        value = value.replace(marker, "");
55    }
56    redact_secret_like_segments(&value)
57        .chars()
58        .take(max_len)
59        .collect()
60}
61
62pub(crate) fn truncate_execution_response(
63    mut response: CodeModeExecutionResponse,
64    max_response_bytes: usize,
65    max_response_tokens: usize,
66    token_estimate_divisor: u32,
67) -> CodeModeExecutionResponse {
68    if response_within_budget(
69        &response,
70        max_response_bytes,
71        max_response_tokens,
72        token_estimate_divisor,
73    ) {
74        return response;
75    }
76    if let Some(result) = response.result.take() {
77        response.result = Some(truncation_marker(&result, token_estimate_divisor));
78    }
79    while !response.logs.is_empty()
80        && !response_within_budget(
81            &response,
82            max_response_bytes,
83            max_response_tokens,
84            token_estimate_divisor,
85        )
86    {
87        response.logs.remove(0);
88    }
89    response
90}
91
92pub(crate) fn response_within_budget(
93    response: &CodeModeExecutionResponse,
94    max_response_bytes: usize,
95    max_response_tokens: usize,
96    token_estimate_divisor: u32,
97) -> bool {
98    serde_json::to_vec(response).is_ok_and(|bytes| {
99        bytes.len() <= max_response_bytes
100            && estimated_tokens(bytes.len(), token_estimate_divisor) <= max_response_tokens.max(1)
101    })
102}
103
104fn truncation_marker(value: &Value, divisor: u32) -> Value {
105    let serialized = serde_json::to_string(value).unwrap_or_else(|_| "null".to_string());
106    json!({
107        "truncated": true,
108        "original_size": serialized.len(),
109        "original_tokens": estimated_tokens(serialized.len(), divisor),
110        "preview": crate::util::utf8_prefix_by_bytes(&serialized, 1024),
111    })
112}
113
114fn estimated_tokens(byte_len: usize, divisor: u32) -> usize {
115    byte_len.div_ceil(divisor.max(1) as usize).max(1)
116}
soma_codemode/truncate.rs

soma_codemode/
truncate.rs
