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