Skip to main content

soma_codemode/
trace.rs

1#![allow(dead_code)]
2
3use serde_json::{json, Map, Value};
4
5use crate::types::CodeModeExecutionResponse;
6
7const REDACTED: &str = "[redacted]";
8const MAX_DEPTH: usize = 16;
9const MAX_COLLECTION_ITEMS: usize = 64;
10const MAX_STRING_CHARS: usize = 512;
11
12#[must_use]
13pub fn redact_trace_value(value: &Value, max_bytes: usize) -> Value {
14    let redacted = redact_value(value, 0);
15    let size = serde_json::to_vec(&redacted).map_or(usize::MAX, |bytes| bytes.len());
16    if size <= max_bytes {
17        redacted
18    } else {
19        json!({
20            "truncated": true,
21            "reason": "redacted_params_exceeded_cap",
22            "original_size_bytes": size,
23            "max_size_bytes": max_bytes,
24        })
25    }
26}
27
28#[must_use]
29pub(crate) fn redact_trace_params(params: &Value, enabled: bool) -> Option<Value> {
30    enabled.then(|| redact_trace_value(params, 4096))
31}
32
33#[must_use]
34pub fn code_mode_execute_trace(response: &CodeModeExecutionResponse) -> Value {
35    let calls = response
36        .calls
37        .iter()
38        .map(|call| {
39            let (namespace, tool) = crate::types::id::split_namespaced_id(&call.id)
40                .map_or(("", call.id.as_str()), |(namespace, tool)| {
41                    (namespace, tool)
42                });
43            json!({
44                "id": call.id,
45                "namespace": namespace,
46                "tool": tool,
47                "params": call
48                    .params
49                    .as_ref()
50                    .map(|params| redact_trace_value(params, 4096)),
51                "ok": call.result.is_some(),
52            })
53        })
54        .collect::<Vec<_>>();
55    let mut trace = Map::new();
56    trace.insert("kind".to_string(), json!("code_mode_execute_trace"));
57    trace.insert("call_count".to_string(), json!(response.calls.len()));
58    trace.insert("calls".to_string(), json!(calls));
59    if let Some(result) = &response.result {
60        trace.insert("result".to_string(), redact_trace_value(result, 4096));
61    }
62    trace.insert(
63        "result_shape".to_string(),
64        response
65            .result
66            .as_ref()
67            .map(compact_result_shape)
68            .unwrap_or_else(|| json!({"type": "undefined"})),
69    );
70    trace.insert("logs_count".to_string(), json!(response.logs.len()));
71    Value::Object(trace)
72}
73
74fn redact_value(value: &Value, depth: usize) -> Value {
75    if depth > MAX_DEPTH {
76        return json!({"truncated": true, "reason": "max_depth"});
77    }
78    match value {
79        Value::Object(object) => Value::Object(
80            object
81                .iter()
82                .take(MAX_COLLECTION_ITEMS)
83                .map(|(key, value)| {
84                    if crate::redact::is_sensitive_key(key) {
85                        (key.clone(), Value::String(REDACTED.to_string()))
86                    } else {
87                        (key.clone(), redact_value(value, depth + 1))
88                    }
89                })
90                .collect(),
91        ),
92        Value::Array(values) => Value::Array(
93            values
94                .iter()
95                .take(MAX_COLLECTION_ITEMS)
96                .map(|value| redact_value(value, depth + 1))
97                .collect(),
98        ),
99        Value::String(text) => {
100            let redacted = crate::truncate::redact_secret_like_segments(text);
101            if redacted != *text {
102                Value::String(redacted)
103            } else if text.chars().count() > MAX_STRING_CHARS {
104                Value::String(format!(
105                    "{}[truncated]",
106                    text.chars().take(MAX_STRING_CHARS).collect::<String>()
107                ))
108            } else {
109                value.clone()
110            }
111        }
112        other => other.clone(),
113    }
114}
115
116fn compact_result_shape(value: &Value) -> Value {
117    let size_bytes = serde_json::to_vec(value).map_or(usize::MAX, |bytes| bytes.len());
118    match value {
119        Value::Null => json!({"type": "null", "size_bytes": size_bytes}),
120        Value::Bool(_) => json!({"type": "boolean", "size_bytes": size_bytes}),
121        Value::Number(_) => json!({"type": "number", "size_bytes": size_bytes}),
122        Value::String(text) => {
123            json!({"type": "string", "size_bytes": size_bytes, "length": text.chars().count()})
124        }
125        Value::Array(values) => {
126            json!({"type": "array", "size_bytes": size_bytes, "length": values.len()})
127        }
128        Value::Object(object) => {
129            let mut keys = object.keys().take(16).cloned().collect::<Vec<_>>();
130            keys.sort();
131            json!({"type": "object", "size_bytes": size_bytes, "keys": keys, "key_count": object.len()})
132        }
133    }
134}