tidaldb/tidal-server/tests/vector_search.rs

// Integration-test exemption (same posture as the other tidal-server tests).
#![allow(clippy::unwrap_used, clippy::cast_possible_truncation)]
//! End-to-end coverage for the m12p1 `POST /vector_search` recall probe.
//!
//! Drives the real standalone handler in-process via `tower::ServiceExt::oneshot`
//! (no TCP bind): seed items + embeddings, then POST a query vector and assert
//! the raw k-NN result — closest-first ordering, `k` honored — plus the boundary
//! 400s (empty vector, dimension mismatch). This pins the surface the
//! `tidal-stress --verify-recall` harness measures against.

use std::sync::Arc;

use axum::{
    body::Body,
    http::{Method, Request, StatusCode},
};
use tidal_server::{router::build_router, state::ServerState};
use tidaldb::TidalDb;
use tower::ServiceExt;

/// Dimensionality of the default schema's `content_vector` slot.
const DIM: usize = 128;

fn make_app() -> axum::Router {
    let (schema, profiles) = tidal_server::config::load_schema(None).unwrap();
    let db = TidalDb::builder()
        .ephemeral()
        .with_schema(schema)
        .with_profiles(profiles)
        .open()
        .unwrap();
    let state = Arc::new(ServerState::new(db));
    build_router(
        state,
        Arc::new(tidal_server::cluster::security::ClusterCreds::unauthenticated()),
    )
}

/// A 128-dim one-hot-ish vector: component `axis` set to `mag`, rest 0 — except
/// we nudge a second axis a hair so no vector is exactly zero-norm.
fn axis_vector(axis: usize, mag: f32) -> Vec<f32> {
    let mut v = vec![0.0_f32; DIM];
    v[axis] = mag;
    v[(axis + 1) % DIM] = 0.01;
    v
}

/// `v` scaled by `s` — a TRUE scalar multiple, i.e. the same direction at a
/// different magnitude.
///
/// Necessary because `axis_vector(axis, mag)` pins a fixed `0.01` on the second
/// axis regardless of `mag`, so varying `mag` there changes the DIRECTION once the
/// vector is normalized, not just its length. A magnitude-independence test built
/// on `axis_vector` would be testing the wrong thing (and would fail for a correct
/// implementation).
fn scaled(v: &[f32], s: f32) -> Vec<f32> {
    v.iter().map(|x| x * s).collect()
}

async fn post_json(app: &axum::Router, uri: &str, body: serde_json::Value) -> StatusCode {
    app.clone()
        .oneshot(
            Request::builder()
                .method(Method::POST)
                .uri(uri)
                .header("Content-Type", "application/json")
                .body(Body::from(serde_json::to_vec(&body).unwrap()))
                .unwrap(),
        )
        .await
        .unwrap()
        .status()
}

async fn post_json_full(
    app: &axum::Router,
    uri: &str,
    body: serde_json::Value,
) -> (StatusCode, serde_json::Value) {
    let response = app
        .clone()
        .oneshot(
            Request::builder()
                .method(Method::POST)
                .uri(uri)
                .header("Content-Type", "application/json")
                .body(Body::from(serde_json::to_vec(&body).unwrap()))
                .unwrap(),
        )
        .await
        .unwrap();
    let status = response.status();
    let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
        .await
        .unwrap();
    let json: serde_json::Value = serde_json::from_slice(&bytes).unwrap_or(serde_json::Value::Null);
    (status, json)
}

async fn seed(app: &axum::Router) {
    // Items on distinct axes: item 1 ~ axis 0, item 3 ~ axis 0 (close to item 1),
    // item 2 ~ axis 64 (far). A query on axis 0 must rank 1 and 3 above 2.
    for (id, v) in [
        (1u64, axis_vector(0, 1.0)),
        (2u64, axis_vector(64, 1.0)),
        (3u64, axis_vector(0, 0.8)),
    ] {
        let s = post_json(
            app,
            "/items",
            serde_json::json!({ "entity_id": id, "metadata": {} }),
        )
        .await;
        assert_eq!(s, StatusCode::CREATED, "item {id}");
        let s = post_json(
            app,
            "/embeddings",
            serde_json::json!({ "entity_id": id, "values": v }),
        )
        .await;
        assert_eq!(s, StatusCode::NO_CONTENT, "embedding {id}");
    }
}

