Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions apps/memos-local-plugin/core/storage/repos/policies.ts
Original file line number Diff line number Diff line change
Expand Up @@ -185,6 +185,7 @@ export function makePoliciesRepo(db: StorageDb) {
where: whereParts.join(" AND "),
params,
hardCap: opts.hardCap,
orderBy: "updated_at DESC, id DESC",
},
);
},
Expand Down
1 change: 1 addition & 0 deletions apps/memos-local-plugin/core/storage/repos/skills.ts
Original file line number Diff line number Diff line change
Expand Up @@ -223,6 +223,7 @@ export function makeSkillsRepo(db: StorageDb) {
where: whereParts.join(" AND "),
params,
hardCap: opts.hardCap,
orderBy: "updated_at DESC, id DESC",
},
);
},
Expand Down
1 change: 1 addition & 0 deletions apps/memos-local-plugin/core/storage/repos/traces.ts
Original file line number Diff line number Diff line change
Expand Up @@ -552,6 +552,7 @@ export function makeTracesRepo(db: StorageDb) {
where: whereParts.join(" AND "),
params,
hardCap: opts.hardCap,
orderBy: "ts DESC, id DESC", // repo-internal constant; validated in scanAndTopK
},
);
},
Expand Down
1 change: 1 addition & 0 deletions apps/memos-local-plugin/core/storage/repos/world_model.ts
Original file line number Diff line number Diff line change
Expand Up @@ -172,6 +172,7 @@ export function makeWorldModelRepo(db: StorageDb) {
vecColumn: "vec",
where,
hardCap: opts.hardCap,
orderBy: "updated_at DESC, id DESC",
});
},

Expand Down
30 changes: 29 additions & 1 deletion apps/memos-local-plugin/core/storage/vector.ts
Original file line number Diff line number Diff line change
Expand Up @@ -185,6 +185,15 @@ export interface VectorScanOptions {
params?: Record<string, unknown>;
/** Optional LIMIT to cap candidates fetched from SQLite. */
hardCap?: number;
/**
* Optional ORDER BY clause (without the "ORDER BY") applied before the
* `hardCap` LIMIT. Without it the bounded candidate window is SQLite's
* arbitrary physical scan prefix, so the globally best vector can be
* excluded purely by physical position. Repos pass their recency column
* (e.g. `ts DESC, id DESC`) so the window is a deterministic, meaningful
* candidate policy — most recent rows first — that existing indexes serve.
*/
orderBy?: string;
}

