Initial commit

2026-08-02 15:26:10 +08:00
parent fcc5d31137
commit eee7fed161
76 changed files with 10686 additions and 3564 deletions
@@ -0,0 +1,18 @@
+[package]
+name = "agentos-memory"
+version.workspace = true
+edition.workspace = true
+license.workspace = true
+repository.workspace = true
+
+[dependencies]
+agentos-core.workspace = true
+rusqlite.workspace = true
+serde.workspace = true
+serde_json.workspace = true
+thiserror.workspace = true
+time.workspace = true
+uuid.workspace = true
+
+[lints]
+workspace = true
@@ -0,0 +1,326 @@
+//! Durable AgentOS state and audit log.
+
+use agentos_core::{AgentId, AgentPrincipal, AuditEvent, WorldlineId, utc_now};
+use rusqlite::{Connection, OptionalExtension, params};
+use serde::{Deserialize, Serialize};
+use serde_json::Value;
+use std::path::{Path, PathBuf};
+use uuid::Uuid;
+
+#[derive(Clone, Debug)]
+pub struct MemoryStore {
+    path: PathBuf,
+}
+
+#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
+pub struct MemoryRecord {
+    pub memory_id: String,
+    pub kind: String,
+    pub content: String,
+    pub importance: f64,
+    pub created_at: String,
+    pub worldline_id: Option<WorldlineId>,
+    pub metadata: Value,
+}
+
+impl MemoryStore {
+    pub fn initialize(path: impl Into<PathBuf>) -> Result<Self, MemoryError> {
+        let path = path.into();
+        if let Some(parent) = path.parent() {
+            std::fs::create_dir_all(parent)?;
+        }
+        let store = Self { path };
+        let connection = store.connect()?;
+        connection.execute_batch(
+            "
+            PRAGMA journal_mode = WAL;
+            PRAGMA foreign_keys = ON;
+            CREATE TABLE IF NOT EXISTS memories (
+                memory_id TEXT PRIMARY KEY,
+                kind TEXT NOT NULL,
+                content TEXT NOT NULL,
+                importance REAL NOT NULL CHECK (importance BETWEEN 0.0 AND 1.0),
+                created_at TEXT NOT NULL,
+                worldline_id TEXT,
+                metadata_json TEXT NOT NULL,
+                state TEXT NOT NULL DEFAULT 'active'
+            );
+            CREATE VIRTUAL TABLE IF NOT EXISTS memory_fts USING fts5(
+                memory_id UNINDEXED,
+                content,
+                tokenize = 'unicode61'
+            );
+            CREATE TABLE IF NOT EXISTS events (
+                sequence INTEGER PRIMARY KEY AUTOINCREMENT,
+                event_id TEXT NOT NULL UNIQUE,
+                event_type TEXT NOT NULL,
+                agent_id TEXT,
+                session_id TEXT,
+                run_id TEXT,
+                worldline_id TEXT,
+                payload_json TEXT NOT NULL,
+                occurred_at TEXT NOT NULL
+            );
+            CREATE TABLE IF NOT EXISTS agent_principals (
+                agent_id TEXT PRIMARY KEY,
+                owner_uid INTEGER NOT NULL,
+                agent_uid INTEGER NOT NULL UNIQUE,
+                agent_gid INTEGER NOT NULL,
+                state TEXT NOT NULL DEFAULT 'active',
+                created_at TEXT NOT NULL,
+                CHECK (agent_uid > 0),
+                CHECK (agent_gid > 0)
+            );
+            ",
+        )?;
+        Ok(store)
+    }
+
+    pub fn path(&self) -> &Path {
+        &self.path
+    }
+
+    pub fn append_event(&self, event: &AuditEvent) -> Result<i64, MemoryError> {
+        let connection = self.connect()?;
+        connection.execute(
+            "INSERT INTO events(
+                event_id, event_type, agent_id, session_id, run_id,
+                worldline_id, payload_json, occurred_at
+             ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)",
+            params![
+                event.event_id.as_str(),
+                event.event_type,
+                event.agent_id.as_ref().map(AgentId::as_str),
+                event.session_id.as_ref().map(ToString::to_string),
+                event.run_id.as_ref().map(ToString::to_string),
+                event.worldline_id.as_ref().map(ToString::to_string),
+                serde_json::to_string(&event.payload)?,
+                event.occurred_at,
+            ],
+        )?;
+        Ok(connection.last_insert_rowid())
+    }
+
+    pub fn bind_principal(&self, principal: &AgentPrincipal) -> Result<(), MemoryError> {
+        principal.validate()?;
+        let connection = self.connect()?;
+        let existing = connection
+            .query_row(
+                "SELECT owner_uid, agent_uid, agent_gid
+                 FROM agent_principals WHERE agent_id = ?1 AND state = 'active'",
+                [principal.agent_id.as_str()],
+                |row| {
+                    Ok((
+                        row.get::<_, u32>(0)?,
+                        row.get::<_, u32>(1)?,
+                        row.get::<_, u32>(2)?,
+                    ))
+                },
+            )
+            .optional()?;
+        if let Some((owner_uid, agent_uid, agent_gid)) = existing {
+            if (owner_uid, agent_uid, agent_gid)
+                == (
+                    principal.owner_uid,
+                    principal.agent_uid,
+                    principal.agent_gid,
+                )
+            {
+                return Ok(());
+            }
+            return Err(MemoryError::PrincipalRebind(principal.agent_id.clone()));
+        }
+        connection.execute(
+            "INSERT INTO agent_principals(
+                agent_id, owner_uid, agent_uid, agent_gid, created_at
+             ) VALUES (?1, ?2, ?3, ?4, ?5)",
+            params![
+                principal.agent_id.as_str(),
+                principal.owner_uid,
+                principal.agent_uid,
+                principal.agent_gid,
+                utc_now(),
+            ],
+        )?;
+        Ok(())
+    }
+
+    pub fn principal(&self, agent_id: &AgentId) -> Result<Option<AgentPrincipal>, MemoryError> {
+        let connection = self.connect()?;
+        Ok(connection
+            .query_row(
+                "SELECT owner_uid, agent_uid, agent_gid
+                 FROM agent_principals WHERE agent_id = ?1 AND state = 'active'",
+                [agent_id.as_str()],
+                |row| {
+                    Ok(AgentPrincipal {
+                        agent_id: agent_id.clone(),
+                        owner_uid: row.get(0)?,
+                        agent_uid: row.get(1)?,
+                        agent_gid: row.get(2)?,
+                    })
+                },
+            )
+            .optional()?)
+    }
+
+    pub fn remember(
+        &self,
+        kind: impl Into<String>,
+        content: impl Into<String>,
+        importance: f64,
+        worldline_id: Option<&WorldlineId>,
+        metadata: Value,
+    ) -> Result<MemoryRecord, MemoryError> {
+        let content = content.into().trim().to_owned();
+        if content.is_empty() {
+            return Err(MemoryError::EmptyMemory);
+        }
+        let record = MemoryRecord {
+            memory_id: Uuid::new_v4().to_string(),
+            kind: kind.into(),
+            content,
+            importance: importance.clamp(0.0, 1.0),
+            created_at: utc_now(),
+            worldline_id: worldline_id.cloned(),
+            metadata,
+        };
+        let mut connection = self.connect()?;
+        let transaction = connection.transaction()?;
+        transaction.execute(
+            "INSERT INTO memories(
+                memory_id, kind, content, importance, created_at,
+                worldline_id, metadata_json
+             ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)",
+            params![
+                record.memory_id,
+                record.kind,
+                record.content,
+                record.importance,
+                record.created_at,
+                record.worldline_id.as_ref().map(ToString::to_string),
+                serde_json::to_string(&record.metadata)?,
+            ],
+        )?;
+        transaction.execute(
+            "INSERT INTO memory_fts(memory_id, content) VALUES (?1, ?2)",
+            params![record.memory_id, record.content],
+        )?;
+        transaction.commit()?;
+        Ok(record)
+    }
+
+    pub fn recall(&self, query: &str, limit: usize) -> Result<Vec<MemoryRecord>, MemoryError> {
+        if limit == 0 {
+            return Ok(Vec::new());
+        }
+        let terms = query
+            .split(|character: char| !character.is_alphanumeric())
+            .filter(|term| term.chars().count() >= 2)
+            .take(12)
+            .map(|term| format!("\"{}\"", term.replace('"', "\"\"")))
+            .collect::<Vec<_>>();
+        let connection = self.connect()?;
+        let sql = if terms.is_empty() {
+            "SELECT m.memory_id, m.kind, m.content, m.importance, m.created_at,
+                    m.worldline_id, m.metadata_json
+             FROM memories m
+             WHERE m.state = 'active'
+             ORDER BY m.importance DESC, m.created_at DESC
+             LIMIT ?1"
+        } else {
+            "SELECT m.memory_id, m.kind, m.content, m.importance, m.created_at,
+                    m.worldline_id, m.metadata_json
+             FROM memory_fts f
+             JOIN memories m ON m.memory_id = f.memory_id
+             WHERE memory_fts MATCH ?1 AND m.state = 'active'
+             ORDER BY bm25(memory_fts), m.importance DESC, m.created_at DESC
+             LIMIT ?2"
+        };
+        let mut statement = connection.prepare(sql)?;
+        let mapper = |row: &rusqlite::Row<'_>| -> rusqlite::Result<MemoryRecord> {
+            let worldline: Option<String> = row.get(5)?;
+            let metadata: String = row.get(6)?;
+            Ok(MemoryRecord {
+                memory_id: row.get(0)?,
+                kind: row.get(1)?,
+                content: row.get(2)?,
+                importance: row.get(3)?,
+                created_at: row.get(4)?,
+                worldline_id: worldline.and_then(|value| WorldlineId::parse(value).ok()),
+                metadata: serde_json::from_str(&metadata).unwrap_or(Value::Null),
+            })
+        };
+        let records = if terms.is_empty() {
+            statement
+                .query_map([i64::try_from(limit).unwrap_or(i64::MAX)], mapper)?
+                .collect::<Result<Vec<_>, _>>()?
+        } else {
+            statement
+                .query_map(
+                    params![terms.join(" OR "), i64::try_from(limit).unwrap_or(i64::MAX)],
+                    mapper,
+                )?
+                .collect::<Result<Vec<_>, _>>()?
+        };
+        Ok(records)
+    }
+
+    fn connect(&self) -> Result<Connection, MemoryError> {
+        let connection = Connection::open(&self.path)?;
+        connection.pragma_update(None, "foreign_keys", "ON")?;
+        connection.busy_timeout(std::time::Duration::from_secs(5))?;
+        Ok(connection)
+    }
+}
+
+#[derive(Debug, thiserror::Error)]
+pub enum MemoryError {
+    #[error("database operation failed: {0}")]
+    Database(#[from] rusqlite::Error),
+    #[error("filesystem operation failed: {0}")]
+    Io(#[from] std::io::Error),
+    #[error("state serialization failed: {0}")]
+    Json(#[from] serde_json::Error),
+    #[error("identity is invalid: {0}")]
+    Identity(#[from] agentos_core::IdentityError),
+    #[error("agent identity cannot be rebound to another Linux principal: {0}")]
+    PrincipalRebind(AgentId),
+    #[error("memory content must not be empty")]
+    EmptyMemory,
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn principal_binding_is_stable_and_memory_is_searchable() {
+        let root = std::env::temp_dir().join(format!("agentos-memory-{}", Uuid::new_v4()));
+        let store = MemoryStore::initialize(root.join("state.sqlite3")).unwrap();
+        let principal = AgentPrincipal {
+            agent_id: AgentId::new(),
+            owner_uid: 1000,
+            agent_uid: 200_001,
+            agent_gid: 200_001,
+        };
+        store.bind_principal(&principal).unwrap();
+        store.bind_principal(&principal).unwrap();
+        assert_eq!(
+            store.principal(&principal.agent_id).unwrap(),
+            Some(principal)
+        );
+
+        store
+            .remember(
+                "fact",
+                "CALCULET runtime uses an NPU bridge",
+                0.8,
+                None,
+                Value::Null,
+            )
+            .unwrap();
+        assert_eq!(store.recall("NPU bridge", 3).unwrap().len(), 1);
+        std::fs::remove_dir_all(root).unwrap();
+    }
+}