#[tokio::test]
async fn vector_search_returns_nearest_closest_first() {
    let app = make_app();
    seed(&app).await;

    let (status, body) = post_json_full(
        &app,
        "/vector_search",
        serde_json::json!({ "vector": axis_vector(0, 1.0), "k": 3 }),
    )
    .await;
    assert_eq!(status, StatusCode::OK, "body: {body}");

    let items = body["items"].as_array().expect("items array");
    assert_eq!(items.len(), 3, "k=3 nearest");

    // Closest-first: an axis-0 query ranks the two axis-0 items (1, 3) above the
    // far axis-64 item (2).
    let ids: Vec<u64> = items
        .iter()
        .map(|it| it["entity_id"].as_u64().unwrap())
        .collect();
    assert_eq!(ids[0], 1, "the exact-axis item is nearest");
    assert!(
        ids[..2].contains(&3),
        "the near-axis item ranks above the far one; got {ids:?}"
    );
    assert_eq!(ids[2], 2, "the far axis-64 item is last");

    // Distances are present and ascending.
    let dists: Vec<f64> = items
        .iter()
        .map(|it| it["distance"].as_f64().unwrap())
        .collect();
    for w in dists.windows(2) {
        assert!(w[0] <= w[1], "distances must ascend: {dists:?}");
    }
}

#[tokio::test]
async fn vector_search_k_defaults_to_ten_and_clamps_to_corpus() {
    let app = make_app();
    seed(&app).await;
    // k omitted → defaults to 10, but only 3 items exist, so 3 come back.
    let (status, body) = post_json_full(
        &app,
        "/vector_search",
        serde_json::json!({ "vector": axis_vector(0, 1.0) }),
    )
    .await;
    assert_eq!(status, StatusCode::OK);
    assert_eq!(body["items"].as_array().unwrap().len(), 3);
}

#[tokio::test]
async fn vector_search_empty_vector_is_400() {
    let app = make_app();
    seed(&app).await;
    let v: Vec<f32> = vec![];
    let status = post_json(&app, "/vector_search", serde_json::json!({ "vector": v })).await;
    assert_eq!(status, StatusCode::BAD_REQUEST);
}

#[tokio::test]
async fn vector_search_dimension_mismatch_is_400() {
    let app = make_app();
    seed(&app).await;
    // 4-dim query against a 128-dim slot → a client error, not a 500.
    let status = post_json(
        &app,
        "/vector_search",
        serde_json::json!({ "vector": vec![0.1f32, 0.2, 0.3, 0.4] }),
    )
    .await;
    assert_eq!(status, StatusCode::BAD_REQUEST);
}

/// An exact-match query must score ~0, not `|q|² − 1`.
///
/// The write path L2-normalizes every stored vector; until 2026-08-30 the read path
/// passed the caller's raw query straight to the index, so the two sides lived in
/// different spaces. With unit `v`, `d = |q|² − 2q·v + 1`. Measured live, a non-unit
/// query returned 591–1174 against a documented `[0.0, 4.0]`.
#[tokio::test]
async fn vector_search_exact_match_scores_zero() {
    let app = make_app();
    seed(&app).await;

    // Item 1's exact direction, scaled x3 so the query is NOT unit-length: the
    // point is that the server owns the normalization. A unit query would pass
    // even against the bug. Must be a true scalar multiple (see `scaled`), or it
    // is a different direction and not an exact match at all.
    let query = scaled(&axis_vector(0, 1.0), 3.0);
    let (status, body) = post_json_full(
        &app,
        "/vector_search",
        serde_json::json!({ "vector": query, "k": 1 }),
    )
    .await;
    assert_eq!(status, StatusCode::OK, "body: {body}");

    let d = body["items"][0]["distance"].as_f64().expect("distance");
    // f16 quantization puts an exact match near 1e-3, not at 0 — assert a tolerance.
    // Equality here would be a flaky test, not a stricter one.
    assert!(
        d < 1e-2,
        "an exact-direction match must score ~0, got {d} (pre-fix this was |q|^2 - 1 ~= 8)"
    );
}

