cookie_agent_protocol/
tool_output.rs

1use std::collections::HashSet;
2
3use schemars::JsonSchema;
4use serde::{Deserialize, Serialize};
5
6use crate::{ArtifactReference, Sha256Digest};
7
8pub const MAX_TOOL_STREAMS: usize = 8;
9pub const MAX_TOOL_STREAM_NAME_BYTES: usize = 64;
10pub const MAX_TOOL_DELTA_BYTES: usize = 64 * 1024;
11pub const MAX_TOOL_DISPLAY_BYTES: usize = 64 * 1024;
12
13pub(crate) fn validate_display(text: &str, maximum: usize) -> bool {
14    text.len() <= maximum
15        && text
16            .chars()
17            .all(|character| !character.is_control() || matches!(character, '\n' | '\t'))
18}
19
20pub(crate) fn deserialize_delta_display<'de, D: serde::Deserializer<'de>>(
21    deserializer: D,
22) -> Result<Option<String>, D::Error> {
23    let display = Option::<String>::deserialize(deserializer)?;
24    if display
25        .as_ref()
26        .is_some_and(|text| !validate_display(text, crate::SafeDisplayText::MAX_BYTES))
27    {
28        return Err(serde::de::Error::custom(
29            "invalid or oversized display delta",
30        ));
31    }
32    Ok(display)
33}
34
35/// The same name grammar is used for declarations, output chunks, and artifact URI suffixes.
36pub fn validate_tool_stream_name(name: &str) -> Result<(), String> {
37    if name.is_empty()
38        || name.len() > MAX_TOOL_STREAM_NAME_BYTES
39        || matches!(name, "." | "..")
40        || !name
41            .bytes()
42            .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'_' | b'-' | b'.'))
43    {
44        return Err(
45            "stream names must be 1..64 URI-safe ASCII characters, excluding . and ..".into(),
46        );
47    }
48    Ok(())
49}
50
51#[derive(Clone, Debug, Default, Deserialize, Eq, JsonSchema, PartialEq, Serialize)]
52#[serde(tag = "kind", rename_all = "snake_case", deny_unknown_fields)]
53pub enum ToolOutputDeclaration {
54    #[default]
55    Single,
56    Named {
57        #[serde(deserialize_with = "deserialize_named_streams")]
58        #[schemars(length(min = 1, max = 8))]
59        streams: Vec<String>,
60    },
61}
62
63fn deserialize_named_streams<'de, D: serde::Deserializer<'de>>(
64    deserializer: D,
65) -> Result<Vec<String>, D::Error> {
66    let declaration = ToolOutputDeclaration::Named {
67        streams: Vec::<String>::deserialize(deserializer)?,
68    };
69    declaration.validate().map_err(serde::de::Error::custom)?;
70    let ToolOutputDeclaration::Named { streams } = declaration else {
71        unreachable!()
72    };
73    Ok(streams)
74}
75
76pub(crate) fn deserialize_stream_name<'de, D: serde::Deserializer<'de>>(
77    deserializer: D,
78) -> Result<String, D::Error> {
79    let name = String::deserialize(deserializer)?;
80    validate_tool_stream_name(&name).map_err(serde::de::Error::custom)?;
81    Ok(name)
82}
83
84impl ToolOutputDeclaration {
85    pub fn validate(&self) -> Result<(), String> {
86        if let Self::Named { streams } = self {
87            if streams.is_empty() || streams.len() > MAX_TOOL_STREAMS {
88                return Err("named output must declare between 1 and 8 streams".into());
89            }
90            let mut names = HashSet::new();
91            for name in streams {
92                validate_tool_stream_name(name)?;
93                if !names.insert(name) {
94                    return Err("output stream names must be unique".into());
95                }
96            }
97        }
98        Ok(())
99    }
100
101    pub fn channels(&self) -> Vec<Option<String>> {
102        match self {
103            Self::Single => vec![None],
104            Self::Named { streams } => streams.iter().cloned().map(Some).collect(),
105        }
106    }
107}
108
109#[derive(Clone, Debug, Deserialize, Eq, JsonSchema, PartialEq, Serialize)]
110#[serde(deny_unknown_fields)]
111pub struct ToolOutputChunk {
112    pub stream: Option<String>,
113    pub text: String,
114}
115
116#[derive(Clone, Debug, Deserialize, Eq, JsonSchema, PartialEq, Serialize)]
117#[serde(tag = "kind", rename_all = "snake_case", deny_unknown_fields)]
118pub enum ToolCompletionOutput {
119    Single { text: String },
120    Named { streams: Vec<ToolOutputChunk> },
121    Streamed,
122}
123
124#[derive(Clone, Debug, Deserialize, Eq, JsonSchema, PartialEq, Serialize)]
125#[serde(deny_unknown_fields)]
126pub struct RetainedToolStream {
127    pub name: Option<String>,
128    pub reference: ArtifactReference,
129    pub sha256: Sha256Digest,
130    pub byte_length: u64,
131    pub line_count: u64,
132    pub truncated: bool,
133    pub next_offset: Option<u64>,
134}
135
136#[derive(Clone, Debug, Deserialize, Eq, JsonSchema, PartialEq, Serialize)]
137#[serde(deny_unknown_fields)]
138pub struct RetainedToolOutput {
139    pub reference: ArtifactReference,
140    pub streams: Vec<RetainedToolStream>,
141    pub incomplete: bool,
142}
143
144#[derive(Clone, Debug, Deserialize, Eq, JsonSchema, PartialEq, Serialize)]
145#[serde(deny_unknown_fields)]
146pub struct ToolOutputManifest {
147    pub streams: Vec<RetainedToolStream>,
148}
149
150impl ToolOutputManifest {
151    pub fn validate(&self) -> Result<(), String> {
152        let names = self
153            .streams
154            .iter()
155            .map(|stream| {
156                stream
157                    .name
158                    .clone()
159                    .ok_or_else(|| "manifest stream must be named".to_owned())
160            })
161            .collect::<Result<Vec<_>, _>>()?;
162        ToolOutputDeclaration::Named { streams: names }.validate()?;
163        for stream in &self.streams {
164            if stream.reference.uri != format!("artifact://sha256/{}", stream.sha256) {
165                return Err("stream reference and digest do not match".into());
166            }
167        }
168        Ok(())
169    }
170}
171
172/// Public artifact routes are deliberately distinct from stored ArtifactReference URIs.
173#[derive(Clone, Debug, Eq, PartialEq)]
174pub struct ArtifactReadPath {
175    pub digest: Sha256Digest,
176    pub stream: Option<String>,
177}
178
179impl ArtifactReadPath {
180    pub fn parse(path: &str) -> Result<Self, String> {
181        let path = path
182            .strip_prefix("artifact://")
183            .ok_or_else(|| "expected artifact:// URI".to_owned())?;
184        let (digest, stream) = path
185            .split_once('/')
186            .map_or((path, None), |(digest, stream)| (digest, Some(stream)));
187        let digest = Sha256Digest::new(digest).map_err(|_| {
188            "artifact URI requires a 64-character lowercase SHA-256 digest".to_owned()
189        })?;
190        if let Some(stream) = stream {
191            validate_tool_stream_name(stream)?;
192        }
193        Ok(Self {
194            digest,
195            stream: stream.map(str::to_owned),
196        })
197    }
198}
199
200#[cfg(test)]
201mod tests {
202    use super::*;
203
204    #[test]
205    fn stream_declarations_share_bounded_uri_safe_names() {
206        for streams in [
207            vec![],
208            vec!["a".into(); 9],
209            vec!["same".into(), "same".into()],
210            vec!["../bad".into()],
211            vec![".".into()],
212            vec!["..".into()],
213            vec!["a/b".into()],
214            vec!["a\\b".into()],
215            vec!["bad?query".into()],
216            vec!["bad#fragment".into()],
217            vec!["bad\n".into()],
218            vec!["x".repeat(65)],
219        ] {
220            let declaration = ToolOutputDeclaration::Named { streams };
221            assert!(declaration.validate().is_err());
222            assert!(
223                serde_json::from_value::<ToolOutputDeclaration>(
224                    serde_json::to_value(declaration).unwrap()
225                )
226                .is_err()
227            );
228        }
229        let declaration = ToolOutputDeclaration::Named {
230            streams: vec!["results".into(), "diagnostics.v2".into(), "EMPTY_1".into()],
231        };
232        declaration.validate().unwrap();
233        assert_eq!(
234            declaration.channels(),
235            vec![
236                Some("results".into()),
237                Some("diagnostics.v2".into()),
238                Some("EMPTY_1".into())
239            ]
240        );
241    }
242
243    #[test]
244    fn public_artifact_uri_parser_is_exact() {
245        let digest = "a".repeat(64);
246        for path in [
247            format!("artifact://{digest}"),
248            format!("artifact://{digest}/diagnostics.v2"),
249        ] {
250            let parsed = ArtifactReadPath::parse(&path).unwrap();
251            assert_eq!(parsed.digest.as_str(), digest);
252        }
253        for path in [
254            digest.clone(),
255            format!("artifact://sha256/{digest}"),
256            format!("artifact://{}", "A".repeat(64)),
257            format!("artifact://{}", "a".repeat(63)),
258            format!("artifact://{digest}/"),
259            format!("artifact://{digest}/../file"),
260            format!("artifact://{digest}/.."),
261            format!("artifact://{digest}/out/more"),
262            format!("artifact://{digest}?query"),
263            format!("artifact://{digest}/out#fragment"),
264            format!("artifact://{digest}/%2e%2e"),
265            format!("artifact://{digest}/bad\\stream"),
266        ] {
267            assert!(ArtifactReadPath::parse(&path).is_err(), "{path}");
268        }
269    }
270}
cookie_agent_protocol/tool_output.rs

cookie_agent_protocol/
tool_output.rs
