openbrain-mcp/tests/e2e_mcp.rs

use serde_json::{json, Value};
use std::process::{Command, Stdio};
use std::time::Duration;
use tokio_postgres::NoTls;
use uuid::Uuid;

fn base_url() -> String {
    std::env::var("OPENBRAIN_E2E_BASE_URL").unwrap_or_else(|_| "http://127.0.0.1:3100".to_string())
}

fn remote_mode() -> bool {
    std::env::var("OPENBRAIN_E2E_REMOTE")
        .map(|v| v == "true" || v == "1")
        .unwrap_or(false)
}

fn api_key() -> Option<String> {
    std::env::var("OPENBRAIN_E2E_API_KEY")
        .ok()
        .or_else(|| std::env::var("OPENBRAIN__AUTH__API_KEYS").ok())
        .map(|keys| keys.split(',').next().unwrap_or("").trim().to_string())
        .filter(|k| !k.is_empty())
}

fn db_url() -> String {
    let host =
        std::env::var("OPENBRAIN__DATABASE__HOST").unwrap_or_else(|_| "localhost".to_string());
    let port = std::env::var("OPENBRAIN__DATABASE__PORT").unwrap_or_else(|_| "5432".to_string());
    let name =
        std::env::var("OPENBRAIN__DATABASE__NAME").unwrap_or_else(|_| "openbrain".to_string());
    let user =
        std::env::var("OPENBRAIN__DATABASE__USER").unwrap_or_else(|_| "openbrain_svc".to_string());
    let password = std::env::var("OPENBRAIN__DATABASE__PASSWORD")
        .unwrap_or_else(|_| "your_secure_password_here".to_string());

    format!("host={host} port={port} dbname={name} user={user} password={password}")
}

async fn ensure_schema() {
    if remote_mode() {
        return;
    }

    let (client, connection) = tokio_postgres::connect(&db_url(), NoTls)
        .await
        .expect("connect to postgres for e2e schema setup");

    tokio::spawn(async move {
        if let Err(e) = connection.await {
            eprintln!("postgres connection error: {e}");
        }
    });

    let vector_exists = client
        .query_one("SELECT to_regtype('vector')::text", &[])
        .await
        .expect("query vector type availability")
        .get::<_, Option<String>>(0)
        .is_some();

    if !vector_exists {
        if let Err(e) = client
            .execute("CREATE EXTENSION IF NOT EXISTS vector", &[])
            .await
        {
            panic!(
                "pgvector extension is not available for this PostgreSQL instance: {e}. \
Install pgvector for your active PostgreSQL major version, then run: CREATE EXTENSION vector;"
            );
        }
    }

    client
        .batch_execute(
            r#"
            CREATE TABLE IF NOT EXISTS memories (
                id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
                auth_scope VARCHAR(255) NOT NULL DEFAULT 'public',
                agent_id VARCHAR(255) NOT NULL,
                content TEXT NOT NULL,
                embedding vector(384) NOT NULL,
                keywords TEXT[] DEFAULT '{}',
                metadata JSONB DEFAULT '{}',
                created_at TIMESTAMPTZ DEFAULT NOW(),
                expires_at TIMESTAMPTZ
            );
            ALTER TABLE memories ADD COLUMN IF NOT EXISTS auth_scope VARCHAR(255) NOT NULL DEFAULT 'public';
            ALTER TABLE memories ADD COLUMN IF NOT EXISTS expires_at TIMESTAMPTZ;
            ALTER TABLE memories ADD COLUMN IF NOT EXISTS tsv tsvector;
            CREATE OR REPLACE FUNCTION memories_tsv_trigger()
            RETURNS trigger
            LANGUAGE plpgsql
            AS $$
            BEGIN
                NEW.tsv :=
                    setweight(to_tsvector('pg_catalog.english', COALESCE(NEW.content, '')), 'A') ||
                    setweight(
                        to_tsvector('pg_catalog.english', COALESCE(array_to_string(NEW.keywords, ' '), '')),
                        'B'
                    );
                RETURN NEW;
            END;
            $$;
            UPDATE memories
            SET tsv =
                setweight(to_tsvector('pg_catalog.english', COALESCE(content, '')), 'A') ||
                setweight(
                    to_tsvector('pg_catalog.english', COALESCE(array_to_string(keywords, ' '), '')),
                    'B'
                )
            WHERE tsv IS NULL;
            DROP TRIGGER IF EXISTS memories_tsv_update ON memories;
            CREATE TRIGGER memories_tsv_update
            BEFORE INSERT OR UPDATE OF content, keywords ON memories
            FOR EACH ROW
            EXECUTE FUNCTION memories_tsv_trigger();
            CREATE INDEX IF NOT EXISTS idx_memories_auth_scope ON memories(auth_scope);
            CREATE INDEX IF NOT EXISTS idx_memories_auth_scope_agent ON memories(auth_scope, agent_id);
            CREATE INDEX IF NOT EXISTS idx_memories_agent ON memories(agent_id);
            CREATE INDEX IF NOT EXISTS idx_memories_embedding ON memories
                USING hnsw (embedding vector_cosine_ops);
            CREATE INDEX IF NOT EXISTS idx_memories_tsv ON memories
                USING GIN (tsv);
            CREATE INDEX IF NOT EXISTS idx_memories_expires_at ON memories (expires_at)
                WHERE expires_at IS NOT NULL;
            "#,
        )
        .await
        .expect("create memories table/indexes for e2e");
}

async fn wait_until_ready(client: &reqwest::Client, base: &str) {
    for _ in 0..60 {
        let resp = client.get(format!("{base}/ready")).send().await;
        if let Ok(resp) = resp {
            if resp.status().is_success() {
                let body: Value = resp.json().await.expect("/ready JSON response");
                if body.get("status").and_then(Value::as_str) == Some("ready") {
                    return;
                }
            }
        }
        tokio::time::sleep(Duration::from_millis(500)).await;
    }

    panic!("Server did not become ready at {base}/ready within timeout");
}

fn apply_request_headers(
    mut req_builder: reqwest::RequestBuilder,
    api_key_override: Option<&str>,
    extra_headers: &[(&str, &str)],
) -> reqwest::RequestBuilder {
    if let Some(key) = api_key_override {
        req_builder = req_builder.header("X-API-Key", key);
    }

    for (name, value) in extra_headers {
        req_builder = req_builder.header(*name, *value);
    }

    req_builder
}

