kcode-dev-tools-chatend 0.1.0

//! Chatend projection for Kennedy's managed development tools.

use std::collections::HashSet;

use anyhow::{Context as _, ensure};
pub use kcode_dev_tools::{ManagedSourceKind, SourceSnapshot};
use kcode_dev_tools::{
    PREVIEW_WRITE_FILE_RUST_BIN_TOOL, PREVIEW_WRITE_FILE_RUST_LIB_TOOL,
    PREVIEW_WRITE_FILE_WEB_LIB_TOOL, WRITE_FILE_FREEFORM_RUST_BIN_TOOL,
    WRITE_FILE_FREEFORM_RUST_LIB_TOOL, WRITE_FILE_FREEFORM_WEB_LIB_TOOL,
};
use kcode_session_history::{
    Session,
    chatend::{BoxContent, BoxId, EventKind, ToolSlotInput},
};
use serde_json::{Value, json};

const RUST_LIB_TOOL_INSTANCE: &str = "managed-rust-libraries";
const WEB_LIB_TOOL_INSTANCE: &str = "managed-web-libraries";
const RUST_BIN_TOOL_INSTANCE: &str = "managed-rust-binaries";
const MAX_CAPTURED_CONTENT_BYTES: usize = 4 * 1024 * 1024;

/// A validated request to capture one complete source file from a provider.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct FreeformWrite {
    kind: ManagedSourceKind,
    name: String,
    path: String,
    update_description: String,
}

impl FreeformWrite {
    /// The prompt acknowledgement that starts raw provider output capture.
    pub fn acknowledgement(&self) -> String {
        format!(
            "Ready. Output the complete contents of {} only, with no Markdown fences or commentary.",
            self.path
        )
    }

    /// The managed source kind selected by the freeform tool name.
    pub fn kind(&self) -> ManagedSourceKind {
        self.kind
    }

    /// The non-mutating development tool used before the captured write.
    pub fn preview_tool(&self) -> &'static str {
        match self.kind {
            ManagedSourceKind::RustLibrary => PREVIEW_WRITE_FILE_RUST_LIB_TOOL,
            ManagedSourceKind::WebLibrary => PREVIEW_WRITE_FILE_WEB_LIB_TOOL,
            ManagedSourceKind::RustBinary => PREVIEW_WRITE_FILE_RUST_BIN_TOOL,
        }
    }

    /// The development tool that commits the captured file contents.
    pub fn write_tool(&self) -> &'static str {
        match self.kind {
            ManagedSourceKind::RustLibrary => WRITE_FILE_FREEFORM_RUST_LIB_TOOL,
            ManagedSourceKind::WebLibrary => WRITE_FILE_FREEFORM_WEB_LIB_TOOL,
            ManagedSourceKind::RustBinary => WRITE_FILE_FREEFORM_RUST_BIN_TOOL,
        }
    }

    /// Resolve the still-open source box this request is allowed to replace.
    pub fn source_box_id(&self, session: &Session) -> anyhow::Result<BoxId> {
        source_box_id(session, self.kind, &self.name).with_context(|| {
            format!(
                "the managed {} box for {:?} is no longer open",
                self.kind.label(),
                self.name
            )
        })
    }

    /// Build the durable application note recorded after the captured write.
    pub fn result_record(&self, ok: bool, result: &str) -> Value {
        json!({
            "tool":self.write_tool(),
            "name":self.name,
            "path":self.path,
            "updateDescription":self.update_description,
            "ok":ok,
            "result":result,
        })
    }

    /// Preserve a top-level provider capture and return its backend arguments.
    pub fn capture(
        &self,
        session: &mut Session,
        recorded_at: &str,
        invocation_box_id: BoxId,
        contents: String,
    ) -> anyhow::Result<Value> {
        let contents = normalize_captured_contents(contents)?;
        session.update_box(
            recorded_at,
            invocation_box_id,
            captured_write_box_content(self, contents.clone()),
        )?;
        session.summarize_box(recorded_at, invocation_box_id, self.summary())?;
        Ok(self.backend_arguments(contents))
    }

    /// Preserve a subagent capture and return its backend arguments.
    pub fn capture_subagent(
        &self,
        session: &mut Session,
        recorded_at: &str,
        contents: String,
    ) -> anyhow::Result<Value> {
        let contents = normalize_captured_contents(contents)?;
        session.record(
            recorded_at,
            EventKind::Note {
                label: "subagent_freeform_write_output".into(),
                value: json!({
                    "tool":self.write_tool(),
                    "name":self.name,
                    "path":self.path,
                    "updateDescription":self.update_description,
                    "contents":contents,
                }),
            },
        )?;
        Ok(self.backend_arguments(contents))
    }

    fn backend_arguments(&self, contents: String) -> Value {
        json!({
            "name":self.name,
            "path":self.path,
            "contents":contents,
        })
    }

    fn summary(&self) -> String {
        format!(
            "Kennedy called write-file on {} in {}, and she describes the update as: {}",
            self.path, self.name, self.update_description
        )
    }
}

