Skip to main content

soma_codemode/
shape.rs

1#![allow(dead_code)]
2
3use serde::{Deserialize, Serialize};
4use serde_json::Value;
5
6use crate::CodeModeResultShapePolicy;
7
8const MIN_SHAPED_RESULT_BYTES: usize = 256;
9
10#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
11pub struct CodeModeResultShapeMetadata {
12    pub policy: CodeModeResultShapePolicy,
13    pub changed: bool,
14    pub truncated: bool,
15    pub original_size_bytes: usize,
16    pub shaped_size_bytes: usize,
17}
18
19#[derive(Debug, Clone, PartialEq)]
20pub(crate) struct ShapedResult {
21    pub(crate) result: Option<Value>,
22    pub(crate) metadata: CodeModeResultShapeMetadata,
23}
24
25pub(crate) fn shape_final_result(
26    result: Option<Value>,
27    policy: CodeModeResultShapePolicy,
28    max_response_bytes: usize,
29    max_response_tokens: usize,
30    token_estimate_divisor: u32,
31) -> ShapedResult {
32    let original_size = result
33        .as_ref()
34        .and_then(|value| serde_json::to_vec(value).ok())
35        .map_or(0, |bytes| bytes.len());
36    match (policy, result) {
37        (CodeModeResultShapePolicy::Off, result) | (_, result @ None) => {
38            unchanged(result, policy, original_size)
39        }
40        (CodeModeResultShapePolicy::Truncate, Some(value)) => shape_truncate(
41            value,
42            policy,
43            original_size,
44            max_response_bytes,
45            max_response_tokens,
46            token_estimate_divisor,
47        ),
48    }
49}
50
51fn unchanged(
52    result: Option<Value>,
53    policy: CodeModeResultShapePolicy,
54    original_size_bytes: usize,
55) -> ShapedResult {
56    ShapedResult {
57        result,
58        metadata: CodeModeResultShapeMetadata {
59            policy,
60            changed: false,
61            truncated: false,
62            original_size_bytes,
63            shaped_size_bytes: original_size_bytes,
64        },
65    }
66}
67
68fn shape_truncate(
69    value: Value,
70    policy: CodeModeResultShapePolicy,
71    original_size_bytes: usize,
72    max_response_bytes: usize,
73    max_response_tokens: usize,
74    token_estimate_divisor: u32,
75) -> ShapedResult {
76    let token_budget_bytes = max_response_tokens
77        .max(1)
78        .saturating_mul(token_estimate_divisor.max(1) as usize);
79    let budget = max_response_bytes
80        .min(token_budget_bytes)
81        .max(MIN_SHAPED_RESULT_BYTES);
82    if original_size_bytes <= budget {
83        return unchanged(Some(value), policy, original_size_bytes);
84    }
85    let serialized = match &value {
86        Value::String(text) => text.clone(),
87        _ => serde_json::to_string_pretty(&value).unwrap_or_else(|_| value.to_string()),
88    };
89    let marker_prefix = format!(
90        "[code mode result truncated]\noriginal_size_bytes={original_size_bytes}, max_size_bytes={budget}\n"
91    );
92    let room = budget.saturating_sub(marker_prefix.len());
93    let marker = format!(
94        "{marker_prefix}{}",
95        crate::util::utf8_prefix_by_bytes(&serialized, room)
96    );
97    let shaped_size = serde_json::to_vec(&Value::String(marker.clone()))
98        .map_or_else(|_| marker.len(), |bytes| bytes.len());
99    ShapedResult {
100        result: Some(Value::String(marker)),
101        metadata: CodeModeResultShapeMetadata {
102            policy,
103            changed: true,
104            truncated: true,
105            original_size_bytes,
106            shaped_size_bytes: shaped_size,
107        },
108    }
109}