async fn call_jsonrpc(client: &reqwest::Client, base: &str, request: Value) -> Value {
    let api_key = api_key();
    call_jsonrpc_with_options(client, base, request, api_key.as_deref(), &[]).await
}

async fn call_jsonrpc_with_options(
    client: &reqwest::Client,
    base: &str,
    request: Value,
    api_key_override: Option<&str>,
    extra_headers: &[(&str, &str)],
) -> Value {
    let req_builder = apply_request_headers(
        client.post(format!("{base}/mcp/message")).json(&request),
        api_key_override,
        extra_headers,
    );

    req_builder
        .send()
        .await
        .expect("JSON-RPC HTTP request")
        .json()
        .await
        .expect("JSON-RPC response body")
}

async fn call_streamable_jsonrpc(
    client: &reqwest::Client,
    base: &str,
    request: Value,
) -> reqwest::Response {
    let api_key = api_key();
    call_streamable_jsonrpc_with_options(client, base, request, api_key.as_deref(), &[]).await
}

async fn call_streamable_jsonrpc_with_options(
    client: &reqwest::Client,
    base: &str,
    request: Value,
    api_key_override: Option<&str>,
    extra_headers: &[(&str, &str)],
) -> reqwest::Response {
    let req_builder = apply_request_headers(
        client
            .post(format!("{base}/mcp"))
            .header("Accept", "application/json, text/event-stream")
            .json(&request),
        api_key_override,
        extra_headers,
    );

    req_builder
        .send()
        .await
        .expect("streamable JSON-RPC HTTP request")
}

/// Make an authenticated GET request to an MCP endpoint
async fn get_mcp_endpoint(client: &reqwest::Client, base: &str, path: &str) -> reqwest::Response {
    let api_key = api_key();
    get_mcp_endpoint_with_options(client, base, path, api_key.as_deref(), &[]).await
}

async fn get_mcp_endpoint_with_options(
    client: &reqwest::Client,
    base: &str,
    path: &str,
    api_key_override: Option<&str>,
    extra_headers: &[(&str, &str)],
) -> reqwest::Response {
    let req_builder = apply_request_headers(
        client.get(format!("{base}{path}")),
        api_key_override,
        extra_headers,
    );

    req_builder.send().await.expect(&format!("GET {path}"))
}

async fn read_sse_event(
    response: &mut reqwest::Response,
    buffer: &mut String,
) -> Option<(Option<String>, String)> {
    loop {
        *buffer = buffer.replace("\r\n", "\n");
        if let Some(idx) = buffer.find("\n\n") {
            let raw_event = buffer[..idx].to_string();
            *buffer = buffer[idx + 2..].to_string();

            let mut event_type = None;
            let mut data_lines = Vec::new();
            for line in raw_event.lines() {
                if let Some(value) = line.strip_prefix("event:") {
                    event_type = Some(value.trim().to_string());
                } else if let Some(value) = line.strip_prefix("data:") {
                    data_lines.push(value.trim_start().to_string());
                }
            }

            return Some((event_type, data_lines.join("\n")));
        }

        let chunk = response.chunk().await.expect("read SSE chunk")?;
        buffer.push_str(std::str::from_utf8(&chunk).expect("SSE chunk should be valid UTF-8"));
    }
}

fn parse_tool_response(tool_name: &str, response: Value) -> Value {
    if let Some(error) = response.get("error") {
        panic!("tools/call for '{tool_name}' failed: {error}");
    }

    let text_payload = response
        .get("result")
        .and_then(|r| r.get("content"))
        .and_then(Value::as_array)
        .and_then(|arr| arr.first())
        .and_then(|item| item.get("text"))
        .and_then(Value::as_str)
        .expect("result.content[0].text payload");

    serde_json::from_str(text_payload).expect("tool text payload to be valid JSON")
}

async fn call_tool(
    client: &reqwest::Client,
    base: &str,
    tool_name: &str,
    arguments: Value,
) -> Value {
    let api_key = api_key();
    call_tool_with_options(client, base, tool_name, arguments, api_key.as_deref(), &[]).await
}

async fn call_tool_with_options(
    client: &reqwest::Client,
    base: &str,
    tool_name: &str,
    arguments: Value,
    api_key_override: Option<&str>,
    extra_headers: &[(&str, &str)],
) -> Value {
    let request = json!({
        "jsonrpc": "2.0",
        "id": Uuid::new_v4().to_string(),
        "method": "tools/call",
        "params": {
            "name": tool_name,
            "arguments": arguments
        }
    });

    let response =
        call_jsonrpc_with_options(client, base, request, api_key_override, extra_headers).await;
    parse_tool_response(tool_name, response)
}

async fn call_tool_streamable_with_options(
    client: &reqwest::Client,
    base: &str,
    tool_name: &str,
    arguments: Value,
    api_key_override: Option<&str>,
    extra_headers: &[(&str, &str)],
) -> Value {
    let request = json!({
        "jsonrpc": "2.0",
        "id": Uuid::new_v4().to_string(),
        "method": "tools/call",
        "params": {
            "name": tool_name,
            "arguments": arguments
        }
    });

    let response = call_streamable_jsonrpc_with_options(
        client,
        base,
        request,
        api_key_override,
        extra_headers,
    )
    .await;
    let response: Value = response
        .json()
        .await
        .expect("streamable tool JSON response body");
    parse_tool_response(tool_name, response)
}

