kcode-kennedy-subagent-context 0.1.2

//! Box-free virtual Ktool state for one Kennedy subagent.

#![forbid(unsafe_code)]

use std::collections::{BTreeMap, BTreeSet};

use kcode_agent_runtime::StateUpdate;
use kcode_dev_tools::{ManagedSourceKind, SourceSnapshot};
use kcode_intelligence_router::AgentProvider;
use kcode_kweb_context::{Context as KwebContext, NodeDraft, ProjectionItem, StagedCreate};

const KWEB_PREFIX: &str = "kweb:";
const MANAGED_PREFIX: &str = "managed:";

/// Current state changes produced by one successful child operation.
#[derive(Clone, Debug, Default, Eq, PartialEq)]
pub struct StateChanges {
    /// Stable replace-or-remove updates consumed by `kcode-agent-runtime`.
    pub updates: Vec<StateUpdate>,
}

impl StateChanges {
    /// Whether no current value changed.
    pub fn is_empty(&self) -> bool {
        self.updates.is_empty()
    }

    /// Complete changed values, in update order, for a result that displays them.
    pub fn display_text(&self) -> String {
        self.updates
            .iter()
            .filter_map(|update| update.text.as_deref())
            .collect::<Vec<_>>()
            .join("\n\n")
    }

    /// Keys for every complete current value returned by [`Self::display_text`].
    pub fn displayed_state_keys(&self) -> Vec<String> {
        self.updates
            .iter()
            .filter(|update| update.text.is_some())
            .map(|update| update.key.clone())
            .collect()
    }
}

/// One applied managed-source snapshot and its stable virtual identity.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct SourceState {
    /// Stable identity for this source kind and project.
    pub key: String,
    /// Complete current state rendered after historical tool results.
    pub text: String,
    /// Replacement update when the state changed.
    pub update: Option<StateUpdate>,
}

/// Ephemeral current state owned exclusively by one subagent run.
#[derive(Clone, Debug)]
pub struct Context {
    initial_sections: Vec<String>,
    kweb: KwebContext,
    states: BTreeMap<String, String>,
    kweb_keys: BTreeSet<String>,
    staged_node_ids: BTreeSet<String>,
}

impl Context {
    /// Creates a child context with provider-appropriate immutable opening sections.
    pub fn new(
        root_node_ids: Vec<String>,
        load_fixed_connections: bool,
        provider: AgentProvider,
        codex_harness_prompt: String,
        selected_node_descriptions: Vec<String>,
    ) -> anyhow::Result<Self> {
        let mut initial_sections = Vec::with_capacity(selected_node_descriptions.len() + 1);
        if provider == AgentProvider::Codex {
            initial_sections.push(codex_harness_prompt);
        }
        initial_sections.extend(selected_node_descriptions);
        Ok(Self {
            initial_sections,
            kweb: KwebContext::with_fixed_connections(root_node_ids, load_fixed_connections)
                .map_err(anyhow::Error::new)?,
            states: BTreeMap::new(),
            kweb_keys: BTreeSet::new(),
            staged_node_ids: BTreeSet::new(),
        })
    }

    /// Immutable provider-specific prompt and selected node descriptions, in model order.
    pub fn initial_sections(&self) -> &[String] {
        &self.initial_sections
    }

    /// The independent Kweb state used by child `LoadNodes` and Kmap tools.
    pub fn kweb(&self) -> &KwebContext {
        &self.kweb
    }

    /// Mutable access for an atomic child Kweb load.
    pub fn kweb_mut(&mut self) -> &mut KwebContext {
        &mut self.kweb
    }

    /// Makes explicitly referenced pending nodes visible to this child.
    pub fn include_staged_nodes(&mut self, ids: impl IntoIterator<Item = String>) {
        self.staged_node_ids.extend(ids);
    }

    /// Reconciles the child's complete Kweb state without touching Chatend.
    pub fn reconcile_kweb(
        &mut self,
        updates: &BTreeMap<String, NodeDraft>,
        creates: &[StagedCreate],
    ) -> anyhow::Result<StateChanges> {
        let visible_updates = updates
            .iter()
            .filter(|(id, _)| self.kweb.contains_full_node(id))
            .map(|(id, update)| (id.clone(), update.clone()))
            .collect();
        let visible_creates = creates
            .iter()
            .filter(|create| self.staged_node_ids.contains(&create.pending_id))
            .cloned()
            .collect::<Vec<_>>();
        let projected = self
            .kweb
            .projection(&visible_updates, &visible_creates)
            .map_err(anyhow::Error::new)?;
        let next_keys = projected
            .iter()
            .map(|item| kweb_key(&item.key))
            .collect::<BTreeSet<_>>();
        let mut changes = StateChanges::default();

        for key in self.kweb_keys.difference(&next_keys) {
            self.states.remove(key);
            changes.updates.push(StateUpdate {
                key: key.clone(),
                text: None,
            });
        }
        for item in projected {
            self.apply_kweb_item(item, &mut changes);
        }
        self.kweb_keys = next_keys;
        Ok(changes)
    }

    /// Installs one current managed-source value in the child context.
    pub fn apply_source_snapshot(&mut self, snapshot: SourceSnapshot) -> SourceState {
        let state = source_state(&snapshot);
        let update = (self.states.get(&state.key) != Some(&state.text)).then(|| StateUpdate {
            key: state.key.clone(),
            text: Some(state.text.clone()),
        });
        self.states.insert(state.key.clone(), state.text.clone());
        SourceState { update, ..state }
    }

    /// Renders a prospective managed-source value without applying it.
    pub fn source_state(&self, snapshot: &SourceSnapshot) -> SourceState {
        let mut state = source_state(snapshot);
        state.update = (self.states.get(&state.key) != Some(&state.text)).then(|| StateUpdate {
            key: state.key.clone(),
            text: Some(state.text.clone()),
        });
        state
    }

