1#![forbid(unsafe_code)]
4
5use std::collections::HashSet;
6
7use anyhow::Context as _;
8use kcode_kennedy_session_kweb_contracts::{
9 DecodedKwebTool, canonical_node_ids, decode as decode_kweb,
10};
11use kcode_session_history::chatend::BoxId;
12use serde_json::Value;
13
14#[derive(Clone, Debug, PartialEq)]
15pub enum DecodedTool {
16 RunSubagent {
17 model: String,
18 reasoning_effort: Option<String>,
19 context_node_ids: Vec<String>,
20 task: String,
21 },
22 EndSession {
23 message: Option<String>,
24 },
25 BoxIds(Vec<BoxId>),
26 SummarizeBox {
27 box_id: BoxId,
28 summary: String,
29 },
30 BoxId(BoxId),
31 LoadNodes(Vec<String>),
32 EmitObject {
33 object_id: String,
34 file_name: Option<String>,
35 },
36 WebSearch {
37 question: String,
38 model: String,
39 },
40 WebFetch(String),
41 StageTelegramGroupMedia(i64),
42 MediaEnrichment {
43 object_id: String,
44 model: String,
45 prompt: String,
46 },
47 GenerateImage {
48 model: String,
49 prompt: String,
50 reference_object_ids: Vec<String>,
51 },
52 ObjectId(String),
53 ConnectNodes(Vec<String>),
54 ConsolidateFanout {
55 parent: String,
56 fanout: Vec<String>,
57 aggregator: String,
58 },
59 SetFixedConnection {
60 parent: String,
61 child: Option<String>,
62 slot: usize,
63 },
64 CreateNode {
65 parents: Vec<String>,
66 owner: String,
67 short_name: String,
68 short_description: String,
69 long_description: String,
70 },
71 UpdateNode {
72 id: String,
73 owner: String,
74 short_name: String,
75 short_description: String,
76 long_description: String,
77 },
78}
79
80impl From<DecodedKwebTool> for DecodedTool {
81 fn from(value: DecodedKwebTool) -> Self {
82 match value {
83 DecodedKwebTool::ConnectNodes(identifiers) => Self::ConnectNodes(identifiers),
84 DecodedKwebTool::ConsolidateFanout {
85 parent,
86 fanout,
87 aggregator,
88 } => Self::ConsolidateFanout {
89 parent,
90 fanout,
91 aggregator,
92 },
93 DecodedKwebTool::SetFixedConnection {
94 parent,
95 child,
96 slot,
97 } => Self::SetFixedConnection {
98 parent,
99 child,
100 slot,
101 },
102 DecodedKwebTool::CreateNode {
103 parents,
104 owner,
105 short_name,
106 short_description,
107 long_description,
108 } => Self::CreateNode {
109 parents,
110 owner,
111 short_name,
112 short_description,
113 long_description,
114 },
115 DecodedKwebTool::UpdateNode {
116 id,
117 owner,
118 short_name,
119 short_description,
120 long_description,
121 } => Self::UpdateNode {
122 id,
123 owner,
124 short_name,
125 short_description,
126 long_description,
127 },
128 }
129 }
130}
131
132#[derive(Clone, Copy, Debug)]
133pub enum ValidationRequest<'a> {
134 Annotation {
135 model: &'a str,
136 media_type: &'a str,
137 },
138 ImageModel(&'a str),
139 TranscriptionModel(&'a str),
140 TranscribableAudio(&'a str),
141 ExtractableDocument {
142 media_type: &'a str,
143 file_name: &'a str,
144 },
145}
146
147#[derive(Clone, Copy, Debug)]
148pub enum ManagedObjectArguments<'a> {
149 RustBinary(&'a Value),
150 WebLibraryAttachment(&'a Value),
151}
152
153pub fn decode(tool: &str, value: &Value) -> anyhow::Result<Option<DecodedTool>> {
154 if let Some(decoded) = decode_kweb(tool, value)? {
155 return Ok(Some(decoded.into()));
156 }
157
158 let decoded = match tool {
159 "RunSubagent" => {
160 exact(
161 value,
162 &["model", "contextNodeIds", "task"],
163 &["reasoningEffort"],
164 )?;
165 DecodedTool::RunSubagent {
166 model: nonempty(value, "model", 128)?,
167 reasoning_effort: value
168 .get("reasoningEffort")
169 .map(|_| nonempty(value, "reasoningEffort", 32))
170 .transpose()?,
171 task: bounded_nonempty(value, "task", 100_000)?,
172 context_node_ids: canonical_node_ids(value, "contextNodeIds", Some(64), false)?,
173 }
174 }
175 "EndSession" => {
176 exact(value, &[], &["message"])?;
177 DecodedTool::EndSession {
178 message: value
179 .get("message")
180 .and_then(Value::as_str)
181 .map(str::to_owned),
182 }
183 }
184 "DehydrateBoxes" | "BoxesIntoObjects" => {
185 exact(value, &["boxIds"], &[])?;
186 DecodedTool::BoxIds(box_ids(value, "boxIds")?)
187 }
188 "SummarizeBox" => {
189 exact(value, &["boxId", "summary"], &[])?;
190 DecodedTool::SummarizeBox {
191 box_id: BoxId(positive_integer(value, "boxId")?),
192 summary: nonempty(value, "summary", 1_000_000)?,
193 }
194 }
195 "HydrateBox" => {
196 exact(value, &["boxId"], &[])?;
197 DecodedTool::BoxId(BoxId(positive_integer(value, "boxId")?))
198 }
199 "LoadNodes" => {
200 exact(value, &["identifiers"], &[])?;
201 DecodedTool::LoadNodes(canonical_node_ids(value, "identifiers", None, true)?)
202 }
203 "EmitObject" => {
204 exact(value, &["objectId"], &["fileName"])?;
205 DecodedTool::EmitObject {
206 object_id: nonempty(value, "objectId", 64)?,
207 file_name: delivery_file_name(value, "fileName")?,
208 }
209 }
210 "WebSearch" => {
211 exact(value, &["question", "model"], &[])?;
212 let model = nonempty(value, "model", 128)?;
213 let question = nonempty(value, "question", 4_000)?;
214 DecodedTool::WebSearch { question, model }
215 }
216 "WebFetch" => {
217 exact(value, &["url"], &[])?;
218 DecodedTool::WebFetch(nonempty(value, "url", 4_096)?)
219 }
220 "StageTelegramGroupMedia" => {
221 exact(value, &["messageId"], &[])?;
222 DecodedTool::StageTelegramGroupMedia(
223 i64::try_from(positive_integer(value, "messageId")?)
224 .context("messageId exceeds Telegram's supported integer range")?,
225 )
226 }
227 "TranscribeAudio" | "AnnotateMedia" => {
228 exact(value, &["objectId", "model", "prompt"], &[])?;
229 let model = nonempty(value, "model", 128)?;
230 let prompt = nonblank(value, "prompt")?;
231 let object_id = nonempty(value, "objectId", 64)?;
232 DecodedTool::MediaEnrichment {
233 object_id,
234 model,
235 prompt,
236 }
237 }
238 "GenerateImage" => {
239 exact(value, &["model", "prompt"], &["referenceObjectIds"])?;
240 let model = nonempty(value, "model", 128)?;
241 validate(ValidationRequest::ImageModel(&model))?;
242 DecodedTool::GenerateImage {
243 model,
244 prompt: bounded_nonempty(value, "prompt", 100_000)?,
245 reference_object_ids: optional_object_ids(value, "referenceObjectIds", 14)?,
246 }
247 }
248 "ExtractDocumentText" => {
249 exact(value, &["objectId"], &[])?;
250 DecodedTool::ObjectId(nonempty(value, "objectId", 64)?)
251 }
252 _ => return Ok(None),
253 };
254 Ok(Some(decoded))
255}
256
257pub fn validate(request: ValidationRequest<'_>) -> anyhow::Result<()> {
258 match request {
259 ValidationRequest::Annotation { model, media_type } => match model {
260 "gpt-5.6" | "gpt-5.6-sol" | "gpt-5.6-terra" | "gpt-5.6-luna" => anyhow::ensure!(
261 media_type.starts_with("image/"),
262 "{model} annotations accept images only"
263 ),
264 "gemini-2.5-flash" | "gemini-3.1-flash-lite" | "gemini-3.1-pro-preview" => {
265 anyhow::ensure!(
266 media_type.starts_with("image/")
267 || media_type.starts_with("audio/")
268 || media_type.starts_with("video/")
269 || media_type == "application/ogg",
270 "{model} annotations accept images, audio, or video only"
271 )
272 }
273 _ => anyhow::bail!("unsupported exact annotation model {model}"),
274 },
275 ValidationRequest::ImageModel(model) => anyhow::ensure!(
276 matches!(model, "gpt-image-2" | "gemini-3-pro-image"),
277 "unsupported exact image model {model}; use gpt-image-2 or gemini-3-pro-image"
278 ),
279 ValidationRequest::TranscriptionModel(model) => anyhow::ensure!(
280 matches!(
281 model,
282 "gpt-4o-transcribe"
283 | "gemini-2.5-flash"
284 | "gemini-3.1-flash-lite"
285 | "gemini-3.1-pro-preview"
286 ),
287 "unsupported exact transcription model {model}"
288 ),
289 ValidationRequest::TranscribableAudio(media_type) => anyhow::ensure!(
290 matches!(
291 media_type,
292 "audio/flac"
293 | "audio/x-flac"
294 | "audio/m4a"
295 | "audio/mp3"
296 | "audio/mp4"
297 | "audio/mpeg"
298 | "audio/mpga"
299 | "audio/ogg"
300 | "audio/opus"
301 | "audio/wav"
302 | "audio/x-wav"
303 | "audio/webm"
304 | "application/ogg"
305 ),
306 "TranscribeAudio accepts a supported FLAC, MP3, MP4, M4A, OGG, WAV, or WebM audio object only"
307 ),
308 ValidationRequest::ExtractableDocument {
309 media_type,
310 file_name,
311 } => {
312 let media_type = normalize_media_type(media_type);
313 let extension = file_name
314 .rsplit_once('.')
315 .map(|(_, extension)| extension.to_ascii_lowercase());
316 let supported_media_type = media_type.starts_with("text/")
317 || matches!(
318 media_type.as_str(),
319 "application/pdf"
320 | "application/msword"
321 | "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
322 | "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
323 | "application/vnd.ms-excel"
324 | "application/vnd.ms-excel.sheet.binary.macroenabled.12"
325 | "application/vnd.oasis.opendocument.spreadsheet"
326 | "application/json"
327 | "application/xml"
328 | "application/yaml"
329 | "application/x-yaml"
330 );
331 anyhow::ensure!(
332 supported_media_type
333 || matches!(
334 extension.as_deref(),
335 Some(
336 "pdf"
337 | "doc"
338 | "docx"
339 | "xlsx"
340 | "xls"
341 | "xlsb"
342 | "ods"
343 | "csv"
344 | "tsv"
345 | "txt"
346 | "md"
347 | "json"
348 | "yaml"
349 | "yml"
350 | "xml"
351 )
352 ),
353 "ExtractDocumentText accepts supported PDF, Word, spreadsheet, and text-family objects only"
354 );
355 }
356 }
357 Ok(())
358}
359
360pub fn decode_managed_objects(request: ManagedObjectArguments<'_>) -> anyhow::Result<Vec<String>> {
361 match request {
362 ManagedObjectArguments::RustBinary(arguments) => {
363 let Some(ids) = arguments.get("objectIds") else {
364 return Ok(Vec::new());
365 };
366 ids.as_array()
367 .context("Rust-binary objectIds must be an array")?
368 .iter()
369 .map(|id| {
370 id.as_str()
371 .map(str::to_owned)
372 .context("Rust-binary objectIds must contain only strings")
373 })
374 .collect()
375 }
376 ManagedObjectArguments::WebLibraryAttachment(arguments) => Ok(vec![
377 arguments
378 .get("objectId")
379 .and_then(Value::as_str)
380 .filter(|id| !id.trim().is_empty())
381 .map(str::to_owned)
382 .context("Web-library attachment objectId must be a nonempty string")?,
383 ]),
384 }
385}
386
387fn exact(value: &Value, required: &[&str], optional: &[&str]) -> anyhow::Result<()> {
388 let map = value
389 .as_object()
390 .context("arguments must be a JSON object")?;
391 let allowed = required
392 .iter()
393 .chain(optional)
394 .copied()
395 .collect::<HashSet<_>>();
396 anyhow::ensure!(
397 required.iter().all(|key| map.contains_key(*key))
398 && map.keys().all(|key| allowed.contains(key.as_str())),
399 "expected exactly: {}{}",
400 required.join(", "),
401 if optional.is_empty() {
402 String::new()
403 } else {
404 format!(" (optional: {})", optional.join(", "))
405 }
406 );
407 Ok(())
408}
409
410fn positive_integer(value: &Value, key: &str) -> anyhow::Result<u64> {
411 value
412 .get(key)
413 .and_then(Value::as_u64)
414 .filter(|value| *value > 0)
415 .with_context(|| format!("{key} must be a positive integer"))
416}
417
418fn box_ids(value: &Value, key: &str) -> anyhow::Result<Vec<BoxId>> {
419 let ids = value
420 .get(key)
421 .and_then(Value::as_array)
422 .with_context(|| format!("{key} must be an array"))?
423 .iter()
424 .map(|value| {
425 value
426 .as_u64()
427 .filter(|value| *value > 0)
428 .map(BoxId)
429 .with_context(|| format!("{key} must contain only positive integers"))
430 })
431 .collect::<anyhow::Result<Vec<_>>>()?;
432 anyhow::ensure!(!ids.is_empty(), "{key} must contain at least one box ID");
433 anyhow::ensure!(
434 ids.iter().copied().collect::<HashSet<_>>().len() == ids.len(),
435 "{key} must not contain duplicate box IDs"
436 );
437 Ok(ids)
438}
439
440fn string(value: &Value, key: &str) -> anyhow::Result<String> {
441 value
442 .get(key)
443 .and_then(Value::as_str)
444 .map(str::to_owned)
445 .with_context(|| format!("{key} must be a string"))
446}
447
448fn nonempty(value: &Value, key: &str, maximum: usize) -> anyhow::Result<String> {
449 let value = string(value, key)?;
450 let trimmed = value.trim();
451 anyhow::ensure!(
452 !trimmed.is_empty() && trimmed.chars().count() <= maximum,
453 "{key} must contain between 1 and {maximum} characters"
454 );
455 Ok(trimmed.into())
456}
457
458fn nonblank(value: &Value, key: &str) -> anyhow::Result<String> {
459 let value = string(value, key)?;
460 anyhow::ensure!(!value.trim().is_empty(), "{key} must not be blank");
461 Ok(value)
462}
463
464fn bounded_nonempty(value: &Value, key: &str, maximum: usize) -> anyhow::Result<String> {
465 let value = string(value, key)?;
466 anyhow::ensure!(
467 !value.trim().is_empty() && value.chars().count() <= maximum,
468 "{key} must contain between 1 and {maximum} characters"
469 );
470 Ok(value)
471}
472
473fn optional_object_ids(value: &Value, key: &str, maximum: usize) -> anyhow::Result<Vec<String>> {
474 let Some(values) = value.get(key) else {
475 return Ok(Vec::new());
476 };
477 let values = values
478 .as_array()
479 .with_context(|| format!("{key} must be an array"))?;
480 anyhow::ensure!(
481 values.len() <= maximum,
482 "{key} must contain at most {maximum} object IDs"
483 );
484 let ids = values
485 .iter()
486 .map(|value| {
487 let id = value
488 .as_str()
489 .with_context(|| format!("{key} entries must be strings"))?;
490 anyhow::ensure!(
491 !id.trim().is_empty() && id.chars().count() <= 64,
492 "{key} entries must contain between 1 and 64 characters"
493 );
494 Ok(id.to_owned())
495 })
496 .collect::<anyhow::Result<Vec<_>>>()?;
497 anyhow::ensure!(
498 ids.iter().collect::<HashSet<_>>().len() == ids.len(),
499 "{key} must not contain duplicate object IDs"
500 );
501 Ok(ids)
502}
503
504fn delivery_file_name(value: &Value, key: &str) -> anyhow::Result<Option<String>> {
505 let Some(file_name) = value.get(key) else {
506 return Ok(None);
507 };
508 let file_name = file_name
509 .as_str()
510 .with_context(|| format!("{key} must be a string"))?;
511 kcode_telegram_session_coordinator::validate_file_name(file_name)?;
512 Ok(Some(file_name.to_owned()))
513}
514
515fn normalize_media_type(value: &str) -> String {
516 value
517 .split(';')
518 .next()
519 .unwrap_or(value)
520 .trim()
521 .to_ascii_lowercase()
522}
523
524#[cfg(test)]
525mod tests {
526 use super::*;
527 use serde_json::json;
528
529 #[test]
530 fn transcribe_preserves_long_nonblank_prompt_and_rejects_optional_fields() {
531 let prompt = "x".repeat(4 * 1024 * 1024 + 1);
532 let decoded = decode(
533 "TranscribeAudio",
534 &json!({
535 "objectId":"pending:1",
536 "model":"gpt-4o-transcribe",
537 "prompt":prompt
538 }),
539 )
540 .unwrap()
541 .unwrap();
542 let DecodedTool::MediaEnrichment { prompt: actual, .. } = decoded else {
543 panic!("wrong decoded variant");
544 };
545 assert_eq!(actual.len(), 4 * 1024 * 1024 + 1);
546 assert!(
547 decode(
548 "TranscribeAudio",
549 &json!({
550 "objectId":"pending:1",
551 "model":"gpt-4o-transcribe",
552 "prompt":"audio",
553 "temperature":0.0
554 }),
555 )
556 .is_err()
557 );
558 }
559
560 #[test]
561 fn image_validation_order_precedes_prompt_and_reference_validation() {
562 let error = decode(
563 "GenerateImage",
564 &json!({
565 "model":"unknown",
566 "prompt":"",
567 "referenceObjectIds":"not-an-array"
568 }),
569 )
570 .unwrap_err()
571 .to_string();
572 assert!(error.contains("unsupported exact image model"));
573 }
574
575 #[test]
576 fn kweb_decode_preserves_character_limits() {
577 let error = decode(
578 "CreateNode",
579 &json!({
580 "parentIdentifiers":["self"],
581 "ownerIdentifier":"self",
582 "shortName":"abc",
583 "shortDescription":"",
584 "longDescription":""
585 }),
586 )
587 .unwrap_err()
588 .to_string();
589 assert!(error.contains("received 3"));
590 }
591}