#[tokio::test]
async fn e2e_store_query_purge_roundtrip() {
    let base = base_url();
    let client = reqwest::Client::builder()
        .timeout(Duration::from_secs(20))
        .build()
        .expect("reqwest client");

    ensure_schema().await;
    wait_until_ready(&client, &base).await;

    let agent_id = format!("e2e-agent-{}", Uuid::new_v4());
    let memory_text = format!(
        "E2E memory {}: user prefers dark theme and vim bindings",
        Uuid::new_v4()
    );

    // Ensure clean slate for this test agent.
    let _ = call_tool(
        &client,
        &base,
        "purge",
        json!({ "agent_id": agent_id, "confirm": true }),
    )
    .await;

    let store_result = call_tool(
        &client,
        &base,
        "store",
        json!({
            "agent_id": agent_id,
            "content": memory_text,
            "metadata": { "source": "e2e-test", "suite": "store-query-purge" }
        }),
    )
    .await;

    assert_eq!(
        store_result.get("success").and_then(Value::as_bool),
        Some(true),
        "store should succeed"
    );

    let query_result = call_tool(
        &client,
        &base,
        "query",
        json!({
            "source_agent_id": agent_id,
            "query": "What are the user's editor preferences?",
            "limit": 5,
            "threshold": 0.0
        }),
    )
    .await;

    let count = query_result
        .get("count")
        .and_then(Value::as_u64)
        .expect("query.count");
    assert!(count >= 1, "query should return at least one stored memory");

    let results = query_result
        .get("results")
        .and_then(Value::as_array)
        .expect("query.results");
    let found_stored_content = results.iter().any(|item| {
        item.get("content")
            .and_then(Value::as_str)
            .map(|content| content == memory_text)
            .unwrap_or(false)
    });
    assert!(
        found_stored_content,
        "query results should include the content stored by this test"
    );

    let purge_result = call_tool(
        &client,
        &base,
        "purge",
        json!({ "agent_id": agent_id, "confirm": true }),
    )
    .await;

    let deleted = purge_result
        .get("deleted")
        .and_then(Value::as_u64)
        .expect("purge.deleted");
    assert!(deleted >= 1, "purge should delete at least one memory");

    let query_after_purge = call_tool(
        &client,
        &base,
        "query",
        json!({
            "source_agent_id": agent_id,
            "query": "dark theme vim bindings",
            "limit": 5,
            "threshold": 0.0
        }),
    )
    .await;

    assert_eq!(
        query_after_purge.get("count").and_then(Value::as_u64),
        Some(0),
        "query after purge should return no memories for this agent"
    );
}

#[tokio::test]
async fn e2e_transport_tools_list_and_unknown_method() {
    let base = base_url();
    let client = reqwest::Client::builder()
        .timeout(Duration::from_secs(20))
        .build()
        .expect("reqwest client");

    wait_until_ready(&client, &base).await;

    let list_response = call_jsonrpc(
        &client,
        &base,
        json!({
            "jsonrpc": "2.0",
            "id": "tools-list-1",
            "method": "tools/list",
            "params": {}
        }),
    )
    .await;

    let tools = list_response
        .get("result")
        .and_then(|r| r.get("tools"))
        .and_then(Value::as_array)
        .expect("tools/list result.tools");

    let tool_names: Vec<&str> = tools
        .iter()
        .filter_map(|t| t.get("name").and_then(Value::as_str))
        .collect();

    assert!(
        tool_names.contains(&"store"),
        "tools/list should include store"
    );
    assert!(
        tool_names.contains(&"query"),
        "tools/list should include query"
    );
    assert!(
        tool_names.contains(&"purge"),
        "tools/list should include purge"
    );

    let unknown_response = call_jsonrpc(
        &client,
        &base,
        json!({
            "jsonrpc": "2.0",
            "id": "unknown-1",
            "method": "not/a/real/method",
            "params": {}
        }),
    )
    .await;

    assert_eq!(
        unknown_response
            .get("error")
            .and_then(|e| e.get("code"))
            .and_then(Value::as_i64),
        Some(-32601),
        "unknown method should return Method Not Found"
    );
}

#[tokio::test]
async fn e2e_streamable_initialize_and_tools_list() {
    let base = base_url();
    let client = reqwest::Client::builder()
        .timeout(Duration::from_secs(20))
        .build()
        .expect("reqwest client");

    wait_until_ready(&client, &base).await;

    let initialize_response: Value = call_streamable_jsonrpc(
        &client,
        &base,
        json!({
            "jsonrpc": "2.0",
            "id": "streamable-init-1",
            "method": "initialize",
            "params": {
                "protocolVersion": "2024-11-05",
                "capabilities": {},
                "clientInfo": {
                    "name": "e2e-client",
                    "version": "0.1.0"
                }
            }
        }),
    )
    .await
    .json()
    .await
    .expect("streamable initialize JSON");

    assert_eq!(
        initialize_response
            .get("result")
            .and_then(|value| value.get("protocolVersion"))
            .and_then(Value::as_str),
        Some("2024-11-05")
    );

    let tools_list_response: Value = call_streamable_jsonrpc(
        &client,
        &base,
        json!({
            "jsonrpc": "2.0",
            "id": "streamable-tools-list-1",
            "method": "tools/list",
            "params": {}
        }),
    )
    .await
    .json()
    .await
    .expect("streamable tools/list JSON");

    assert!(
        tools_list_response
            .get("result")
            .and_then(|value| value.get("tools"))
            .and_then(Value::as_array)
            .map(|tools| !tools.is_empty())
            .unwrap_or(false),
        "streamable /mcp tools/list should return tool definitions"
    );
}

#[tokio::test]
async fn e2e_streamable_get_returns_405() {
    let base = base_url();
    let client = reqwest::Client::builder()
        .timeout(Duration::from_secs(20))
        .build()
        .expect("reqwest client");

    wait_until_ready(&client, &base).await;

    let mut request = client
        .get(format!("{base}/mcp"))
        .header("Accept", "text/event-stream");

    if let Some(key) = api_key() {
        request = request.header("X-API-Key", key);
    }

    let response = request.send().await.expect("GET /mcp");
    assert_eq!(
        response.status(),
        reqwest::StatusCode::METHOD_NOT_ALLOWED,
        "streamable GET /mcp should explicitly return 405 when standalone SSE streams are not offered"
    );
}

#[tokio::test]
async fn e2e_purge_requires_confirm_flag() {
    let base = base_url();
    let client = reqwest::Client::builder()
        .timeout(Duration::from_secs(20))
        .build()
        .expect("reqwest client");

    wait_until_ready(&client, &base).await;

    let response = call_jsonrpc(
        &client,
        &base,
        json!({
            "jsonrpc": "2.0",
            "id": "purge-confirm-1",
            "method": "tools/call",
            "params": {
                "name": "purge",
                "arguments": {
                    "agent_id": format!("e2e-agent-{}", Uuid::new_v4()),
                    "confirm": false
                }
            }
        }),
    )
    .await;

    let error_message = response
        .get("error")
        .and_then(|e| e.get("message"))
        .and_then(Value::as_str)
        .expect("purge without confirm should return JSON-RPC error");

    assert!(
        error_message.contains("confirm: true") || error_message.contains("confirm"),
        "purge error should explain confirmation requirement"
    );
}

