kimetsu_brain/
delivery.rs

1//! Final serving boundary. Retrieval/reranking decides relevance; this module
2//! admits only whole chosen capsules that fit the serialized delivery budget.
3use super::{ContextCapsule, memory_revision_bindings};
4use serde_json::{Value, json};
5
6pub struct Delivery {
7    pub payload: Value,
8    pub capsules: Vec<ContextCapsule>,
9}
10
11/// Conservative tokenizer-independent bound: one token per UTF-8 byte, including
12/// the MCP content envelope and both layers of JSON escaping. This is an upper
13/// bound for byte-based tokenizers, not a measured model tokenizer count. JSON-RPC
14/// request IDs/framing are transport-only and are not included.
15pub fn serialized_output_tokens(payload: &Value) -> u32 {
16    let bytes = json!({"content": [{"type": "text", "text": payload.to_string()}]})
17        .to_string()
18        .len();
19    u32::try_from(bytes).unwrap_or(u32::MAX)
20}
21
22fn account(payload: &mut Value) -> u32 {
23    payload["used_tokens"] = json!(0);
24    loop {
25        let bound = serialized_output_tokens(payload);
26        if payload["used_tokens"].as_u64() == Some(u64::from(bound)) {
27            return bound;
28        }
29        payload["used_tokens"] = json!(bound);
30    }
31}
32
33pub fn compact_capsules(capsules: &[ContextCapsule]) -> Vec<Value> {
34    capsules
35        .iter()
36        .map(|c| {
37            json!({
38                "id": c.id, "kind": c.kind, "summary": c.summary,
39                "expansion_handle": c.expansion_handle, "score": c.score,
40            })
41        })
42        .collect()
43}
44
45/// `render` must rebuild all text and counts from this slice, including duplicated
46/// playbook text. Rejected candidates never enter the render callback. If even an
47/// empty envelope cannot fit, return an explicit error with its true bound (which
48/// can exceed the requested tiny budget); never report success or delivered IDs.
49pub fn fit_json(
50    mut capsules: Vec<ContextCapsule>,
51    budget: u32,
52    render: impl Fn(&[ContextCapsule]) -> Value,
53) -> Delivery {
54    loop {
55        let mut payload = render(&capsules);
56        payload["budget_tokens"] = json!(budget);
57        payload["token_accounting"] = json!("utf8_byte_upper_bound");
58        if account(&mut payload) <= budget {
59            return Delivery { payload, capsules };
60        }
61        if capsules.pop().is_none() {
62            let mut payload = json!({"ok": false, "error": "budget_too_small",
63                "budget_tokens": budget, "capsules": [], "capsule_count": 0,
64                "token_accounting": "utf8_byte_upper_bound"});
65            account(&mut payload);
66            return Delivery { payload, capsules };
67        }
68    }
69}
70
71/// Add optional framing only if the complete payload still fits. Evidence wins
72/// over warm-start hints; callers must invoke this before logging the exposure.
73pub fn add_optional_field(delivery: &mut Delivery, key: &str, value: Value, budget: u32) {
74    let mut payload = delivery.payload.clone();
75    payload[key] = value;
76    if account(&mut payload) <= budget {
77        delivery.payload = payload;
78    }
79}
80
81/// Exposure is built from the final delivered slice, never by re-reading current
82/// claims. An empty revision map explicitly means unbound, not legacy attribution.
83pub fn injected_payload(capsules: &[ContextCapsule], used_tokens: u32) -> Value {
84    json!({
85        "memory_ids": capsules.iter().filter_map(|c| c.expansion_handle.strip_prefix("memory:")).collect::<Vec<_>>(),
86        "memory_revisions": memory_revision_bindings(capsules),
87        "capsule_handles": capsules.iter().map(|c| c.expansion_handle.as_str()).collect::<Vec<_>>(),
88        "capsule_count": capsules.len(), "used_tokens": used_tokens,
89    })
90}
91
92#[cfg(test)]
93mod tests {
94    use super::*;
95
96    fn capsule(id: &str, text: &str) -> ContextCapsule {
97        let mut c = ContextCapsule::wire_minimal(text.into(), "memory".into(), 0.9);
98        c.id = id.into();
99        c.expansion_handle = format!("memory:{id}");
100        c.claim_revision = Some(format!("revision-{id}"));
101        c
102    }
103    fn render(c: &[ContextCapsule]) -> Value {
104        json!({"ok":true,"capsules":compact_capsules(c),"capsule_count":c.len()})
105    }
106
107    #[test]
108    fn final_serialization_bounds_unicode_identifiers_and_escaping() {
109        for text in [
110            "字".repeat(1000),
111            "no_space_identifier".repeat(1000),
112            "\"\\\n".repeat(1000),
113        ] {
114            let delivery = fit_json(
115                vec![capsule("kept", "short"), capsule("dropped", &text)],
116                800,
117                render,
118            );
119            assert_eq!(delivery.capsules.len(), 1);
120            assert_eq!(
121                delivery.payload["used_tokens"].as_u64(),
122                Some(u64::from(serialized_output_tokens(&delivery.payload)))
123            );
124            assert!(serialized_output_tokens(&delivery.payload) <= 800);
125            let event = injected_payload(&delivery.capsules, 0);
126            assert_eq!(event["memory_ids"], json!(["kept"]));
127            assert_eq!(event["memory_revisions"], json!({"kept":"revision-kept"}));
128        }
129    }
130
131    #[test]
132    fn tiny_budget_reports_actual_error_cost_and_no_exposure() {
133        let delivery = fit_json(vec![capsule("secret", "secret")], 1, render);
134        assert_eq!(delivery.payload["error"], "budget_too_small");
135        assert!(delivery.payload["used_tokens"].as_u64().unwrap() > 1);
136        assert!(delivery.capsules.is_empty());
137        assert_eq!(injected_payload(&[], 0)["memory_revisions"], json!({}));
138    }
139
140    #[test]
141    fn repeated_playbook_text_is_included_in_final_bound() {
142        let delivery = fit_json(
143            vec![capsule("large", &"x".repeat(500))],
144            900,
145            |c| json!({"capsules":compact_capsules(c),"playbook_markdown":c.iter().map(|c|c.summary.as_str()).collect::<Vec<_>>().join("\n")}),
146        );
147        assert!(delivery.capsules.is_empty());
148        assert_eq!(delivery.payload["playbook_markdown"], "");
149        assert!(serialized_output_tokens(&delivery.payload) <= 900);
150    }
151}
kimetsu_brain/delivery.rs

kimetsu_brain/
delivery.rs
