xiaoxixi 06811160f0 docs: add memory system implementation plan

2026-05-07 23:02:47 +08:00

40 KiB

Raw Blame History

Memory System Implementation Plan

For execution: Use superpowers:executing-plans to implement this plan task-by-task.

Goal: Add a memory system to PicoBot that persists knowledge (facts/preferences) and timeline (compressed conversation summaries) via SQLite+FTS5, unified with the existing ContextCompressor.

Architecture: Extend existing Storage module with a memories table + FTS5 virtual table. MemoryManager wraps CRUD operations. ContextCompressor gets a MemoryManager handle to auto-write timeline entries on compression. SystemPromptBuilder injects FTS5-retrieved knowledge into every system prompt. Three new agent tools (memory_store, memory_recall, memory_forget) registered in create_default_tools().

Tech Stack: Rust 2024 edition, sqlx (SQLite + FTS5), serde_json, anyhow, uuid, chrono

Task 1: Memory Type Definitions

Files:

Create: src/memory/types.rs
Modify: src/lib.rs:14 (add pub mod memory;)

Step 1: Write module declaration

Add to src/lib.rs after line 11 (pub mod observability;):

pub mod memory;

Step 2: Create memory types file

Write src/memory/types.rs:

use serde::{Deserialize, Serialize};

/// Memory categories for the memory system.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum MemoryCategory {
    /// Long-term facts, preferences, patterns, insights (merged fact+insight).
    Knowledge,
    /// Conversation summaries produced by context compression.
    Timeline,
}

impl MemoryCategory {
    pub fn as_str(&self) -> &str {
        match self {
            Self::Knowledge => "knowledge",
            Self::Timeline => "timeline",
        }
    }

    pub fn from_str(s: &str) -> Option<Self> {
        match s {
            "knowledge" => Some(Self::Knowledge),
            "timeline" => Some(Self::Timeline),
            _ => None,
        }
    }
}

/// A single memory entry stored in SQLite.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MemoryEntry {
    pub id: String,
    /// Semantic identifier (e.g. "user_prefers_rust").
    pub key: String,
    /// The memory content.
    pub content: String,
    /// Category: knowledge or timeline.
    pub category: MemoryCategory,
    /// Importance score 0.0–1.0 (default 0.5).
    pub importance: f64,
    /// Associated session ID (optional).
    pub session_id: Option<String>,
    /// RFC 3339 creation timestamp.
    pub created_at: String,
    /// RFC 3339 last update timestamp.
    pub updated_at: String,
}

/// Result from an LLM consolidation call.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ConsolidationResult {
    /// New or updated knowledge entries extracted from conversation.
    pub facts: Vec<ConsolidationFact>,
    /// Summary entry for timeline (formatted as "[YYYY-MM-DD HH:MM] text...").
    pub timeline: Option<String>,
    /// Keys of existing memories that should be invalidated.
    pub invalidations: Vec<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ConsolidationFact {
    pub key: String,
    pub content: String,
    pub importance: f64,
}

Step 3: Create module root file

Write src/memory/mod.rs:

pub mod types;

pub use types::{ConsolidationFact, ConsolidationResult, MemoryCategory, MemoryEntry};

Step 4: Verify compilation

Run: cargo build 2>&1 | head -20 Expected: compiles without errors (types module loads).

Task 2: MemoryConfig + Config Extension

Files:

Modify: src/config/mod.rs

Step 1: Read current GatewayConfig struct for context

Already known from prior analysis — GatewayConfig has host, port, session_ttl_hours, etc.

Step 2: Add MemoryConfig struct

Insert after the GatewayConfig impl block (around line 211):

#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct MemoryConfig {
    /// Master switch for the memory system.
    #[serde(default)]
    pub enabled: bool,
    /// Provider name for consolidation LLM calls (key in `providers`).
    #[serde(default)]
    pub consolidation_provider: Option<String>,
    /// Model name for consolidation LLM calls (key in `models`).
    #[serde(default)]
    pub consolidation_model: Option<String>,
    /// Max knowledge entries injected into system prompt per turn.
    #[serde(default = "default_recall_limit")]
    pub recall_limit: usize,
    /// Number of turns without consolidation before forcing one.
    #[serde(default = "default_consolidation_turn_threshold")]
    pub consolidation_turn_threshold: usize,
    /// Idle minutes before triggering consolidation (for async channels).
    #[serde(default = "default_idle_consolidation_minutes")]
    pub idle_consolidation_minutes: u64,
    /// Days before timeline entries are auto-cleaned.
    #[serde(default = "default_timeline_retention_days")]
    pub timeline_retention_days: u64,
    /// Consecutive consolidation failures before degrading to raw archive.
    #[serde(default = "default_max_failures_before_degrade")]
    pub max_failures_before_degrade: usize,
}