export interface ScanRow {
Expand All @@ -208,11 +217,26 @@ export interface ScanRow {
*/
export const DEFAULT_SCAN_HARD_CAP = 5_000;

/**
* `orderBy` is interpolated into SQL, so it must be a repo-internal constant.
* This allowlist (identifier[, identifier]… each optionally ASC/DESC) rejects
* anything else at the boundary, so a future caller passing request-derived
* input fails loudly instead of opening SQL injection.
*/
const SAFE_ORDER_BY_RE =
/^[a-z_][a-z0-9_]*(\s+(asc|desc))?(,\s*[a-z_][a-z0-9_]*(\s+(asc|desc))?)*$/i;

/**
* Stream rows from `table`, decode vectors, and run top-K cosine against
* `query`. `selectExtra` lets callers bring along columns that will surface in
* `VectorHit.meta`.
*
* Bounded-scan semantics: the `hardCap` LIMIT bounds how many rows enter
* cosine ranking. Without `orderBy` that window is SQLite's arbitrary
* physical scan prefix; repos pass a deterministic recency order (e.g.
* `ts DESC, id DESC`) so the window is a defined candidate policy — the most
* recent qualifying rows — instead of a physical accident (#2233).
*
* Streaming: we use `.iterate()` (not `.all()`) so at most one row's
* BLOB is decoded at a time. The top-K min-heap keeps only `k`
* vectors of state, so peak RSS is O(k * dim) regardless of how many
Expand All @@ -229,12 +253,16 @@ export function scanAndTopK<TMeta = undefined>(
): Array<VectorHit<string, TMeta>> {
if (k <= 0 || query.length === 0) return [];

const { vecColumn, norm2Column, where, params, hardCap } = opts;
const { vecColumn, norm2Column, where, params, hardCap, orderBy } = opts;
if (orderBy !== undefined && !SAFE_ORDER_BY_RE.test(orderBy)) {
throw new Error(`scanAndTopK: unsafe orderBy value: ${JSON.stringify(orderBy)}`);
}
const cap = hardCap ?? DEFAULT_SCAN_HARD_CAP;
const cols = ["id", vecColumn, ...(norm2Column ? [norm2Column] : []), ...selectExtra];
const sql = [
`SELECT ${cols.join(", ")} FROM ${table}`,
where ? `WHERE ${where}` : "",
orderBy ? `ORDER BY ${orderBy}` : "",
`LIMIT ${cap}`,
]
.filter(Boolean)
Expand Down
77 changes: 77 additions & 0 deletions apps/memos-local-plugin/tests/unit/storage/vector-stream.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -226,4 +226,81 @@ describe("scanAndTopK — streaming rewrite (#2076)", () => {
db.close();
}
});

it("orderBy makes the hardCap window a deterministic recency policy, not a physical prefix (#2233)", () => {
// Same table shape `traces`/`policies` vector search uses: a recency
// column next to the vector BLOB.
const query = vec([1, 0]);
const makeRows = () => {
// Two perfect matches: an OLD one and a NEW one; 18 orthogonal fillers.
// Which one a bounded window sees is entirely decided by its ordering
// policy — the bug was that the policy was "whatever SQLite's physical
// scan happens to visit first".
const rows = [
{ id: "r-new", v: vec([1, 0]), ts: 1_000 },
{ id: "r-old", v: vec([1, 0]), ts: 1 },
];
for (let i = 0; i < 18; i++) {
rows.push({ id: `r-filler-${i}`, v: vec([0, 1]), ts: 2 + i });
}
return rows;
};

// With the recency order the repos now pass, a cap-1 window is
// deterministically the newest row — the perfect match with ts=1000 —
// regardless of the table's physical row layout.
for (const layout of ["insertion", "reversed"] as const) {
const db = new Database(":memory:");
db.exec(`CREATE TABLE bench (id TEXT PRIMARY KEY, vec BLOB, ts INTEGER NOT NULL);`);
try {
const insert = db.prepare("INSERT INTO bench (id, vec, ts) VALUES (?, ?, ?)");
const rows = makeRows();
const ordered = layout === "reversed" ? [...rows].reverse() : rows;
for (const r of ordered) insert.run(r.id, encodeVector(r.v), r.ts);

const hits = scanAndTopK(
// eslint-disable-next-line @typescript-eslint/no-explicit-any
db as any,
"bench",
[],
query,
1,
{ vecColumn: "vec", where: "vec IS NOT NULL", hardCap: 1, orderBy: "ts DESC, id DESC" },
);
expect(hits).toHaveLength(1);
expect(hits[0]!.id).toBe("r-new");
expect(hits[0]!.score).toBeCloseTo(1, 5);
} finally {
db.close();
}
}

// Without orderBy the cap-1 window is whichever row the planner visits
// first — layout-dependent by definition, so there is nothing stable to
// assert about the winner; the point of the option is that it no
// longer matters.
});

it("rejects an orderBy that is not a repo-internal column list (SQL-injection boundary)", () => {
const db = openTinyVecDb();
try {
const inject = "ts DESC; DROP TABLE bench --";
expect(() =>
scanAndTopK(
// eslint-disable-next-line @typescript-eslint/no-explicit-any
db as any,
"bench",
[],
vec([1, 0]),
1,
{ vecColumn: "vec", where: "vec IS NOT NULL", orderBy: inject },
),
).toThrow(/unsafe orderBy/);
// The table is untouched — the guard fired before any SQL ran.
const n = db.prepare("SELECT COUNT(*) AS n FROM bench").get() as { n: number };
expect(n.n).toBe(0);
} finally {
db.close();
}
});
});