Skip to main content

kcode_dev_tools_chatend/
lib.rs

1//! Chatend projection for Kennedy's managed development tools.
2
3use 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/// A validated request to capture one complete source file from a provider.
24#[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    /// The prompt acknowledgement that starts raw provider output capture.
34    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    /// The managed source kind selected by the freeform tool name.
42    pub fn kind(&self) -> ManagedSourceKind {
43        self.kind
44    }
45
46    /// The non-mutating development tool used before the captured write.
47    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    /// The development tool that commits the captured file contents.
56    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    /// Resolve the still-open source box this request is allowed to replace.
65    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    /// Build the durable application note recorded after the captured write.
76    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    /// Preserve a top-level provider capture and return its backend arguments.
88    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    /// Preserve a subagent capture and return its backend arguments.
106    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
145/// Parse a managed freeform-write call and ensure that its source is open.
146///
147/// Unrelated tool names return `None` without inspecting their arguments.
148pub 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
181/// Find the active stable Chatend box for one managed project.
182pub 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
198/// Create or revise the stable Chatend source box represented by a snapshot.
199pub 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 &current.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}