PicoBot/src/storage/tests.rs

use super::migrations::has_column;
use super::*;
use crate::bus::SYSTEM_CONTEXT_AGENT_PROMPT;
use crate::domain::messages::ToolCall;

const TEST_CHANNEL: &str = "test-channel";

#[test]
fn test_persistent_session_id_for_cli_and_channel() {
    assert_eq!(persistent_session_id("cli", "abc"), "abc");
    // 幂等：已带前缀的 chat_id 会被清理，不会累积前缀
    assert_eq!(persistent_session_id("websocket", "websocket:abc"), "abc");
    assert_eq!(
        persistent_session_id("websocket", "websocket:websocket:abc"),
        "abc"
    );
    assert_eq!(
        persistent_session_id(TEST_CHANNEL, "abc"),
        "test-channel:abc"
    );
    // 其他通道也幂等
    assert_eq!(
        persistent_session_id(TEST_CHANNEL, "test-channel:abc"),
        "test-channel:abc"
    );
}

#[test]
fn test_session_store_roundtrip_and_lifecycle() {
    let store = SessionStore::in_memory().unwrap();

    let session = store.create_cli_session(Some("demo")).unwrap();
    assert_eq!(session.title, "demo");
    assert_eq!(session.channel_name, "cli");
    assert_eq!(session.chat_id, session.id);
    assert_eq!(session.message_count, 0);
    assert_eq!(session.user_turn_count, 0);
    assert_eq!(session.agent_prompt_reinjection_count, 0);

    let first = ChatMessage::user("hello");
    let second = ChatMessage::assistant("world");
    store.append_message(&session.id, &first).unwrap();
    store.append_message(&session.id, &second).unwrap();

    let stored = store.get_session(&session.id).unwrap().unwrap();
    assert_eq!(stored.message_count, 2);
    assert!(stored.archived_at.is_none());
    assert_eq!(stored.user_turn_count, 1);
    assert_eq!(stored.agent_prompt_reinjection_count, 0);

    let messages = store.load_messages(&session.id).unwrap();
    assert_eq!(messages.len(), 2);
    assert_eq!(messages[0].role, "user");
    assert_eq!(messages[0].content, "hello");
    assert_eq!(messages[1].role, "assistant");
    assert_eq!(messages[1].content, "world");

    store.rename_session(&session.id, "renamed").unwrap();
    let renamed = store.get_session(&session.id).unwrap().unwrap();
    assert_eq!(renamed.title, "renamed");

    store.archive_session(&session.id).unwrap();
    let archived = store.get_session(&session.id).unwrap().unwrap();
    assert!(archived.archived_at.is_some());

    let active_only = store.list_sessions("cli", false).unwrap();
    assert!(active_only.is_empty());

    let including_archived = store.list_sessions("cli", true).unwrap();
    assert_eq!(including_archived.len(), 1);

    store.clear_messages(&session.id).unwrap();
    let cleared = store.load_messages(&session.id).unwrap();
    assert!(cleared.is_empty());
    let cleared_session = store.get_session(&session.id).unwrap().unwrap();
    assert_eq!(cleared_session.message_count, 0);
    assert_eq!(cleared_session.user_turn_count, 0);
    assert_eq!(cleared_session.agent_prompt_reinjection_count, 0);

    store.delete_session(&session.id).unwrap();
    assert!(store.get_session(&session.id).unwrap().is_none());
}

#[test]
fn test_ensure_channel_session_is_stable() {
    let store = SessionStore::in_memory().unwrap();

    let first = store
        .ensure_channel_session(TEST_CHANNEL, "chat-1")
        .unwrap();
    let second = store
        .ensure_channel_session(TEST_CHANNEL, "chat-1")
        .unwrap();

    assert_eq!(first.id, second.id);
    assert_eq!(first.chat_id, "chat-1");
    assert_eq!(second.channel_name, TEST_CHANNEL);
}