fn default_recall_limit() -> usize { 5 }
fn default_consolidation_turn_threshold() -> usize { 3 }
fn default_idle_consolidation_minutes() -> u64 { 10 }
fn default_timeline_retention_days() -> u64 { 90 }
fn default_max_failures_before_degrade() -> usize { 3 }

Step 3: Add memory field to Config struct

Add after the channels field (around line 50-51 in Config):

    #[serde(default)]
    pub memory: MemoryConfig,

Step 4: Verify compilation

Run: cargo build 2>&1 | head -20 Expected: compiles. MemoryConfig default values are correct.

Task 3: SQLite Schema — memories Table + FTS5

Files:

Create: src/storage/memory.rs
Modify: src/storage/mod.rs

Step 1: Add module declaration in storage/mod.rs

Add after line 4 (pub mod scheduler;):

pub mod memory;

Step 2: Add memories table creation to init_schema()

In src/storage/mod.rs, after the messages index creation (around line 86), add:

        sqlx::query(
            r#"
            CREATE TABLE IF NOT EXISTS memories (
                id          TEXT PRIMARY KEY,
                key         TEXT NOT NULL UNIQUE,
                content     TEXT NOT NULL,
                category    TEXT NOT NULL DEFAULT 'knowledge',
                importance  REAL NOT NULL DEFAULT 0.5,
                session_id  TEXT,
                created_at  TEXT NOT NULL,
                updated_at  TEXT NOT NULL
            )
            "#,
        )
        .execute(&self.pool)
        .await?;

        // FTS5 virtual table for full-text search on memories
        sqlx::query(
            r#"
            CREATE VIRTUAL TABLE IF NOT EXISTS memory_fts USING fts5(
                key,
                content,
                content=memories,
                content_rowid=rowid
            )
            "#,
        )
        .execute(&self.pool)
        .await?;

Step 3: Add last_consolidated_at column to sessions (migration)

After the sessions table creation, add an ALTER TABLE migration:

        // Migration: add last_consolidated_at column if not exists
        sqlx::query(
            r#"
            ALTER TABLE sessions ADD COLUMN last_consolidated_at INTEGER
            "#,
        )
        .execute(&self.pool)
        .await
        .ok(); // Ignore error if column already exists

Step 4: Verify compilation

Run: cargo build 2>&1 | head -20 Expected: compiles. Schema initialization includes memories table and FTS5.

Task 4: Storage CRUD for Memories

Files:

Modify: src/storage/memory.rs

Step 1: Write memory CRUD operations

use sqlx::{Pool, Sqlite, Row};
use std::path::Path;

use crate::memory::{MemoryCategory, MemoryEntry};

use super::StorageError;

impl super::Storage {
    /// Store or update a memory entry (upsert by key).
    pub async fn upsert_memory(&self, entry: &MemoryEntry) -> Result<(), StorageError> {
        let category_str = entry.category.as_str();
        sqlx::query(
            r#"
            INSERT INTO memories (id, key, content, category, importance, session_id, created_at, updated_at)
            VALUES (?, ?, ?, ?, ?, ?, ?, ?)
            ON CONFLICT(key) DO UPDATE SET
                content = excluded.content,
                category = excluded.category,
                importance = excluded.importance,
                session_id = excluded.session_id,
                updated_at = excluded.updated_at
            "#,
        )
        .bind(&entry.id)
        .bind(&entry.key)
        .bind(&entry.content)
        .bind(category_str)
        .bind(entry.importance)
        .bind(&entry.session_id)
        .bind(&entry.created_at)
        .bind(&entry.updated_at)
        .execute(&self.pool)
        .await?;
        Ok(())
    }

    /// Delete a memory entry by key.
    pub async fn delete_memory(&self, key: &str) -> Result<(), StorageError> {
        sqlx::query("DELETE FROM memories WHERE key = ?")
            .bind(key)
            .execute(&self.pool)
            .await?;
        Ok(())
    }

