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}