#[tokio::test]
async fn e2e_query_shares_memories_across_agent_ids() {
    let base = base_url();
    let client = reqwest::Client::builder()
        .timeout(Duration::from_secs(20))
        .build()
        .expect("reqwest client");

    ensure_schema().await;
    wait_until_ready(&client, &base).await;

    let agent_a = format!("e2e-agent-a-{}", Uuid::new_v4());
    let agent_b = format!("e2e-agent-b-{}", Uuid::new_v4());
    let a_text = format!("A {} prefers dark mode", Uuid::new_v4());
    let b_text = format!("B {} prefers light mode", Uuid::new_v4());

    let _ = call_tool(
        &client,
        &base,
        "purge",
        json!({ "agent_id": agent_a, "confirm": true }),
    )
    .await;
    let _ = call_tool(
        &client,
        &base,
        "purge",
        json!({ "agent_id": agent_b, "confirm": true }),
    )
    .await;

    let _ = call_tool(
        &client,
        &base,
        "store",
        json!({ "agent_id": agent_a, "content": a_text, "metadata": {"suite": "agent-isolation"} }),
    )
    .await;

    let _ = call_tool(
        &client,
        &base,
        "store",
        json!({ "agent_id": agent_b, "content": b_text, "metadata": {"suite": "agent-isolation"} }),
    )
    .await;

    let shared_query = call_tool(
        &client,
        &base,
        "query",
        json!({
            "query": "mode preference",
            "limit": 10,
            "threshold": 0.0
        }),
    )
    .await;

    let results = shared_query
        .get("results")
        .and_then(Value::as_array)
        .expect("query results");

    let has_a = results.iter().any(|item| {
        item.get("content")
            .and_then(Value::as_str)
            .map(|s| s == a_text)
            .unwrap_or(false)
    });
    let has_b = results.iter().any(|item| {
        item.get("content")
            .and_then(Value::as_str)
            .map(|s| s == b_text)
            .unwrap_or(false)
    });

    assert!(has_a, "shared query should include agent A memory");
    assert!(has_b, "shared query should include agent B memory");

    let filtered_query = call_tool(
        &client,
        &base,
        "query",
        json!({
            "source_agent_id": agent_a,
            "query": "mode preference",
            "limit": 10,
            "threshold": 0.0
        }),
    )
    .await;

    let filtered_results = filtered_query
        .get("results")
        .and_then(Value::as_array)
        .expect("filtered query results");
    let filtered_has_a = filtered_results.iter().any(|item| {
        item.get("content")
            .and_then(Value::as_str)
            .map(|s| s == a_text)
            .unwrap_or(false)
    });
    let filtered_has_b = filtered_results.iter().any(|item| {
        item.get("content")
            .and_then(Value::as_str)
            .map(|s| s == b_text)
            .unwrap_or(false)
    });

    assert!(
        filtered_has_a,
        "filtered query should include agent A memory"
    );
    assert!(
        !filtered_has_b,
        "filtered query should exclude agent B memory"
    );

    let _ = call_tool(
        &client,
        &base,
        "purge",
        json!({ "agent_id": agent_a, "confirm": true }),
    )
    .await;
    let _ = call_tool(
        &client,
        &base,
        "purge",
        json!({ "agent_id": agent_b, "confirm": true }),
    )
    .await;
}

#[tokio::test]
async fn e2e_initialize_contract() {
    let base = base_url();
    let client = reqwest::Client::builder()
        .timeout(Duration::from_secs(20))
        .build()
        .expect("reqwest client");

    let response = call_jsonrpc(
        &client,
        &base,
        json!({
            "jsonrpc": "2.0",
            "id": "init-1",
            "method": "initialize",
            "params": {}
        }),
    )
    .await;

    let result = response.get("result").expect("initialize result");
    assert_eq!(
        result.get("protocolVersion").and_then(Value::as_str),
        Some("2024-11-05")
    );
    assert_eq!(
        result
            .get("serverInfo")
            .and_then(|v| v.get("name"))
            .and_then(Value::as_str),
        Some("openbrain-mcp")
    );
}

#[tokio::test]
async fn e2e_initialized_notification_is_accepted() {
    let base = base_url();
    let client = reqwest::Client::builder()
        .timeout(Duration::from_secs(20))
        .build()
        .expect("reqwest client");

    wait_until_ready(&client, &base).await;

    let mut request = client.post(format!("{base}/mcp/message")).json(&json!({
        "jsonrpc": "2.0",
        "method": "notifications/initialized",
        "params": {}
    }));

    if let Some(key) = api_key() {
        request = request.header("X-API-Key", key);
    }

    let response = request
        .send()
        .await
        .expect("initialized notification request");
    assert_eq!(
        response.status(),
        reqwest::StatusCode::ACCEPTED,
        "notifications/initialized should be accepted without a JSON-RPC response body"
    );
}