#[test]
fn test_assistant_tool_calls_roundtrip() {
    let store = SessionStore::in_memory().unwrap();
    let session = store.create_cli_session(Some("tools")).unwrap();

    let assistant = ChatMessage::assistant_with_tool_calls(
        "calling tool",
        vec![ToolCall {
            id: "call_1".to_string(),
            name: "calculator".to_string(),
            arguments: serde_json::json!({ "expression": "3*7" }),
        }],
    );

    store.append_message(&session.id, &assistant).unwrap();

    let messages = store.load_messages(&session.id).unwrap();
    assert_eq!(messages.len(), 1);
    assert_eq!(messages[0].role, "assistant");
    assert_eq!(messages[0].tool_calls.as_ref().unwrap().len(), 1);
    assert_eq!(messages[0].tool_calls.as_ref().unwrap()[0].id, "call_1");
    assert_eq!(
        messages[0].tool_calls.as_ref().unwrap()[0].name,
        "calculator"
    );
}

#[test]
fn test_assistant_reasoning_content_roundtrip() {
    let store = SessionStore::in_memory().unwrap();
    let session = store.create_cli_session(Some("reasoning")).unwrap();

    let assistant = ChatMessage::assistant_with_reasoning("final answer", "hidden reasoning");

    store.append_message(&session.id, &assistant).unwrap();

    let messages = store.load_messages(&session.id).unwrap();
    assert_eq!(messages.len(), 1);
    assert_eq!(messages[0].content, "final answer");
    assert_eq!(
        messages[0].reasoning_content.as_deref(),
        Some("hidden reasoning")
    );
}

#[test]
fn test_schema_migration_adds_user_turn_and_reinjection_columns() {
    let tmp = std::env::temp_dir().join(format!("picobot_test_mig2_{}.db", uuid::Uuid::new_v4()));
    let conn = Connection::open(&tmp).unwrap();
    conn.execute_batch(
        "
        CREATE TABLE sessions (
            id TEXT PRIMARY KEY,
            title TEXT NOT NULL,
            channel_name TEXT NOT NULL,
            chat_id TEXT NOT NULL,
            summary TEXT,
            created_at INTEGER NOT NULL,
            updated_at INTEGER NOT NULL,
            last_active_at INTEGER NOT NULL,
            archived_at INTEGER,
            deleted_at INTEGER,
            message_count INTEGER NOT NULL DEFAULT 0
        );

        CREATE TABLE messages (
            id TEXT PRIMARY KEY,
            session_id TEXT NOT NULL,
            seq INTEGER NOT NULL,
            role TEXT NOT NULL,
            content TEXT NOT NULL,
            media_refs_json TEXT NOT NULL,
            tool_call_id TEXT,
            tool_name TEXT,
            tool_calls_json TEXT,
            created_at INTEGER NOT NULL,
            FOREIGN KEY(session_id) REFERENCES sessions(id) ON DELETE CASCADE,
            UNIQUE(session_id, seq)
        );
        ",
    )
    .unwrap();

    let path_str = tmp.to_string_lossy().to_string();
    let store = SessionStore::from_connection(conn, &path_str).unwrap();
    let session = store.create_cli_session(Some("migrated")).unwrap();
    assert_eq!(session.user_turn_count, 0);
    assert_eq!(session.agent_prompt_reinjection_count, 0);
}

#[test]
fn test_schema_migration_adds_reasoning_content_column_to_messages() {
    let tmp = std::env::temp_dir().join(format!("picobot_test_mig_{}.db", uuid::Uuid::new_v4()));
    let conn = Connection::open(&tmp).unwrap();
    conn.execute_batch(
        "
        CREATE TABLE sessions (
            id TEXT PRIMARY KEY,
            title TEXT NOT NULL,
            channel_name TEXT NOT NULL,
            chat_id TEXT NOT NULL,
            summary TEXT,
            created_at INTEGER NOT NULL,
            updated_at INTEGER NOT NULL,
            last_active_at INTEGER NOT NULL,
            archived_at INTEGER,
            deleted_at INTEGER,
            message_count INTEGER NOT NULL DEFAULT 0
        );

        CREATE TABLE messages (
            id TEXT PRIMARY KEY,
            session_id TEXT NOT NULL,
            seq INTEGER NOT NULL,
            role TEXT NOT NULL,
            content TEXT NOT NULL,
            media_refs_json TEXT NOT NULL,
            tool_call_id TEXT,
            tool_name TEXT,
            tool_calls_json TEXT,
            created_at INTEGER NOT NULL,
            FOREIGN KEY(session_id) REFERENCES sessions(id) ON DELETE CASCADE,
            UNIQUE(session_id, seq)
        );
        ",
    )
    .unwrap();

    let path_str = tmp.to_string_lossy().to_string();
    let _store = SessionStore::from_connection(conn, &path_str).unwrap();
    let conn = _store.pool.get().unwrap();

    assert!(has_column(&conn, "messages", "reasoning_content").unwrap());
}