    /// Search memories by keyword using FTS5.
    /// Falls back to LIKE query if FTS5 returns no results.
    pub async fn search_memories(
        &self,
        query: &str,
        category: Option<&MemoryCategory>,
        limit: usize,
    ) -> Result<Vec<MemoryEntry>, StorageError> {
        // Build FTS5 query: wrap each word in quotes and join with OR
        let fts_query = query
            .split_whitespace()
            .map(|w| format!("\"{}\"", w.replace('"', "")))
            .collect::<Vec<_>>()
            .join(" OR ");

        let category_filter = category.map(|c| c.as_str());

        // Try FTS5 first
        let rows = sqlx::query(
            r#"
            SELECT m.id, m.key, m.content, m.category, m.importance,
                   m.session_id, m.created_at, m.updated_at
            FROM memory_fts f
            JOIN memories m ON f.rowid = m.rowid
            WHERE memory_fts MATCH ? AND (? IS NULL OR m.category = ?)
            ORDER BY rank
            LIMIT ?
            "#,
        )
        .bind(&fts_query)
        .bind(category_filter)
        .bind(category_filter)
        .bind(limit as i64)
        .fetch_all(&self.pool)
        .await?;

        let mut entries = parse_memory_rows(&rows)?;

        // Fallback to LIKE if FTS5 returned nothing
        if entries.is_empty() {
            let like_pattern = format!("%{}%", query.replace('%', "").replace('_', ""));
            let rows = sqlx::query(
                r#"
                SELECT id, key, content, category, importance,
                       session_id, created_at, updated_at
                FROM memories
                WHERE (key LIKE ? OR content LIKE ?)
                  AND (? IS NULL OR category = ?)
                ORDER BY importance DESC, updated_at DESC
                LIMIT ?
                "#,
            )
            .bind(&like_pattern)
            .bind(&like_pattern)
            .bind(category_filter)
            .bind(category_filter)
            .bind(limit as i64)
            .fetch_all(&self.pool)
            .await?;

            entries = parse_memory_rows(&rows)?;
        }

        Ok(entries)
    }

    /// Retrieve memories within a time range.
    pub async fn search_memories_by_time(
        &self,
        since: i64,
        until: i64,
        category: Option<&MemoryCategory>,
        limit: usize,
    ) -> Result<Vec<MemoryEntry>, StorageError> {
        let category_filter = category.map(|c| c.as_str());
        let since_str = chrono::DateTime::from_timestamp_millis(since)
            .unwrap_or_default()
            .to_rfc3339();
        let until_str = chrono::DateTime::from_timestamp_millis(until)
            .unwrap_or_default()
            .to_rfc3339();

        let rows = sqlx::query(
            r#"
            SELECT id, key, content, category, importance,
                   session_id, created_at, updated_at
            FROM memories
            WHERE created_at >= ? AND created_at <= ?
              AND (? IS NULL OR category = ?)
            ORDER BY created_at DESC
            LIMIT ?
            "#,
        )
        .bind(&since_str)
        .bind(&until_str)
        .bind(category_filter)
        .bind(category_filter)
        .bind(limit as i64)
        .fetch_all(&self.pool)
        .await?;

        parse_memory_rows(&rows)
    }

    /// Delete old timeline entries beyond retention period.
    pub async fn cleanup_old_timelines(&self, retention_days: u64) -> Result<u64, StorageError> {
        let cutoff = chrono::Utc::now() - chrono::Duration::days(retention_days as i64);
        let cutoff_str = cutoff.to_rfc3339();

        let result = sqlx::query(
            "DELETE FROM memories WHERE category = 'timeline' AND created_at < ?",
        )
        .bind(&cutoff_str)
        .execute(&self.pool)
        .await?;

        Ok(result.rows_affected())
    }
}

fn parse_memory_rows(
    rows: &[sqlx::sqlite::SqliteRow],
) -> Result<Vec<MemoryEntry>, StorageError> {
    rows.iter()
        .map(|row| {
            Ok(MemoryEntry {
                id: row.try_get("id")?,
                key: row.try_get("key")?,
                content: row.try_get("content")?,
                category: MemoryCategory::from_str(&row.try_get::<String, _>("category")?)
                    .unwrap_or(MemoryCategory::Knowledge),
                importance: row.try_get::<f64, _>("importance")?,
                session_id: row.try_get::<Option<String>, _>("session_id")?,
                created_at: row.try_get("created_at")?,
                updated_at: row.try_get("updated_at")?,
            })
        })
        .collect()
}

Step 2: Verify compilation

Run: cargo build 2>&1 | head -30 Expected: compiles. Fix any import errors for chrono.

Task 5: MemoryManager API

Files:

Modify: src/memory/mod.rs

Step 1: Write MemoryManager

Replace src/memory/mod.rs content:

pub mod types;

use std::sync::Arc;
use uuid::Uuid;

use crate::storage::Storage;
pub use types::{ConsolidationFact, ConsolidationResult, MemoryCategory, MemoryEntry};

/// MemoryManager provides high-level memory operations.
/// Wraps the Storage SQLite layer with semantic methods.
#[derive(Clone)]
pub struct MemoryManager {
    storage: Arc<Storage>,
}

impl MemoryManager {
    pub fn new(storage: Arc<Storage>) -> Self {
        Self { storage }
    }

    /// Store or update a memory entry. Generates timestamp and UUID.
    pub async fn store(
        &self,
        key: &str,
        content: &str,
        category: MemoryCategory,
        session_id: Option<&str>,
        importance: Option<f64>,
    ) -> Result<(), crate::storage::StorageError> {
        let now = chrono::Utc::now().to_rfc3339();
        let entry = MemoryEntry {
            id: Uuid::new_v4().to_string(),
            key: key.to_string(),
            content: content.to_string(),
            category,
            importance: importance.unwrap_or(0.5),
            session_id: session_id.map(|s| s.to_string()),
            created_at: now.clone(),
            updated_at: now,
        };
        self.storage.upsert_memory(&entry).await
    }

