kcode_kennedy_subagent_context/
lib.rs

1//! Box-free virtual Ktool state for one Kennedy subagent.
2
3#![forbid(unsafe_code)]
4
5use std::collections::{BTreeMap, BTreeSet};
6
7use kcode_agent_runtime::StateUpdate;
8use kcode_dev_tools::{ManagedSourceKind, SourceSnapshot};
9use kcode_intelligence_router::AgentProvider;
10use kcode_kweb_context::{Context as KwebContext, NodeDraft, ProjectionItem, StagedCreate};
11
12const KWEB_PREFIX: &str = "kweb:";
13const MANAGED_PREFIX: &str = "managed:";
14
15/// Current state changes produced by one successful child operation.
16#[derive(Clone, Debug, Default, Eq, PartialEq)]
17pub struct StateChanges {
18    /// Stable replace-or-remove updates consumed by `kcode-agent-runtime`.
19    pub updates: Vec<StateUpdate>,
20}
21
22impl StateChanges {
23    /// Whether no current value changed.
24    pub fn is_empty(&self) -> bool {
25        self.updates.is_empty()
26    }
27
28    /// Complete changed values, in update order, for a result that displays them.
29    pub fn display_text(&self) -> String {
30        self.updates
31            .iter()
32            .filter_map(|update| update.text.as_deref())
33            .collect::<Vec<_>>()
34            .join("\n\n")
35    }
36
37    /// Keys for every complete current value returned by [`Self::display_text`].
38    pub fn displayed_state_keys(&self) -> Vec<String> {
39        self.updates
40            .iter()
41            .filter(|update| update.text.is_some())
42            .map(|update| update.key.clone())
43            .collect()
44    }
45}
46
47/// One applied managed-source snapshot and its stable virtual identity.
48#[derive(Clone, Debug, Eq, PartialEq)]
49pub struct SourceState {
50    /// Stable identity for this source kind and project.
51    pub key: String,
52    /// Complete current state rendered after historical tool results.
53    pub text: String,
54    /// Replacement update when the state changed.
55    pub update: Option<StateUpdate>,
56}
57
58/// Ephemeral current state owned exclusively by one subagent run.
59#[derive(Clone, Debug)]
60pub struct Context {
61    initial_sections: Vec<String>,
62    kweb: KwebContext,
63    states: BTreeMap<String, String>,
64    kweb_keys: BTreeSet<String>,
65    staged_node_ids: BTreeSet<String>,
66}
67
68impl Context {
69    /// Creates a child context with provider-appropriate immutable opening sections.
70    pub fn new(
71        root_node_ids: Vec<String>,
72        load_fixed_connections: bool,
73        provider: AgentProvider,
74        codex_harness_prompt: String,
75        selected_node_descriptions: Vec<String>,
76    ) -> anyhow::Result<Self> {
77        let mut initial_sections = Vec::with_capacity(selected_node_descriptions.len() + 1);
78        if provider == AgentProvider::Codex {
79            initial_sections.push(codex_harness_prompt);
80        }
81        initial_sections.extend(selected_node_descriptions);
82        Ok(Self {
83            initial_sections,
84            kweb: KwebContext::with_fixed_connections(root_node_ids, load_fixed_connections)
85                .map_err(anyhow::Error::new)?,
86            states: BTreeMap::new(),
87            kweb_keys: BTreeSet::new(),
88            staged_node_ids: BTreeSet::new(),
89        })
90    }
91
92    /// Immutable provider-specific prompt and selected node descriptions, in model order.
93    pub fn initial_sections(&self) -> &[String] {
94        &self.initial_sections
95    }
96
97    /// The independent Kweb state used by child `LoadNodes` and Kmap tools.
98    pub fn kweb(&self) -> &KwebContext {
99        &self.kweb
100    }
101
102    /// Mutable access for an atomic child Kweb load.
103    pub fn kweb_mut(&mut self) -> &mut KwebContext {
104        &mut self.kweb
105    }
106
107    /// Makes explicitly referenced pending nodes visible to this child.
108    pub fn include_staged_nodes(&mut self, ids: impl IntoIterator<Item = String>) {
109        self.staged_node_ids.extend(ids);
110    }
111
112    /// Reconciles the child's complete Kweb state without touching Chatend.
113    pub fn reconcile_kweb(
114        &mut self,
115        updates: &BTreeMap<String, NodeDraft>,
116        creates: &[StagedCreate],
117    ) -> anyhow::Result<StateChanges> {
118        let visible_updates = updates
119            .iter()
120            .filter(|(id, _)| self.kweb.contains_full_node(id))
121            .map(|(id, update)| (id.clone(), update.clone()))
122            .collect();
123        let visible_creates = creates
124            .iter()
125            .filter(|create| self.staged_node_ids.contains(&create.pending_id))
126            .cloned()
127            .collect::<Vec<_>>();
128        let projected = self
129            .kweb
130            .projection(&visible_updates, &visible_creates)
131            .map_err(anyhow::Error::new)?;
132        let next_keys = projected
133            .iter()
134            .map(|item| kweb_key(&item.key))
135            .collect::<BTreeSet<_>>();
136        let mut changes = StateChanges::default();
137
138        for key in self.kweb_keys.difference(&next_keys) {
139            self.states.remove(key);
140            changes.updates.push(StateUpdate {
141                key: key.clone(),
142                text: None,
143            });
144        }
145        for item in projected {
146            self.apply_kweb_item(item, &mut changes);
147        }
148        self.kweb_keys = next_keys;
149        Ok(changes)
150    }
151
152    /// Installs one current managed-source value in the child context.
153    pub fn apply_source_snapshot(&mut self, snapshot: SourceSnapshot) -> SourceState {
154        let state = source_state(&snapshot);
155        let update = (self.states.get(&state.key) != Some(&state.text)).then(|| StateUpdate {
156            key: state.key.clone(),
157            text: Some(state.text.clone()),
158        });
159        self.states.insert(state.key.clone(), state.text.clone());
160        SourceState { update, ..state }
161    }
162
163    /// Renders a prospective managed-source value without applying it.
164    pub fn source_state(&self, snapshot: &SourceSnapshot) -> SourceState {
165        let mut state = source_state(snapshot);
166        state.update = (self.states.get(&state.key) != Some(&state.text)).then(|| StateUpdate {
167            key: state.key.clone(),
168            text: Some(state.text.clone()),
169        });
170        state
171    }
172
173    /// Whether this exact source kind and project is open in the child.
174    pub fn source_is_open(&self, kind: ManagedSourceKind, name: &str) -> bool {
175        self.states.contains_key(&managed_key(kind, name))
176    }
177
178    fn apply_kweb_item(&mut self, item: ProjectionItem, changes: &mut StateChanges) {
179        let key = kweb_key(&item.key);
180        let text = format!("Current {}:\n{}", item.name, item.text);
181        if self.states.get(&key) != Some(&text) {
182            self.states.insert(key.clone(), text.clone());
183            changes.updates.push(StateUpdate {
184                key,
185                text: Some(text),
186            });
187        }
188    }
189}
190
191fn source_state(snapshot: &SourceSnapshot) -> SourceState {
192    SourceState {
193        key: managed_key(snapshot.kind, &snapshot.name),
194        text: format!(
195            "Current Managed {} {}:\n{}",
196            snapshot.kind.label(),
197            snapshot.name,
198            snapshot.text
199        ),
200        update: None,
201    }
202}
203
204fn kweb_key(logical: &str) -> String {
205    format!("{KWEB_PREFIX}{logical}")
206}
207
208fn managed_key(kind: ManagedSourceKind, name: &str) -> String {
209    let kind = match kind {
210        ManagedSourceKind::RustLibrary => "rust-library",
211        ManagedSourceKind::WebLibrary => "web-library",
212        ManagedSourceKind::RustBinary => "rust-binary",
213    };
214    format!("{MANAGED_PREFIX}{kind}:{name}")
215}
216
217#[cfg(test)]
218mod tests {
219    use super::*;
220    use kcode_kweb_context::Node;
221
222    fn node(description: &str) -> Node {
223        Node {
224            id: "AAAAAAAE".into(),
225            short_name: "Root".into(),
226            short_description: "Summary".into(),
227            long_description: description.into(),
228            owner: "self".into(),
229            fixed_connections: Vec::new(),
230            recent_connections: Vec::new(),
231            objects: Vec::new(),
232            last_modified_by: "test".into(),
233            last_modified_at: None,
234        }
235    }
236
237    fn empty_context() -> Context {
238        Context::new(
239            vec!["AAAAAAAE".into()],
240            false,
241            AgentProvider::OpenAi,
242            "Codex only".into(),
243            Vec::new(),
244        )
245        .unwrap()
246    }
247
248    #[test]
249    fn codex_is_the_only_provider_that_receives_the_harness_prompt() {
250        for (provider, expected) in [
251            (
252                AgentProvider::Codex,
253                vec!["Codex only", "first node", "second node"],
254            ),
255            (AgentProvider::OpenAi, vec!["first node", "second node"]),
256            (AgentProvider::Gemini, vec!["first node", "second node"]),
257        ] {
258            let context = Context::new(
259                vec!["AAAAAAAE".into()],
260                false,
261                provider,
262                "Codex only".into(),
263                vec!["first node".into(), "second node".into()],
264            )
265            .unwrap();
266            assert_eq!(context.initial_sections(), expected);
267        }
268    }
269
270    #[test]
271    fn kweb_changes_reuse_stable_keys_without_boxes() {
272        let mut context = empty_context();
273        context
274            .kweb_mut()
275            .apply_load(node("old"), Vec::new())
276            .unwrap();
277        let first = context.reconcile_kweb(&BTreeMap::new(), &[]).unwrap();
278        assert_eq!(first.updates[0].key, "kweb:AAAAAAAE");
279        assert!(first.display_text().contains("old"));
280
281        context.kweb_mut().refresh([node("new")]).unwrap();
282        let second = context.reconcile_kweb(&BTreeMap::new(), &[]).unwrap();
283        assert_eq!(second.updates.len(), 1);
284        assert_eq!(second.updates[0].key, "kweb:AAAAAAAE");
285        assert!(second.updates[0].text.as_deref().unwrap().contains("new"));
286    }
287
288    #[test]
289    fn managed_sources_are_isolated_by_kind_and_name() {
290        let mut context = empty_context();
291        let snapshot = SourceSnapshot {
292            kind: ManagedSourceKind::RustLibrary,
293            name: "example".into(),
294            text: "source".into(),
295        };
296        let applied = context.apply_source_snapshot(snapshot.clone());
297        assert_eq!(applied.key, "managed:rust-library:example");
298        assert!(applied.update.is_some());
299        assert!(context.source_is_open(snapshot.kind, &snapshot.name));
300        assert!(context.apply_source_snapshot(snapshot).update.is_none());
301        assert!(!context.source_is_open(ManagedSourceKind::WebLibrary, "example"));
302    }
303
304    #[test]
305    fn kweb_projection_hides_unloaded_parent_plan_state() {
306        let mut context = empty_context();
307        context
308            .kweb_mut()
309            .apply_load(node("loaded"), Vec::new())
310            .unwrap();
311        let unrelated_id = "AAAAAAAI".to_owned();
312        let updates = BTreeMap::from([(
313            unrelated_id,
314            NodeDraft {
315                short_name: "Hidden".into(),
316                short_description: "Hidden".into(),
317                long_description: "unrelated parent update".into(),
318                owner: "self".into(),
319                fixed_connections: Vec::new(),
320                recent_connections: Vec::new(),
321                objects: Vec::new(),
322            },
323        )]);
324        let hidden = StagedCreate {
325            pending_id: "pending:hidden".into(),
326            data: updates.values().next().unwrap().clone(),
327        };
328
329        let changes = context
330            .reconcile_kweb(&updates, std::slice::from_ref(&hidden))
331            .unwrap();
332        assert!(!changes.display_text().contains("unrelated parent update"));
333        assert!(!changes.display_text().contains("pending:hidden"));
334
335        context.include_staged_nodes([hidden.pending_id.clone()]);
336        let changes = context.reconcile_kweb(&updates, &[hidden]).unwrap();
337        assert!(changes.display_text().contains("pending:hidden"));
338    }
339}
kcode_kennedy_subagent_context/lib.rs

kcode_kennedy_subagent_context/
lib.rs