#[test]
fn test_compact_active_history_rebuilds_active_segment_with_delta_messages() {
    let store = SessionStore::in_memory().unwrap();
    let session = store.create_cli_session(Some("compact-history")).unwrap();

    let agent_prompt =
        ChatMessage::system_with_context("agent", Some(SYSTEM_CONTEXT_AGENT_PROMPT.to_string()));
    let seed_messages = vec![
        agent_prompt.clone(),
        ChatMessage::user("u1"),
        ChatMessage::assistant("a1"),
        ChatMessage::user("u2"),
        ChatMessage::assistant("a2"),
        ChatMessage::user("u3"),
        ChatMessage::assistant("a3"),
        ChatMessage::user("u4"),
        ChatMessage::assistant("a4"),
    ];

    for message in &seed_messages {
        store.append_message(&session.id, message).unwrap();
    }

    let snapshot_end_seq = store
        .get_session(&session.id)
        .unwrap()
        .unwrap()
        .message_count;
    let preserved_messages = store.load_messages(&session.id).unwrap()[3..].to_vec();
    let preserved_system_messages = vec![agent_prompt];

    store
        .append_message(&session.id, &ChatMessage::user("u5"))
        .unwrap();
    store
        .append_message(&session.id, &ChatMessage::assistant("a5"))
        .unwrap();

    let summary_message = ChatMessage::system("[Compressed History]\n\nsummary");
    let compacted = store
        .compact_active_history(
            &session.id,
            snapshot_end_seq,
            &preserved_system_messages,
            &summary_message,
            &preserved_messages,
        )
        .unwrap();

    assert!(compacted);

    let active_messages = store.load_messages(&session.id).unwrap();
    assert_eq!(active_messages.len(), 10);
    assert_eq!(active_messages[0].role, "system");
    assert_eq!(active_messages[0].content, "agent");
    assert_eq!(
        active_messages[0].system_context.as_deref(),
        Some(SYSTEM_CONTEXT_AGENT_PROMPT)
    );
    assert_eq!(active_messages[1].role, "system");
    assert_eq!(
        active_messages[1].content,
        "[Compressed History]\n\nsummary"
    );
    assert_eq!(active_messages[2].content, "u2");
    assert_eq!(active_messages[3].content, "a2");
    assert_eq!(active_messages[8].content, "u5");
    assert_eq!(active_messages[9].content, "a5");

    let stored = store.get_session(&session.id).unwrap().unwrap();
    assert_eq!(stored.user_turn_count, 4);

    let all_messages = store.load_all_messages(&session.id).unwrap();
    assert_eq!(all_messages.len(), 10);
}

#[test]
fn test_mark_agent_prompt_reinjected_increments_counter() {
    let store = SessionStore::in_memory().unwrap();
    let session = store.create_cli_session(Some("prompt")).unwrap();

    store.mark_agent_prompt_reinjected(&session.id).unwrap();
    store.mark_agent_prompt_reinjected(&session.id).unwrap();

    let stored = store.get_session(&session.id).unwrap().unwrap();
    assert_eq!(stored.agent_prompt_reinjection_count, 2);
}

#[test]
fn test_tool_result_roundtrip() {
    let store = SessionStore::in_memory().unwrap();
    let session = store.create_cli_session(Some("tool-result")).unwrap();

    let tool_message = ChatMessage::tool("call_9", "write", "saved to /tmp/output.txt");
    store.append_message(&session.id, &tool_message).unwrap();

    let messages = store.load_messages(&session.id).unwrap();
    assert_eq!(messages.len(), 1);
    assert_eq!(messages[0].role, "tool");
    assert_eq!(messages[0].content, "saved to /tmp/output.txt");
    assert_eq!(messages[0].tool_call_id.as_deref(), Some("call_9"));
    assert_eq!(messages[0].tool_name.as_deref(), Some("write"));
    assert!(messages[0].tool_calls.is_none());
}

