1use std::collections::HashSet;
4
5use anyhow::{Context as _, ensure};
6pub use kcode_dev_tools::{ManagedSourceKind, SourceSnapshot};
7use kcode_dev_tools::{
8 PREVIEW_WRITE_FILE_RUST_BIN_TOOL, PREVIEW_WRITE_FILE_RUST_LIB_TOOL,
9 PREVIEW_WRITE_FILE_WEB_LIB_TOOL, WRITE_FILE_FREEFORM_RUST_BIN_TOOL,
10 WRITE_FILE_FREEFORM_RUST_LIB_TOOL, WRITE_FILE_FREEFORM_WEB_LIB_TOOL,
11};
12use kcode_session_history::{
13 Session,
14 chatend::{BoxContent, BoxId, EventKind, ToolSlotInput},
15};
16use serde_json::{Value, json};
17
18const RUST_LIB_TOOL_INSTANCE: &str = "managed-rust-libraries";
19const WEB_LIB_TOOL_INSTANCE: &str = "managed-web-libraries";
20const RUST_BIN_TOOL_INSTANCE: &str = "managed-rust-binaries";
21const MAX_CAPTURED_CONTENT_BYTES: usize = 4 * 1024 * 1024;
22
23#[derive(Clone, Debug, Eq, PartialEq)]
25pub struct FreeformWrite {
26 kind: ManagedSourceKind,
27 name: String,
28 path: String,
29 update_description: String,
30}
31
32impl FreeformWrite {
33 pub fn acknowledgement(&self) -> String {
35 format!(
36 "Ready. Output the complete contents of {} only, with no Markdown fences or commentary.",
37 self.path
38 )
39 }
40
41 pub fn kind(&self) -> ManagedSourceKind {
43 self.kind
44 }
45
46 pub fn name(&self) -> &str {
48 &self.name
49 }
50
51 pub fn preview_tool(&self) -> &'static str {
53 match self.kind {
54 ManagedSourceKind::RustLibrary => PREVIEW_WRITE_FILE_RUST_LIB_TOOL,
55 ManagedSourceKind::WebLibrary => PREVIEW_WRITE_FILE_WEB_LIB_TOOL,
56 ManagedSourceKind::RustBinary => PREVIEW_WRITE_FILE_RUST_BIN_TOOL,
57 }
58 }
59
60 pub fn write_tool(&self) -> &'static str {
62 match self.kind {
63 ManagedSourceKind::RustLibrary => WRITE_FILE_FREEFORM_RUST_LIB_TOOL,
64 ManagedSourceKind::WebLibrary => WRITE_FILE_FREEFORM_WEB_LIB_TOOL,
65 ManagedSourceKind::RustBinary => WRITE_FILE_FREEFORM_RUST_BIN_TOOL,
66 }
67 }
68
69 pub fn source_box_id(&self, session: &Session) -> anyhow::Result<BoxId> {
71 source_box_id(session, self.kind, &self.name).with_context(|| {
72 format!(
73 "the managed {} box for {:?} is no longer open",
74 self.kind.label(),
75 self.name
76 )
77 })
78 }
79
80 pub fn result_record(&self, ok: bool, result: &str) -> Value {
82 json!({
83 "tool":self.write_tool(),
84 "name":self.name,
85 "path":self.path,
86 "updateDescription":self.update_description,
87 "ok":ok,
88 "result":result,
89 })
90 }
91
92 pub fn capture(
94 &self,
95 session: &mut Session,
96 recorded_at: &str,
97 invocation_box_id: BoxId,
98 contents: String,
99 ) -> anyhow::Result<Value> {
100 let contents = normalize_captured_contents(contents)?;
101 session.update_box(
102 recorded_at,
103 invocation_box_id,
104 captured_write_box_content(self, contents.clone()),
105 )?;
106 session.summarize_box(recorded_at, invocation_box_id, self.summary())?;
107 Ok(self.backend_arguments(contents))
108 }
109
110 pub fn capture_subagent(
112 &self,
113 session: &mut Session,
114 recorded_at: &str,
115 contents: String,
116 ) -> anyhow::Result<Value> {
117 let contents = normalize_captured_contents(contents)?;
118 session.record(
119 recorded_at,
120 EventKind::Note {
121 label: "subagent_freeform_write_output".into(),
122 value: json!({
123 "tool":self.write_tool(),
124 "name":self.name,
125 "path":self.path,
126 "updateDescription":self.update_description,
127 "contents":contents,
128 }),
129 },
130 )?;
131 Ok(self.backend_arguments(contents))
132 }
133
134 fn backend_arguments(&self, contents: String) -> Value {
135 json!({
136 "name":self.name,
137 "path":self.path,
138 "contents":contents,
139 })
140 }
141
142 fn summary(&self) -> String {
143 format!(
144 "Kennedy called write-file on {} in {}, and she describes the update as: {}",
145 self.path, self.name, self.update_description
146 )
147 }
148}
149
150pub fn prepare_freeform_write(
154 session: &Session,
155 tool_name: &str,
156 arguments: &Value,
157) -> anyhow::Result<Option<FreeformWrite>> {
158 let Some(request) = decode_freeform_write(tool_name, arguments)? else {
159 return Ok(None);
160 };
161 ensure!(
162 source_box_id(session, request.kind, &request.name).is_some(),
163 "{} {:?} is not open in this Kennedy session. Call {} first.",
164 request.kind.label(),
165 request.name,
166 request.kind.open_tool()
167 );
168 Ok(Some(request))
169}
170
171pub fn decode_freeform_write(
177 tool_name: &str,
178 arguments: &Value,
179) -> anyhow::Result<Option<FreeformWrite>> {
180 let Some(kind) = freeform_kind(tool_name) else {
181 return Ok(None);
182 };
183 validate_exact_arguments(arguments, &["name", "path", "updateDescription"])?;
184 let request = FreeformWrite {
185 kind,
186 name: nonempty_string(arguments, "name", 255)?,
187 path: nonempty_string(arguments, "path", 4_096)?,
188 update_description: nonempty_string(arguments, "updateDescription", 4_000)?,
189 };
190 ensure!(
191 !request.path.contains(['\r', '\n']),
192 "path must contain exactly one line"
193 );
194 ensure!(
195 !request.update_description.contains(['\r', '\n']),
196 "updateDescription must contain exactly one line"
197 );
198 Ok(Some(request))
199}
200
201pub fn source_box_id(session: &Session, kind: ManagedSourceKind, name: &str) -> Option<BoxId> {
203 session
204 .state()
205 .tools
206 .get(tool_instance(kind))?
207 .slots
208 .iter()
209 .find_map(|slot| {
210 if slot.retired {
211 return None;
212 }
213 let state = session.state().box_state(slot.box_id)?;
214 (logical_name(kind, state, &slot.slot) == name).then_some(slot.box_id)
215 })
216}
217
218pub fn apply_snapshot(
220 session: &mut Session,
221 recorded_at: &str,
222 snapshot: SourceSnapshot,
223) -> anyhow::Result<BoxId> {
224 let kind = snapshot.kind;
225 let current = session
226 .state()
227 .tools
228 .get(tool_instance(kind))
229 .cloned()
230 .unwrap_or_default();
231 let mut selected_slot = None;
232 let mut slots = Vec::with_capacity(current.slots.len() + 1);
233 let mut used_slots = current
234 .slots
235 .iter()
236 .map(|slot| slot.slot.clone())
237 .collect::<HashSet<_>>();
238
239 for slot in ¤t.slots {
240 let state = session
241 .state()
242 .box_state(slot.box_id)
243 .with_context(|| format!("managed {} slot box is missing", kind.label()))?;
244 let selected = selected_slot.is_none()
245 && !slot.retired
246 && logical_name(kind, state, &slot.slot) == snapshot.name;
247 if selected {
248 selected_slot = Some(slot.slot.clone());
249 slots.push(ToolSlotInput {
250 slot: slot.slot.clone(),
251 name: format!("Managed {} {}", kind.label(), snapshot.name),
252 content: source_box_content(kind, &snapshot),
253 retired: false,
254 });
255 } else {
256 slots.push(ToolSlotInput {
257 slot: slot.slot.clone(),
258 name: state.name.clone(),
259 content: state.canonical.content.clone(),
260 retired: slot.retired,
261 });
262 }
263 }
264
265 let selected_slot = selected_slot.unwrap_or_else(|| {
266 let slot = unique_slot(&snapshot.name, &mut used_slots);
267 slots.push(ToolSlotInput {
268 slot: slot.clone(),
269 name: format!("Managed {} {}", kind.label(), snapshot.name),
270 content: source_box_content(kind, &snapshot),
271 retired: false,
272 });
273 slot
274 });
275
276 session.apply_tool_slots(recorded_at, tool_instance(kind), slots)?;
277 session
278 .state()
279 .tools
280 .get(tool_instance(kind))
281 .and_then(|tool| {
282 tool.slots
283 .iter()
284 .find(|slot| slot.slot == selected_slot && !slot.retired)
285 })
286 .map(|slot| slot.box_id)
287 .with_context(|| format!("managed {} box was not installed", kind.label()))
288}
289
290fn freeform_kind(tool_name: &str) -> Option<ManagedSourceKind> {
291 match tool_name {
292 WRITE_FILE_FREEFORM_RUST_LIB_TOOL => Some(ManagedSourceKind::RustLibrary),
293 WRITE_FILE_FREEFORM_WEB_LIB_TOOL => Some(ManagedSourceKind::WebLibrary),
294 WRITE_FILE_FREEFORM_RUST_BIN_TOOL => Some(ManagedSourceKind::RustBinary),
295 _ => None,
296 }
297}
298
299fn tool_instance(kind: ManagedSourceKind) -> &'static str {
300 match kind {
301 ManagedSourceKind::RustLibrary => RUST_LIB_TOOL_INSTANCE,
302 ManagedSourceKind::WebLibrary => WEB_LIB_TOOL_INSTANCE,
303 ManagedSourceKind::RustBinary => RUST_BIN_TOOL_INSTANCE,
304 }
305}
306
307fn metadata_key(kind: ManagedSourceKind) -> &'static str {
308 match kind {
309 ManagedSourceKind::RustLibrary => "managedRustLibrary",
310 ManagedSourceKind::WebLibrary => "managedWebLibrary",
311 ManagedSourceKind::RustBinary => "managedRustBinary",
312 }
313}
314
315fn logical_name(
316 kind: ManagedSourceKind,
317 state: &kcode_session_history::chatend::BoxState,
318 fallback: &str,
319) -> String {
320 state
321 .canonical
322 .content
323 .metadata
324 .get(metadata_key(kind))
325 .and_then(Value::as_str)
326 .unwrap_or(fallback)
327 .to_owned()
328}
329
330fn source_box_content(kind: ManagedSourceKind, snapshot: &SourceSnapshot) -> BoxContent {
331 BoxContent {
332 text: snapshot.text.clone(),
333 objects: Vec::new(),
334 metadata: json!({metadata_key(kind):snapshot.name}),
335 }
336}
337
338fn captured_write_box_content(request: &FreeformWrite, contents: String) -> BoxContent {
339 BoxContent {
340 text: contents,
341 objects: Vec::new(),
342 metadata: json!({
343 "capturedFreeformOutput":true,
344 "toolName":request.write_tool(),
345 "arguments":{
346 "name":request.name,
347 "path":request.path,
348 "updateDescription":request.update_description,
349 },
350 }),
351 }
352}
353
354fn normalize_captured_contents(mut contents: String) -> anyhow::Result<String> {
355 let needs_final_newline = !contents.ends_with('\n');
356 let normalized_len = contents
357 .len()
358 .checked_add(usize::from(needs_final_newline))
359 .context("captured contents length overflow")?;
360 ensure!(
361 normalized_len <= MAX_CAPTURED_CONTENT_BYTES,
362 "normalized captured contents must not exceed {MAX_CAPTURED_CONTENT_BYTES} bytes"
363 );
364 if needs_final_newline {
365 contents.push('\n');
366 }
367 Ok(contents)
368}
369
370fn unique_slot(logical: &str, used: &mut HashSet<String>) -> String {
371 if used.insert(logical.to_owned()) {
372 return logical.to_owned();
373 }
374 let mut generation = 2_u64;
375 loop {
376 let candidate = format!("{logical}#generation-{generation}");
377 if used.insert(candidate.clone()) {
378 return candidate;
379 }
380 generation += 1;
381 }
382}
383
384fn validate_exact_arguments(value: &Value, required: &[&str]) -> anyhow::Result<()> {
385 let map = value
386 .as_object()
387 .context("arguments must be a JSON object")?;
388 ensure!(
389 required.iter().all(|key| map.contains_key(*key)) && map.len() == required.len(),
390 "expected exactly: {}",
391 required.join(", ")
392 );
393 Ok(())
394}
395
396fn nonempty_string(value: &Value, key: &str, max: usize) -> anyhow::Result<String> {
397 let value = value
398 .get(key)
399 .and_then(Value::as_str)
400 .with_context(|| format!("{key} must be a string"))?;
401 let trimmed = value.trim();
402 ensure!(
403 !trimmed.is_empty() && trimmed.chars().count() <= max,
404 "{key} must contain between 1 and {max} characters"
405 );
406 Ok(trimmed.into())
407}
408
409#[cfg(test)]
410mod tests {
411 use std::{
412 sync::atomic::{AtomicU64, Ordering},
413 time::{SystemTime, UNIX_EPOCH},
414 };
415
416 use kcode_dev_tools::{WRITE_FILE_FREEFORM_RUST_LIB_TOOL, WRITE_FILE_FREEFORM_WEB_LIB_TOOL};
417 use kcode_session_history::chatend::{BoxContent, BoxOwner, Representation, SessionKind};
418 use kcode_session_history::{Config, NewSession, SessionHistory};
419
420 use super::*;
421
422 static NEXT_SESSION_SECOND: AtomicU64 = AtomicU64::new(0);
423
424 fn test_session(label: &str) -> (std::path::PathBuf, Session) {
425 let root = std::env::temp_dir().join(format!(
426 "kcode-dev-tools-chatend-{label}-{}-{}",
427 std::process::id(),
428 SystemTime::now()
429 .duration_since(UNIX_EPOCH)
430 .unwrap()
431 .as_nanos()
432 ));
433 let history = SessionHistory::open(Config {
434 directory: root.join("sessions"),
435 completed_list: root.join("completed.jsonl"),
436 provider_cost_compatibility: None,
437 })
438 .unwrap();
439 let second = NEXT_SESSION_SECOND.fetch_add(1, Ordering::Relaxed);
440 assert!(second < 60);
441 let session = history
442 .create_session(NewSession {
443 kind: SessionKind::Conversation,
444 created_at: format!("2026-07-31T00:00:{second:02}Z"),
445 effective_context_tokens: 10_000,
446 channel: Value::Null,
447 })
448 .unwrap();
449 (root, session)
450 }
451
452 #[test]
453 fn snapshots_keep_one_stable_box_per_kind_and_project() {
454 let (root, mut session) = test_session("stable-boxes");
455 let rust_box = apply_snapshot(
456 &mut session,
457 "t1",
458 SourceSnapshot {
459 kind: ManagedSourceKind::RustLibrary,
460 name: "shared-name".into(),
461 text: "old Rust source".into(),
462 },
463 )
464 .unwrap();
465 let web_box = apply_snapshot(
466 &mut session,
467 "t2",
468 SourceSnapshot {
469 kind: ManagedSourceKind::WebLibrary,
470 name: "shared-name".into(),
471 text: "Web source".into(),
472 },
473 )
474 .unwrap();
475 let binary_box = apply_snapshot(
476 &mut session,
477 "t3",
478 SourceSnapshot {
479 kind: ManagedSourceKind::RustBinary,
480 name: "shared-name".into(),
481 text: "binary source".into(),
482 },
483 )
484 .unwrap();
485 let same_rust_box = apply_snapshot(
486 &mut session,
487 "t4",
488 SourceSnapshot {
489 kind: ManagedSourceKind::RustLibrary,
490 name: "shared-name".into(),
491 text: "new Rust source".into(),
492 },
493 )
494 .unwrap();
495
496 assert_eq!(same_rust_box, rust_box);
497 assert_ne!(rust_box, web_box);
498 assert_ne!(rust_box, binary_box);
499 assert_ne!(web_box, binary_box);
500 assert_eq!(session.state().tools[RUST_LIB_TOOL_INSTANCE].slots.len(), 1);
501 assert_eq!(session.state().tools[WEB_LIB_TOOL_INSTANCE].slots.len(), 1);
502 assert_eq!(session.state().tools[RUST_BIN_TOOL_INSTANCE].slots.len(), 1);
503 assert_eq!(
504 session
505 .state()
506 .box_state(rust_box)
507 .unwrap()
508 .canonical
509 .content
510 .text,
511 "new Rust source"
512 );
513 std::fs::remove_dir_all(root).unwrap();
514 }
515
516 #[test]
517 fn snapshot_updates_preserve_representation_choices() {
518 let (root, mut session) = test_session("representation");
519 let box_id = apply_snapshot(
520 &mut session,
521 "t1",
522 SourceSnapshot {
523 kind: ManagedSourceKind::RustLibrary,
524 name: "summary-lib".into(),
525 text: "old canonical source".into(),
526 },
527 )
528 .unwrap();
529 session
530 .summarize_box("t2", box_id, "Kennedy's retained library summary")
531 .unwrap();
532 apply_snapshot(
533 &mut session,
534 "t3",
535 SourceSnapshot {
536 kind: ManagedSourceKind::RustLibrary,
537 name: "summary-lib".into(),
538 text: "new canonical source".into(),
539 },
540 )
541 .unwrap();
542
543 let state = session.state().box_state(box_id).unwrap();
544 assert_eq!(state.canonical.content.text, "new canonical source");
545 assert!(state.stale());
546 assert!(matches!(
547 state.representation,
548 Representation::Summarized { .. }
549 ));
550 assert!(
551 session
552 .state()
553 .render()
554 .contains("Kennedy's retained library summary")
555 );
556 assert!(!session.state().render().contains("new canonical source"));
557 std::fs::remove_dir_all(root).unwrap();
558 }
559
560 #[test]
561 fn freeform_capture_is_exact_except_for_one_missing_final_newline() {
562 let (root, mut session) = test_session("freeform-capture");
563 apply_snapshot(
564 &mut session,
565 "t1",
566 SourceSnapshot {
567 kind: ManagedSourceKind::RustLibrary,
568 name: "example-lib".into(),
569 text: "existing source".into(),
570 },
571 )
572 .unwrap();
573 let request = prepare_freeform_write(
574 &session,
575 WRITE_FILE_FREEFORM_RUST_LIB_TOOL,
576 &json!({
577 "name":"example-lib",
578 "path":"src/lib.rs",
579 "updateDescription":"Preserved raw Rust source",
580 }),
581 )
582 .unwrap()
583 .unwrap();
584 let invocation = session
585 .create_box(
586 "t2",
587 "Kennedy tool call",
588 BoxOwner::Kennedy,
589 BoxContent::text("call"),
590 )
591 .unwrap();
592 let arguments = request
593 .capture(
594 &mut session,
595 "t3",
596 invocation,
597 "\n//! leading newline\npub fn quote() -> &'static str { \"raw\\\\text\" }".into(),
598 )
599 .unwrap();
600
601 let exact = "\n//! leading newline\npub fn quote() -> &'static str { \"raw\\\\text\" }\n";
602 let state = session.state().box_state(invocation).unwrap();
603 assert_eq!(state.canonical.content.text, exact);
604 assert_eq!(arguments["contents"], exact);
605 assert_eq!(
606 state.canonical.content.metadata["toolName"],
607 request.write_tool()
608 );
609 assert!(session.state().render().contains(
610 "Kennedy called write-file on src/lib.rs in example-lib, and she describes the update as: Preserved raw Rust source"
611 ));
612 assert!(!session.state().render().contains("raw\\\\text"));
613 std::fs::remove_dir_all(root).unwrap();
614 }
615
616 #[test]
617 fn captured_contents_enforce_normalized_byte_limit_before_persistence() {
618 let (root, mut session) = test_session("capture-limit");
619 apply_snapshot(
620 &mut session,
621 "t1",
622 SourceSnapshot {
623 kind: ManagedSourceKind::RustLibrary,
624 name: "bounded-lib".into(),
625 text: "existing source".into(),
626 },
627 )
628 .unwrap();
629 let request = prepare_freeform_write(
630 &session,
631 WRITE_FILE_FREEFORM_RUST_LIB_TOOL,
632 &json!({
633 "name":"bounded-lib",
634 "path":"src/lib.rs",
635 "updateDescription":"Bound captured source",
636 }),
637 )
638 .unwrap()
639 .unwrap();
640
641 let accepted_invocation = session
642 .create_box(
643 "t2",
644 "Accepted Kennedy tool call",
645 BoxOwner::Kennedy,
646 BoxContent::text("accepted call"),
647 )
648 .unwrap();
649 let arguments = request
650 .capture(
651 &mut session,
652 "t3",
653 accepted_invocation,
654 "a".repeat(MAX_CAPTURED_CONTENT_BYTES - 1),
655 )
656 .unwrap();
657 let accepted = arguments["contents"].as_str().unwrap();
658 assert_eq!(accepted.len(), MAX_CAPTURED_CONTENT_BYTES);
659 assert!(accepted.ends_with('\n'));
660
661 let rejected_invocation = session
662 .create_box(
663 "t4",
664 "Rejected Kennedy tool call",
665 BoxOwner::Kennedy,
666 BoxContent::text("unchanged call"),
667 )
668 .unwrap();
669 assert!(
670 request
671 .capture(
672 &mut session,
673 "t5",
674 rejected_invocation,
675 "b".repeat(MAX_CAPTURED_CONTENT_BYTES),
676 )
677 .is_err()
678 );
679 assert_eq!(
680 session
681 .state()
682 .box_state(rejected_invocation)
683 .unwrap()
684 .canonical
685 .content
686 .text,
687 "unchanged call"
688 );
689 std::fs::remove_dir_all(root).unwrap();
690 }
691
692 #[test]
693 fn freeform_metadata_is_strict_and_kind_specific() {
694 let (root, mut session) = test_session("freeform-validation");
695 apply_snapshot(
696 &mut session,
697 "t1",
698 SourceSnapshot {
699 kind: ManagedSourceKind::WebLibrary,
700 name: "example-ui".into(),
701 text: "existing source".into(),
702 },
703 )
704 .unwrap();
705 let valid = json!({
706 "name":"example-ui",
707 "path":"index.js",
708 "updateDescription":"Replace the entry module",
709 });
710 let request = prepare_freeform_write(&session, WRITE_FILE_FREEFORM_WEB_LIB_TOOL, &valid)
711 .unwrap()
712 .unwrap();
713 assert_eq!(request.kind(), ManagedSourceKind::WebLibrary);
714 assert_eq!(request.write_tool(), WRITE_FILE_FREEFORM_WEB_LIB_TOOL);
715
716 let mut extra = valid.clone();
717 extra["contents"] = json!("not accepted in the Ktool call");
718 assert!(
719 prepare_freeform_write(&session, WRITE_FILE_FREEFORM_WEB_LIB_TOOL, &extra).is_err()
720 );
721 let mut multiline_path = valid.clone();
722 multiline_path["path"] = json!("index.js\nanother");
723 assert!(
724 prepare_freeform_write(&session, WRITE_FILE_FREEFORM_WEB_LIB_TOOL, &multiline_path)
725 .is_err()
726 );
727 let mut multiline_description = valid;
728 multiline_description["updateDescription"] = json!("line one\nline two");
729 assert!(
730 prepare_freeform_write(
731 &session,
732 WRITE_FILE_FREEFORM_WEB_LIB_TOOL,
733 &multiline_description,
734 )
735 .is_err()
736 );
737 assert!(
738 prepare_freeform_write(&session, "unrelated-tool", &Value::Null)
739 .unwrap()
740 .is_none()
741 );
742 std::fs::remove_dir_all(root).unwrap();
743 }
744
745 #[test]
746 fn box_free_decode_does_not_require_session_history() {
747 let request = decode_freeform_write(
748 WRITE_FILE_FREEFORM_RUST_LIB_TOOL,
749 &json!({
750 "name":"box-free",
751 "path":"src/lib.rs",
752 "updateDescription":"Replace the source",
753 }),
754 )
755 .unwrap()
756 .unwrap();
757
758 assert_eq!(request.kind(), ManagedSourceKind::RustLibrary);
759 assert_eq!(request.name(), "box-free");
760 }
761}