    /// Whether this exact source kind and project is open in the child.
    pub fn source_is_open(&self, kind: ManagedSourceKind, name: &str) -> bool {
        self.states.contains_key(&managed_key(kind, name))
    }

    fn apply_kweb_item(&mut self, item: ProjectionItem, changes: &mut StateChanges) {
        let key = kweb_key(&item.key);
        let text = format!("Current {}:\n{}", item.name, item.text);
        if self.states.get(&key) != Some(&text) {
            self.states.insert(key.clone(), text.clone());
            changes.updates.push(StateUpdate {
                key,
                text: Some(text),
            });
        }
    }
}

fn source_state(snapshot: &SourceSnapshot) -> SourceState {
    SourceState {
        key: managed_key(snapshot.kind, &snapshot.name),
        text: format!(
            "Current Managed {} {}:\n{}",
            snapshot.kind.label(),
            snapshot.name,
            snapshot.text
        ),
        update: None,
    }
}

fn kweb_key(logical: &str) -> String {
    format!("{KWEB_PREFIX}{logical}")
}

fn managed_key(kind: ManagedSourceKind, name: &str) -> String {
    let kind = match kind {
        ManagedSourceKind::RustLibrary => "rust-library",
        ManagedSourceKind::WebLibrary => "web-library",
        ManagedSourceKind::RustBinary => "rust-binary",
    };
    format!("{MANAGED_PREFIX}{kind}:{name}")
}

#[cfg(test)]
mod tests {
    use super::*;
    use kcode_kweb_context::Node;

    fn node(description: &str) -> Node {
        Node {
            id: "AAAAAAAE".into(),
            short_name: "Root".into(),
            short_description: "Summary".into(),
            long_description: description.into(),
            owner: "self".into(),
            fixed_connections: Vec::new(),
            recent_connections: Vec::new(),
            objects: Vec::new(),
            last_modified_by: "test".into(),
            last_modified_at: None,
        }
    }

    fn empty_context() -> Context {
        Context::new(
            vec!["AAAAAAAE".into()],
            false,
            AgentProvider::OpenAi,
            "Codex only".into(),
            Vec::new(),
        )
        .unwrap()
    }

    #[test]
    fn codex_is_the_only_provider_that_receives_the_harness_prompt() {
        for (provider, expected) in [
            (
                AgentProvider::Codex,
                vec!["Codex only", "first node", "second node"],
            ),
            (AgentProvider::OpenAi, vec!["first node", "second node"]),
            (AgentProvider::Gemini, vec!["first node", "second node"]),
        ] {
            let context = Context::new(
                vec!["AAAAAAAE".into()],
                false,
                provider,
                "Codex only".into(),
                vec!["first node".into(), "second node".into()],
            )
            .unwrap();
            assert_eq!(context.initial_sections(), expected);
        }
    }

    #[test]
    fn kweb_changes_reuse_stable_keys_without_boxes() {
        let mut context = empty_context();
        context
            .kweb_mut()
            .apply_load(node("old"), Vec::new())
            .unwrap();
        let first = context.reconcile_kweb(&BTreeMap::new(), &[]).unwrap();
        assert_eq!(first.updates[0].key, "kweb:AAAAAAAE");
        assert!(first.display_text().contains("old"));

        context.kweb_mut().refresh([node("new")]).unwrap();
        let second = context.reconcile_kweb(&BTreeMap::new(), &[]).unwrap();
        assert_eq!(second.updates.len(), 1);
        assert_eq!(second.updates[0].key, "kweb:AAAAAAAE");
        assert!(second.updates[0].text.as_deref().unwrap().contains("new"));
    }

    #[test]
    fn managed_sources_are_isolated_by_kind_and_name() {
        let mut context = empty_context();
        let snapshot = SourceSnapshot {
            kind: ManagedSourceKind::RustLibrary,
            name: "example".into(),
            text: "source".into(),
        };
        let applied = context.apply_source_snapshot(snapshot.clone());
        assert_eq!(applied.key, "managed:rust-library:example");
        assert!(applied.update.is_some());
        assert!(context.source_is_open(snapshot.kind, &snapshot.name));
        assert!(context.apply_source_snapshot(snapshot).update.is_none());
        assert!(!context.source_is_open(ManagedSourceKind::WebLibrary, "example"));
    }

    #[test]
    fn kweb_projection_hides_unloaded_parent_plan_state() {
        let mut context = empty_context();
        context
            .kweb_mut()
            .apply_load(node("loaded"), Vec::new())
            .unwrap();
        let unrelated_id = "AAAAAAAI".to_owned();
        let updates = BTreeMap::from([(
            unrelated_id,
            NodeDraft {
                short_name: "Hidden".into(),
                short_description: "Hidden".into(),
                long_description: "unrelated parent update".into(),
                owner: "self".into(),
                fixed_connections: Vec::new(),
                recent_connections: Vec::new(),
                objects: Vec::new(),
            },
        )]);
        let hidden = StagedCreate {
            pending_id: "pending:hidden".into(),
            data: updates.values().next().unwrap().clone(),
        };

        let changes = context
            .reconcile_kweb(&updates, std::slice::from_ref(&hidden))
            .unwrap();
        assert!(!changes.display_text().contains("unrelated parent update"));
        assert!(!changes.display_text().contains("pending:hidden"));

        context.include_staged_nodes([hidden.pending_id.clone()]);
        let changes = context.reconcile_kweb(&updates, &[hidden]).unwrap();
        assert!(changes.display_text().contains("pending:hidden"));
    }
}