#[test]
fn test_skill_events_roundtrip() {
    let store = SessionStore::in_memory().unwrap();
    let session = store.create_cli_session(Some("skill-events")).unwrap();

    store
        .append_skill_event(None, "discovered", None, &serde_json::json!({"count": 2}))
        .unwrap();
    store
        .append_skill_event(
            Some(&session.id),
            "activated",
            Some("code-review"),
            &serde_json::json!({"source": "project"}),
        )
        .unwrap();

    let global_events = store.list_skill_events(None).unwrap();
    assert_eq!(global_events.len(), 1);
    assert_eq!(global_events[0].event_type, "discovered");
    assert_eq!(global_events[0].payload["count"], 2);

    let session_events = store.list_skill_events(Some(&session.id)).unwrap();
    assert_eq!(session_events.len(), 1);
    assert_eq!(session_events[0].event_type, "activated");
    assert_eq!(session_events[0].skill_name.as_deref(), Some("code-review"));
    assert_eq!(session_events[0].payload["source"], "project");
}

#[test]
fn test_memory_roundtrip_with_source_fields() {
    let store = SessionStore::in_memory().unwrap();

    let saved = store
        .put_memory(&MemoryUpsert {
            scope_kind: "user".to_string(),
            scope_key: format!("{}:user-1", TEST_CHANNEL),
            namespace: "user".to_string(),
            memory_key: "language".to_string(),
            content: "Rust".to_string(),
            source_type: "message".to_string(),
            source_session_id: Some(format!("{}:chat-1", TEST_CHANNEL)),
            source_message_id: Some("msg-1".to_string()),
            source_message_seq: Some(7),
            source_channel_name: Some(TEST_CHANNEL.to_string()),
            source_chat_id: Some("chat-1".to_string()),
        })
        .unwrap();

    assert_eq!(saved.content, "Rust");
    assert_eq!(saved.source_type, "message");
    assert_eq!(
        saved.source_session_id.as_deref(),
        Some("test-channel:chat-1")
    );
    assert_eq!(saved.source_message_id.as_deref(), Some("msg-1"));
    assert_eq!(saved.source_message_seq, Some(7));

    let fetched = store
        .get_memory("user", "test-channel:user-1", "user", "language")
        .unwrap()
        .unwrap();
    assert_eq!(fetched.id, saved.id);
    assert_eq!(fetched.source_chat_id.as_deref(), Some("chat-1"));
}

#[test]
fn test_memory_fts_tracks_upsert_and_delete() {
    let store = SessionStore::in_memory().unwrap();

    store
        .put_memory(&MemoryUpsert {
            scope_kind: "user".to_string(),
            scope_key: format!("{}:user-1", TEST_CHANNEL),
            namespace: "user".to_string(),
            memory_key: "editor".to_string(),
            content: "Prefers rust-analyzer and cargo test output".to_string(),
            source_type: "message".to_string(),
            source_session_id: Some(format!("{}:chat-2", TEST_CHANNEL)),
            source_message_id: Some("msg-2".to_string()),
            source_message_seq: Some(3),
            source_channel_name: Some(TEST_CHANNEL.to_string()),
            source_chat_id: Some("chat-2".to_string()),
        })
        .unwrap();

    let hits = store
        .search_memories("user", "test-channel:user-1", "rust-analyzer", None, 10)
        .unwrap();
    assert_eq!(hits.len(), 1);
    assert_eq!(hits[0].memory_key, "editor");

    store
        .put_memory(&MemoryUpsert {
            scope_kind: "user".to_string(),
            scope_key: format!("{}:user-1", TEST_CHANNEL),
            namespace: "user".to_string(),
            memory_key: "editor".to_string(),
            content: "Prefers clippy diagnostics".to_string(),
            source_type: "message".to_string(),
            source_session_id: Some(format!("{}:chat-3", TEST_CHANNEL)),
            source_message_id: Some("msg-3".to_string()),
            source_message_seq: Some(4),
            source_channel_name: Some(TEST_CHANNEL.to_string()),
            source_chat_id: Some("chat-3".to_string()),
        })
        .unwrap();

    let old_hits = store
        .search_memories("user", "test-channel:user-1", "rust-analyzer", None, 10)
        .unwrap();
    assert!(old_hits.is_empty());

    let new_hits = store
        .search_memories("user", "test-channel:user-1", "clippy", None, 10)
        .unwrap();
    assert_eq!(new_hits.len(), 1);

    let deleted = store
        .delete_memory("user", "test-channel:user-1", "user", "editor")
        .unwrap();
    assert!(deleted);

    let hits_after_delete = store
        .search_memories("user", "test-channel:user-1", "clippy", None, 10)
        .unwrap();
    assert!(hits_after_delete.is_empty());
}