#[tokio::test]
async fn e2e_sse_session_routes_posted_response() {
    let base = base_url();
    let client = reqwest::Client::builder()
        .timeout(Duration::from_secs(20))
        .build()
        .expect("reqwest client");

    wait_until_ready(&client, &base).await;

    let mut sse_request = client
        .get(format!("{base}/mcp/sse"))
        .header("Accept", "text/event-stream");

    if let Some(key) = api_key() {
        sse_request = sse_request.header("X-API-Key", key);
    }

    let mut sse_response = sse_request.send().await.expect("GET /mcp/sse");
    assert_eq!(sse_response.status(), reqwest::StatusCode::OK);
    assert!(
        sse_response
            .headers()
            .get(reqwest::header::CONTENT_TYPE)
            .and_then(|value| value.to_str().ok())
            .map(|value| value.starts_with("text/event-stream"))
            .unwrap_or(false),
        "SSE endpoint should return text/event-stream"
    );

    let mut buffer = String::new();
    let (event_type, endpoint) = tokio::time::timeout(
        Duration::from_secs(10),
        read_sse_event(&mut sse_response, &mut buffer),
    )
    .await
    .expect("timed out waiting for SSE endpoint event")
    .expect("SSE endpoint event");

    assert_eq!(event_type.as_deref(), Some("endpoint"));
    assert!(
        endpoint.contains("/mcp/message?sessionId="),
        "endpoint event should advertise a session-specific message URL"
    );

    let post_url = if endpoint.starts_with("http://") || endpoint.starts_with("https://") {
        endpoint
    } else {
        format!("{base}{endpoint}")
    };

    let mut post_request = client.post(post_url).json(&json!({
        "jsonrpc": "2.0",
        "id": "sse-tools-list-1",
        "method": "tools/list",
        "params": {}
    }));

    if let Some(key) = api_key() {
        post_request = post_request.header("X-API-Key", key);
    }

    let post_response = post_request.send().await.expect("POST session message");
    assert_eq!(
        post_response.status(),
        reqwest::StatusCode::ACCEPTED,
        "session-bound POST should be accepted and routed over SSE"
    );

    let (event_type, payload) = tokio::time::timeout(
        Duration::from_secs(10),
        read_sse_event(&mut sse_response, &mut buffer),
    )
    .await
    .expect("timed out waiting for SSE message event")
    .expect("SSE message event");

    assert_eq!(event_type.as_deref(), Some("message"));

    let message: Value = serde_json::from_str(&payload).expect("SSE payload should be valid JSON");
    assert_eq!(
        message.get("id").and_then(Value::as_str),
        Some("sse-tools-list-1")
    );
    assert!(
        message
            .get("result")
            .and_then(|value| value.get("tools"))
            .and_then(Value::as_array)
            .map(|tools| !tools.is_empty())
            .unwrap_or(false),
        "SSE-routed tools/list response should include tool definitions"
    );
}

#[tokio::test]
async fn e2e_health_endpoints() {
    let base = base_url();
    let client = reqwest::Client::builder()
        .timeout(Duration::from_secs(20))
        .build()
        .expect("reqwest client");

    // Root health endpoint - no auth required
    let root_health: Value = client
        .get(format!("{base}/health"))
        .send()
        .await
        .expect("GET /health")
        .json()
        .await
        .expect("/health JSON");

    assert_eq!(
        root_health.get("status").and_then(Value::as_str),
        Some("ok"),
        "/health should report server liveness"
    );

    // MCP health endpoint - requires auth if enabled
    let mcp_health: Value = get_mcp_endpoint(&client, &base, "/mcp/health")
        .await
        .json()
        .await
        .expect("/mcp/health JSON");

    assert_eq!(
        mcp_health.get("status").and_then(Value::as_str),
        Some("healthy"),
        "/mcp/health should report MCP transport health"
    );
    assert_eq!(
        mcp_health.get("server").and_then(Value::as_str),
        Some("openbrain-mcp")
    );
}

#[tokio::test]
async fn e2e_store_requires_content() {
    let base = base_url();
    let client = reqwest::Client::builder()
        .timeout(Duration::from_secs(20))
        .build()
        .expect("reqwest client");

    wait_until_ready(&client, &base).await;

    let response = call_jsonrpc(
        &client,
        &base,
        json!({
            "jsonrpc": "2.0",
            "id": "store-missing-content-1",
            "method": "tools/call",
            "params": {
                "name": "store",
                "arguments": {
                    "agent_id": format!("e2e-agent-{}", Uuid::new_v4()),
                    "metadata": {"suite": "validation"}
                }
            }
        }),
    )
    .await;

    let message = response
        .get("error")
        .and_then(|e| e.get("message"))
        .and_then(Value::as_str)
        .expect("store missing content should return an error message");

    assert!(
        message.contains("Missing required parameter: content"),
        "store validation should mention missing content"
    );
}

#[tokio::test]
async fn e2e_auth_rejection_without_key() {
    // This test only runs when auth is expected to be enabled
    let auth_enabled = std::env::var("OPENBRAIN__AUTH__ENABLED")
        .map(|v| v == "true")
        .unwrap_or(false);

    if !auth_enabled {
        println!("Skipping auth rejection test - OPENBRAIN__AUTH__ENABLED is not true");
        return;
    }

    let base = base_url();
    let client = reqwest::Client::builder()
        .timeout(Duration::from_secs(20))
        .build()
        .expect("reqwest client");

    // Make request WITHOUT API key
    let response = client
        .post(format!("{base}/mcp/message"))
        .json(&json!({
            "jsonrpc": "2.0",
            "id": "auth-test-1",
            "method": "tools/list",
            "params": {}
        }))
        .send()
        .await
        .expect("HTTP request");

    assert_eq!(
        response.status().as_u16(),
        401,
        "Request without API key should return 401 Unauthorized"
    );
}

fn pick_free_port() -> u16 {
    std::net::TcpListener::bind("127.0.0.1:0")
        .expect("bind ephemeral port")
        .local_addr()
        .expect("local addr")
        .port()
}

async fn wait_for_status(url: &str, expected_status: reqwest::StatusCode) {
    let client = reqwest::Client::builder()
        .timeout(Duration::from_secs(2))
        .build()
        .expect("reqwest client");

    for _ in 0..80 {
        if let Ok(resp) = client.get(url).send().await {
            if resp.status() == expected_status {
                return;
            }
        }
        tokio::time::sleep(Duration::from_millis(250)).await;
    }

    panic!("Timed out waiting for {url} to return status {expected_status}");
}

fn spawn_auth_enabled_test_server(port: u16, test_key: &str) -> std::process::Child {
    Command::new(env!("CARGO_BIN_EXE_openbrain-mcp"))
        .current_dir(env!("CARGO_MANIFEST_DIR"))
        .env("OPENBRAIN__SERVER__PORT", port.to_string())
        .env("OPENBRAIN__AUTH__ENABLED", "true")
        .env("OPENBRAIN__AUTH__API_KEYS", test_key)
        .stdout(Stdio::null())
        .stderr(Stdio::null())
        .spawn()
        .expect("spawn openbrain-mcp for auth-enabled e2e test")
}

