1use crate::id::{IdentityError, ToolName, MAX_IDENTITY_BYTES};
6use crate::limits::InputLimits;
7use serde::{Deserialize, Serialize};
8use thiserror::Error;
9
10#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
17pub struct CanonicalInput {
18 messages: Vec<CanonicalMessage>,
19}
20
21impl CanonicalInput {
22 pub fn try_new(
24 messages: Vec<CanonicalMessage>,
25 limits: &InputLimits,
26 ) -> Result<Self, InputValidationError> {
27 if messages.is_empty() {
28 return Err(InputValidationError::EmptyMessages);
29 }
30 if messages.len() > limits.max_messages {
31 return Err(InputValidationError::TooManyMessages {
32 count: messages.len(),
33 max: limits.max_messages,
34 });
35 }
36
37 let mut aggregate_text = 0usize;
38 let mut seen_tool_call_ids: Vec<String> = Vec::new();
39
40 for (index, msg) in messages.iter().enumerate() {
41 msg.validate(limits, index, &mut aggregate_text, &mut seen_tool_call_ids)?;
42 }
43
44 if aggregate_text > limits.max_aggregate_text_bytes {
45 return Err(InputValidationError::AggregateTextTooLarge {
46 bytes: aggregate_text,
47 max: limits.max_aggregate_text_bytes,
48 });
49 }
50
51 Ok(Self { messages })
52 }
53
54 pub fn messages(&self) -> &[CanonicalMessage] {
56 &self.messages
57 }
58
59 pub fn into_messages(self) -> Vec<CanonicalMessage> {
61 self.messages
62 }
63}
64
65#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
67pub enum CanonicalMessage {
68 System {
70 content: Vec<TextPart>,
72 name: Option<String>,
74 },
75 User {
77 content: Vec<TextPart>,
79 name: Option<String>,
81 },
82 Assistant {
84 content: Vec<TextPart>,
86 tool_calls: Vec<CanonicalAssistantToolCall>,
88 },
89 Tool {
91 tool_call_id: String,
93 content: Vec<TextPart>,
95 },
96}
97
98#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
100pub struct TextPart {
101 text: String,
102}
103
104impl TextPart {
105 pub fn try_new(
107 text: impl Into<String>,
108 max_bytes: usize,
109 ) -> Result<Self, InputValidationError> {
110 let text = text.into();
111 if text.is_empty() {
112 return Err(InputValidationError::EmptyTextPart);
113 }
114 if text.len() > max_bytes {
115 return Err(InputValidationError::TextPartTooLarge {
116 bytes: text.len(),
117 max: max_bytes,
118 });
119 }
120 Ok(Self { text })
121 }
122
123 pub fn text(&self) -> &str {
125 &self.text
126 }
127}
128
129#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
131pub struct CanonicalAssistantToolCall {
132 pub tool_call_id: String,
134 pub tool_name: ToolName,
136 pub arguments: serde_json::Value,
138}
139
140impl CanonicalMessage {
141 fn validate(
142 &self,
143 limits: &InputLimits,
144 index: usize,
145 aggregate_text: &mut usize,
146 seen_tool_call_ids: &mut Vec<String>,
147 ) -> Result<(), InputValidationError> {
148 match self {
149 Self::System { content, name } | Self::User { content, name } => {
150 validate_name(name, limits)?;
151 validate_text_parts(content, limits, true, aggregate_text)?;
152 }
153 Self::Assistant {
154 content,
155 tool_calls,
156 } => {
157 if content.is_empty() && tool_calls.is_empty() {
158 return Err(InputValidationError::EmptyAssistant { index });
159 }
160 validate_text_parts(content, limits, false, aggregate_text)?;
161 if tool_calls.len() > limits.max_tool_calls {
162 return Err(InputValidationError::TooManyToolCalls {
163 count: tool_calls.len(),
164 max: limits.max_tool_calls,
165 });
166 }
167 for call in tool_calls {
168 validate_tool_call_id(&call.tool_call_id, limits)?;
169 if seen_tool_call_ids.iter().any(|id| id == &call.tool_call_id) {
170 return Err(InputValidationError::DuplicateToolCallId {
171 id: call.tool_call_id.clone(),
172 });
173 }
174 seen_tool_call_ids.push(call.tool_call_id.clone());
175 validate_json_value(
176 &call.arguments,
177 limits.max_json_depth,
178 limits.max_tool_argument_bytes,
179 )?;
180 let _ = call.tool_name.as_str(); }
182 }
183 Self::Tool {
184 tool_call_id,
185 content,
186 } => {
187 validate_tool_call_id(tool_call_id, limits)?;
188 if !seen_tool_call_ids.iter().any(|id| id == tool_call_id) {
189 return Err(InputValidationError::UnknownToolCallId {
190 id: tool_call_id.clone(),
191 });
192 }
193 validate_text_parts(content, limits, true, aggregate_text)?;
194 }
195 }
196 Ok(())
197 }
198}
199
200fn validate_name(name: &Option<String>, limits: &InputLimits) -> Result<(), InputValidationError> {
201 if let Some(n) = name {
202 if n.is_empty() {
203 return Err(InputValidationError::EmptyName);
204 }
205 if n.len() > limits.max_name_bytes {
206 return Err(InputValidationError::NameTooLong {
207 bytes: n.len(),
208 max: limits.max_name_bytes,
209 });
210 }
211 if n.chars().any(|c| c.is_control()) {
212 return Err(InputValidationError::ControlCharacter);
213 }
214 }
215 Ok(())
216}
217
218fn validate_tool_call_id(id: &str, limits: &InputLimits) -> Result<(), InputValidationError> {
219 if id.is_empty() {
220 return Err(InputValidationError::EmptyToolCallId);
221 }
222 if id.len() > limits.max_tool_call_id_bytes {
223 return Err(InputValidationError::ToolCallIdTooLong {
224 bytes: id.len(),
225 max: limits.max_tool_call_id_bytes,
226 });
227 }
228 if id.chars().any(|c| c.is_control()) {
229 return Err(InputValidationError::ControlCharacter);
230 }
231 Ok(())
232}
233
234fn validate_text_parts(
235 parts: &[TextPart],
236 limits: &InputLimits,
237 require_non_empty: bool,
238 aggregate_text: &mut usize,
239) -> Result<(), InputValidationError> {
240 if require_non_empty && parts.is_empty() {
241 return Err(InputValidationError::EmptyTextParts);
242 }
243 if parts.len() > limits.max_content_parts {
244 return Err(InputValidationError::TooManyContentParts {
245 count: parts.len(),
246 max: limits.max_content_parts,
247 });
248 }
249 for p in parts {
250 if p.text.is_empty() {
251 return Err(InputValidationError::EmptyTextPart);
252 }
253 if p.text.len() > limits.max_text_part_bytes {
254 return Err(InputValidationError::TextPartTooLarge {
255 bytes: p.text.len(),
256 max: limits.max_text_part_bytes,
257 });
258 }
259 *aggregate_text = aggregate_text.saturating_add(p.text.len());
260 }
261 Ok(())
262}
263
264fn validate_json_value(
265 value: &serde_json::Value,
266 max_depth: u32,
267 max_bytes: usize,
268) -> Result<(), InputValidationError> {
269 let depth = json_depth(value);
270 if depth > max_depth {
271 return Err(InputValidationError::JsonTooDeep {
272 depth,
273 max: max_depth,
274 });
275 }
276 let encoded = serde_json::to_vec(value).map_err(|_| InputValidationError::JsonEncodeFailed)?;
277 if encoded.len() > max_bytes {
278 return Err(InputValidationError::ToolArgumentsTooLarge {
279 bytes: encoded.len(),
280 max: max_bytes,
281 });
282 }
283 Ok(())
284}
285
286fn json_depth(value: &serde_json::Value) -> u32 {
287 match value {
288 serde_json::Value::Array(items) => 1 + items.iter().map(json_depth).max().unwrap_or(0),
289 serde_json::Value::Object(map) => 1 + map.values().map(json_depth).max().unwrap_or(0),
290 _ => 1,
291 }
292}
293
294pub fn user_text_input(text: impl Into<String>) -> Result<CanonicalInput, InputValidationError> {
296 let limits = InputLimits::default();
297 let part = TextPart::try_new(text, limits.max_text_part_bytes)?;
298 CanonicalInput::try_new(
299 vec![CanonicalMessage::User {
300 content: vec![part],
301 name: None,
302 }],
303 &limits,
304 )
305}
306
307#[derive(Clone, Debug, Error, PartialEq, Eq)]
309pub enum InputValidationError {
310 #[error("canonical input requires at least one message")]
312 EmptyMessages,
313 #[error("message count {count} exceeds max {max}")]
315 TooManyMessages {
316 count: usize,
318 max: usize,
320 },
321 #[error("aggregate text bytes {bytes} exceeds max {max}")]
323 AggregateTextTooLarge {
324 bytes: usize,
326 max: usize,
328 },
329 #[error("message requires at least one text part")]
331 EmptyTextParts,
332 #[error("text part must be non-empty")]
334 EmptyTextPart,
335 #[error("text part bytes {bytes} exceeds max {max}")]
337 TextPartTooLarge {
338 bytes: usize,
340 max: usize,
342 },
343 #[error("content part count {count} exceeds max {max}")]
345 TooManyContentParts {
346 count: usize,
348 max: usize,
350 },
351 #[error("assistant message at index {index} is empty")]
353 EmptyAssistant {
354 index: usize,
356 },
357 #[error("tool call count {count} exceeds max {max}")]
359 TooManyToolCalls {
360 count: usize,
362 max: usize,
364 },
365 #[error("duplicate tool_call_id {id}")]
367 DuplicateToolCallId {
368 id: String,
370 },
371 #[error("tool message references unknown tool_call_id {id}")]
373 UnknownToolCallId {
374 id: String,
376 },
377 #[error("tool_call_id must be non-empty")]
379 EmptyToolCallId,
380 #[error("tool_call_id bytes {bytes} exceeds max {max}")]
382 ToolCallIdTooLong {
383 bytes: usize,
385 max: usize,
387 },
388 #[error("message name must be non-empty when present")]
390 EmptyName,
391 #[error("name bytes {bytes} exceeds max {max}")]
393 NameTooLong {
394 bytes: usize,
396 max: usize,
398 },
399 #[error("input string must not contain control characters")]
401 ControlCharacter,
402 #[error("JSON depth {depth} exceeds max {max}")]
404 JsonTooDeep {
405 depth: u32,
407 max: u32,
409 },
410 #[error("tool argument bytes {bytes} exceeds max {max}")]
412 ToolArgumentsTooLarge {
413 bytes: usize,
415 max: usize,
417 },
418 #[error("JSON encode failed")]
420 JsonEncodeFailed,
421 #[error(transparent)]
423 Identity(#[from] IdentityError),
424}
425
426const _: usize = MAX_IDENTITY_BYTES;
428
429#[cfg(test)]
430mod tests {
431 use super::*;
432 use crate::id::ToolName;
433
434 #[test]
435 fn requires_messages_and_text() {
436 let limits = InputLimits::default();
437 assert!(CanonicalInput::try_new(vec![], &limits).is_err());
438 let empty_user = CanonicalMessage::User {
439 content: vec![],
440 name: None,
441 };
442 assert!(CanonicalInput::try_new(vec![empty_user], &limits).is_err());
443 }
444
445 #[test]
446 fn tool_must_reference_prior_assistant_call() {
447 let limits = InputLimits::default();
448 let part = TextPart::try_new("ok", limits.max_text_part_bytes).unwrap();
449 let bad = CanonicalMessage::Tool {
450 tool_call_id: "missing".into(),
451 content: vec![part],
452 };
453 assert!(matches!(
454 CanonicalInput::try_new(vec![bad], &limits),
455 Err(InputValidationError::UnknownToolCallId { .. })
456 ));
457 }
458
459 #[test]
460 fn historical_tool_round_trip_ok() {
461 let limits = InputLimits::default();
462 let call = CanonicalAssistantToolCall {
463 tool_call_id: "c1".into(),
464 tool_name: ToolName::try_new("search").unwrap(),
465 arguments: serde_json::json!({"q": "x"}),
466 };
467 let messages = vec![
468 CanonicalMessage::User {
469 content: vec![TextPart::try_new("hi", limits.max_text_part_bytes).unwrap()],
470 name: None,
471 },
472 CanonicalMessage::Assistant {
473 content: vec![],
474 tool_calls: vec![call],
475 },
476 CanonicalMessage::Tool {
477 tool_call_id: "c1".into(),
478 content: vec![TextPart::try_new("result", limits.max_text_part_bytes).unwrap()],
479 },
480 ];
481 let input = CanonicalInput::try_new(messages, &limits).unwrap();
482 let json = serde_json::to_string(&input).unwrap();
483 let back: CanonicalInput = serde_json::from_str(&json).unwrap();
484 assert_eq!(input, back);
485 }
486}