/// Parse a managed freeform-write call and ensure that its source is open.
///
/// Unrelated tool names return `None` without inspecting their arguments.
pub fn prepare_freeform_write(
    session: &Session,
    tool_name: &str,
    arguments: &Value,
) -> anyhow::Result<Option<FreeformWrite>> {
    let Some(kind) = freeform_kind(tool_name) else {
        return Ok(None);
    };
    validate_exact_arguments(arguments, &["name", "path", "updateDescription"])?;
    let request = FreeformWrite {
        kind,
        name: nonempty_string(arguments, "name", 255)?,
        path: nonempty_string(arguments, "path", 4_096)?,
        update_description: nonempty_string(arguments, "updateDescription", 4_000)?,
    };
    ensure!(
        !request.path.contains(['\r', '\n']),
        "path must contain exactly one line"
    );
    ensure!(
        !request.update_description.contains(['\r', '\n']),
        "updateDescription must contain exactly one line"
    );
    ensure!(
        source_box_id(session, kind, &request.name).is_some(),
        "{} {:?} is not open in this Kennedy session. Call {} first.",
        kind.label(),
        request.name,
        kind.open_tool()
    );
    Ok(Some(request))
}

/// Find the active stable Chatend box for one managed project.
pub fn source_box_id(session: &Session, kind: ManagedSourceKind, name: &str) -> Option<BoxId> {
    session
        .state()
        .tools
        .get(tool_instance(kind))?
        .slots
        .iter()
        .find_map(|slot| {
            if slot.retired {
                return None;
            }
            let state = session.state().box_state(slot.box_id)?;
            (logical_name(kind, state, &slot.slot) == name).then_some(slot.box_id)
        })
}