    /// Search memories by keyword query. Returns entries sorted by relevance.
    pub async fn recall(
        &self,
        query: &str,
        limit: usize,
        category: Option<MemoryCategory>,
    ) -> Result<Vec<MemoryEntry>, crate::storage::StorageError> {
        self.storage.search_memories(query, category.as_ref(), limit).await
    }

    /// Search memories by time range (Unix milliseconds).
    pub async fn recall_by_time(
        &self,
        since: i64,
        until: i64,
        limit: usize,
        category: Option<MemoryCategory>,
    ) -> Result<Vec<MemoryEntry>, crate::storage::StorageError> {
        self.storage.search_memories_by_time(since, until, category.as_ref(), limit).await
    }

    /// Delete a memory entry by key.
    pub async fn forget(&self, key: &str) -> Result<(), crate::storage::StorageError> {
        self.storage.delete_memory(key).await
    }

    /// Check if the memory system has any entries (for testing/health check).
    pub async fn is_empty(&self) -> Result<bool, crate::storage::StorageError> {
        self.recall("*", 1, None).await.map(|r| r.is_empty())
    }
}

Step 2: Verify compilation

Run: cargo build 2>&1 | head -20 Expected: compiles. MemoryManager is usable.

Task 6: Session last_consolidated_at Field

Files:

Modify: src/session/session.rs
Modify: src/storage/session.rs

Step 1: Add field to Session struct

In src/session/session.rs, add to the Session struct (after routing_info):

    /// Timestamp (Unix ms) of the last consolidation.
    /// Messages before this time have been compressed into memory.
    pub last_consolidated_at: Option<i64>,

Step 2: Initialize in Session constructor

Find the Session::new() or Session::create() method. Add to its initialization:

            last_consolidated_at: None,

Step 3: Update session load/save to include the field

In src/storage/session.rs, find the session reading/writing queries and ensure last_consolidated_at is included:

Read query: add last_consolidated_at to SELECT and to the struct mapping. Write query: add last_consolidated_at to INSERT/UPDATE.

Step 4: Verify compilation

Run: cargo build 2>&1 | head -20 Expected: compiles.

Task 7: ContextCompressor Memory Integration

Files:

Modify: src/agent/context_compressor.rs

Step 1: Add memory field to ContextCompressor

Add after provider field (around line 51):

    /// Memory manager handle (optional). When set, compressed
    /// context summaries are persisted as timeline memory entries.
    memory: Option<Arc<crate::memory::MemoryManager>>,

Step 2: Add with_memory() constructor method

Add after with_config() method (around line 77):

    /// Attach a memory manager to persist compressed summaries.
    pub fn with_memory(mut self, memory: Arc<crate::memory::MemoryManager>) -> Self {
        self.memory = Some(memory);
        self
    }

Step 3: Initialize memory field in constructors

In new():

            memory: None,

In with_config():

            memory: None,

Step 4: Write timeline entry during compress_once

In the compress_once() method, after generating the summary, add:

        // Persist compressed summary as timeline memory entry
        if let Some(ref mm) = self.memory {
            let ts = chrono::Utc::now().format("%Y-%m-%d %H:%M").to_string();
            let timeline_content = format!("[{}] Compressed {} conversation segments:\n{}",
                ts, segment_count, summary);
            let key = format!("ctx_compressed_{}", uuid::Uuid::new_v4());
            // Fire-and-forget: don't block compression on memory write
            let mm_clone = mm.clone();
            let session_id = session_id.cloned();
            tokio::spawn(async move {
                if let Err(e) = mm_clone.store(
                    &key,
                    &timeline_content,
                    crate::memory::MemoryCategory::Timeline,
                    session_id.as_deref(),
                    Some(0.3),
                ).await {
                    tracing::warn!(error = %e, "Failed to store compressed context as timeline");
                }
            });
        }

Note: this requires compress_once() to accept an Option<&str> session_id parameter, or the ContextCompressor to track it.

Step 5: Add session_id tracking to ContextCompressor

Add a field for the current session_id that gets updated via a setter:

    /// Current session ID for timeline memory writes.
    session_id: Option<String>,

Add setter:

    pub fn set_session_id(&mut self, id: Option<String>) {
        self.session_id = id;
    }

Initialize in constructors:

            session_id: None,

Step 6: Verify compilation

Run: cargo build 2>&1 | head -30 Expected: compiles. Import fixes for uuid may be needed.

Task 8: SystemPromptBuilder Memory Context Injection

Files:

Modify: src/agent/system_prompt.rs

Step 1: Add MemoryContextSection

After CrossChannelSection impl (before Helper Functions):

/// Injects relevant knowledge memories into the system prompt.
pub struct MemoryContextSection {
    memory: Option<Arc<crate::memory::MemoryManager>>,
    recall_limit: usize,
}