#[test]
fn test_memory_search_matches_memory_key_field() {
    let store = SessionStore::in_memory().unwrap();

    store
        .put_memory(&MemoryUpsert {
            scope_kind: "user".to_string(),
            scope_key: format!("{}:user-1", TEST_CHANNEL),
            namespace: "user".to_string(),
            memory_key: "email_folder_preference".to_string(),
            content: "用户提到邮件时默认查看代收邮箱。".to_string(),
            source_type: "message".to_string(),
            source_session_id: Some(format!("{}:chat-8", TEST_CHANNEL)),
            source_message_id: Some("msg-8".to_string()),
            source_message_seq: Some(8),
            source_channel_name: Some(TEST_CHANNEL.to_string()),
            source_chat_id: Some("chat-8".to_string()),
        })
        .unwrap();

    let hits = store
        .search_memories(
            "user",
            "test-channel:user-1",
            "email_folder_preference",
            None,
            10,
        )
        .unwrap();

    assert_eq!(hits.len(), 1);
    assert_eq!(hits[0].memory_key, "email_folder_preference");
}

#[test]
fn test_search_memories_any_matches_multiple_keywords_once() {
    let store = SessionStore::in_memory().unwrap();

    store
        .put_memory(&MemoryUpsert {
            scope_kind: "user".to_string(),
            scope_key: format!("{}:user-1", TEST_CHANNEL),
            namespace: "user".to_string(),
            memory_key: "editor".to_string(),
            content: "Prefers rust-analyzer and cargo test output".to_string(),
            source_type: "message".to_string(),
            source_session_id: Some(format!("{}:chat-2", TEST_CHANNEL)),
            source_message_id: Some("msg-2".to_string()),
            source_message_seq: Some(3),
            source_channel_name: Some(TEST_CHANNEL.to_string()),
            source_chat_id: Some("chat-2".to_string()),
        })
        .unwrap();

    store
        .put_memory(&MemoryUpsert {
            scope_kind: "user".to_string(),
            scope_key: format!("{}:user-1", TEST_CHANNEL),
            namespace: "episodic".to_string(),
            memory_key: "quality".to_string(),
            content: "Tracks clippy warnings before release".to_string(),
            source_type: "message".to_string(),
            source_session_id: Some(format!("{}:chat-3", TEST_CHANNEL)),
            source_message_id: Some("msg-3".to_string()),
            source_message_seq: Some(4),
            source_channel_name: Some(TEST_CHANNEL.to_string()),
            source_chat_id: Some("chat-3".to_string()),
        })
        .unwrap();

    let hits = store
        .search_memories_any(
            "user",
            "test-channel:user-1",
            &["rust-analyzer".to_string(), "clippy".to_string()],
            None,
            10,
        )
        .unwrap();

    assert_eq!(hits.len(), 2);
    assert!(hits.iter().any(|memory| memory.memory_key == "editor"));
    assert!(hits.iter().any(|memory| memory.memory_key == "quality"));
}