/// Create or revise the stable Chatend source box represented by a snapshot.
pub fn apply_snapshot(
    session: &mut Session,
    recorded_at: &str,
    snapshot: SourceSnapshot,
) -> anyhow::Result<BoxId> {
    let kind = snapshot.kind;
    let current = session
        .state()
        .tools
        .get(tool_instance(kind))
        .cloned()
        .unwrap_or_default();
    let mut selected_slot = None;
    let mut slots = Vec::with_capacity(current.slots.len() + 1);
    let mut used_slots = current
        .slots
        .iter()
        .map(|slot| slot.slot.clone())
        .collect::<HashSet<_>>();

    for slot in &current.slots {
        let state = session
            .state()
            .box_state(slot.box_id)
            .with_context(|| format!("managed {} slot box is missing", kind.label()))?;
        let selected = selected_slot.is_none()
            && !slot.retired
            && logical_name(kind, state, &slot.slot) == snapshot.name;
        if selected {
            selected_slot = Some(slot.slot.clone());
            slots.push(ToolSlotInput {
                slot: slot.slot.clone(),
                name: format!("Managed {} {}", kind.label(), snapshot.name),
                content: source_box_content(kind, &snapshot),
                retired: false,
            });
        } else {
            slots.push(ToolSlotInput {
                slot: slot.slot.clone(),
                name: state.name.clone(),
                content: state.canonical.content.clone(),
                retired: slot.retired,
            });
        }
    }

    let selected_slot = selected_slot.unwrap_or_else(|| {
        let slot = unique_slot(&snapshot.name, &mut used_slots);
        slots.push(ToolSlotInput {
            slot: slot.clone(),
            name: format!("Managed {} {}", kind.label(), snapshot.name),
            content: source_box_content(kind, &snapshot),
            retired: false,
        });
        slot
    });

    session.apply_tool_slots(recorded_at, tool_instance(kind), slots)?;
    session
        .state()
        .tools
        .get(tool_instance(kind))
        .and_then(|tool| {
            tool.slots
                .iter()
                .find(|slot| slot.slot == selected_slot && !slot.retired)
        })
        .map(|slot| slot.box_id)
        .with_context(|| format!("managed {} box was not installed", kind.label()))
}

fn freeform_kind(tool_name: &str) -> Option<ManagedSourceKind> {
    match tool_name {
        WRITE_FILE_FREEFORM_RUST_LIB_TOOL => Some(ManagedSourceKind::RustLibrary),
        WRITE_FILE_FREEFORM_WEB_LIB_TOOL => Some(ManagedSourceKind::WebLibrary),
        WRITE_FILE_FREEFORM_RUST_BIN_TOOL => Some(ManagedSourceKind::RustBinary),
        _ => None,
    }
}

fn tool_instance(kind: ManagedSourceKind) -> &'static str {
    match kind {
        ManagedSourceKind::RustLibrary => RUST_LIB_TOOL_INSTANCE,
        ManagedSourceKind::WebLibrary => WEB_LIB_TOOL_INSTANCE,
        ManagedSourceKind::RustBinary => RUST_BIN_TOOL_INSTANCE,
    }
}

fn metadata_key(kind: ManagedSourceKind) -> &'static str {
    match kind {
        ManagedSourceKind::RustLibrary => "managedRustLibrary",
        ManagedSourceKind::WebLibrary => "managedWebLibrary",
        ManagedSourceKind::RustBinary => "managedRustBinary",
    }
}

fn logical_name(
    kind: ManagedSourceKind,
    state: &kcode_session_history::chatend::BoxState,
    fallback: &str,
) -> String {
    state
        .canonical
        .content
        .metadata
        .get(metadata_key(kind))
        .and_then(Value::as_str)
        .unwrap_or(fallback)
        .to_owned()
}

fn source_box_content(kind: ManagedSourceKind, snapshot: &SourceSnapshot) -> BoxContent {
    BoxContent {
        text: snapshot.text.clone(),
        objects: Vec::new(),
        metadata: json!({metadata_key(kind):snapshot.name}),
    }
}

fn captured_write_box_content(request: &FreeformWrite, contents: String) -> BoxContent {
    BoxContent {
        text: contents,
        objects: Vec::new(),
        metadata: json!({
            "capturedFreeformOutput":true,
            "toolName":request.write_tool(),
            "arguments":{
                "name":request.name,
                "path":request.path,
                "updateDescription":request.update_description,
            },
        }),
    }
}

fn normalize_captured_contents(mut contents: String) -> anyhow::Result<String> {
    let needs_final_newline = !contents.ends_with('\n');
    let normalized_len = contents
        .len()
        .checked_add(usize::from(needs_final_newline))
        .context("captured contents length overflow")?;
    ensure!(
        normalized_len <= MAX_CAPTURED_CONTENT_BYTES,
        "normalized captured contents must not exceed {MAX_CAPTURED_CONTENT_BYTES} bytes"
    );
    if needs_final_newline {
        contents.push('\n');
    }
    Ok(contents)
}