impl MemoryContextSection {
    pub fn new(memory: Option<Arc<crate::memory::MemoryManager>>, recall_limit: usize) -> Self {
        Self { memory, recall_limit }
    }
}

impl PromptSection for MemoryContextSection {
    fn name(&self) -> &str {
        "memory_context"
    }

    fn build(&self, ctx: &PromptContext<'_>) -> String {
        let Some(ref memory) = self.memory else {
            return String::new();
        };

        // We cannot do async in a sync trait method.
        // Instead, the caller should pre-fetch memories and pass them via PromptContext.
        // For now, return empty — memory context is injected before building the prompt.
        String::new()
    }
}

Step 2: Add memory_context field to PromptContext

pub struct PromptContext<'a> {
    pub workspace_dir: &'a Path,
    pub model_name: &'a str,
    pub tools: &'a ToolRegistry,
    pub session_id: Option<&'a str>,
    /// Pre-fetched memory context string to inject.
    pub memory_context: Option<&'a str>,
}

Step 3: Add MemorySection to default sections

In SystemPromptBuilder::with_defaults(), add:

    Box::new(MemorySection),

Step 4: Write MemorySection that renders the pre-fetched context

pub struct MemorySection;

impl PromptSection for MemorySection {
    fn name(&self) -> &str {
        "memory"
    }

    fn build(&self, ctx: &PromptContext<'_>) -> String {
        match ctx.memory_context {
            Some(context) if !context.is_empty() => {
                format!("## 记忆上下文\n\n{}", context)
            }
            _ => String::new(),
        }
    }
}

Step 5: Update build_system_prompt() signature

Add memory_context: Option<&str> parameter:

pub fn build_system_prompt(
    workspace_dir: &Path,
    model_name: &str,
    tools: &ToolRegistry,
    session_id: Option<&str>,
    memory_context: Option<&str>,
) -> String {
    let ctx = PromptContext {
        workspace_dir,
        model_name,
        tools,
        session_id,
        memory_context,
    };
    SystemPromptBuilder::with_defaults().build(&ctx)
}

Step 6: Update callers of build_system_prompt()

In src/session/session.rs, find all calls to build_system_prompt() and add None as the memory_context parameter (will be wired later).

Step 7: Verify compilation

Run: cargo build 2>&1 | head -30 Expected: compiles. Fix all call sites.

Task 9: Agent Memory Tools

Files:

Create: src/tools/memory_store.rs
Create: src/tools/memory_recall.rs
Create: src/tools/memory_forget.rs
Modify: src/tools/mod.rs

Step 1: Write memory_store tool

src/tools/memory_store.rs:

use super::traits::{Tool, ToolResult};
use async_trait::async_trait;
use serde_json::json;
use std::sync::Arc;

use crate::memory::{MemoryCategory, MemoryManager};

pub struct MemoryStoreTool {
    memory: Arc<MemoryManager>,
}

impl MemoryStoreTool {
    pub fn new(memory: Arc<MemoryManager>) -> Self {
        Self { memory }
    }
}

#[async_trait]
impl Tool for MemoryStoreTool {
    fn name(&self) -> &str { "memory_store" }

    fn description(&self) -> &str {
        "Store a fact, preference, or insight into long-term memory. \
         Use this when the user shares important information you should remember. \
         Provide a descriptive key (e.g., 'user_prefers_python', 'project_auth_approach') \
         and the full content to remember."
    }

    fn read_only(&self) -> bool { false }

    fn parameters_schema(&self) -> serde_json::Value {
        json!({
            "type": "object",
            "properties": {
                "key": {
                    "type": "string",
                    "description": "Semantic identifier for this memory (e.g., 'user_language_pref'). Unique key."
                },
                "content": {
                    "type": "string",
                    "description": "The full content of the memory entry."
                },
                "category": {
                    "type": "string",
                    "enum": ["knowledge", "timeline"],
                    "description": "Memory category. Use 'knowledge' for facts/preferences/insights, 'timeline' for conversation summaries."
                },
                "importance": {
                    "type": "number",
                    "description": "Importance score 0.0-1.0. Higher = more important. Use 0.8+ for critical facts, 0.5 for general info."
                }
            },
            "required": ["key", "content"]
        })
    }

    async fn execute(&self, args: serde_json::Value) -> anyhow::Result<ToolResult> {
        let key = args.get("key").and_then(|v| v.as_str())
            .ok_or_else(|| anyhow::anyhow!("Missing required parameter: key"))?;

        let content = args.get("content").and_then(|v| v.as_str())
            .ok_or_else(|| anyhow::anyhow!("Missing required parameter: content"))?;

        let category = args.get("category")
            .and_then(|v| v.as_str())
            .and_then(MemoryCategory::from_str)
            .unwrap_or(MemoryCategory::Knowledge);

        let importance = args.get("importance")
            .and_then(|v| v.as_f64());

        self.memory.store(key, content, category, None, importance).await?;

        Ok(ToolResult {
            success: true,
            output: format!("Memory stored: {}", key),
            error: None,
        })
    }
}

