mcp-pocket-memory 0.4.0

use std::{
    collections::BTreeSet,
    env,
    path::{Path, PathBuf},
    sync::Arc,
    time::Duration,
};

use ax_cache::Cache;
use rmcp::{
    ErrorData as McpError, ServerHandler,
    handler::server::{
        router::tool::ToolRouter,
        wrapper::{Json, Parameters},
    },
    model::{ServerCapabilities, ServerInfo},
    schemars, tool, tool_handler, tool_router,
};
use serde::{Deserialize, Serialize};

use crate::db::{Database, Memory};

const FIND_CACHE_CAPACITY_ENV: &str = "POCKET_MEMORY_FIND_CACHE_CAPACITY";
const FIND_CACHE_TTL_SECS_ENV: &str = "POCKET_MEMORY_FIND_CACHE_TTL_SECS";
const DEFAULT_FIND_CACHE_CAPACITY: usize = 512;
const DEFAULT_FIND_CACHE_TTL_SECS: u64 = 30;

#[derive(Debug, Deserialize, schemars::JsonSchema)]
pub struct UpsertMemoryRequest {
    #[schemars(
        description = "Stable kebab-case identifier. Reuse the exact same key to update existing memory. Examples: 'error-handling-pattern', 'db-pool-config', 'frontend-state-rule'"
    )]
    pub key: String,
    #[schemars(
        description = "Repository name this memory belongs to (e.g., 'mcp-pocket-memory'). REQUIRED for project-specific knowledge. Use null ONLY for universal/cross-repo knowledge."
    )]
    pub repo: Option<String>,
    #[schemars(
        description = "Optional topic tags for retrieval. Use 1-10 short lowercase tags. Standard tags such as architecture, api, auth, bug, convention, decision, dependency, deployment, performance, preference, requirement, security, testing, unicode, workaround are recommended, but custom tags like edge-case or long-text are allowed."
    )]
    pub tags: Option<Vec<String>>,
    #[schemars(
        description = "Concise but complete durable knowledge. Include WHAT, WHY, and WHEN. 1-3 sentences. Store only cross-session knowledge: decisions, conventions, bugs, constraints, preferences. DO NOT store current file contents or temporary analysis."
    )]
    pub content: String,
}

#[derive(Debug, Deserialize, schemars::JsonSchema)]
pub struct DeleteMemoryRequest {
    #[schemars(description = "Exact key of the memory to permanently delete.")]
    pub key: String,
}

#[derive(Debug, Deserialize, schemars::JsonSchema)]
pub struct FindMemoryRequest {
    #[schemars(
        description = "Optional broad keyword text. Split into words and matched against memory key, content, repo, and tags. Use a few terms when you do not know the exact wording."
    )]
    pub keyword: Option<String>,
    #[schemars(
        description = "Repository name filter. ALWAYS provide this when working in a known repository. This is the highest-priority filter."
    )]
    pub repo: Option<String>,
    #[schemars(
        description = "Optional topic tags. Multiple tags use AND logic (all must match). Use to narrow context quickly. Custom tags are allowed."
    )]
    pub tags: Option<Vec<String>>,
    #[schemars(description = "Maximum results. Values clamped to 1-50. Default 10.")]
    pub limit: Option<usize>,
}

#[derive(Debug, Serialize, schemars::JsonSchema)]
pub struct UpsertMemoryResponse {
    pub memory: Memory,
}

#[derive(Debug, Serialize, schemars::JsonSchema)]
pub struct DeleteMemoryResponse {
    pub key: String,
    pub deleted: bool,
}

#[derive(Debug, Serialize, schemars::JsonSchema)]
pub struct FindMemoryResponse {
    pub memories: Vec<Memory>,
}

#[derive(Clone)]
pub struct PocketMemoryServer {
    db_path: PathBuf,
    find_cache: Arc<Cache<String, Vec<Memory>>>,
    find_cache_ttl: Duration,
    #[allow(dead_code)]
    tool_router: ToolRouter<Self>,
}

impl PocketMemoryServer {
    pub fn new(db_path: impl AsRef<Path>) -> Self {
        let cache_capacity = find_cache_capacity();
        Self {
            db_path: db_path.as_ref().to_owned(),
            find_cache: Arc::new(Cache::new(cache_capacity)),
            find_cache_ttl: find_cache_ttl(),
            tool_router: Self::tool_router(),
        }
    }

    fn db(&self) -> Result<Database, McpError> {
        Database::open(&self.db_path).map_err(map_db_error)
    }
}