fn unique_slot(logical: &str, used: &mut HashSet<String>) -> String {
    if used.insert(logical.to_owned()) {
        return logical.to_owned();
    }
    let mut generation = 2_u64;
    loop {
        let candidate = format!("{logical}#generation-{generation}");
        if used.insert(candidate.clone()) {
            return candidate;
        }
        generation += 1;
    }
}

fn validate_exact_arguments(value: &Value, required: &[&str]) -> anyhow::Result<()> {
    let map = value
        .as_object()
        .context("arguments must be a JSON object")?;
    ensure!(
        required.iter().all(|key| map.contains_key(*key)) && map.len() == required.len(),
        "expected exactly: {}",
        required.join(", ")
    );
    Ok(())
}

fn nonempty_string(value: &Value, key: &str, max: usize) -> anyhow::Result<String> {
    let value = value
        .get(key)
        .and_then(Value::as_str)
        .with_context(|| format!("{key} must be a string"))?;
    let trimmed = value.trim();
    ensure!(
        !trimmed.is_empty() && trimmed.chars().count() <= max,
        "{key} must contain between 1 and {max} characters"
    );
    Ok(trimmed.into())
}

#[cfg(test)]
mod tests {
    use std::{
        sync::atomic::{AtomicU64, Ordering},
        time::{SystemTime, UNIX_EPOCH},
    };

    use kcode_dev_tools::{WRITE_FILE_FREEFORM_RUST_LIB_TOOL, WRITE_FILE_FREEFORM_WEB_LIB_TOOL};
    use kcode_session_history::chatend::{BoxContent, BoxOwner, Representation, SessionKind};
    use kcode_session_history::{Config, NewSession, SessionHistory};

    use super::*;

    static NEXT_SESSION_SECOND: AtomicU64 = AtomicU64::new(0);

    fn test_session(label: &str) -> (std::path::PathBuf, Session) {
        let root = std::env::temp_dir().join(format!(
            "kcode-dev-tools-chatend-{label}-{}-{}",
            std::process::id(),
            SystemTime::now()
                .duration_since(UNIX_EPOCH)
                .unwrap()
                .as_nanos()
        ));
        let history = SessionHistory::open(Config {
            directory: root.join("sessions"),
            completed_list: root.join("completed.jsonl"),
        })
        .unwrap();
        let second = NEXT_SESSION_SECOND.fetch_add(1, Ordering::Relaxed);
        assert!(second < 60);
        let session = history
            .create_session(NewSession {
                kind: SessionKind::Conversation,
                created_at: format!("2026-07-31T00:00:{second:02}Z"),
                effective_context_tokens: 10_000,
                channel: Value::Null,
            })
            .unwrap();
        (root, session)
    }

    #[test]
    fn snapshots_keep_one_stable_box_per_kind_and_project() {
        let (root, mut session) = test_session("stable-boxes");
        let rust_box = apply_snapshot(
            &mut session,
            "t1",
            SourceSnapshot {
                kind: ManagedSourceKind::RustLibrary,
                name: "shared-name".into(),
                text: "old Rust source".into(),
            },
        )
        .unwrap();
        let web_box = apply_snapshot(
            &mut session,
            "t2",
            SourceSnapshot {
                kind: ManagedSourceKind::WebLibrary,
                name: "shared-name".into(),
                text: "Web source".into(),
            },
        )
        .unwrap();
        let binary_box = apply_snapshot(
            &mut session,
            "t3",
            SourceSnapshot {
                kind: ManagedSourceKind::RustBinary,
                name: "shared-name".into(),
                text: "binary source".into(),
            },
        )
        .unwrap();
        let same_rust_box = apply_snapshot(
            &mut session,
            "t4",
            SourceSnapshot {
                kind: ManagedSourceKind::RustLibrary,
                name: "shared-name".into(),
                text: "new Rust source".into(),
            },
        )
        .unwrap();

        assert_eq!(same_rust_box, rust_box);
        assert_ne!(rust_box, web_box);
        assert_ne!(rust_box, binary_box);
        assert_ne!(web_box, binary_box);
        assert_eq!(session.state().tools[RUST_LIB_TOOL_INSTANCE].slots.len(), 1);
        assert_eq!(session.state().tools[WEB_LIB_TOOL_INSTANCE].slots.len(), 1);
        assert_eq!(session.state().tools[RUST_BIN_TOOL_INSTANCE].slots.len(), 1);
        assert_eq!(
            session
                .state()
                .box_state(rust_box)
                .unwrap()
                .canonical
                .content
                .text,
            "new Rust source"
        );
        std::fs::remove_dir_all(root).unwrap();
    }

