rac_engine/
budget.rs

1//! Per-response character budget (ADR-033) — a port of `src/asdecided/mcp/budget.py`
2//! (ORACLE-NEXT revision, which adds the `items` rule). Shared by the CLI
3//! retrieve surface (`commands::cmd_retrieve`) and the decided-mcp server.
4//!
5//! The budget unit is CHARACTERS (Python `len` of the serialized string —
6//! Unicode code points, not bytes) of the payload serialized as
7//! `json.dumps(payload, ensure_ascii=False)` with DEFAULT separators — i.e.
8//! `", "` / `": "` WITH spaces (`pyjson::dumps_compact`). The docstring in the
9//! oracle says "no spaces"; the code does not pass `separators`, so the wire
10//! truth is *with* spaces (PORT-CONTRACT.d/10 §2) — port the code, not the
11//! comment.
12//!
13//! Truncation is deterministic and whole-item wherever a repeated collection is
14//! involved. Every successful payload is either reduced to the configured
15//! character budget or replaced with a small explicit budget error; an
16//! over-budget success is never returned.
17
18use crate::pyjson::dumps_compact;
19use serde_json::{json, Map, Value};
20
21pub const DEFAULT_BUDGET: i64 = 10_000;
22/// Smallest supported configured budget. This leaves room for a structured
23/// budget error when fixed response fields alone cannot fit.
24pub const MIN_BUDGET: i64 = 128;
25
26pub const MARKER_TRUNCATED: &str = "truncated";
27pub const MARKER_OMITTED: &str = "omitted";
28pub const MARKER_HINT: &str = "hint";
29
30pub const HINT_SEARCH: &str = "Narrow the query or request a specific artifact ID.";
31pub const HINT_RELATED: &str = "Request the artifact directly, or narrow what you are changing.";
32pub const HINT_CONTENT: &str =
33    "Request a more specific artifact, or read the file directly for the full content.";
34pub const HINT_SUMMARY: &str = "The repository summary exceeds the response budget; raise the \
35server budget to see the full overview.";
36pub const HINT_RETRIEVE: &str = "Lower top_k, raise the budget, or narrow the task.";
37pub const BUDGET_ERROR: &str = "response_budget_exceeded";
38pub const BUDGET_ERROR_HINT: &str = "Raise the response budget or narrow the request.";
39
40type TruncationStrategy = fn(&Value, i64) -> (Value, bool);
41
42pub fn valid_configured_budget(budget: i64) -> bool {
43    budget >= MIN_BUDGET
44}
45
46pub fn validate_call_budget(budget: i64) -> Result<(), String> {
47    if budget > 0 && budget < MIN_BUDGET {
48        return Err(format!(
49            "Requested response budget {budget} is below the minimum supported budget of {MIN_BUDGET} characters."
50        ));
51    }
52    Ok(())
53}
54
55/// `len(text)` in Python — code points, not bytes.
56pub fn char_len(s: &str) -> i64 {
57    s.chars().count() as i64
58}
59
60/// `text[:stop]` with Python slice semantics (negative stop trims the tail;
61/// the truncators only pass non-negative stops, but the retrieve excerpt
62/// share can go negative).
63pub fn py_slice_to(s: &str, stop: i64) -> String {
64    let n = char_len(s);
65    let stop = if stop < 0 { (n + stop).max(0) } else { stop.min(n) };
66    s.chars().take(stop as usize).collect()
67}
68
69fn length(payload: &Value) -> i64 {
70    char_len(&dumps_compact(payload))
71}
72
73/// `budget.serialize(payload, budget)`.
74pub fn serialize(payload: &Value, budget: i64) -> String {
75    let text = dumps_compact(payload);
76    if char_len(&text) <= budget {
77        return text;
78    }
79    let truncated = truncate(payload, budget);
80    let text = dumps_compact(&truncated);
81    if char_len(&text) <= budget {
82        return text;
83    }
84    budget_error(budget)
85}
86
87fn truncate(payload: &Value, budget: i64) -> Value {
88    let Some(obj) = payload.as_object() else {
89        return json!({"error": BUDGET_ERROR, "hint": BUDGET_ERROR_HINT});
90    };
91    let mut candidate = Value::Object(obj.clone());
92
93    // A shape can require more than one reduction (for example, a deep
94    // relationship response has both incoming and neighborhood collections).
95    // Apply strategies in a stable order, then remove optional fixed fields as
96    // a last resort before returning the explicit error from `serialize`.
97    let strategies: [TruncationStrategy; 9] = [
98        |value, limit| truncate_list_strategy(value, "matches", limit, HINT_SEARCH),
99        |value, limit| truncate_items_strategy(value, limit),
100        |value, limit| truncate_content_strategy(value, limit),
101        |value, limit| truncate_related_strategy(value, limit),
102        |value, limit| truncate_list_strategy(value, "decisions", limit, HINT_SEARCH),
103        |value, limit| truncate_list_strategy(value, "attention", limit, HINT_SUMMARY),
104        |value, limit| truncate_list_strategy(value, "neighborhood", limit, HINT_RELATED),
105        |value, limit| truncate_outgoing_strategy(value, limit),
106        truncate_optional_strategy,
107    ];
108    for strategy in strategies {
109        if char_len(&dumps_compact(&candidate)) <= budget {
110            break;
111        }
112        let (next, changed) = strategy(&candidate, budget);
113        candidate = next;
114        if !changed {
115            continue;
116        }
117    }
118    candidate
119}
120
121fn budget_error(budget: i64) -> String {
122    let full = dumps_compact(&json!({
123        "error": BUDGET_ERROR,
124        "hint": BUDGET_ERROR_HINT,
125    }));
126    if char_len(&full) <= budget {
127        return full;
128    }
129    let short = dumps_compact(&json!({"error": BUDGET_ERROR}));
130    if char_len(&short) <= budget {
131        return short;
132    }
133    if budget >= 2 {
134        return "{}".to_string();
135    }
136    if budget == 1 {
137        return "0".to_string();
138    }
139    String::new()
140}
141
142/// A copy of `payload` with `key` replaced by `kept` and the marker added.
143/// `IndexMap::insert` keeps an existing key's position, matching Python
144/// dict-update semantics (a `truncated` key already present — the
145/// `get_related` edge-overflow marker — is overwritten in place).
146fn with_marker(payload: &Value, key: &str, kept: Vec<Value>, omitted: i64, hint: &str) -> Value {
147    let mut marked: Map<String, Value> = payload.as_object().expect("object").clone();
148    marked.insert(key.to_string(), Value::Array(kept));
149    marked.insert(MARKER_TRUNCATED.to_string(), json!(true));
150    marked.insert(MARKER_OMITTED.to_string(), json!(omitted));
151    marked.insert(MARKER_HINT.to_string(), json!(hint));
152    Value::Object(marked)
153}
154
155fn existing_omitted(payload: &Value) -> i64 {
156    payload
157        .get(MARKER_OMITTED)
158        .and_then(Value::as_i64)
159        .unwrap_or(0)
160}
161
162fn truncate_list_strategy(
163    payload: &Value,
164    key: &str,
165    budget: i64,
166    hint: &str,
167) -> (Value, bool) {
168    let Some(items) = payload.get(key).and_then(Value::as_array) else {
169        return (payload.clone(), false);
170    };
171    if items.is_empty() {
172        return (payload.clone(), false);
173    }
174    let items: Vec<Value> = payload[key].as_array().cloned().unwrap_or_default();
175    let total = items.len() as i64;
176    let mut kept = items;
177    while !kept.is_empty() {
178        let candidate = with_marker(payload, key, kept.clone(), total - kept.len() as i64, hint);
179        if length(&candidate) <= budget {
180            return (candidate, true);
181        }
182        kept.pop();
183    }
184    (with_marker(payload, key, Vec::new(), total, hint), true)
185}
186
187fn truncate_items_strategy(payload: &Value, budget: i64) -> (Value, bool) {
188    if payload
189        .get("items")
190        .and_then(Value::as_array)
191        .is_none_or(|items| items.is_empty())
192    {
193        return (payload.clone(), false);
194    }
195    (truncate_items(payload, budget), true)
196}
197
198fn truncate_content_strategy(payload: &Value, budget: i64) -> (Value, bool) {
199    let Some(content) = payload.get("content").and_then(Value::as_str) else {
200        return (payload.clone(), false);
201    };
202    if content.is_empty() {
203        return (payload.clone(), false);
204    }
205    (truncate_content(payload, budget), true)
206}
207
208fn truncate_content(payload: &Value, budget: i64) -> Value {
209    let content = payload["content"].as_str().unwrap_or("").to_string();
210    let total = char_len(&content);
211    let with_content = |kept: String, omitted: i64| -> Value {
212        let mut marked: Map<String, Value> = payload.as_object().expect("object").clone();
213        marked.insert("content".to_string(), json!(kept));
214        marked.insert(MARKER_TRUNCATED.to_string(), json!(true));
215        marked.insert(MARKER_OMITTED.to_string(), json!(omitted));
216        marked.insert(MARKER_HINT.to_string(), json!(HINT_CONTENT));
217        Value::Object(marked)
218    };
219    let (mut lo, mut hi) = (0i64, total);
220    let mut best = 0i64;
221    while lo <= hi {
222        let mid = (lo + hi).div_euclid(2);
223        let candidate = with_content(py_slice_to(&content, mid), total - mid);
224        if length(&candidate) <= budget {
225            best = mid;
226            lo = mid + 1;
227        } else {
228            hi = mid - 1;
229        }
230    }
231    with_content(py_slice_to(&content, best), total - best)
232}
233
234/// The retrieve `items` rule (ADR-113): excerpt-first, then whole-item.
235fn truncate_items(payload: &Value, budget: i64) -> Value {
236    let items: Vec<Value> = payload["items"].as_array().cloned().unwrap_or_default();
237    let total = items.len() as i64;
238    let mut kept = items;
239    while !kept.is_empty() {
240        let omitted = total - kept.len() as i64;
241        let candidate = with_marker(payload, "items", kept.clone(), omitted, HINT_RETRIEVE);
242        if length(&candidate) <= budget {
243            return candidate;
244        }
245        // Trim the last kept item's excerpt before dropping it entirely.
246        let mut last = kept
247            .last()
248            .and_then(Value::as_object)
249            .cloned()
250            .unwrap_or_default();
251        let excerpt: String = last
252            .get("excerpt")
253            .and_then(Value::as_str)
254            .unwrap_or("")
255            .to_string();
256        let (mut lo, mut hi) = (0i64, char_len(&excerpt));
257        let mut best: Option<i64> = None;
258        while lo <= hi {
259            let mid = (lo + hi).div_euclid(2);
260            last.insert("excerpt".to_string(), json!(py_slice_to(&excerpt, mid)));
261            let mut trial_items: Vec<Value> = kept[..kept.len() - 1].to_vec();
262            trial_items.push(Value::Object(last.clone()));
263            let trial = with_marker(payload, "items", trial_items, omitted, HINT_RETRIEVE);
264            if length(&trial) <= budget {
265                best = Some(mid);
266                lo = mid + 1;
267            } else {
268                hi = mid - 1;
269            }
270        }
271        if let Some(best) = best {
272            last.insert("excerpt".to_string(), json!(py_slice_to(&excerpt, best)));
273            let mut final_items: Vec<Value> = kept[..kept.len() - 1].to_vec();
274            final_items.push(Value::Object(last));
275            return with_marker(payload, "items", final_items, omitted, HINT_RETRIEVE);
276        }
277        kept.pop();
278    }
279    with_marker(payload, "items", Vec::new(), total, HINT_RETRIEVE)
280}
281
282fn truncate_related_strategy(payload: &Value, budget: i64) -> (Value, bool) {
283    let has_incoming = payload.get("incoming").is_some_and(Value::is_array);
284    let has_neighborhood = payload.get("neighborhood").is_some_and(Value::is_array);
285    if !has_incoming && !has_neighborhood {
286        return (payload.clone(), false);
287    }
288    let mut candidate = payload.clone();
289    let mut omitted = existing_omitted(payload);
290    let mut changed = false;
291    for key in ["incoming", "neighborhood"] {
292        let Some(items) = candidate.get(key).and_then(Value::as_array).cloned() else {
293            continue;
294        };
295        let mut kept = items;
296        while length(&candidate) > budget && !kept.is_empty() {
297            kept.pop();
298            omitted += 1;
299            changed = true;
300            let mut marked = candidate.as_object().expect("object").clone();
301            marked.insert(key.to_string(), Value::Array(kept.clone()));
302            marked.insert(MARKER_TRUNCATED.to_string(), json!(true));
303            marked.insert(MARKER_OMITTED.to_string(), json!(omitted));
304            marked.insert(MARKER_HINT.to_string(), json!(HINT_RELATED));
305            candidate = Value::Object(marked);
306        }
307        if length(&candidate) <= budget {
308            return (candidate, changed);
309        }
310    }
311    (candidate, changed)
312}
313
314fn truncate_outgoing_strategy(payload: &Value, budget: i64) -> (Value, bool) {
315    let Some(outgoing) = payload.get("outgoing").and_then(Value::as_object) else {
316        return (payload.clone(), false);
317    };
318    if !outgoing.values().any(Value::is_array) {
319        return (payload.clone(), false);
320    }
321    let mut candidate = payload.clone();
322    let mut omitted = existing_omitted(payload);
323    let mut changed = false;
324    while length(&candidate) > budget {
325        let Some((section, targets)) = candidate
326            .get("outgoing")
327            .and_then(Value::as_object)
328            .and_then(|map| {
329                map.iter()
330                    .rev()
331                    .find(|(_, value)| value.as_array().is_some_and(|items| !items.is_empty()))
332            })
333        else {
334            break;
335        };
336        let mut updated = candidate.as_object().expect("object").clone();
337        let mut outgoing = updated
338            .get("outgoing")
339            .and_then(Value::as_object)
340            .cloned()
341            .expect("outgoing object");
342        let mut kept = targets.as_array().cloned().expect("outgoing targets");
343        kept.pop();
344        if kept.is_empty() {
345            outgoing.remove(section);
346        } else {
347            outgoing.insert(section.clone(), Value::Array(kept));
348        }
349        omitted += 1;
350        changed = true;
351        updated.insert("outgoing".to_string(), Value::Object(outgoing));
352        updated.insert(MARKER_TRUNCATED.to_string(), json!(true));
353        updated.insert(MARKER_OMITTED.to_string(), json!(omitted));
354        updated.insert(MARKER_HINT.to_string(), json!(HINT_RELATED));
355        candidate = Value::Object(updated);
356    }
357    (candidate, changed)
358}
359
360fn truncate_optional_strategy(payload: &Value, budget: i64) -> (Value, bool) {
361    let Some(object) = payload.as_object() else {
362        return (payload.clone(), false);
363    };
364    // These fields are derived context, not the artifact identity itself. Drop
365    // them in a fixed order only after whole-item collection truncation has
366    // been exhausted. The marker tells the caller that context was omitted.
367    const OPTIONAL: [&str; 10] = [
368        "provenance",
369        "evidence",
370        "outgoing",
371        "incoming",
372        "neighborhood",
373        "attention",
374        "completeness",
375        "relationships",
376        "health",
377        "validation_status",
378    ];
379    let mut candidate = object.clone();
380    let mut changed = false;
381    for key in OPTIONAL {
382        if candidate.contains_key(key) && length(&Value::Object(candidate.clone())) > budget {
383            candidate.remove(key);
384            candidate.insert(MARKER_TRUNCATED.to_string(), json!(true));
385            candidate.insert(MARKER_OMITTED.to_string(), json!(existing_omitted(payload)));
386            candidate.insert(MARKER_HINT.to_string(), json!(HINT_RELATED));
387            changed = true;
388        }
389        if length(&Value::Object(candidate.clone())) <= budget {
390            break;
391        }
392    }
393    (Value::Object(candidate), changed)
394}
395
396#[cfg(test)]
397mod tests {
398    use super::*;
399
400    fn repeated(prefix: &str, count: usize) -> Vec<Value> {
401        (0..count)
402            .map(|index| json!({"id": format!("{prefix}-{index:04}"), "path": format!("{prefix}/{index}.md")}))
403            .collect()
404    }
405
406    #[test]
407    fn summary_attention_is_truncated_to_the_budget() {
408        let payload = json!({
409            "schema_version": "1",
410            "directory": "decisions",
411            "recursive": true,
412            "attention": repeated("attention", 256),
413            "health": {"score": 1.0}
414        });
415        let text = serialize(&payload, 512);
416        assert!(char_len(&text) <= 512, "{} characters", char_len(&text));
417        let value: Value = serde_json::from_str(&text).expect("budget result is JSON");
418        assert_eq!(value[MARKER_TRUNCATED], json!(true));
419        assert!(value[MARKER_OMITTED].as_i64().unwrap_or(0) > 0);
420        assert!(value["attention"].as_array().unwrap().len() < 256);
421    }
422
423    #[test]
424    fn related_incoming_and_neighborhood_are_truncated_deterministically() {
425        let payload = json!({
426            "schema_version": "1",
427            "id": "ADR-001",
428            "depth": 3,
429            "incoming": repeated("incoming", 128),
430            "neighborhood": repeated("neighbor", 128)
431        });
432        let first = serialize(&payload, 512);
433        let second = serialize(&payload, 512);
434        assert_eq!(first, second);
435        assert!(char_len(&first) <= 512, "{} characters", char_len(&first));
436        let value: Value = serde_json::from_str(&first).expect("budget result is JSON");
437        assert_eq!(value[MARKER_TRUNCATED], json!(true));
438        assert!(value[MARKER_OMITTED].as_i64().unwrap_or(0) > 0);
439    }
440
441    #[test]
442    fn outgoing_relationship_targets_are_truncated_deterministically() {
443        let payload = json!({
444            "schema_version": "1",
445            "id": "ADR-001",
446            "depth": 3,
447            "outgoing": {
448                "related decisions": repeated("decision", 96),
449                "related requirements": repeated("requirement", 96)
450            }
451        });
452        let first = serialize(&payload, 512);
453        let second = serialize(&payload, 512);
454        assert_eq!(first, second);
455        assert!(char_len(&first) <= 512, "{} characters", char_len(&first));
456        let value: Value = serde_json::from_str(&first).expect("budget result is JSON");
457        assert_eq!(value[MARKER_TRUNCATED], json!(true));
458        assert!(value[MARKER_OMITTED].as_i64().unwrap_or(0) > 0);
459    }
460
461    #[test]
462    fn fixed_fields_return_an_explicit_error_instead_of_an_oversized_success() {
463        let payload = json!({"query": "x".repeat(20_000)});
464        let text = serialize(&payload, MIN_BUDGET);
465        assert!(char_len(&text) <= MIN_BUDGET);
466        let value: Value = serde_json::from_str(&text).expect("budget error is JSON");
467        assert_eq!(value["error"], json!(BUDGET_ERROR));
468    }
469
470    #[test]
471    fn configured_and_per_call_minimums_are_explicit() {
472        assert!(valid_configured_budget(MIN_BUDGET));
473        assert!(!valid_configured_budget(MIN_BUDGET - 1));
474        assert!(validate_call_budget(0).is_ok());
475        assert!(validate_call_budget(MIN_BUDGET).is_ok());
476        assert!(validate_call_budget(MIN_BUDGET - 1).is_err());
477    }
478}
rac_engine/budget.rs

rac_engine/
budget.rs