Step 2: Write memory_recall tool

src/tools/memory_recall.rs:

use super::traits::{Tool, ToolResult};
use async_trait::async_trait;
use serde_json::json;
use std::sync::Arc;

use crate::memory::{MemoryCategory, MemoryManager};

pub struct MemoryRecallTool {
    memory: Arc<MemoryManager>,
}

impl MemoryRecallTool {
    pub fn new(memory: Arc<MemoryManager>) -> Self {
        Self { memory }
    }
}

#[async_trait]
impl Tool for MemoryRecallTool {
    fn name(&self) -> &str { "memory_recall" }

    fn description(&self) -> &str {
        "Search and retrieve entries from long-term memory. \
         Use this to recall previously stored facts, preferences, or conversation history. \
         Supports keyword search and optional time-range filtering."
    }

    fn read_only(&self) -> bool { true }

    fn parameters_schema(&self) -> serde_json::Value {
        json!({
            "type": "object",
            "properties": {
                "query": {
                    "type": "string",
                    "description": "Search query — keywords to match against memory keys and content."
                },
                "category": {
                    "type": "string",
                    "enum": ["knowledge", "timeline"],
                    "description": "Filter by memory category. Omit to search all categories."
                },
                "since": {
                    "type": "integer",
                    "description": "Start of time range (Unix milliseconds)."
                },
                "until": {
                    "type": "integer",
                    "description": "End of time range (Unix milliseconds)."
                },
                "limit": {
                    "type": "integer",
                    "description": "Max results to return (default 10)."
                }
            },
            "required": ["query"]
        })
    }

    async fn execute(&self, args: serde_json::Value) -> anyhow::Result<ToolResult> {
        let query = args.get("query").and_then(|v| v.as_str())
            .ok_or_else(|| anyhow::anyhow!("Missing required parameter: query"))?;

        let category = args.get("category")
            .and_then(|v| v.as_str())
            .and_then(MemoryCategory::from_str);

        let limit = args.get("limit")
            .and_then(|v| v.as_u64())
            .unwrap_or(10) as usize;

        let entries = if args.get("since").is_some() || args.get("until").is_some() {
            let since = args.get("since").and_then(|v| v.as_i64()).unwrap_or(0);
            let until = args.get("until").and_then(|v| v.as_i64())
                .unwrap_or(chrono::Utc::now().timestamp_millis());
            self.memory.recall_by_time(since, until, limit, category).await?
        } else {
            self.memory.recall(query, limit, category).await?
        };

        if entries.is_empty() {
            return Ok(ToolResult {
                success: true,
                output: "No matching memories found.".to_string(),
                error: None,
            });
        }

        let formatted = entries.iter()
            .map(|e| format!("- {} [{}] [importance: {:.1}]: {}",
                e.key, e.category.as_str(), e.importance, e.content))
            .collect::<Vec<_>>()
            .join("\n");

        Ok(ToolResult {
            success: true,
            output: format!("Found {} memories:\n{}", entries.len(), formatted),
            error: None,
        })
    }
}

Step 3: Write memory_forget tool

src/tools/memory_forget.rs:

use super::traits::{Tool, ToolResult};
use async_trait::async_trait;
use serde_json::json;
use std::sync::Arc;

use crate::memory::MemoryManager;

pub struct MemoryForgetTool {
    memory: Arc<MemoryManager>,
}

impl MemoryForgetTool {
    pub fn new(memory: Arc<MemoryManager>) -> Self {
        Self { memory }
    }
}

#[async_trait]
impl Tool for MemoryForgetTool {
    fn name(&self) -> &str { "memory_forget" }

    fn description(&self) -> &str {
        "Delete a memory entry by its key. Use this when information is outdated, \
         incorrect, or the user asks to forget something."
    }

    fn read_only(&self) -> bool { false }

    fn parameters_schema(&self) -> serde_json::Value {
        json!({
            "type": "object",
            "properties": {
                "key": {
                    "type": "string",
                    "description": "The key of the memory entry to delete."
                }
            },
            "required": ["key"]
        })
    }

    async fn execute(&self, args: serde_json::Value) -> anyhow::Result<ToolResult> {
        let key = args.get("key").and_then(|v| v.as_str())
            .ok_or_else(|| anyhow::anyhow!("Missing required parameter: key"))?;

        self.memory.forget(key).await?;

        Ok(ToolResult {
            success: true,
            output: format!("Memory deleted: {}", key),
            error: None,
        })
    }
}

Step 4: Register tools in create_default_tools()

In src/tools/mod.rs, add module declarations:

pub mod memory_store;
pub mod memory_recall;
pub mod memory_forget;

Add pub use:

pub use memory_store::MemoryStoreTool;
pub use memory_recall::MemoryRecallTool;
pub use memory_forget::MemoryForgetTool;

Modify create_default_tools() signature to accept Option<Arc<MemoryManager>>:

pub fn create_default_tools(
    skills_loader: Arc<SkillsLoader>,
    memory: Option<Arc<crate::memory::MemoryManager>>,
) -> ToolRegistry {
    // ... existing tools ...

    // Register memory tools if memory system is available
    if let Some(mm) = memory {
        registry.register(MemoryStoreTool::new(mm.clone()));
        registry.register(MemoryRecallTool::new(mm.clone()));
        registry.register(MemoryForgetTool::new(mm.clone()));
    }

    registry
}

Step 5: Update callers of create_default_tools()

In src/session/session.rs, find the call to create_default_tools() and pass None for now (will be wired in GatewayState task).

Step 6: Verify compilation

Run: cargo build 2>&1 | head -30 Expected: compiles. Fix import issues.

Task 10: GatewayState Initialization — Wiring

Files:

Modify: src/gateway/mod.rs

Step 1: Initialize MemoryManager in GatewayState::new()

After the storage initialization (around line 56), add:

        // Initialize MemoryManager if memory system is enabled
        let memory_manager = if config.memory.enabled {
            let mm = Arc::new(crate::memory::MemoryManager::new(storage.clone()));
            tracing::info!("Memory system enabled");
            Some(mm)
        } else {
            None
        };

Step 2: Pass memory_manager to create_default_tools()

Update the call to create_default_tools() in SessionManager::new() to accept memory_manager.clone().

Note: SessionManager::new() currently calls create_default_tools() internally. You'll need to propagate the memory_manager parameter through SessionManager::new().

In src/session/session.rs, update SessionManager::new():

impl SessionManager {
    pub fn new(
        session_ttl_hours: u32,
        provider_config: LLMProviderConfig,
        storage: Arc<Storage>,
        bus: Arc<MessageBus>,
        memory_manager: Option<Arc<crate::memory::MemoryManager>>,
    ) -> Result<Arc<Self>, SessionError> {
        // ... existing init ...

        let skills_loader = Arc::new(SkillsLoader::new());
        let tools = create_default_tools(skills_loader, memory_manager.clone());

        // Store memory_manager on SessionManager for later use
        // (add memory_manager field to SessionManager struct)
    }
}

Step 3: Attach MemoryManager to ContextCompressor in Session creation

In the Session creation path within SessionManager, when creating the ContextCompressor:

    let mut compressor = ContextCompressor::new(provider.clone(), token_limit);
    if let Some(ref mm) = memory_manager {
        compressor = compressor.with_memory(mm.clone());
    }

Step 4: Wire memory context into system prompt building

In handle_message() (session.rs), before building the system prompt:

    // Fetch memory context
    let memory_context = if let Some(ref mm) = memory_manager {
        match mm.recall(&user_message, 5, Some(MemoryCategory::Knowledge)).await {
            Ok(entries) if !entries.is_empty() => {
                Some(entries.iter()
                    .map(|e| format!("- {}: {}", e.key, e.content))
                    .collect::<Vec<_>>()
                    .join("\n"))
            }
            _ => None,
        }
    } else {
        None
    };

    let system_prompt = build_system_prompt(
        &workspace_dir,
        &model_name,
        &tools,
        session_id.as_deref(),
        memory_context.as_deref(),
    );

Step 5: Verify compilation

Run: cargo build 2>&1 | head -40 Expected: compiles. Fix all wiring issues. This is the most integration-heavy task.

Task 11: Unit Tests

Files:

Create: src/memory/types.rs (add tests module)
Modify: src/storage/memory.rs (add tests module)
Modify: src/memory/mod.rs (add tests module)

Step 1: Type tests

In src/memory/types.rs, add:

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

    #[test]
    fn test_memory_category_as_str() {
        assert_eq!(MemoryCategory::Knowledge.as_str(), "knowledge");
        assert_eq!(MemoryCategory::Timeline.as_str(), "timeline");
    }

    #[test]
    fn test_memory_category_from_str() {
        assert_eq!(MemoryCategory::from_str("knowledge"), Some(MemoryCategory::Knowledge));
        assert_eq!(MemoryCategory::from_str("timeline"), Some(MemoryCategory::Timeline));
        assert_eq!(MemoryCategory::from_str("invalid"), None);
    }
}

Run: cargo test --lib memory::types::tests Expected: 2 tests pass.

Step 2: MemoryManager integration test (with real Storage)