#[tokio::test]
async fn e2e_auth_enabled_accepts_test_key() {
    if remote_mode() {
        println!("Skipping local auth spawn test in OPENBRAIN_E2E_REMOTE mode");
        return;
    }

    ensure_schema().await;

    let port = pick_free_port();
    let base = format!("http://127.0.0.1:{port}");
    let test_key = "e2e-test-key-123";

    let mut server = spawn_auth_enabled_test_server(port, test_key);

    wait_for_status(&format!("{base}/ready"), reqwest::StatusCode::OK).await;

    let client = reqwest::Client::builder()
        .timeout(Duration::from_secs(20))
        .build()
        .expect("reqwest client");

    let request = json!({
        "jsonrpc": "2.0",
        "id": "auth-enabled-1",
        "method": "tools/list",
        "params": {}
    });

    let unauthorized = client
        .post(format!("{base}/mcp/message"))
        .json(&request)
        .send()
        .await
        .expect("unauthorized request");
    assert_eq!(
        unauthorized.status(),
        reqwest::StatusCode::UNAUTHORIZED,
        "request without key should be rejected when auth is enabled"
    );

    let authorized: Value = client
        .post(format!("{base}/mcp/message"))
        .header("X-API-Key", test_key)
        .json(&request)
        .send()
        .await
        .expect("authorized request")
        .json()
        .await
        .expect("authorized JSON response");

    assert!(
        authorized.get("error").is_none(),
        "valid key should not return JSON-RPC error"
    );
    assert!(
        authorized
            .get("result")
            .and_then(|r| r.get("tools"))
            .and_then(Value::as_array)
            .map(|tools| !tools.is_empty())
            .unwrap_or(false),
        "authorized tools/list should return tool definitions"
    );

    let bearer_authorized: Value = client
        .post(format!("{base}/mcp/message"))
        .header("Authorization", format!("Bearer {test_key}"))
        .json(&request)
        .send()
        .await
        .expect("bearer-authorized request")
        .json()
        .await
        .expect("bearer-authorized JSON response");

    assert!(
        bearer_authorized.get("error").is_none(),
        "valid bearer token should not return JSON-RPC error"
    );
    assert!(
        bearer_authorized
            .get("result")
            .and_then(|r| r.get("tools"))
            .and_then(Value::as_array)
            .map(|tools| !tools.is_empty())
            .unwrap_or(false),
        "authorized bearer tools/list should return tool definitions"
    );

    let _ = server.kill();
    let _ = server.wait();
}

#[tokio::test]
async fn e2e_same_token_shares_memories_across_agent_ids_and_agent_types() {
    if remote_mode() {
        println!("Skipping local auth spawn test in OPENBRAIN_E2E_REMOTE mode");
        return;
    }

    ensure_schema().await;

    let port = pick_free_port();
    let base = format!("http://127.0.0.1:{port}");
    let test_key = format!("e2e-shared-token-{}", Uuid::new_v4());
    let mut server = spawn_auth_enabled_test_server(port, &test_key);

    wait_for_status(&format!("{base}/ready"), reqwest::StatusCode::OK).await;

    let client = reqwest::Client::builder()
        .timeout(Duration::from_secs(20))
        .build()
        .expect("reqwest client");

    let alpha_agent_id = "agent-alpha";
    let beta_agent_id = "agent-beta";
    let query_agent_id = "agent-gamma";
    let alpha_text = format!("Shared token alpha memory {}", Uuid::new_v4());
    let beta_text = format!("Shared token beta memory {}", Uuid::new_v4());

    let _ = call_tool_with_options(
        &client,
        &base,
        "store",
        json!({
            "content": alpha_text,
            "metadata": { "suite": "shared-token-cross-agent", "transport": "legacy" }
        }),
        Some(test_key.as_str()),
        &[
            ("X-Agent-ID", alpha_agent_id),
            ("X-Agent-Type", "agent-zero"),
        ],
    )
    .await;

    let _ = call_tool_streamable_with_options(
        &client,
        &base,
        "store",
        json!({
            "content": beta_text,
            "metadata": { "suite": "shared-token-cross-agent", "transport": "streamable" }
        }),
        Some(test_key.as_str()),
        &[("X-Agent-ID", beta_agent_id), ("X-Agent-Type", "codex")],
    )
    .await;

    let shared_query = call_tool_streamable_with_options(
        &client,
        &base,
        "query",
        json!({
            "query": "shared token memory",
            "limit": 10,
            "threshold": 0.0
        }),
        Some(test_key.as_str()),
        &[
            ("X-Agent-ID", query_agent_id),
            ("X-Agent-Type", "claude-code"),
        ],
    )
    .await;

    let results = shared_query["results"]
        .as_array()
        .expect("shared query results");

    let has_alpha = results.iter().any(|item| {
        item.get("content")
            .and_then(Value::as_str)
            .map(|content| content == alpha_text)
            .unwrap_or(false)
    });
    let has_beta = results.iter().any(|item| {
        item.get("content")
            .and_then(Value::as_str)
            .map(|content| content == beta_text)
            .unwrap_or(false)
    });
    let source_ids: Vec<&str> = results
        .iter()
        .filter_map(|item| item.get("agent_id").and_then(Value::as_str))
        .collect();

    assert!(has_alpha, "shared query should return alpha agent memory");
    assert!(has_beta, "shared query should return beta agent memory");
    assert!(
        source_ids.contains(&alpha_agent_id),
        "results should preserve alpha source agent provenance"
    );
    assert!(
        source_ids.contains(&beta_agent_id),
        "results should preserve beta source agent provenance"
    );

    let filtered_query = call_tool_with_options(
        &client,
        &base,
        "query",
        json!({
            "source_agent_id": alpha_agent_id,
            "query": "shared token memory",
            "limit": 10,
            "threshold": 0.0
        }),
        Some(test_key.as_str()),
        &[
            ("X-Agent-ID", query_agent_id),
            ("X-Agent-Type", "claude-code"),
        ],
    )
    .await;

    let filtered_results = filtered_query["results"]
        .as_array()
        .expect("filtered query results");
    assert!(
        filtered_results.iter().any(|item| {
            item.get("content")
                .and_then(Value::as_str)
                .map(|content| content == alpha_text)
                .unwrap_or(false)
        }),
        "source_agent_id filter should retain alpha memory"
    );
    assert!(
        !filtered_results.iter().any(|item| {
            item.get("content")
                .and_then(Value::as_str)
                .map(|content| content == beta_text)
                .unwrap_or(false)
        }),
        "source_agent_id filter should exclude beta memory"
    );

    let _ = call_tool_with_options(
        &client,
        &base,
        "purge",
        json!({ "confirm": true }),
        Some(test_key.as_str()),
        &[
            ("X-Agent-ID", query_agent_id),
            ("X-Agent-Type", "claude-code"),
        ],
    )
    .await;

    let _ = server.kill();
    let _ = server.wait();
}