#[tool_router]
impl PocketMemoryServer {
    #[tool(
        name = "upsert_memory",
        title = "Upsert Memory",
        description = "Store durable, cross-session knowledge. CALL THIS IMMEDIATELY when you learn:
- A project convention or rule (e.g., 'use anyhow for errors', 'naming convention for tables')
- An architecture decision and its rationale
- A recurring bug and its workaround
- A user preference or constraint
- A dependency version requirement or compatibility issue

DO NOT store: current file contents, temporary analysis, or information retrievable from code.

Parameters:
- key: Stable kebab-case identifier. Reuse exact key to update. Examples: 'error-handling-pattern', 'db-connection-pool-size', 'frontend-state-management'
- repo: REQUIRED if memory is project-specific. Use the repository name (e.g., 'mcp-pocket-memory'). Use null ONLY for universal knowledge.
- tags: Optional. Use 1-10 short tags when useful. Prefer consistent tags, but custom tags such as edge-case, security, unicode, and long-text are allowed.
- content: Concise but complete. Include WHAT the decision/rule is, WHY it exists, and WHEN it applies. 1-3 sentences.",
        annotations(
            title = "Upsert Memory",
            read_only_hint = false,
            destructive_hint = false,
            idempotent_hint = false,
            open_world_hint = false
        )
    )]
    fn upsert_memory(
        &self,
        Parameters(request): Parameters<UpsertMemoryRequest>,
    ) -> Result<Json<UpsertMemoryResponse>, McpError> {
        let key = clean_key(request.key)?;
        let repo = clean_optional(request.repo);
        let tags = clean_tags(request.tags)?;
        let content = clean_required("content", request.content)?;
        let memory = self
            .db()?
            .upsert_memory(&key, repo.as_deref(), &tags, &content)
            .map_err(map_db_error)?;
        self.find_cache.clear();

        Ok(Json(UpsertMemoryResponse { memory }))
    }

    #[tool(
        name = "delete_memory",
        title = "Delete Memory",
        description = "Permanently remove a memory by its exact key. Use ONLY when:
- The information is factually wrong and misleading
- The decision/convention has been explicitly revoked by the user
- The memory is obsolete (e.g., old bug fixed, old dependency removed)

DO NOT use this to correct a memory — use upsert_memory with the same key to update instead.",
        annotations(
            title = "Delete Memory",
            read_only_hint = false,
            destructive_hint = true,
            idempotent_hint = true,
            open_world_hint = false
        )
    )]
    fn delete_memory(
        &self,
        Parameters(request): Parameters<DeleteMemoryRequest>,
    ) -> Result<Json<DeleteMemoryResponse>, McpError> {
        let key = clean_key(request.key)?;
        let deleted = self.db()?.delete_memory(&key).map_err(map_db_error)?;
        self.find_cache.clear();

        Ok(Json(DeleteMemoryResponse { key, deleted }))
    }

    #[tool(
        name = "find_memory",
        title = "Find Memory",
        description = "Retrieve durable memory BEFORE acting or answering. Always call this first. It can also be called without keyword, repo, or tags to load recent pocket-memory context.

Search strategy (apply in order):
1. If you know the current repo: ALWAYS include repo filter first.
2. If looking for a specific topic: add relevant tags (e.g., ['auth'] for authentication).
3. If you only know rough words: add keyword with a few terms. The search matches any term across key, content, repo, and tags.
4. If you do not know the repo or terms yet: call without filters to inspect recent memory context.

Parameter guidelines:
- repo: The current working repository name. This is the most important filter.
- tags: Use to narrow scope. Multiple tags use AND logic (all must match).
- keyword: Use broad word search across key, content, repo, and tags. Multiple words are searched broadly.
- limit: Default 10. Increase only if you need broad context.",
        annotations(
            title = "Find Memory",
            read_only_hint = true,
            destructive_hint = false,
            idempotent_hint = true,
            open_world_hint = false
        )
    )]
    fn find_memory(
        &self,
        Parameters(request): Parameters<FindMemoryRequest>,
    ) -> Result<Json<FindMemoryResponse>, McpError> {
        let keyword = request
            .keyword
            .map(|keyword| keyword.trim().to_owned())
            .filter(|keyword| !keyword.is_empty());
        let repo = clean_optional(request.repo);
        let tags = clean_tags(request.tags)?;

        let limit = request.limit.unwrap_or(10);
        let cache_key = find_cache_key(keyword.as_deref(), repo.as_deref(), &tags, limit);
        if let Some(memories) = self.find_cache.get(&cache_key) {
            return Ok(Json(FindMemoryResponse { memories }));
        }

        let memories = self
            .db()?
            .find_memories(keyword.as_deref(), repo.as_deref(), &tags, limit)
            .map_err(map_db_error)?;
        self.find_cache
            .insert_with_ttl(cache_key, memories.clone(), self.find_cache_ttl);

        Ok(Json(FindMemoryResponse { memories }))
    }
}

