ainl_memory/
mission_query.rs

1//! Mission subgraph query helpers (graph-native mission substrate).
2
3use crate::edge_labels::{HAS_ASSERTION, HAS_FEATURE, HANDED_OFF_BY, PROGRESS_FOR};
4use crate::node::AinlMemoryNode;
5use crate::snapshot::SnapshotEdge;
6use crate::store::GraphStore;
7use ainl_contracts::{Assertion, Feature, Handoff, Mission, MissionState};
8use serde::{Deserialize, Serialize};
9use std::collections::HashSet;
10use uuid::Uuid;
11
12/// Lightweight mission row for orchestrator prompts and API list views.
13#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
14pub struct MissionSummary {
15    pub mission_id: String,
16    pub state: MissionState,
17    pub objective_preview: String,
18    pub milestone_count: usize,
19}
20
21/// Mission node plus connected features, assertions, handoffs, and edge rows.
22#[derive(Debug, Clone)]
23pub struct MissionSubgraph {
24    pub mission: Mission,
25    pub mission_node_id: Uuid,
26    pub features: Vec<Feature>,
27    pub assertions: Vec<Assertion>,
28    pub handoffs: Vec<Handoff>,
29    pub edges: Vec<SnapshotEdge>,
30}
31
32fn node_matches_agent(node: &AinlMemoryNode, agent_id: &str) -> bool {
33    node.agent_id.is_empty() || node.agent_id == agent_id
34}
35
36fn objective_preview(objective_md: &str, max_chars: usize) -> String {
37    let flat: String = objective_md
38        .lines()
39        .map(str::trim)
40        .filter(|l| !l.is_empty())
41        .collect::<Vec<_>>()
42        .join(" ");
43    if flat.chars().count() <= max_chars {
44        flat
45    } else {
46        flat.chars().take(max_chars).collect::<String>() + "…"
47    }
48}
49
50fn find_mission_node(
51    store: &dyn GraphStore,
52    agent_id: &str,
53    mission_id: &str,
54) -> Result<AinlMemoryNode, String> {
55    for node in store.find_by_type("mission")? {
56        if !node_matches_agent(&node, agent_id) {
57            continue;
58        }
59        if node
60            .mission()
61            .is_some_and(|m| m.mission_id.as_str() == mission_id)
62        {
63            return Ok(node);
64        }
65    }
66    Err(format!("mission not found: {mission_id}"))
67}
68
69/// Active missions for this agent (not `Completed` / `Cancelled`).
70pub fn find_active_missions(
71    store: &dyn GraphStore,
72    agent_id: &str,
73) -> Result<Vec<MissionSummary>, String> {
74    let mut out = Vec::new();
75    for node in store.find_by_type("mission")? {
76        if !node_matches_agent(&node, agent_id) {
77            continue;
78        }
79        let Some(mission) = node.mission() else {
80            continue;
81        };
82        if matches!(
83            mission.state,
84            MissionState::Completed | MissionState::Cancelled
85        ) {
86            continue;
87        }
88        out.push(MissionSummary {
89            mission_id: mission.mission_id.as_str().to_string(),
90            state: mission.state,
91            objective_preview: objective_preview(&mission.objective_md, 160),
92            milestone_count: mission.milestone_ids.len(),
93        });
94    }
95    out.sort_by(|a, b| a.mission_id.cmp(&b.mission_id));
96    Ok(out)
97}
98
99/// Load mission record + connected DAG nodes in one sweep.
100pub fn mission_subgraph(
101    store: &dyn GraphStore,
102    agent_id: &str,
103    mission_id: &str,
104) -> Result<MissionSubgraph, String> {
105    let mission_node = find_mission_node(store, agent_id, mission_id)?;
106    let mission = mission_node
107        .mission()
108        .cloned()
109        .ok_or_else(|| format!("invalid mission node payload for {mission_id}"))?;
110
111    let feature_nodes = store.walk_edges(mission_node.id, HAS_FEATURE)?;
112    let assertion_nodes = store.walk_edges(mission_node.id, HAS_ASSERTION)?;
113
114    let features = feature_nodes
115        .iter()
116        .filter_map(|n| n.feature().cloned())
117        .collect::<Vec<_>>();
118    let assertions = assertion_nodes
119        .iter()
120        .filter_map(|n| n.assertion().cloned())
121        .collect::<Vec<_>>();
122
123    let mut handoffs = Vec::new();
124    let mut handoff_ids = HashSet::new();
125    for feat_node in &feature_nodes {
126        for handoff_node in store.walk_edges(feat_node.id, HANDED_OFF_BY)? {
127            if handoff_ids.insert(handoff_node.id) {
128                if let Some(h) = handoff_node.handoff() {
129                    handoffs.push(h.clone());
130                }
131            }
132        }
133    }
134
135    Ok(MissionSubgraph {
136        mission,
137        mission_node_id: mission_node.id,
138        features,
139        assertions,
140        handoffs,
141        edges: Vec::new(),
142    })
143}
144
145/// Count `Episode --PROGRESS_FOR--> Feature` edges for features in this mission since `since_ts`.
146pub fn count_progress_events_since(
147    store: &dyn GraphStore,
148    agent_id: &str,
149    mission_id: &str,
150    since_ts: i64,
151) -> Result<usize, String> {
152    let subgraph = mission_subgraph(store, agent_id, mission_id)?;
153    if subgraph.features.is_empty() {
154        return Ok(0);
155    }
156    let feature_ids: HashSet<String> = subgraph
157        .features
158        .iter()
159        .map(|f| f.feature_id.as_str().to_string())
160        .collect();
161
162    let mut count = 0usize;
163    for node in store.query_episodes_since(since_ts, 10_000)? {
164        for target in store.walk_edges(node.id, PROGRESS_FOR)? {
165            if target
166                .feature()
167                .is_some_and(|f| feature_ids.contains(f.feature_id.as_str()))
168            {
169                count += 1;
170            }
171        }
172    }
173    Ok(count)
174}
175
176/// Features filtered by milestone label within a mission subgraph.
177pub fn find_features_by_milestone(
178    store: &dyn GraphStore,
179    agent_id: &str,
180    mission_id: &str,
181    milestone: &str,
182) -> Result<Vec<Feature>, String> {
183    Ok(mission_subgraph(store, agent_id, mission_id)?
184        .features
185        .into_iter()
186        .filter(|f| f.milestone.as_deref() == Some(milestone))
187        .collect())
188}
189
190/// Assertions filtered by milestone label within a mission subgraph.
191pub fn find_assertions_for_milestone(
192    store: &dyn GraphStore,
193    agent_id: &str,
194    mission_id: &str,
195    milestone: &str,
196) -> Result<Vec<Assertion>, String> {
197    Ok(mission_subgraph(store, agent_id, mission_id)?
198        .assertions
199        .into_iter()
200        .filter(|a| a.milestone.as_deref() == Some(milestone))
201        .collect())
202}
203
204#[cfg(test)]
205mod tests {
206    use super::*;
207    use crate::node::{AinlMemoryNode, AinlNodeType, MemoryCategory};
208    use crate::GraphMemory;
209    use ainl_contracts::{FeatureId, MissionCapabilityFlags, MissionId};
210    use chrono::Utc;
211
212    fn write_mission(memory: &GraphMemory, agent_id: &str, mid: &str) -> Uuid {
213        let mission = Mission {
214            mission_id: MissionId(mid.into()),
215            objective_md: "Build the mission substrate".into(),
216            state: MissionState::Running,
217            milestone_ids: vec!["m1".into()],
218            mission_root: None,
219            created_at: Utc::now(),
220            last_orchestrator_turn_at: None,
221            capability_flags: MissionCapabilityFlags::default(),
222        };
223        let mut node = AinlMemoryNode {
224            id: Uuid::new_v4(),
225            memory_category: MemoryCategory::Mission,
226            importance_score: 0.5,
227            agent_id: agent_id.to_string(),
228            project_id: None,
229            node_type: AinlNodeType::Mission { mission },
230            edges: Vec::new(),
231            plugin_data: None,
232        };
233        let id = node.id;
234        memory.write_node(&node).expect("write mission");
235        id
236    }
237
238    fn write_feature(memory: &GraphMemory, agent_id: &str, fid: &str, milestone: Option<&str>) -> AinlMemoryNode {
239        let feature = Feature {
240            feature_id: FeatureId(fid.into()),
241            description: format!("feature {fid}"),
242            status: ainl_contracts::FeatureStatus::Pending,
243            milestone: milestone.map(str::to_string),
244            skill_name: None,
245            touches_files: vec!["src/lib.rs".into()],
246            preconditions: vec![],
247            expected_behavior: vec![],
248            verification_steps: vec![],
249            fulfills: vec![],
250            snapshot: None,
251        };
252        let node = AinlMemoryNode {
253            id: Uuid::new_v4(),
254            memory_category: MemoryCategory::Feature,
255            importance_score: 0.5,
256            agent_id: agent_id.to_string(),
257            project_id: None,
258            node_type: AinlNodeType::Feature { feature },
259            edges: Vec::new(),
260            plugin_data: None,
261        };
262        memory.write_node(&node).expect("write feature");
263        node
264    }
265
266    #[test]
267    fn mission_subgraph_and_active_missions_roundtrip() {
268        let dir = tempfile::tempdir().expect("tempdir");
269        let db = dir.path().join("mission_query.db");
270        let memory = GraphMemory::new(&db).expect("graph memory");
271        let agent = "agent-mission-q";
272
273        let mission_id = write_mission(&memory, agent, "mission-1");
274        let feat = write_feature(&memory, agent, "f1", Some("m1"));
275        let mut mission_node = memory
276            .store()
277            .read_node(mission_id)
278            .expect("read")
279            .expect("node");
280        mission_node.add_edge(feat.id, HAS_FEATURE);
281        memory.write_node(&mission_node).expect("edge mission->feature");
282
283        let store = memory.store();
284        let active = find_active_missions(store, agent).expect("active");
285        assert_eq!(active.len(), 1);
286        assert_eq!(active[0].mission_id, "mission-1");
287
288        let subgraph = mission_subgraph(store, agent, "mission-1").expect("subgraph");
289        assert_eq!(subgraph.features.len(), 1);
290        assert_eq!(subgraph.features[0].feature_id.as_str(), "f1");
291
292        let by_m = find_features_by_milestone(store, agent, "mission-1", "m1").expect("by milestone");
293        assert_eq!(by_m.len(), 1);
294    }
295
296    #[test]
297    fn mission_migration_idempotent_and_fts_indexed() {
298        let dir = tempfile::tempdir().expect("tempdir");
299        let db = dir.path().join("mission_migrate.db");
300        let memory = GraphMemory::new(&db).expect("open v1");
301        let mission_node_id = write_mission(&memory, "agent-a", "m-fts");
302        let _feat = write_feature(&memory, "agent-a", "f-fts", Some("m1"));
303        drop(memory);
304
305        let memory2 = GraphMemory::new(&db).expect("reopen migrates idempotently");
306        let hits = memory2
307            .search_all_nodes_fts("agent-a", "substrate", None, 10)
308            .expect("fts");
309        assert!(
310            hits.iter().any(|n| n.id == mission_node_id),
311            "mission nodes should be FTS-indexed"
312        );
313        let _ = GraphMemory::new(&db).expect("third open is idempotent");
314    }
315
316    #[test]
317    fn mission_subgraph_query_under_budget_on_fixture() {
318        use std::time::Instant;
319        let dir = tempfile::tempdir().expect("tempdir");
320        let db = dir.path().join("mission_perf.db");
321        let memory = GraphMemory::new(&db).expect("open");
322        let agent = "perf-agent";
323        let _mid = write_mission(&memory, agent, "perf-m1");
324        for i in 0..50 {
325            write_feature(&memory, agent, &format!("f{i}"), Some("m1"));
326        }
327        let store = memory.store();
328        let start = Instant::now();
329        for _ in 0..20 {
330            let _ = mission_subgraph(store, agent, "perf-m1").expect("subgraph");
331        }
332        let elapsed_ms = start.elapsed().as_millis();
333        assert!(
334            elapsed_ms < 500,
335            "20 mission_subgraph queries took {elapsed_ms}ms (budget 500ms on ~50-feature fixture)"
336        );
337    }
338}
ainl_memory/mission_query.rs

ainl_memory/
mission_query.rs
