Skip to main content

ainl_memory/
lib.rs

1//! AINL Memory - Graph-based agent memory substrate
2//!
3//! **Graph-as-memory for AI agents. Execution IS the memory.**
4//!
5//! AINL Memory implements agent memory as an execution graph. Every agent turn,
6//! tool call, and delegation becomes a typed graph node. No separate retrieval
7//! layer—the graph itself is the memory.
8//!
9//! # Quick Start
10//!
11//! ```no_run
12//! use ainl_memory::GraphMemory;
13//! use std::path::Path;
14//!
15//! let memory = GraphMemory::new(Path::new("memory.db")).unwrap();
16//!
17//! // Record an episode
18//! memory.write_episode(
19//!     vec!["file_read".to_string(), "agent_delegate".to_string()],
20//!     Some("agent-B".to_string()),
21//!     None,
22//! ).unwrap();
23//!
24//! // Recall recent episodes
25//! let recent = memory.recall_recent(100).unwrap();
26//! ```
27//!
28//! # Architecture
29//!
30//! AINL Memory is designed as infrastructure that any agent framework can adopt:
31//! - Zero dependencies on specific agent runtimes
32//! - Simple trait-based API via `GraphStore`
33//! - Bring your own storage backend
34//!
35//! ## Graph store: query, export, validation (since 0.1.4-alpha)
36//!
37//! - **[`SqliteGraphStore`]**: SQLite backend with **`PRAGMA foreign_keys = ON`**, `FOREIGN KEY` constraints
38//!   on `ainl_graph_edges`, one-time migration for legacy DBs (see [CHANGELOG.md](../CHANGELOG.md)).
39//! - **[`GraphQuery`]**: `store.query(agent_id)` — agent-scoped SQL helpers (episodes, lineage, tags, …).
40//! - **Snapshots**: [`AgentGraphSnapshot`], [`SnapshotEdge`], [`SNAPSHOT_SCHEMA_VERSION`];
41//!   [`SqliteGraphStore::export_graph`] / [`SqliteGraphStore::import_graph`] (strict vs repair via
42//!   `allow_dangling_edges`).
43//! - **Validation**: [`GraphValidationReport`], [`DanglingEdgeDetail`]; [`SqliteGraphStore::validate_graph`]
44//!   for agent-scoped semantics beyond raw FK enforcement.
45//! - **[`GraphMemory`]** forwards the above where hosts should not reach past the facade (see impl block).
46//!
47//! ## Node Types
48//!
49//! - **Episode**: What happened during an agent turn (tool calls, delegations)
50//! - **Semantic**: Facts learned with confidence scores
51//! - **Procedural**: Reusable compiled workflow patterns
52//! - **Persona**: Agent traits learned over time
53//! - **Runtime state** (`RuntimeStateNode`, `node_type = runtime_state`): Optional persisted session
54//!   counters and persona snapshot JSON for **ainl-runtime** (see [`GraphMemory::read_runtime_state`] /
55//!   [`GraphMemory::write_runtime_state`]).
56//! - **Trajectory** (`TrajectoryNode`): execution traces for replay / learning.
57//! - **Failure** (`FailureNode`): typed failures (e.g. loop guard) with optional FTS search
58//!   ([`GraphMemory::search_failures_for_agent`]).
59
60pub mod anchored_summary;
61pub mod edge_labels;
62pub mod mission_query;
63pub mod node;
64pub mod pattern_promotion;
65pub mod query;
66pub mod snapshot;
67pub mod store;
68mod trajectory_persist;
69pub mod trajectory_table;
70
71pub use anchored_summary::{anchored_summary_id, ANCHORED_SUMMARY_TAG};
72
73pub use trajectory_persist::{
74    persist_trajectory_coarse_tools, persist_trajectory_for_episode, trajectory_env_enabled,
75};
76
77pub use node::{
78    AinlEdge, AinlMemoryNode, AinlNodeKind, AinlNodeType, EpisodicNode, FailureNode,
79    MemoryCategory, PersonaLayer, PersonaNode, PersonaSource, ProceduralNode, ProcedureType,
80    RuntimeStateNode, SemanticNode, Sentiment, StrengthEvent, TrajectoryNode,
81};
82pub use mission_query::{MissionSubgraph, MissionSummary};
83pub use query::{
84    count_by_topic_cluster, find_high_confidence_facts, find_patterns, find_strong_traits,
85    recall_by_procedure_type, recall_by_topic_cluster, recall_contradictions,
86    recall_delta_by_relevance, recall_episodes_by_conversation, recall_episodes_with_signal,
87    recall_flagged_episodes, recall_low_success_procedures, recall_recent, recall_strength_history,
88    recall_task_scoped_episodes, walk_from, GraphQuery,
89};
90pub use snapshot::{
91    AgentGraphSnapshot, DanglingEdgeDetail, GraphValidationReport, SnapshotEdge,
92    SNAPSHOT_SCHEMA_VERSION,
93};
94pub use store::{GraphStore, GraphValidationError, SnapshotImportError, SqliteGraphStore};
95pub use trajectory_table::TrajectoryDetailRecord;
96
97use ainl_contracts::{
98    ProcedureArtifact, ProcedureLifecycle, ProcedureReuseOutcome, ProcedureStepKind,
99    TrajectoryOutcome,
100};
101use uuid::Uuid;
102
103/// High-level graph memory API - the main entry point for AINL memory.
104///
105/// Wraps a GraphStore implementation with a simplified 5-method API.
106pub struct GraphMemory {
107    store: SqliteGraphStore,
108}
109
110fn score_procedure_artifact(
111    artifact: &ProcedureArtifact,
112    intent: &str,
113    available_tools: &[String],
114) -> f32 {
115    let haystack = format!(
116        "{} {} {}",
117        artifact.title.to_ascii_lowercase(),
118        artifact.intent.to_ascii_lowercase(),
119        artifact.summary.to_ascii_lowercase()
120    );
121    let tokens = intent
122        .split(|c: char| !c.is_ascii_alphanumeric())
123        .filter(|token| token.len() >= 3)
124        .map(str::to_ascii_lowercase)
125        .collect::<Vec<_>>();
126    let intent_score = if tokens.is_empty() {
127        0.0
128    } else {
129        tokens
130            .iter()
131            .filter(|token| haystack.contains(token.as_str()))
132            .count() as f32
133            / tokens.len() as f32
134    };
135    let tool_score = if artifact.required_tools.is_empty() {
136        0.2
137    } else {
138        artifact
139            .required_tools
140            .iter()
141            .filter(|tool| available_tools.iter().any(|available| available == *tool))
142            .count() as f32
143            / artifact.required_tools.len() as f32
144    };
145    ((intent_score * 0.55) + (tool_score * 0.30) + (artifact.fitness.clamp(0.0, 1.0) * 0.15))
146        .clamp(0.0, 1.0)
147}
148
149impl GraphMemory {
150    /// Create a new graph memory at the given database path.
151    ///
152    /// This will create the database file if it doesn't exist, and
153    /// ensure the AINL graph schema is initialized.
154    pub fn new(db_path: &std::path::Path) -> Result<Self, String> {
155        let store = SqliteGraphStore::open(db_path)?;
156        Ok(Self { store })
157    }
158
159    /// Create from an existing SQLite connection (for integration with existing memory pools)
160    pub fn from_connection(conn: rusqlite::Connection) -> Result<Self, String> {
161        let store = SqliteGraphStore::from_connection(conn)?;
162        Ok(Self { store })
163    }
164
165    /// Wrap an already-open [`SqliteGraphStore`] (for hosts that manage connections externally).
166    pub fn from_sqlite_store(store: SqliteGraphStore) -> Self {
167        Self { store }
168    }
169
170    /// Write an episode node (what happened during an agent turn).
171    ///
172    /// # Arguments
173    /// * `tool_calls` - List of tools executed during this turn
174    /// * `delegation_to` - Agent ID this turn delegated to (if any)
175    /// * `trace_event` - Optional orchestration trace event (serialized JSON)
176    ///
177    /// # Returns
178    /// The ID of the created episode node
179    pub fn write_episode(
180        &self,
181        tool_calls: Vec<String>,
182        delegation_to: Option<String>,
183        trace_event: Option<serde_json::Value>,
184    ) -> Result<Uuid, String> {
185        let turn_id = Uuid::new_v4();
186        let timestamp = chrono::Utc::now().timestamp();
187
188        let node =
189            AinlMemoryNode::new_episode(turn_id, timestamp, tool_calls, delegation_to, trace_event);
190
191        let node_id = node.id;
192        self.store.write_node(&node)?;
193        Ok(node_id)
194    }
195
196    /// Write a semantic fact (learned information with confidence).
197    ///
198    /// # Arguments
199    /// * `fact` - The fact in natural language
200    /// * `confidence` - Confidence score (0.0-1.0)
201    /// * `source_turn_id` - Turn ID that generated this fact
202    ///
203    /// # Returns
204    /// The ID of the created semantic node
205    pub fn write_fact(
206        &self,
207        fact: String,
208        confidence: f32,
209        source_turn_id: Uuid,
210    ) -> Result<Uuid, String> {
211        let node = AinlMemoryNode::new_fact(fact, confidence, source_turn_id);
212        let node_id = node.id;
213        self.store.write_node(&node)?;
214        Ok(node_id)
215    }
216
217    /// Store a procedural pattern (compiled workflow).
218    ///
219    /// # Arguments
220    /// * `pattern_name` - Name/identifier for the pattern
221    /// * `compiled_graph` - Binary representation of the compiled graph
222    ///
223    /// # Returns
224    /// The ID of the created procedural node
225    pub fn store_pattern(
226        &self,
227        pattern_name: String,
228        compiled_graph: Vec<u8>,
229    ) -> Result<Uuid, String> {
230        let node = AinlMemoryNode::new_pattern(pattern_name, compiled_graph);
231        let node_id = node.id;
232        self.store.write_node(&node)?;
233        Ok(node_id)
234    }
235
236    /// Store a procedural pattern derived from a live tool sequence (heuristic extraction).
237    ///
238    /// This path treats the row as **curated** (prompt-eligible) so a single write is visible in
239    /// suggested-procedure style recall; use [`Self::write_node`] with a hand-built node if you
240    /// need candidate-only semantics.
241    pub fn write_procedural(
242        &self,
243        pattern_name: &str,
244        tool_sequence: Vec<String>,
245        confidence: f32,
246    ) -> Result<Uuid, String> {
247        let mut node = AinlMemoryNode::new_procedural_tools(
248            pattern_name.to_string(),
249            tool_sequence,
250            confidence,
251        );
252        if let AinlNodeType::Procedural { ref mut procedural } = node.node_type {
253            procedural.pattern_observation_count = procedural
254                .pattern_observation_count
255                .max(crate::pattern_promotion::DEFAULT_MIN_OBSERVATIONS);
256            let floor = crate::pattern_promotion::DEFAULT_FITNESS_FLOOR;
257            if let Some(f) = procedural.fitness {
258                procedural.fitness = Some(f.max(floor));
259            } else {
260                procedural.fitness = Some(floor);
261            }
262            procedural.prompt_eligible = true;
263        }
264        let node_id = node.id;
265        self.store.write_node(&node)?;
266        Ok(node_id)
267    }
268
269    /// Store a portable procedure artifact as a procedural graph node.
270    ///
271    /// The canonical JSON artifact is stored in `compiled_graph` so older graph consumers can
272    /// ignore it safely, while new consumers can recall and deserialize validated procedure
273    /// artifacts without adding a separate table.
274    pub fn write_procedure_artifact(&self, artifact: &ProcedureArtifact) -> Result<Uuid, String> {
275        self.write_procedure_artifact_for_agent("", artifact)
276    }
277
278    /// Store a portable procedure artifact for a specific agent.
279    pub fn write_procedure_artifact_for_agent(
280        &self,
281        agent_id: &str,
282        artifact: &ProcedureArtifact,
283    ) -> Result<Uuid, String> {
284        let artifact_json = serde_json::to_vec(artifact).map_err(|e| e.to_string())?;
285        let tool_sequence = artifact
286            .steps
287            .iter()
288            .filter_map(|step| match &step.kind {
289                ProcedureStepKind::ToolCall { tool, .. } => Some(tool.clone()),
290                _ => None,
291            })
292            .collect::<Vec<_>>();
293        let mut node = AinlMemoryNode::new_pattern(artifact.id.clone(), artifact_json);
294        node.agent_id = agent_id.to_string();
295        if let AinlNodeType::Procedural { ref mut procedural } = node.node_type {
296            procedural.tool_sequence = tool_sequence;
297            procedural.confidence = Some(artifact.fitness.clamp(0.0, 1.0));
298            procedural.fitness = Some(artifact.fitness.clamp(0.0, 1.0));
299            procedural.pattern_observation_count = artifact.observation_count;
300            procedural.prompt_eligible = matches!(
301                artifact.lifecycle,
302                ProcedureLifecycle::Validated | ProcedureLifecycle::Promoted
303            );
304            procedural.label = artifact.id.clone();
305            procedural.trigger_conditions = vec![artifact.intent.clone()];
306        }
307        let node_id = node.id;
308        self.store.write_node(&node)?;
309        Ok(node_id)
310    }
311
312    /// Update an existing procedural node for `artifact.id`, or write a new one if no node exists.
313    pub fn upsert_procedure_artifact_for_agent(
314        &self,
315        agent_id: &str,
316        artifact: &ProcedureArtifact,
317    ) -> Result<Uuid, String> {
318        for mut node in self.store.find_by_type("procedural")? {
319            if node.agent_id != agent_id {
320                continue;
321            }
322            let Some(procedural) = node.procedural() else {
323                continue;
324            };
325            let matches_id = procedural.label == artifact.id
326                || serde_json::from_slice::<ProcedureArtifact>(&procedural.compiled_graph)
327                    .map(|existing| existing.id == artifact.id)
328                    .unwrap_or(false);
329            if !matches_id {
330                continue;
331            }
332            let artifact_json = serde_json::to_vec(artifact).map_err(|e| e.to_string())?;
333            let tool_sequence = artifact
334                .steps
335                .iter()
336                .filter_map(|step| match &step.kind {
337                    ProcedureStepKind::ToolCall { tool, .. } => Some(tool.clone()),
338                    _ => None,
339                })
340                .collect::<Vec<_>>();
341            if let AinlNodeType::Procedural { ref mut procedural } = node.node_type {
342                procedural.compiled_graph = artifact_json;
343                procedural.tool_sequence = tool_sequence;
344                procedural.confidence = Some(artifact.fitness.clamp(0.0, 1.0));
345                procedural.fitness = Some(artifact.fitness.clamp(0.0, 1.0));
346                procedural.pattern_observation_count = artifact.observation_count;
347                procedural.prompt_eligible = matches!(
348                    artifact.lifecycle,
349                    ProcedureLifecycle::Validated | ProcedureLifecycle::Promoted
350                );
351                procedural.label = artifact.id.clone();
352                procedural.trigger_conditions = vec![artifact.intent.clone()];
353            }
354            let node_id = node.id;
355            self.store.write_node(&node)?;
356            return Ok(node_id);
357        }
358        self.write_procedure_artifact_for_agent(agent_id, artifact)
359    }
360
361    /// Recall portable procedure artifacts previously stored with [`Self::write_procedure_artifact`].
362    pub fn recall_procedure_artifacts(&self) -> Result<Vec<ProcedureArtifact>, String> {
363        let mut out = Vec::new();
364        for node in self.store.find_by_type("procedural")? {
365            let Some(procedural) = node.procedural() else {
366                continue;
367            };
368            if !procedural.prompt_eligible || procedural.compiled_graph.is_empty() {
369                continue;
370            }
371            if let Ok(artifact) =
372                serde_json::from_slice::<ProcedureArtifact>(&procedural.compiled_graph)
373            {
374                out.push(artifact);
375            }
376        }
377        Ok(out)
378    }
379
380    /// Search validated/promoted procedure artifacts by intent text and required tool overlap.
381    pub fn search_procedure_artifacts_for_agent(
382        &self,
383        agent_id: &str,
384        intent: &str,
385        available_tools: &[String],
386        limit: usize,
387    ) -> Result<Vec<ProcedureArtifact>, String> {
388        let mut scored = Vec::new();
389        for node in self.store.find_by_type("procedural")? {
390            if node.agent_id != agent_id {
391                continue;
392            }
393            let Some(procedural) = node.procedural() else {
394                continue;
395            };
396            if !procedural.prompt_eligible || procedural.compiled_graph.is_empty() {
397                continue;
398            }
399            let Ok(artifact) =
400                serde_json::from_slice::<ProcedureArtifact>(&procedural.compiled_graph)
401            else {
402                continue;
403            };
404            if matches!(artifact.lifecycle, ProcedureLifecycle::Deprecated) {
405                continue;
406            }
407            let score = score_procedure_artifact(&artifact, intent, available_tools);
408            if score > 0.0 {
409                scored.push((score, artifact));
410            }
411        }
412        scored.sort_by(|a, b| b.0.partial_cmp(&a.0).unwrap_or(std::cmp::Ordering::Equal));
413        Ok(scored
414            .into_iter()
415            .take(limit)
416            .map(|(_, artifact)| artifact)
417            .collect())
418    }
419
420    /// Update artifact lifecycle and fitness after an attempted reuse.
421    pub fn record_procedure_reuse_outcome_for_agent(
422        &self,
423        agent_id: &str,
424        outcome: &ProcedureReuseOutcome,
425    ) -> Result<Uuid, String> {
426        let mut artifacts =
427            self.search_procedure_artifacts_for_agent(agent_id, "", &[], usize::MAX)?;
428        let Some(mut artifact) = artifacts
429            .drain(..)
430            .find(|artifact| artifact.id == outcome.procedure_id)
431        else {
432            return Err(format!(
433                "procedure artifact not found: {}",
434                outcome.procedure_id
435            ));
436        };
437        artifact.observation_count = artifact.observation_count.saturating_add(1);
438        let delta = match outcome.outcome {
439            TrajectoryOutcome::Success => 0.04,
440            TrajectoryOutcome::PartialSuccess => 0.01,
441            TrajectoryOutcome::Failure => -0.08,
442            TrajectoryOutcome::Aborted => -0.12,
443        };
444        artifact.fitness = (artifact.fitness + delta).clamp(0.0, 1.0);
445        if let Some(failure_id) = outcome.failure_id.as_ref() {
446            if !artifact
447                .source_failure_ids
448                .iter()
449                .any(|id| id == failure_id)
450            {
451                artifact.source_failure_ids.push(failure_id.clone());
452            }
453        }
454        self.upsert_procedure_artifact_for_agent(agent_id, &artifact)
455    }
456
457    /// Write a graph edge between nodes (e.g. episode timeline `follows`).
458    pub fn write_edge(&self, source: Uuid, target: Uuid, rel: &str) -> Result<(), String> {
459        self.store.insert_graph_edge(source, target, rel)
460    }
461
462    /// Recall recent episodes (within the last N seconds).
463    ///
464    /// # Arguments
465    /// * `seconds_ago` - Only return episodes from the last N seconds
466    ///
467    /// # Returns
468    /// Vector of episode nodes, most recent first
469    pub fn recall_recent(&self, seconds_ago: i64) -> Result<Vec<AinlMemoryNode>, String> {
470        let since = chrono::Utc::now().timestamp() - seconds_ago;
471        self.store.query_episodes_since(since, 100)
472    }
473
474    /// Recall nodes of a specific kind written in the last `seconds_ago` seconds.
475    pub fn recall_by_type(
476        &self,
477        kind: AinlNodeKind,
478        seconds_ago: i64,
479    ) -> Result<Vec<AinlMemoryNode>, String> {
480        let since = chrono::Utc::now().timestamp() - seconds_ago;
481        self.store
482            .query_nodes_by_type_since(kind.as_str(), since, 500)
483    }
484
485    /// Find a recent procedural (tool-sequence) row for this agent whose `tool_sequence` matches
486    /// `tool_sequence` (per-element trim). Returns the **newest** match if several exist
487    /// (e.g. legacy duplicates before merge).
488    pub fn find_procedural_by_tool_sequence(
489        &self,
490        agent_id: &str,
491        tool_sequence: &[String],
492    ) -> Result<Option<AinlMemoryNode>, String> {
493        let norm: Vec<String> = tool_sequence.iter().map(|s| s.trim().to_string()).collect();
494        if norm.is_empty() {
495            return Ok(None);
496        }
497        let nodes = self.recall_by_type(AinlNodeKind::Procedural, 60 * 60 * 24 * 365 * 5)?;
498        for n in nodes {
499            if n.agent_id != agent_id {
500                continue;
501            }
502            let AinlNodeType::Procedural { ref procedural } = n.node_type else {
503                continue;
504            };
505            if procedural.tool_sequence.len() != norm.len() {
506                continue;
507            }
508            let same = procedural
509                .tool_sequence
510                .iter()
511                .zip(norm.iter())
512                .all(|(a, b)| a.trim() == b.trim());
513            if same {
514                // `recall_by_type` is most-recent first; first hit is the canonical row to update.
515                return Ok(Some(n));
516            }
517        }
518        Ok(None)
519    }
520
521    /// Write a persona trait node.
522    pub fn write_persona(
523        &self,
524        trait_name: &str,
525        strength: f32,
526        learned_from: Vec<Uuid>,
527    ) -> Result<Uuid, String> {
528        let node = AinlMemoryNode::new_persona(trait_name.to_string(), strength, learned_from);
529        let node_id = node.id;
530        self.store.write_node(&node)?;
531        Ok(node_id)
532    }
533
534    /// Get direct access to the underlying store for advanced queries
535    pub fn store(&self) -> &dyn GraphStore {
536        &self.store
537    }
538
539    /// SQLite backing store (for components such as `ainl-graph-extractor` that require concrete SQL access).
540    pub fn sqlite_store(&self) -> &SqliteGraphStore {
541        &self.store
542    }
543
544    /// [`SqliteGraphStore::validate_graph`] for the same backing database (checkpoint / boot gate).
545    pub fn validate_graph(&self, agent_id: &str) -> Result<GraphValidationReport, String> {
546        self.store.validate_graph(agent_id)
547    }
548
549    /// [`SqliteGraphStore::export_graph`].
550    pub fn export_graph(&self, agent_id: &str) -> Result<AgentGraphSnapshot, String> {
551        self.store.export_graph(agent_id)
552    }
553
554    /// [`SqliteGraphStore::import_graph`] — use `allow_dangling_edges: false` for normal loads; `true` only for repair.
555    pub fn import_graph(
556        &mut self,
557        snapshot: &AgentGraphSnapshot,
558        allow_dangling_edges: bool,
559    ) -> Result<(), String> {
560        self.store.import_graph(snapshot, allow_dangling_edges)
561    }
562
563    /// [`SqliteGraphStore::agent_subgraph_edges`].
564    pub fn agent_subgraph_edges(&self, agent_id: &str) -> Result<Vec<SnapshotEdge>, String> {
565        self.store.agent_subgraph_edges(agent_id)
566    }
567
568    /// [`SqliteGraphStore::write_node_with_edges`] (transactional; fails if embedded edge targets are missing).
569    pub fn write_node_with_edges(&mut self, node: &AinlMemoryNode) -> Result<(), String> {
570        self.store.write_node_with_edges(node)
571    }
572
573    /// [`SqliteGraphStore::insert_graph_edge_checked`].
574    pub fn insert_graph_edge_checked(
575        &self,
576        from_id: Uuid,
577        to_id: Uuid,
578        label: &str,
579    ) -> Result<(), String> {
580        self.store.insert_graph_edge_checked(from_id, to_id, label)
581    }
582
583    /// Read persisted [`RuntimeStateNode`] for `agent_id` (most recent row).
584    pub fn read_runtime_state(&self, agent_id: &str) -> Result<Option<RuntimeStateNode>, String> {
585        self.store.read_runtime_state(agent_id)
586    }
587
588    /// Upsert persisted [`RuntimeStateNode`] for the given agent (stable node id per `agent_id`).
589    pub fn write_runtime_state(&self, state: &RuntimeStateNode) -> Result<(), String> {
590        self.store.write_runtime_state(state)
591    }
592
593    /// Write a fully constructed node (additive API for callers that set extended metadata).
594    pub fn write_node(&self, node: &AinlMemoryNode) -> Result<(), String> {
595        self.store.write_node(node)
596    }
597
598    /// Insert a detailed trajectory row (see [`SqliteGraphStore::insert_trajectory_detail`]).
599    pub fn insert_trajectory_detail(&self, row: &TrajectoryDetailRecord) -> Result<(), String> {
600        self.store.insert_trajectory_detail(row)
601    }
602
603    /// Recent trajectory detail rows for an agent (see [`SqliteGraphStore::list_trajectories_for_agent`]).
604    pub fn list_trajectories_for_agent(
605        &self,
606        agent_id: &str,
607        limit: usize,
608        since_timestamp: Option<i64>,
609    ) -> Result<Vec<TrajectoryDetailRecord>, String> {
610        self.store
611            .list_trajectories_for_agent(agent_id, limit, since_timestamp)
612    }
613
614    /// How many `ainl_trajectories` detail rows would be removed by
615    /// [`Self::prune_trajectory_details_before`] (same `before_recorded_at` semantics).
616    pub fn count_trajectory_details_before(
617        &self,
618        agent_id: &str,
619        before_recorded_at: i64,
620    ) -> Result<usize, String> {
621        self.store
622            .count_trajectory_details_before(agent_id, before_recorded_at)
623    }
624
625    /// Remove persisted trajectory **detail** rows with `recorded_at` **strictly before** `before_recorded_at` (seconds).
626    ///
627    /// This targets the `ainl_trajectories` table only. Graph `Trajectory` nodes and cross-links are not
628    /// deleted here; use exports / graph tooling if you need a full-store consistency pass after pruning.
629    pub fn prune_trajectory_details_before(
630        &self,
631        agent_id: &str,
632        before_recorded_at: i64,
633    ) -> Result<usize, String> {
634        self.store
635            .delete_trajectory_details_before(agent_id, before_recorded_at)
636    }
637
638    /// Search persisted [`FailureNode`] rows for an agent (FTS5 over `ainl_failures_fts`).
639    pub fn search_failures_for_agent(
640        &self,
641        agent_id: &str,
642        query: &str,
643        limit: usize,
644    ) -> Result<Vec<AinlMemoryNode>, String> {
645        self.store
646            .search_failures_fts_for_agent(agent_id, query, limit)
647    }
648
649    /// Full-graph FTS5 search (`ainl_nodes_fts`); see [`SqliteGraphStore::search_all_nodes_fts_for_agent`].
650    pub fn search_all_nodes_fts(
651        &self,
652        agent_id: &str,
653        query: &str,
654        project_id: Option<&str>,
655        limit: usize,
656    ) -> Result<Vec<AinlMemoryNode>, String> {
657        self.store
658            .search_all_nodes_fts_for_agent(agent_id, query, project_id, limit)
659    }
660}
661
662#[cfg(test)]
663mod tests {
664    use super::*;
665
666    #[test]
667    fn test_graph_memory_api() {
668        let temp_dir = std::env::temp_dir();
669        let db_path = temp_dir.join("ainl_lib_test.db");
670        let _ = std::fs::remove_file(&db_path);
671
672        let memory = GraphMemory::new(&db_path).expect("Failed to create memory");
673
674        // Write an episode
675        let episode_id = memory
676            .write_episode(
677                vec!["file_read".to_string(), "agent_delegate".to_string()],
678                Some("agent-B".to_string()),
679                None,
680            )
681            .expect("Failed to write episode");
682
683        assert_ne!(episode_id, Uuid::nil());
684
685        // Write a fact
686        let fact_id = memory
687            .write_fact(
688                "User prefers concise responses".to_string(),
689                0.85,
690                episode_id,
691            )
692            .expect("Failed to write fact");
693
694        assert_ne!(fact_id, Uuid::nil());
695
696        // Recall recent episodes
697        let recent = memory.recall_recent(60).expect("Failed to recall");
698        assert_eq!(recent.len(), 1);
699
700        // Verify the episode content
701        if let AinlNodeType::Episode { episodic } = &recent[0].node_type {
702            assert_eq!(episodic.delegation_to, Some("agent-B".to_string()));
703            assert_eq!(episodic.tool_calls.len(), 2);
704        } else {
705            panic!("Wrong node type");
706        }
707    }
708
709    #[test]
710    fn test_store_pattern() {
711        let temp_dir = std::env::temp_dir();
712        let db_path = temp_dir.join("ainl_lib_test_pattern.db");
713        let _ = std::fs::remove_file(&db_path);
714
715        let memory = GraphMemory::new(&db_path).expect("Failed to create memory");
716
717        let pattern_id = memory
718            .store_pattern("research_workflow".to_string(), vec![1, 2, 3, 4])
719            .expect("Failed to store pattern");
720
721        assert_ne!(pattern_id, Uuid::nil());
722
723        // Query it back
724        let patterns = find_patterns(memory.store(), "research").expect("Query failed");
725        assert_eq!(patterns.len(), 1);
726    }
727
728    /// End-to-end: `Failure` graph row + `ainl_failures_fts` sync + `search_failures_for_agent`.
729    #[test]
730    fn failure_write_and_fts_search_roundtrip() {
731        let dir = tempfile::tempdir().expect("tempdir");
732        let db_path = dir.path().join("ainl_failure_fts_smoke.db");
733        let memory = GraphMemory::new(&db_path).expect("graph memory");
734        let agent_id = "agent-smoke-fts";
735
736        let mut node = AinlMemoryNode::new_loop_guard_failure(
737            "block",
738            Some("shell_exec"),
739            "repeated identical tool invocation blocked by loop guard",
740            Some("session-xyz"),
741        );
742        node.agent_id = agent_id.to_string();
743        let nid = node.id;
744        memory.write_node(&node).expect("write failure node");
745
746        let hits = memory
747            .search_failures_for_agent(agent_id, "loop", 10)
748            .expect("search loop");
749        assert_eq!(hits.len(), 1, "expected one FTS hit for token 'loop'");
750        assert_eq!(hits[0].id, nid);
751        assert!(
752            matches!(&hits[0].node_type, AinlNodeType::Failure { .. }),
753            "expected Failure node type"
754        );
755
756        let hits2 = memory
757            .search_failures_for_agent(agent_id, "shell_exec", 10)
758            .expect("search tool name");
759        assert_eq!(hits2.len(), 1);
760        assert_eq!(hits2[0].id, nid);
761
762        let empty = memory
763            .search_failures_for_agent(agent_id, "   ", 10)
764            .expect("whitespace-only query");
765        assert!(empty.is_empty());
766
767        let wrong_agent = memory
768            .search_failures_for_agent("other-agent", "loop", 10)
769            .expect("wrong agent id");
770        assert!(wrong_agent.is_empty());
771    }
772
773    /// Full-graph `ainl_nodes_fts` — semantic fact is searchable, not only failures.
774    #[test]
775    fn all_nodes_fts_write_and_search_roundtrip() {
776        let dir = tempfile::tempdir().expect("tempdir");
777        let db_path = dir.path().join("ainl_all_nodes_fts.db");
778        let memory = GraphMemory::new(&db_path).expect("graph memory");
779        let agent_id = "agent-fts-all";
780        let mut node =
781            AinlMemoryNode::new_fact("unique-fts-violet-cat-42".into(), 0.8, Uuid::new_v4());
782        node.agent_id = agent_id.to_string();
783        let nid = node.id;
784        memory.write_node(&node).expect("write fact");
785
786        let hits = memory
787            .search_all_nodes_fts(agent_id, "violet", None, 10)
788            .expect("search");
789        assert_eq!(hits.len(), 1, "expected one all-nodes FTS hit");
790        assert_eq!(hits[0].id, nid);
791    }
792
793    #[test]
794    fn tool_execution_failure_write_and_fts_search_roundtrip() {
795        let dir = tempfile::tempdir().expect("tempdir");
796        let db_path = dir.path().join("ainl_tool_failure_fts.db");
797        let memory = GraphMemory::new(&db_path).expect("graph memory");
798        let agent_id = "agent-tool-ft";
799
800        let mut node = AinlMemoryNode::new_tool_execution_failure(
801            "file_read",
802            "ENOENT: no such file or directory",
803            Some("sess-tool-1"),
804        );
805        node.agent_id = agent_id.to_string();
806        let nid = node.id;
807        memory.write_node(&node).expect("write tool failure node");
808
809        let hits = memory
810            .search_failures_for_agent(agent_id, "ENOENT", 10)
811            .expect("search ENOENT");
812        assert_eq!(hits.len(), 1);
813        assert_eq!(hits[0].id, nid);
814
815        let src_hits = memory
816            .search_failures_for_agent(agent_id, "tool_runner", 10)
817            .expect("search source");
818        assert_eq!(src_hits.len(), 1);
819        assert_eq!(src_hits[0].id, nid);
820    }
821
822    #[test]
823    fn agent_loop_precheck_failure_write_and_fts_search_roundtrip() {
824        let dir = tempfile::tempdir().expect("tempdir");
825        let db_path = dir.path().join("ainl_precheck_failure_fts.db");
826        let memory = GraphMemory::new(&db_path).expect("graph memory");
827        let agent_id = "agent-precheck-ft";
828
829        let mut node = AinlMemoryNode::new_agent_loop_precheck_failure(
830            "param_validation",
831            "file_write",
832            "missing required field: path",
833            Some("sess-pv-1"),
834        );
835        node.agent_id = agent_id.to_string();
836        let nid = node.id;
837        memory.write_node(&node).expect("write precheck failure");
838
839        let hits = memory
840            .search_failures_for_agent(agent_id, "param_validation", 10)
841            .expect("search kind");
842        assert_eq!(hits.len(), 1);
843        assert_eq!(hits[0].id, nid);
844
845        let hits2 = memory
846            .search_failures_for_agent(agent_id, "agent_loop", 10)
847            .expect("search agent_loop prefix");
848        assert_eq!(hits2.len(), 1);
849    }
850
851    #[test]
852    fn ainl_runtime_graph_validation_failure_write_and_fts_search_roundtrip() {
853        let dir = tempfile::tempdir().expect("tempdir");
854        let db_path = dir.path().join("ainl_graph_validation_failure_fts.db");
855        let memory = GraphMemory::new(&db_path).expect("graph memory");
856        let agent_id = "agent-graph-val-ft";
857
858        let mut node = AinlMemoryNode::new_ainl_runtime_graph_validation_failure(
859            "graph validation failed before turn: dangling edges …",
860            Some("sess-gv-1"),
861        );
862        node.agent_id = agent_id.to_string();
863        let nid = node.id;
864        memory
865            .write_node(&node)
866            .expect("write graph validation failure");
867
868        let hits = memory
869            .search_failures_for_agent(agent_id, "graph_validation", 10)
870            .expect("search source label");
871        assert_eq!(hits.len(), 1);
872        assert_eq!(hits[0].id, nid);
873
874        let hits2 = memory
875            .search_failures_for_agent(agent_id, "dangling", 10)
876            .expect("search message body");
877        assert_eq!(hits2.len(), 1);
878    }
879
880    #[test]
881    fn trajectory_detail_prune_before_drops_only_old_rows() {
882        use ainl_contracts::{TrajectoryOutcome, TrajectoryStep};
883
884        let dir = tempfile::tempdir().expect("tempdir");
885        let db_path = dir.path().join("ainl_traj_prune.db");
886        let memory = GraphMemory::new(&db_path).expect("graph memory");
887        let agent = "agent-traj-prune";
888        let ep_old = memory
889            .write_episode(vec![], None, None)
890            .expect("episode for old traj");
891        let ep_new = memory
892            .write_episode(vec![], None, None)
893            .expect("episode for new traj");
894        let mk_step = |sid: &str| TrajectoryStep {
895            step_id: sid.to_string(),
896            timestamp_ms: 0,
897            adapter: "a".into(),
898            operation: "o".into(),
899            inputs_preview: None,
900            outputs_preview: None,
901            duration_ms: 1,
902            success: true,
903            error: None,
904            vitals: None,
905            freshness_at_step: None,
906            frame_vars: None,
907            tool_telemetry: None,
908        };
909        let r_old = TrajectoryDetailRecord {
910            id: Uuid::new_v4(),
911            episode_id: ep_old,
912            graph_trajectory_node_id: None,
913            agent_id: agent.to_string(),
914            session_id: "s-old".into(),
915            project_id: None,
916            recorded_at: 100,
917            outcome: TrajectoryOutcome::Success,
918            ainl_source_hash: None,
919            duration_ms: 1,
920            steps: vec![mk_step("1")],
921            frame_vars: None,
922            fitness_delta: None,
923        };
924        let r_new = TrajectoryDetailRecord {
925            id: Uuid::new_v4(),
926            episode_id: ep_new,
927            graph_trajectory_node_id: None,
928            agent_id: agent.to_string(),
929            session_id: "s-new".into(),
930            project_id: None,
931            recorded_at: 200,
932            outcome: TrajectoryOutcome::Success,
933            ainl_source_hash: None,
934            duration_ms: 1,
935            steps: vec![mk_step("2")],
936            frame_vars: None,
937            fitness_delta: None,
938        };
939        memory.insert_trajectory_detail(&r_old).expect("insert old");
940        memory.insert_trajectory_detail(&r_new).expect("insert new");
941        let before = memory
942            .list_trajectories_for_agent(agent, 10, None)
943            .expect("list");
944        assert_eq!(before.len(), 2);
945        let removed = memory
946            .prune_trajectory_details_before(agent, 200)
947            .expect("prune");
948        assert_eq!(removed, 1);
949        let after = memory
950            .list_trajectories_for_agent(agent, 10, None)
951            .expect("list after");
952        assert_eq!(after.len(), 1);
953        assert_eq!(after[0].recorded_at, 200);
954    }
955
956    #[test]
957    fn stores_and_recalls_validated_procedure_artifact() {
958        use ainl_contracts::{
959            ProcedureArtifact, ProcedureArtifactFormat, ProcedureLifecycle, ProcedureStep,
960            ProcedureStepKind, ProcedureVerification, LEARNER_SCHEMA_VERSION,
961        };
962
963        let tmp = tempfile::tempdir().unwrap();
964        let memory = GraphMemory::new(&tmp.path().join("memory.db")).unwrap();
965        let artifact = ProcedureArtifact {
966            schema_version: LEARNER_SCHEMA_VERSION,
967            id: "proc:test".into(),
968            title: "Test Procedure".into(),
969            intent: "test intent".into(),
970            summary: "summary".into(),
971            required_tools: vec!["file_read".into()],
972            required_adapters: vec![],
973            inputs: vec![],
974            outputs: vec![],
975            preconditions: vec![],
976            steps: vec![ProcedureStep {
977                step_id: "s1".into(),
978                title: "Read".into(),
979                kind: ProcedureStepKind::ToolCall {
980                    tool: "file_read".into(),
981                    args_schema: serde_json::json!({"type":"object"}),
982                },
983                rationale: None,
984            }],
985            verification: ProcedureVerification::default(),
986            known_failures: vec![],
987            recovery: vec![],
988            source_trajectory_ids: vec![],
989            source_failure_ids: vec![],
990            fitness: 0.9,
991            observation_count: 3,
992            lifecycle: ProcedureLifecycle::Validated,
993            render_targets: vec![ProcedureArtifactFormat::PromptOnly],
994        };
995        memory.write_procedure_artifact(&artifact).unwrap();
996        let recalled = memory.recall_procedure_artifacts().unwrap();
997        assert_eq!(recalled, vec![artifact]);
998    }
999
1000    #[test]
1001    fn searches_and_updates_procedure_reuse_fitness() {
1002        use ainl_contracts::{
1003            ProcedureArtifact, ProcedureArtifactFormat, ProcedureLifecycle, ProcedureReuseOutcome,
1004            ProcedureStep, ProcedureStepKind, ProcedureVerification, TrajectoryOutcome,
1005            LEARNER_SCHEMA_VERSION,
1006        };
1007
1008        let tmp = tempfile::tempdir().unwrap();
1009        let memory = GraphMemory::new(&tmp.path().join("memory.db")).unwrap();
1010        let artifact = ProcedureArtifact {
1011            schema_version: LEARNER_SCHEMA_VERSION,
1012            id: "proc:review".into(),
1013            title: "Review PR".into(),
1014            intent: "review pull request".into(),
1015            summary: "review code changes safely".into(),
1016            required_tools: vec!["file_read".into(), "shell_exec".into()],
1017            required_adapters: vec![],
1018            inputs: vec![],
1019            outputs: vec![],
1020            preconditions: vec![],
1021            steps: vec![ProcedureStep {
1022                step_id: "s1".into(),
1023                title: "Read".into(),
1024                kind: ProcedureStepKind::ToolCall {
1025                    tool: "file_read".into(),
1026                    args_schema: serde_json::json!({"type":"object"}),
1027                },
1028                rationale: None,
1029            }],
1030            verification: ProcedureVerification::default(),
1031            known_failures: vec![],
1032            recovery: vec![],
1033            source_trajectory_ids: vec![],
1034            source_failure_ids: vec![],
1035            fitness: 0.6,
1036            observation_count: 3,
1037            lifecycle: ProcedureLifecycle::Promoted,
1038            render_targets: vec![ProcedureArtifactFormat::PromptOnly],
1039        };
1040        memory
1041            .write_procedure_artifact_for_agent("agent-search", &artifact)
1042            .unwrap();
1043        let hits = memory
1044            .search_procedure_artifacts_for_agent(
1045                "agent-search",
1046                "please review this pull request",
1047                &["file_read".into(), "shell_exec".into()],
1048                5,
1049            )
1050            .unwrap();
1051        assert_eq!(hits.len(), 1);
1052        assert_eq!(hits[0].id, "proc:review");
1053
1054        memory
1055            .record_procedure_reuse_outcome_for_agent(
1056                "agent-search",
1057                &ProcedureReuseOutcome {
1058                    procedure_id: "proc:review".into(),
1059                    outcome: TrajectoryOutcome::Failure,
1060                    failure_id: Some("failure-x".into()),
1061                    notes: None,
1062                },
1063            )
1064            .unwrap();
1065        let updated = memory
1066            .search_procedure_artifacts_for_agent("agent-search", "review pull request", &[], 5)
1067            .unwrap();
1068        assert_eq!(updated[0].observation_count, 4);
1069        assert!(updated[0].fitness < 0.6);
1070        assert!(updated[0].source_failure_ids.contains(&"failure-x".into()));
1071    }
1072}