rpi_agent/
tool_program_results.rs

1//! Invocation-owned finalized results. Request disposal is deliberately independent.
2use super::*;
3use tokio::sync::OwnedSemaphorePermit;
4
5pub(super) const TREE_BYTES: usize = 64 * 1024 * 1024;
6#[derive(Default)]
7pub(super) struct Store(HashMap<String, Retained>);
8struct Retained {
9    native: crate::AgentToolResult,
10    is_error: bool,
11    encoded: Box<str>,
12    _bytes: OwnedSemaphorePermit,
13}
14impl Store {
15    pub(super) fn clear(&mut self) {
16        self.0.clear();
17    }
18}
19// JSON's encoded punctuation can be much smaller than its native Value/map
20// nodes. Include a conservative per-node allocation estimate, not just strings.
21fn value_overhead(value: &Value) -> Option<usize> {
22    let add = |sum: usize, child: &Value| sum.checked_add(value_overhead(child)?);
23    match value {
24        Value::Array(values) => values.iter().try_fold(128, add),
25        Value::Object(values) => values.values().try_fold(128, add),
26        _ => Some(128),
27    }
28}
29impl ToolProgram {
30    pub(super) fn retain_result(
31        &self,
32        id: &str,
33        native: crate::AgentToolResult,
34        is_error: bool,
35        encoded: String,
36    ) -> Result<Value, String> {
37        // Charge encoded bytes plus a native payload estimate and structural overhead.
38        // The shared permit follows ownership and is returned on release/finish/drop.
39        let charge = encoded
40            .len()
41            .checked_mul(2)
42            .and_then(|n| n.checked_add(4096))
43            .and_then(|n| n.checked_add(value_overhead(&native.details)?))
44            .and_then(|n| n.checked_add(native.content.len().checked_mul(256)?))
45            .filter(|n| *n <= TREE_BYTES)
46            .ok_or("ToolProgram result reference byte quota exceeded")?;
47        let mut state = self.state.lock().unwrap();
48        if state.results.0.len() >= 16 {
49            return Err("ToolProgram retained result reference limit exceeded".into());
50        }
51        let bytes = self
52            .result_bytes
53            .clone()
54            .try_acquire_many_owned(charge as u32)
55            .map_err(|_| "ToolProgram invocation-tree result reference byte quota exceeded")?;
56        let manifest = json!({"requestId":id,"status":"done","isError":is_error,
57            "resultReference":{"version":1,"id":id,"byteLength":encoded.len(),
58                "contentBlocks":native.content.len(),"isError":is_error,"terminate":native.terminate}});
59        state.results.0.insert(
60            id.to_owned(),
61            Retained {
62                native,
63                is_error,
64                encoded: encoded.into_boxed_str(),
65                _bytes: bytes,
66            },
67        );
68        Ok(manifest)
69    }
70    pub(super) fn read_result(&self, args: &Value) -> Result<Value, String> {
71        let id = args["resultId"].as_str().ok_or("resultId is required")?;
72        let offset = integer(args, "offset", 0, 0, u64::MAX)?;
73        let limit = integer(args, "limit", 65536, 4, 65536)? as usize;
74        let state = self.state.lock().unwrap();
75        if self.signal.is_cancelled() {
76            return Err("parent invocation is no longer active".into());
77        }
78        let retained = state
79            .results
80            .0
81            .get(id)
82            .ok_or("unknown result reference for this invocation")?;
83        let text = &retained.encoded;
84        let offset = usize::try_from(offset).map_err(|_| "offset is past EOF")?;
85        if offset > text.len() || !text.is_char_boundary(offset) {
86            return Err("offset must be a UTF-8 boundary at or before EOF".into());
87        }
88        let mut end = offset.saturating_add(limit).min(text.len());
89        while !text.is_char_boundary(end) {
90            end -= 1;
91        }
92        let reply = json!({"resultId":id,"offset":offset,"data":&text[offset..end],
93            "byteLength":text.len(),"nextOffset":if end < text.len() {Some(end)} else {None}});
94        // Even worst-case JSON escaping of 64 KiB stays below the frame limit.
95        bounded(&reply)?;
96        Ok(reply)
97    }
98    pub(super) fn release_result(&self, args: &Value) -> Result<Value, String> {
99        let id = args["resultId"].as_str().ok_or("resultId is required")?;
100        self.state
101            .lock()
102            .unwrap()
103            .results
104            .0
105            .remove(id)
106            .ok_or("unknown result reference for this invocation")?;
107        Ok(json!({"resultId":id,"status":"released"}))
108    }
109    /// Resolve native authority only within this owning program's active lifetime.
110    /// Usage/names remain in the receipt for inspection but are merged by host
111    /// accounting, never by the attachment. No transcript calls are synthesized.
112    pub fn resolve_result_reference(
113        &self,
114        id: &str,
115    ) -> Result<(crate::AgentToolResult, bool), String> {
116        let state = self.state.lock().unwrap();
117        if self.signal.is_cancelled() {
118            return Err("parent invocation is no longer active".into());
119        }
120        let retained = state
121            .results
122            .0
123            .get(id)
124            .ok_or("unknown result reference for this invocation")?;
125        let mut result = retained.native.clone();
126        result.usage = None;
127        result.added_tool_names.clear();
128        Ok((result, retained.is_error))
129    }
130}
131
132const CONTENT_BLOCKS: usize = 256;
133const CONTENT_BYTES: usize = 32 * 1024 * 1024;
134
135/// Fallible serialized-byte sink: no payload JSON/string allocation.
136#[derive(Default)]
137struct ContentBudget {
138    blocks: usize,
139    encoded_bytes: usize,
140}
141impl ContentBudget {
142    fn block(&mut self, block: &TextContentOrImage) -> Result<(), String> {
143        if self.blocks == CONTENT_BLOCKS {
144            return Err("content exceeds 256 blocks".into());
145        }
146        // n blocks require n-1 commas and two brackets.
147        let mut counter = ByteCounter {
148            bytes: 0,
149            remaining: CONTENT_BYTES
150                .checked_sub(self.encoded_bytes + self.blocks + 2)
151                .ok_or("content array exceeds 32 MiB")?,
152        };
153        match block {
154            TextContentOrImage::Text(text) => serde_json::to_writer(&mut counter, text),
155            TextContentOrImage::Image(image) => serde_json::to_writer(&mut counter, image),
156        }
157        .map_err(|e| e.to_string())?;
158        self.encoded_bytes += counter.bytes;
159        self.blocks += 1;
160        Ok(())
161    }
162}
163struct ByteCounter {
164    bytes: usize,
165    remaining: usize,
166}
167impl std::io::Write for ByteCounter {
168    fn write(&mut self, bytes: &[u8]) -> std::io::Result<usize> {
169        if bytes.len() > self.remaining {
170            return Err(std::io::Error::other("content array exceeds 32 MiB"));
171        }
172        self.bytes += bytes.len();
173        self.remaining -= bytes.len();
174        Ok(bytes.len())
175    }
176    fn flush(&mut self) -> std::io::Result<()> {
177        Ok(())
178    }
179}
180fn visit_selected(
181    blocks: &[TextContentOrImage],
182    indices: Option<&[usize]>,
183    mut visit: impl FnMut(&TextContentOrImage) -> Result<(), String>,
184) -> Result<(), String> {
185    let count = indices.map_or(blocks.len(), <[usize]>::len);
186    if count > CONTENT_BLOCKS {
187        return Err("content selection exceeds 256 blocks".into());
188    }
189    for position in 0..count {
190        let index = indices.map_or(position, |indices| indices[position]);
191        visit(
192            blocks
193                .get(index)
194                .ok_or("content index is past native content")?,
195        )?;
196    }
197    Ok(())
198}
199impl ToolProgram {
200    fn check_content_live(&self) -> Result<(), String> {
201        if self.signal.is_cancelled() {
202            Err("parent invocation is no longer active".into())
203        } else {
204            Ok(())
205        }
206    }
207    pub(super) fn measure_result_content(&self, args: &Value) -> Result<Value, String> {
208        self.check_content_live()?;
209        let id = args["resultId"].as_str().ok_or("resultId is required")?;
210        let indices = match args.get("indices") {
211            None | Some(Value::Null) => None,
212            Some(Value::Array(values)) if values.len() <= CONTENT_BLOCKS => Some(
213                values
214                    .iter()
215                    .map(|v| {
216                        v.as_u64()
217                            .and_then(|n| usize::try_from(n).ok())
218                            .ok_or_else(|| "indices must contain nonnegative integers".to_owned())
219                    })
220                    .collect::<Result<Vec<_>, _>>()?,
221            ),
222            _ => return Err("indices must be null or an array of at most 256 integers".into()),
223        };
224        let state = self.state.lock().unwrap();
225        let retained = state
226            .results
227            .0
228            .get(id)
229            .ok_or("unknown result reference for this invocation")?;
230        let mut budget = ContentBudget::default();
231        visit_selected(&retained.native.content, indices.as_deref(), |block| {
232            self.check_content_live()?;
233            budget.block(block)
234        })?;
235        self.check_content_live()?;
236        let reply = json!({"resultId":id,"contentBlocks":budget.blocks,"encodedBytes":budget.encoded_bytes});
237        bounded(&reply)?;
238        Ok(reply)
239    }
240    /// Resolve from this invocation's authentic native store. Count each block
241    /// before cloning it, including repeated selections, to bound expansion.
242    /// The serialized array limit is not an RSS ceiling or a retained-tree charge.
243    pub fn resolve_content_parts(
244        &self,
245        parts: Vec<ToolProgramContentPart>,
246    ) -> Result<Vec<TextContentOrImage>, String> {
247        self.check_content_live()?;
248        if parts.len() > CONTENT_BLOCKS {
249            return Err("content exceeds 256 parts".into());
250        }
251        let state = self.state.lock().unwrap();
252        let mut budget = ContentBudget::default();
253        let mut content = Vec::new();
254        for part in parts {
255            self.check_content_live()?;
256            match part {
257                ToolProgramContentPart::Inline(blocks) => {
258                    for block in blocks {
259                        self.check_content_live()?;
260                        budget.block(&block)?;
261                        content.push(block);
262                    }
263                }
264                ToolProgramContentPart::Reference { id, indices } => {
265                    let retained = state
266                        .results
267                        .0
268                        .get(&id)
269                        .ok_or("unknown result reference for this invocation")?;
270                    visit_selected(&retained.native.content, indices.as_deref(), |block| {
271                        self.check_content_live()?;
272                        budget.block(block)?;
273                        content.push(block.clone());
274                        Ok(())
275                    })?;
276                }
277            }
278        }
279        self.check_content_live()?;
280        Ok(content)
281    }
282}
rpi_agent/tool_program_results.rs

rpi_agent/
tool_program_results.rs