#[test]
fn test_memory_scope_listing_and_full_scope_read() {
    let store = SessionStore::in_memory().unwrap();

    store
        .put_memory(&MemoryUpsert {
            scope_kind: "user".to_string(),
            scope_key: format!("{}:user-2", TEST_CHANNEL),
            namespace: "user".to_string(),
            memory_key: "style".to_string(),
            content: "偏好简洁表达".to_string(),
            source_type: "message".to_string(),
            source_session_id: Some(format!("{}:chat-2", TEST_CHANNEL)),
            source_message_id: Some("msg-2".to_string()),
            source_message_seq: Some(2),
            source_channel_name: Some(TEST_CHANNEL.to_string()),
            source_chat_id: Some("chat-2".to_string()),
        })
        .unwrap();
    store
        .put_memory(&MemoryUpsert {
            scope_kind: "user".to_string(),
            scope_key: format!("{}:user-1", TEST_CHANNEL),
            namespace: "user".to_string(),
            memory_key: "work".to_string(),
            content: "用户在做AI产品".to_string(),
            source_type: "message".to_string(),
            source_session_id: Some(format!("{}:chat-1", TEST_CHANNEL)),
            source_message_id: Some("msg-1".to_string()),
            source_message_seq: Some(1),
            source_channel_name: Some(TEST_CHANNEL.to_string()),
            source_chat_id: Some("chat-1".to_string()),
        })
        .unwrap();
    store
        .put_memory(&MemoryUpsert {
            scope_kind: "user".to_string(),
            scope_key: format!("{}:user-1", TEST_CHANNEL),
            namespace: "patterns".to_string(),
            memory_key: "workflow".to_string(),
            content: "习惯先问方案再要代码".to_string(),
            source_type: "message".to_string(),
            source_session_id: Some(format!("{}:chat-1", TEST_CHANNEL)),
            source_message_id: Some("msg-3".to_string()),
            source_message_seq: Some(3),
            source_channel_name: Some(TEST_CHANNEL.to_string()),
            source_chat_id: Some("chat-1".to_string()),
        })
        .unwrap();

    let scope_keys = store.list_memory_scope_keys("user").unwrap();
    assert_eq!(
        scope_keys,
        vec![
            "test-channel:user-1".to_string(),
            "test-channel:user-2".to_string()
        ]
    );

    let full_scope = store
        .list_memories_for_scope("user", "test-channel:user-1")
        .unwrap();
    assert_eq!(full_scope.len(), 2);
    assert!(
        full_scope
            .iter()
            .all(|memory| memory.scope_key == "test-channel:user-1")
    );
    assert!(full_scope.iter().any(|memory| memory.memory_key == "work"));
    assert!(
        full_scope
            .iter()
            .any(|memory| memory.memory_key == "workflow")
    );
}

#[test]
fn test_scheduler_job_roundtrip_and_runtime_update() {
    let store = SessionStore::in_memory().unwrap();

    let saved = store
        .upsert_scheduler_job(&SchedulerJobUpsert {
            id: "heartbeat".to_string(),
            kind: "outbound_message".to_string(),
            schedule: serde_json::json!({
                "type": "interval",
                "seconds": 300,
                "startup_delay_secs": 10,
            }),
            interval_secs: 300,
            startup_delay_secs: 10,
            target: serde_json::json!({
                "channel": "test-channel",
                "chat_id": "oc_demo",
            }),
            payload: serde_json::json!({
                "content": "heartbeat",
            }),
            enabled: true,
            state: SchedulerJobState::Scheduled,
            last_status: None,
            last_error: None,
            run_count: 0,
            max_runs: Some(3),
            last_fired_at: None,
            next_fire_at: Some(1_700_000_000_000),
            paused_at: None,
            completed_at: None,
        })
        .unwrap();

    assert_eq!(saved.id, "heartbeat");
    assert_eq!(saved.kind, "outbound_message");
    assert_eq!(saved.state, SchedulerJobState::Scheduled);
    assert_eq!(saved.max_runs, Some(3));

    store
        .update_scheduler_job_runtime(
            "heartbeat",
            SchedulerJobState::Completed,
            Some(SchedulerJobStatus::Ok),
            None,
            1,
            Some(1_700_000_000_000),
            None,
            None,
            Some(1_700_000_000_100),
        )
        .unwrap();

    let fetched = store.get_scheduler_job("heartbeat").unwrap().unwrap();
    assert_eq!(fetched.state, SchedulerJobState::Completed);
    assert_eq!(fetched.last_status, Some(SchedulerJobStatus::Ok));
    assert_eq!(fetched.run_count, 1);
    assert_eq!(fetched.completed_at, Some(1_700_000_000_100));
}