/// Every distance must fall inside the range documented at
/// `tidal/src/storage/vector/mod.rs:49`.
///
/// This is the assertion whose absence let 591–1174 ship unnoticed: the existing
/// ordering test proves distances ASCEND, which stayed true the whole time.
#[tokio::test]
async fn vector_search_distances_within_documented_range() {
    let app = make_app();
    seed(&app).await;

    // A range of query magnitudes, since the bug scaled distances by |q|^2 — one
    // magnitude could coincidentally land in range.
    for mag in [0.5_f32, 1.0, 3.0, 10.0] {
        let (status, body) = post_json_full(
            &app,
            "/vector_search",
            serde_json::json!({ "vector": scaled(&axis_vector(0, 1.0), mag), "k": 3 }),
        )
        .await;
        assert_eq!(status, StatusCode::OK, "mag {mag}, body: {body}");

        let dists: Vec<f64> = body["items"]
            .as_array()
            .expect("items")
            .iter()
            .map(|it| it["distance"].as_f64().expect("distance"))
            .collect();
        assert!(
            dists.iter().all(|d| (0.0..=4.0).contains(d)),
            "mag {mag}: outside documented [0.0, 4.0] (storage/vector/mod.rs:49): {dists:?}"
        );
    }
}

/// A zero-norm query has no direction, so "nearest by cosine" is undefined for it.
/// Reject it as caller error rather than returning an arbitrary ranking.
#[tokio::test]
async fn vector_search_rejects_zero_query() {
    let app = make_app();
    seed(&app).await;

    let (status, body) = post_json_full(
        &app,
        "/vector_search",
        serde_json::json!({ "vector": vec![0.0_f32; DIM], "k": 3 }),
    )
    .await;
    assert_eq!(
        status,
        StatusCode::BAD_REQUEST,
        "a zero query must be rejected, not answered with an arbitrary ranking; body: {body}"
    );
}

/// Normalizing the query must NOT change ranking — it is a contract fix, not a
/// behavioral one.
///
/// `|q|²` and `1` are constant across candidates, so ordering was already correct
/// cosine order. This is the regression guard: if a future change to the read path
/// alters ranking, this fails even though the range assertions still pass.
#[tokio::test]
async fn vector_search_ordering_unchanged_by_normalization() {
    let app = make_app();
    seed(&app).await;

    // The same direction at four magnitudes must produce the IDENTICAL id order.
    // Pre-fix this held too (the bug was scale-only), so this test passes before
    // AND after — which is exactly what makes it a guard rather than a symptom test.
    let mut orders = Vec::new();
    for mag in [0.5_f32, 1.0, 3.0, 10.0] {
        let (status, body) = post_json_full(
            &app,
            "/vector_search",
            serde_json::json!({ "vector": scaled(&axis_vector(0, 1.0), mag), "k": 3 }),
        )
        .await;
        assert_eq!(status, StatusCode::OK, "mag {mag}");
        orders.push(
            body["items"]
                .as_array()
                .expect("items")
                .iter()
                .map(|it| it["entity_id"].as_u64().expect("entity_id"))
                .collect::<Vec<u64>>(),
        );
    }
    assert_eq!(orders[0], vec![1, 3, 2], "golden order for an axis-0 query");
    for (i, o) in orders.iter().enumerate() {
        assert_eq!(
            o, &orders[0],
            "query magnitude must not affect ranking (index {i}): {orders:?}"
        );
    }
}