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 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() {
    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(),
                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()
            );
            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);
            "#,
        )
        .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");
}

async fn call_jsonrpc(client: &reqwest::Client, base: &str, request: Value) -> Value {
    let mut req_builder = client
        .post(format!("{base}/mcp/message"))
        .json(&request);
    
    // Add API key header if available
    if let Some(key) = api_key() {
        req_builder = req_builder.header("X-API-Key", key);
    }
    
    req_builder
        .send()
        .await
        .expect("JSON-RPC HTTP request")
        .json()
        .await
        .expect("JSON-RPC response body")
}

/// Make an authenticated GET request to an MCP endpoint
async fn get_mcp_endpoint(client: &reqwest::Client, base: &str, path: &str) -> reqwest::Response {
    let mut req_builder = client.get(format!("{base}{path}"));
    
    if let Some(key) = api_key() {
        req_builder = req_builder.header("X-API-Key", key);
    }
    
    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"));
    }
}

async fn call_tool(
    client: &reqwest::Client,
    base: &str,
    tool_name: &str,
    arguments: Value,
) -> 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(client, base, request).await;

    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")
}

#[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!({
            "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!({
            "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_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_isolated_by_agent_id() {
    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 query_a = call_tool(
        &client,
        &base,
        "query",
        json!({
            "agent_id": agent_a,
            "query": "mode preference",
            "limit": 10,
            "threshold": 0.0
        }),
    )
    .await;

    let results = query_a
        .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, "agent A query should include agent A memory");
    assert!(!has_b, "agent A query must not include 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}");
}

#[tokio::test]
async fn e2e_auth_enabled_accepts_test_key() {
    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 = 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");

    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();
}