#[tool_handler]
impl ServerHandler for PocketMemoryServer {
    fn get_info(&self) -> ServerInfo {
        ServerInfo::new(ServerCapabilities::builder().enable_tools().build())
            .with_instructions(
                "You are working with a durable memory system. Follow this workflow STRICTLY:

1. BEFORE starting any task: call find_memory with repo=<current_repo> to load existing context.
2. If the current repo is unknown, call find_memory without filters or with a few broad keyword terms to inspect recent context.
3. BEFORE answering questions about project conventions, past decisions, or bugs: call find_memory first.
4. AFTER learning something that should survive across sessions (conventions, decisions, bugs, preferences, constraints): call upsert_memory immediately.
5. NEVER store transient information (current file contents, temporary thoughts, or information easily found in code).
6. For tags, prefer consistent short lowercase tags. Standard tags are recommended, but custom tags are allowed when they describe the memory better.
7. Keys must be kebab-case and descriptive (e.g., 'auth-jwt-strategy', 'db-migration-policy').",
            )
    }
}

fn clean_required(field: &'static str, value: String) -> Result<String, McpError> {
    let value = value.trim().to_owned();
    if value.is_empty() {
        return Err(McpError::invalid_params(
            format!("{field} must not be empty"),
            None,
        ));
    }

    Ok(value)
}

fn clean_key(value: String) -> Result<String, McpError> {
    let value = clean_required("key", value)?;
    if !is_kebab_case(&value) {
        return Err(McpError::invalid_params(
            "key must be kebab-case using lowercase letters, numbers, and single hyphens",
            None,
        ));
    }

    Ok(value)
}

fn clean_optional(value: Option<String>) -> Option<String> {
    value
        .map(|value| value.trim().to_owned())
        .filter(|value| !value.is_empty())
}

fn clean_tags(value: Option<Vec<String>>) -> Result<Vec<String>, McpError> {
    let tags = value
        .unwrap_or_default()
        .into_iter()
        .map(normalize_tag)
        .filter(|tag| !tag.is_empty())
        .collect::<BTreeSet<_>>()
        .into_iter()
        .collect::<Vec<_>>();

    if tags.len() > 10 {
        return Err(McpError::invalid_params(
            "tags must include at most 10 values",
            None,
        ));
    }

    if let Some(tag) = tags.iter().find(|tag| !is_valid_tag(tag)) {
        return Err(McpError::invalid_params(
            format!("invalid tag: {tag}"),
            None,
        ));
    }

    Ok(tags)
}

fn find_cache_key(
    keyword: Option<&str>,
    repo: Option<&str>,
    tags: &[String],
    limit: usize,
) -> String {
    format!(
        "keyword={}|repo={}|tags={}|limit={}",
        keyword.unwrap_or_default().trim().to_lowercase(),
        repo.unwrap_or_default().trim().to_lowercase(),
        tags.join(","),
        limit.clamp(1, 50),
    )
}

fn find_cache_capacity() -> usize {
    env::var(FIND_CACHE_CAPACITY_ENV)
        .ok()
        .and_then(|value| value.parse::<usize>().ok())
        .filter(|value| *value > 0)
        .unwrap_or(DEFAULT_FIND_CACHE_CAPACITY)
}

fn find_cache_ttl() -> Duration {
    let secs = env::var(FIND_CACHE_TTL_SECS_ENV)
        .ok()
        .and_then(|value| value.parse::<u64>().ok())
        .unwrap_or(DEFAULT_FIND_CACHE_TTL_SECS);

    Duration::from_secs(secs)
}

fn normalize_tag(tag: String) -> String {
    let mut normalized = String::new();
    let mut previous_hyphen = false;

    for ch in tag.trim().to_lowercase().chars() {
        let ch = match ch {
            'a'..='z' | '0'..='9' => ch,
            '-' | '_' | ' ' | '\t' => '-',
            _ => continue,
        };

        if ch == '-' {
            if normalized.is_empty() || previous_hyphen {
                continue;
            }
            previous_hyphen = true;
        } else {
            previous_hyphen = false;
        }

        normalized.push(ch);
    }

    normalized.trim_matches('-').to_owned()
}

fn is_kebab_case(value: &str) -> bool {
    let mut previous_hyphen = false;
    let mut has_char = false;

    for ch in value.chars() {
        match ch {
            'a'..='z' | '0'..='9' => {
                previous_hyphen = false;
                has_char = true;
            }
            '-' if has_char && !previous_hyphen => {
                previous_hyphen = true;
            }
            _ => return false,
        }
    }

    has_char && !previous_hyphen
}

fn is_valid_tag(tag: &str) -> bool {
    !tag.is_empty()
        && tag.len() <= 64
        && tag
            .chars()
            .all(|ch| matches!(ch, 'a'..='z' | '0'..='9' | '-'))
        && !tag.starts_with('-')
        && !tag.ends_with('-')
        && !tag.contains("--")
}

fn map_db_error(error: rusqlite::Error) -> McpError {
    McpError::internal_error(format!("database error: {error}"), None)
}