In src/memory/mod.rs, add:

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::Arc;
    use tempfile::tempdir;

    async fn setup_memory_manager() -> (Arc<MemoryManager>, tempfile::TempDir) {
        let dir = tempdir().unwrap();
        let db_path = dir.path().join("test.db");
        let storage = Arc::new(crate::storage::Storage::new(&db_path).await.unwrap());
        let mm = Arc::new(MemoryManager::new(storage));
        (mm, dir)
    }

    #[tokio::test]
    async fn test_store_and_recall() {
        let (mm, _dir) = setup_memory_manager().await;

        mm.store("test_key", "This is a test memory", MemoryCategory::Knowledge, None, Some(0.8))
            .await
            .unwrap();

        let results = mm.recall("test memory", 10, None).await.unwrap();
        assert_eq!(results.len(), 1);
        assert_eq!(results[0].key, "test_key");
        assert_eq!(results[0].content, "This is a test memory");
        assert!((results[0].importance - 0.8).abs() < 0.01);
    }

    #[tokio::test]
    async fn test_upsert_overwrites() {
        let (mm, _dir) = setup_memory_manager().await;

        mm.store("dup_key", "original", MemoryCategory::Knowledge, None, None)
            .await.unwrap();
        mm.store("dup_key", "updated", MemoryCategory::Knowledge, None, Some(0.9))
            .await.unwrap();

        let results = mm.recall("updated", 10, None).await.unwrap();
        assert_eq!(results.len(), 1);
        assert_eq!(results[0].content, "updated");
    }

    #[tokio::test]
    async fn test_forget() {
        let (mm, _dir) = setup_memory_manager().await;

        mm.store("to_delete", "will be deleted", MemoryCategory::Knowledge, None, None)
            .await.unwrap();
        mm.forget("to_delete").await.unwrap();

        let results = mm.recall("deleted", 10, None).await.unwrap();
        assert!(results.is_empty());
    }

    #[tokio::test]
    async fn test_category_filter() {
        let (mm, _dir) = setup_memory_manager().await;

        mm.store("knowledge_1", "fact content", MemoryCategory::Knowledge, None, None)
            .await.unwrap();
        mm.store("timeline_1", "summary content", MemoryCategory::Timeline, None, None)
            .await.unwrap();

        let know_results = mm.recall("content", 10, Some(MemoryCategory::Knowledge)).await.unwrap();
        assert_eq!(know_results.len(), 1);
        assert_eq!(know_results[0].key, "knowledge_1");

        let time_results = mm.recall("content", 10, Some(MemoryCategory::Timeline)).await.unwrap();
        assert_eq!(time_results.len(), 1);
        assert_eq!(time_results[0].key, "timeline_1");
    }
}

Run: cargo test --lib memory::tests Expected: 4 tests pass (store+recall, upsert, forget, category filter).

Step 3: SystemPromptBuilder memory_context test

In src/agent/system_prompt.rs tests module, add:

    #[test]
    fn test_memory_section_with_context() {
        let temp_dir = std::env::temp_dir();
        let tools = ToolRegistry::new();

        let ctx = PromptContext {
            workspace_dir: &temp_dir,
            model_name: "test",
            tools: &tools,
            session_id: None,
            memory_context: Some("- user_pref: Prefers Rust"),
        };

        let prompt = SystemPromptBuilder::with_defaults().build(&ctx);
        assert!(prompt.contains("## 记忆上下文"));
        assert!(prompt.contains("Prefers Rust"));
    }

    #[test]
    fn test_memory_section_without_context() {
        let temp_dir = std::env::temp_dir();
        let tools = ToolRegistry::new();

        let ctx = PromptContext {
            workspace_dir: &temp_dir,
            model_name: "test",
            tools: &tools,
            session_id: None,
            memory_context: None,
        };

        let prompt = SystemPromptBuilder::with_defaults().build(&ctx);
        assert!(!prompt.contains("## 记忆上下文"));
    }

Run: cargo test --lib agent::system_prompt::tests Expected: all tests pass including new memory tests.

Task 12: Full Build and Test

Step 1: Full build

Run: cargo build 2>&1 Expected: zero errors, zero warnings.

Step 2: All unit tests

Run: cargo test --lib Expected: all tests pass.

Step 3: Integration tests

Run: cargo test -- --ignored (requires test.env with API keys) Expected: existing integration tests still pass.

Step 4: Commit

git add -A && git commit -m "feat: add memory system with FTS5 search and context compression integration"

40 KiB Raw Blame History Unescape Escape

Memory System Implementation Plan

Task 1: Memory Type Definitions

Task 2: MemoryConfig + Config Extension

Task 3: SQLite Schema — memories Table + FTS5

Task 4: Storage CRUD for Memories

Task 5: MemoryManager API

Task 6: Session last_consolidated_at Field

Task 7: ContextCompressor Memory Integration

Task 8: SystemPromptBuilder Memory Context Injection

Task 9: Agent Memory Tools

Task 10: GatewayState Initialization — Wiring

Task 11: Unit Tests

Task 12: Full Build and Test

40 KiB

Raw Blame History