    #[test]
    fn snapshot_updates_preserve_representation_choices() {
        let (root, mut session) = test_session("representation");
        let box_id = apply_snapshot(
            &mut session,
            "t1",
            SourceSnapshot {
                kind: ManagedSourceKind::RustLibrary,
                name: "summary-lib".into(),
                text: "old canonical source".into(),
            },
        )
        .unwrap();
        session
            .summarize_box("t2", box_id, "Kennedy's retained library summary")
            .unwrap();
        apply_snapshot(
            &mut session,
            "t3",
            SourceSnapshot {
                kind: ManagedSourceKind::RustLibrary,
                name: "summary-lib".into(),
                text: "new canonical source".into(),
            },
        )
        .unwrap();

        let state = session.state().box_state(box_id).unwrap();
        assert_eq!(state.canonical.content.text, "new canonical source");
        assert!(state.stale());
        assert!(matches!(
            state.representation,
            Representation::Summarized { .. }
        ));
        assert!(
            session
                .state()
                .render()
                .contains("Kennedy's retained library summary")
        );
        assert!(!session.state().render().contains("new canonical source"));
        std::fs::remove_dir_all(root).unwrap();
    }

    #[test]
    fn freeform_capture_is_exact_except_for_one_missing_final_newline() {
        let (root, mut session) = test_session("freeform-capture");
        apply_snapshot(
            &mut session,
            "t1",
            SourceSnapshot {
                kind: ManagedSourceKind::RustLibrary,
                name: "example-lib".into(),
                text: "existing source".into(),
            },
        )
        .unwrap();
        let request = prepare_freeform_write(
            &session,
            WRITE_FILE_FREEFORM_RUST_LIB_TOOL,
            &json!({
                "name":"example-lib",
                "path":"src/lib.rs",
                "updateDescription":"Preserved raw Rust source",
            }),
        )
        .unwrap()
        .unwrap();
        let invocation = session
            .create_box(
                "t2",
                "Kennedy tool call",
                BoxOwner::Kennedy,
                BoxContent::text("call"),
            )
            .unwrap();
        let arguments = request
            .capture(
                &mut session,
                "t3",
                invocation,
                "\n//! leading newline\npub fn quote() -> &'static str { \"raw\\\\text\" }".into(),
            )
            .unwrap();

        let exact = "\n//! leading newline\npub fn quote() -> &'static str { \"raw\\\\text\" }\n";
        let state = session.state().box_state(invocation).unwrap();
        assert_eq!(state.canonical.content.text, exact);
        assert_eq!(arguments["contents"], exact);
        assert_eq!(
            state.canonical.content.metadata["toolName"],
            request.write_tool()
        );
        assert!(session.state().render().contains(
            "Kennedy called write-file on src/lib.rs in example-lib, and she describes the update as: Preserved raw Rust source"
        ));
        assert!(!session.state().render().contains("raw\\\\text"));
        std::fs::remove_dir_all(root).unwrap();
    }