// =============================================================================
// Batch Store Tests (Issue #12)
// =============================================================================

#[tokio::test]
async fn e2e_batch_store_basic() -> anyhow::Result<()> {
    let base = base_url();
    let client = reqwest::Client::builder()
        .timeout(Duration::from_secs(20))
        .build()
        .expect("reqwest client");

    ensure_schema().await;
    wait_until_ready(&client, &base).await;

    let agent = format!("batch_{}", uuid::Uuid::new_v4());
    let _ = call_tool(
        &client,
        &base,
        "purge",
        json!({ "agent_id": agent.clone(), "confirm": true }),
    )
    .await;

    let result = call_tool(
        &client,
        &base,
        "batch_store",
        serde_json::json!({
            "agent_id": agent.clone(),
            "entries": [
                {"content": "Fact alpha for batch test"},
                {"content": "Fact beta for batch test"},
                {"content": "Fact gamma for batch test"}
            ]
        }),
    )
    .await;

    let _ = call_tool(
        &client,
        &base,
        "purge",
        json!({ "agent_id": agent, "confirm": true }),
    )
    .await;

    assert!(result["success"].as_bool().unwrap_or(false));
    assert_eq!(result["count"].as_i64().unwrap_or(0), 3);
    Ok(())
}

#[tokio::test]
async fn e2e_batch_store_empty_rejected() -> anyhow::Result<()> {
    let base = base_url();
    let client = reqwest::Client::builder()
        .timeout(Duration::from_secs(20))
        .build()
        .expect("reqwest client");

    wait_until_ready(&client, &base).await;

    let response = call_jsonrpc(
        &client,
        &base,
        json!({
            "jsonrpc": "2.0",
            "id": "batch-empty-1",
            "method": "tools/call",
            "params": {
                "name": "batch_store",
                "arguments": {
                    "entries": []
                }
            }
        }),
    )
    .await;

    assert!(
        response.get("error").is_some(),
        "empty batch_store should return an error"
    );
    Ok(())
}

#[tokio::test]
async fn e2e_batch_store_exceeds_max() -> anyhow::Result<()> {
    let base = base_url();
    let client = reqwest::Client::builder()
        .timeout(Duration::from_secs(20))
        .build()
        .expect("reqwest client");

    wait_until_ready(&client, &base).await;

    let entries: Vec<Value> = (0..51)
        .map(|i| serde_json::json!({"content": format!("Entry {}", i)}))
        .collect();
    let response = call_jsonrpc(
        &client,
        &base,
        json!({
            "jsonrpc": "2.0",
            "id": "batch-too-large-1",
            "method": "tools/call",
            "params": {
                "name": "batch_store",
                "arguments": {
                    "entries": entries
                }
            }
        }),
    )
    .await;

    assert!(
        response.get("error").is_some(),
        "oversized batch_store should return an error"
    );
    Ok(())
}

#[tokio::test]
async fn e2e_batch_store_missing_content() -> anyhow::Result<()> {
    let base = base_url();
    let client = reqwest::Client::builder()
        .timeout(Duration::from_secs(20))
        .build()
        .expect("reqwest client");

    wait_until_ready(&client, &base).await;

    let response = call_jsonrpc(
        &client,
        &base,
        json!({
            "jsonrpc": "2.0",
            "id": "batch-missing-content-1",
            "method": "tools/call",
            "params": {
                "name": "batch_store",
                "arguments": {
                    "entries": [{"content": "Valid entry"}, {"metadata": {}}]
                }
            }
        }),
    )
    .await;

    assert!(
        response.get("error").is_some(),
        "missing batch entry content should return an error"
    );
    Ok(())
}

#[tokio::test]
async fn e2e_batch_store_appears_in_tools() -> anyhow::Result<()> {
    let base = base_url();
    let client = reqwest::Client::builder()
        .timeout(Duration::from_secs(20))
        .build()
        .expect("reqwest client");

    wait_until_ready(&client, &base).await;

    let response = call_jsonrpc(
        &client,
        &base,
        json!({
            "jsonrpc": "2.0",
            "id": "batch-tools-list-1",
            "method": "tools/list",
            "params": {}
        }),
    )
    .await;

    let names: Vec<&str> = response
        .get("result")
        .and_then(|value| value.get("tools"))
        .and_then(Value::as_array)
        .expect("tools/list result.tools")
        .iter()
        .filter_map(|t| t.get("name").and_then(Value::as_str))
        .collect();
    assert!(names.contains(&"batch_store"));
    Ok(())
}

#[tokio::test]
async fn e2e_existing_store_unchanged() -> anyhow::Result<()> {
    let base = base_url();
    let client = reqwest::Client::builder()
        .timeout(Duration::from_secs(20))
        .build()
        .expect("reqwest client");

    ensure_schema().await;
    wait_until_ready(&client, &base).await;

    let agent = format!("compat_{}", uuid::Uuid::new_v4());
    let _ = call_tool(
        &client,
        &base,
        "purge",
        json!({ "agent_id": agent.clone(), "confirm": true }),
    )
    .await;

    let result = call_tool(
        &client,
        &base,
        "store",
        serde_json::json!({
            "agent_id": agent.clone(),
            "content": "Original store still works"
        }),
    )
    .await;

    let _ = call_tool(
        &client,
        &base,
        "purge",
        json!({ "agent_id": agent, "confirm": true }),
    )
    .await;

    assert!(result["success"].as_bool().unwrap_or(false));
    assert_eq!(result["deduplicated"].as_bool(), Some(false));
    Ok(())
}

// =============================================================================
// Deduplication Tests (Issue #14)
// =============================================================================