#[cfg(test)]
mod tests {
    use super::*;
    use rmcp::handler::server::wrapper::Parameters;
    use std::sync::Mutex;

    static ENV_LOCK: Mutex<()> = Mutex::new(());

    #[test]
    fn clean_key_accepts_kebab_case() {
        assert_eq!(
            clean_key("repo-context-rule".into()).unwrap(),
            "repo-context-rule"
        );
        assert_eq!(
            clean_key("auth-v2-policy".into()).unwrap(),
            "auth-v2-policy"
        );
    }

    #[test]
    fn clean_key_rejects_non_kebab_case() {
        assert!(clean_key("RepoContextRule".into()).is_err());
        assert!(clean_key("repo_context_rule".into()).is_err());
        assert!(clean_key("repo--context".into()).is_err());
    }

    #[test]
    fn clean_tags_accepts_optional_custom_tags() {
        assert_eq!(
            clean_tags(Some(vec![
                "Convention".into(),
                "edge case".into(),
                "long_text".into(),
                "security".into(),
            ]))
            .unwrap(),
            vec!["convention", "edge-case", "long-text", "security"]
        );
        assert!(clean_tags(None).unwrap().is_empty());
        assert!(
            clean_tags(Some(vec![
                "bug".into(),
                "auth".into(),
                "testing".into(),
                "decision".into(),
                "security".into(),
                "unicode".into(),
                "edge-case".into(),
                "long-text".into(),
                "convention".into(),
                "preference".into(),
                "performance".into(),
            ]),)
            .is_err()
        );
    }

    #[test]
    fn find_cache_key_is_normalized() {
        assert_eq!(
            find_cache_key(
                Some(" Gateway Auth "),
                Some(" Test Repo "),
                &["auth".into(), "decision".into()],
                99,
            ),
            "keyword=gateway auth|repo=test repo|tags=auth,decision|limit=50"
        );
    }

    #[test]
    fn upsert_clears_find_cache() {
        let db_path = std::env::temp_dir().join(format!(
            "mcp-pocket-memory-cache-{}.sqlite3",
            std::process::id()
        ));
        let _ = std::fs::remove_file(&db_path);
        Database::open(&db_path).unwrap();
        let server = PocketMemoryServer::new(&db_path);

        let before = server
            .find_memory(Parameters(FindMemoryRequest {
                keyword: None,
                repo: Some("test-repo".into()),
                tags: None,
                limit: Some(10),
            }))
            .unwrap();
        assert!(before.0.memories.is_empty());

        server
            .upsert_memory(Parameters(UpsertMemoryRequest {
                key: "cache-invalidation-test".into(),
                repo: Some("test-repo".into()),
                tags: Some(vec!["cache".into()]),
                content: "cache invalidation should expose new memory".into(),
            }))
            .unwrap();

        let after = server
            .find_memory(Parameters(FindMemoryRequest {
                keyword: None,
                repo: Some("test-repo".into()),
                tags: None,
                limit: Some(10),
            }))
            .unwrap();

        let _ = std::fs::remove_file(&db_path);
        let _ = std::fs::remove_file(db_path.with_extension("sqlite3-wal"));
        let _ = std::fs::remove_file(db_path.with_extension("sqlite3-shm"));

        assert_eq!(after.0.memories.len(), 1);
        assert_eq!(after.0.memories[0].key, "cache-invalidation-test");
    }

    #[test]
    fn find_cache_config_uses_env_overrides() {
        let _lock = ENV_LOCK.lock().unwrap();
        unsafe {
            env::set_var(FIND_CACHE_CAPACITY_ENV, "42");
            env::set_var(FIND_CACHE_TTL_SECS_ENV, "7");
        }

        assert_eq!(find_cache_capacity(), 42);
        assert_eq!(find_cache_ttl(), Duration::from_secs(7));

        unsafe {
            env::remove_var(FIND_CACHE_CAPACITY_ENV);
            env::remove_var(FIND_CACHE_TTL_SECS_ENV);
        }
    }

    #[test]
    fn find_cache_config_falls_back_for_invalid_env() {
        let _lock = ENV_LOCK.lock().unwrap();
        unsafe {
            env::set_var(FIND_CACHE_CAPACITY_ENV, "0");
            env::set_var(FIND_CACHE_TTL_SECS_ENV, "invalid");
        }

        assert_eq!(find_cache_capacity(), DEFAULT_FIND_CACHE_CAPACITY);
        assert_eq!(
            find_cache_ttl(),
            Duration::from_secs(DEFAULT_FIND_CACHE_TTL_SECS)
        );

        unsafe {
            env::remove_var(FIND_CACHE_CAPACITY_ENV);
            env::remove_var(FIND_CACHE_TTL_SECS_ENV);
        }
    }
}