    #[test]
    fn captured_contents_enforce_normalized_byte_limit_before_persistence() {
        let (root, mut session) = test_session("capture-limit");
        apply_snapshot(
            &mut session,
            "t1",
            SourceSnapshot {
                kind: ManagedSourceKind::RustLibrary,
                name: "bounded-lib".into(),
                text: "existing source".into(),
            },
        )
        .unwrap();
        let request = prepare_freeform_write(
            &session,
            WRITE_FILE_FREEFORM_RUST_LIB_TOOL,
            &json!({
                "name":"bounded-lib",
                "path":"src/lib.rs",
                "updateDescription":"Bound captured source",
            }),
        )
        .unwrap()
        .unwrap();

        let accepted_invocation = session
            .create_box(
                "t2",
                "Accepted Kennedy tool call",
                BoxOwner::Kennedy,
                BoxContent::text("accepted call"),
            )
            .unwrap();
        let arguments = request
            .capture(
                &mut session,
                "t3",
                accepted_invocation,
                "a".repeat(MAX_CAPTURED_CONTENT_BYTES - 1),
            )
            .unwrap();
        let accepted = arguments["contents"].as_str().unwrap();
        assert_eq!(accepted.len(), MAX_CAPTURED_CONTENT_BYTES);
        assert!(accepted.ends_with('\n'));

        let rejected_invocation = session
            .create_box(
                "t4",
                "Rejected Kennedy tool call",
                BoxOwner::Kennedy,
                BoxContent::text("unchanged call"),
            )
            .unwrap();
        assert!(
            request
                .capture(
                    &mut session,
                    "t5",
                    rejected_invocation,
                    "b".repeat(MAX_CAPTURED_CONTENT_BYTES),
                )
                .is_err()
        );
        assert_eq!(
            session
                .state()
                .box_state(rejected_invocation)
                .unwrap()
                .canonical
                .content
                .text,
            "unchanged call"
        );
        std::fs::remove_dir_all(root).unwrap();
    }

    #[test]
    fn freeform_metadata_is_strict_and_kind_specific() {
        let (root, mut session) = test_session("freeform-validation");
        apply_snapshot(
            &mut session,
            "t1",
            SourceSnapshot {
                kind: ManagedSourceKind::WebLibrary,
                name: "example-ui".into(),
                text: "existing source".into(),
            },
        )
        .unwrap();
        let valid = json!({
            "name":"example-ui",
            "path":"index.js",
            "updateDescription":"Replace the entry module",
        });
        let request = prepare_freeform_write(&session, WRITE_FILE_FREEFORM_WEB_LIB_TOOL, &valid)
            .unwrap()
            .unwrap();
        assert_eq!(request.kind(), ManagedSourceKind::WebLibrary);
        assert_eq!(request.write_tool(), WRITE_FILE_FREEFORM_WEB_LIB_TOOL);

        let mut extra = valid.clone();
        extra["contents"] = json!("not accepted in the Ktool call");
        assert!(
            prepare_freeform_write(&session, WRITE_FILE_FREEFORM_WEB_LIB_TOOL, &extra).is_err()
        );
        let mut multiline_path = valid.clone();
        multiline_path["path"] = json!("index.js\nanother");
        assert!(
            prepare_freeform_write(&session, WRITE_FILE_FREEFORM_WEB_LIB_TOOL, &multiline_path)
                .is_err()
        );
        let mut multiline_description = valid;
        multiline_description["updateDescription"] = json!("line one\nline two");
        assert!(
            prepare_freeform_write(
                &session,
                WRITE_FILE_FREEFORM_WEB_LIB_TOOL,
                &multiline_description,
            )
            .is_err()
        );
        assert!(
            prepare_freeform_write(&session, "unrelated-tool", &Value::Null)
                .unwrap()
                .is_none()
        );
        std::fs::remove_dir_all(root).unwrap();
    }
}