#[tokio::test]
async fn e2e_store_deduplicates_and_merges_metadata() -> anyhow::Result<()> {
    let base = base_url();
    let client = reqwest::Client::builder()
        .timeout(Duration::from_secs(20))
        .build()
        .expect("reqwest client");

    ensure_schema().await;
    wait_until_ready(&client, &base).await;

    let agent = format!("dedup_{}", uuid::Uuid::new_v4());
    let content = format!(
        "Dedup fact {} prefers concise replies",
        uuid::Uuid::new_v4()
    );
    let _ = call_tool(
        &client,
        &base,
        "purge",
        json!({ "agent_id": agent.clone(), "confirm": true }),
    )
    .await;

    let first = call_tool(
        &client,
        &base,
        "store",
        json!({
            "agent_id": agent.clone(),
            "content": content.clone(),
            "metadata": {
                "source": "first",
                "keep": true,
                "override": "old"
            }
        }),
    )
    .await;

    assert_eq!(first["deduplicated"].as_bool(), Some(false));

    let first_query = call_tool(
        &client,
        &base,
        "query",
        json!({
            "source_agent_id": agent.clone(),
            "query": content.clone(),
            "limit": 5,
            "threshold": 0.0
        }),
    )
    .await;
    let first_created_at = first_query["results"]
        .as_array()
        .and_then(|items| items.first())
        .and_then(|item| item.get("created_at"))
        .and_then(Value::as_str)
        .expect("first created_at")
        .to_string();

    tokio::time::sleep(Duration::from_millis(1100)).await;

    let second = call_tool(
        &client,
        &base,
        "store",
        json!({
            "agent_id": agent.clone(),
            "content": content.clone(),
            "metadata": {
                "override": "new",
                "second": true
            }
        }),
    )
    .await;

    assert_eq!(second["deduplicated"].as_bool(), Some(true));
    assert_eq!(second["id"], first["id"]);

    let query = call_tool(
        &client,
        &base,
        "query",
        json!({
            "source_agent_id": agent.clone(),
            "query": content.clone(),
            "limit": 5,
            "threshold": 0.0
        }),
    )
    .await;

    assert_eq!(query["count"].as_u64(), Some(1));
    let stored = query["results"]
        .as_array()
        .and_then(|items| items.first())
        .expect("dedup query result");

    assert_eq!(stored["metadata"]["source"], "first");
    assert_eq!(stored["metadata"]["keep"], true);
    assert_eq!(stored["metadata"]["override"], "new");
    assert_eq!(stored["metadata"]["second"], true);

    let second_created_at = stored["created_at"].as_str().expect("second created_at");
    assert!(
        second_created_at > first_created_at.as_str(),
        "deduplicated write should refresh created_at"
    );

    let _ = call_tool(
        &client,
        &base,
        "purge",
        json!({ "agent_id": agent, "confirm": true }),
    )
    .await;

    Ok(())
}

#[tokio::test]
async fn e2e_store_dedup_is_agent_scoped() -> anyhow::Result<()> {
    let base = base_url();
    let client = reqwest::Client::builder()
        .timeout(Duration::from_secs(20))
        .build()
        .expect("reqwest client");

    ensure_schema().await;
    wait_until_ready(&client, &base).await;

    let agent_a = format!("dedup_scope_a_{}", uuid::Uuid::new_v4());
    let agent_b = format!("dedup_scope_b_{}", uuid::Uuid::new_v4());
    let content = format!("Shared cross-agent fact {}", uuid::Uuid::new_v4());

    let _ = call_tool(
        &client,
        &base,
        "purge",
        json!({ "agent_id": agent_a.clone(), "confirm": true }),
    )
    .await;
    let _ = call_tool(
        &client,
        &base,
        "purge",
        json!({ "agent_id": agent_b.clone(), "confirm": true }),
    )
    .await;

    let first = call_tool(
        &client,
        &base,
        "store",
        json!({
            "agent_id": agent_a.clone(),
            "content": content.clone()
        }),
    )
    .await;
    let second = call_tool(
        &client,
        &base,
        "store",
        json!({
            "agent_id": agent_b.clone(),
            "content": content.clone()
        }),
    )
    .await;

    assert_eq!(first["deduplicated"].as_bool(), Some(false));
    assert_eq!(second["deduplicated"].as_bool(), Some(false));
    assert_ne!(first["id"], second["id"]);

    let _ = call_tool(
        &client,
        &base,
        "purge",
        json!({ "agent_id": agent_a, "confirm": true }),
    )
    .await;
    let _ = call_tool(
        &client,
        &base,
        "purge",
        json!({ "agent_id": agent_b, "confirm": true }),
    )
    .await;

    Ok(())
}

#[tokio::test]
async fn e2e_batch_store_deduplicates_within_batch() -> anyhow::Result<()> {
    let base = base_url();
    let client = reqwest::Client::builder()
        .timeout(Duration::from_secs(20))
        .build()
        .expect("reqwest client");

    ensure_schema().await;
    wait_until_ready(&client, &base).await;

    let agent = format!("batch_dedup_{}", uuid::Uuid::new_v4());
    let content = format!("Batch dedup fact {}", uuid::Uuid::new_v4());
    let _ = call_tool(
        &client,
        &base,
        "purge",
        json!({ "agent_id": agent.clone(), "confirm": true }),
    )
    .await;

    let result = call_tool(
        &client,
        &base,
        "batch_store",
        json!({
            "agent_id": agent.clone(),
            "entries": [
                {
                    "content": content.clone(),
                    "metadata": { "source": "first", "keep": "yes" }
                },
                {
                    "content": content.clone(),
                    "metadata": { "source": "second", "merged": "yes" }
                }
            ]
        }),
    )
    .await;

    let results = result["results"].as_array().expect("batch results");
    assert_eq!(result["count"].as_u64(), Some(2));
    assert_eq!(results[0]["deduplicated"].as_bool(), Some(false));
    assert_eq!(results[0]["status"], "stored");
    assert_eq!(results[1]["deduplicated"].as_bool(), Some(true));
    assert_eq!(results[1]["status"], "deduplicated");
    assert_eq!(results[0]["id"], results[1]["id"]);

    let query = call_tool(
        &client,
        &base,
        "query",
        json!({
            "source_agent_id": agent.clone(),
            "query": content.clone(),
            "limit": 5,
            "threshold": 0.0
        }),
    )
    .await;

    assert_eq!(query["count"].as_u64(), Some(1));
    let stored = query["results"]
        .as_array()
        .and_then(|items| items.first())
        .expect("batch dedup query result");
    assert_eq!(stored["metadata"]["source"], "second");
    assert_eq!(stored["metadata"]["keep"], "yes");
    assert_eq!(stored["metadata"]["merged"], "yes");

    let _ = call_tool(
        &client,
        &base,
        "purge",
        json!({ "agent_id": agent, "confirm": true }),
    )
    .await;

    Ok(())
}