mirror of
https://github.com/badlogic/pi-mono.git
synced 2026-11-06 23:44:36 +00:00
fix(agent): harden bounded branch queries
This commit is contained in:
Christian Klotz
parent
9b75872017
commit
081d7fbfca
5 changed files with 443 additions and 25 deletions
|
|
@ -9,6 +9,12 @@ export interface CachedBranch {
|
|||
leafSeq: number;
|
||||
}
|
||||
|
||||
export interface CachedBranchQuery {
|
||||
stopAtType?: SessionEntryRow["type"];
|
||||
stopAtId?: string;
|
||||
order?: "newestFirst" | "oldestFirst";
|
||||
}
|
||||
|
||||
export async function readCachedBranch(
|
||||
db: SqliteDatabase,
|
||||
sessionId: string,
|
||||
|
|
@ -23,6 +29,78 @@ export async function readCachedBranch(
|
|||
return { branchId: membership.branch_id, leafSeq: membership.entry_seq };
|
||||
}
|
||||
|
||||
export async function isCachedBranchValid(
|
||||
db: SqliteDatabase,
|
||||
sessionId: string,
|
||||
branch: CachedBranch,
|
||||
leafId: string,
|
||||
startSeq = 0,
|
||||
): Promise<boolean> {
|
||||
const result = await db
|
||||
.prepare(
|
||||
`WITH path AS (
|
||||
SELECT
|
||||
b.entry_id,
|
||||
b.entry_seq,
|
||||
e.id AS stored_entry_id,
|
||||
e.parent_id,
|
||||
LAG(b.entry_id) OVER (ORDER BY b.entry_seq) AS previous_entry_id
|
||||
FROM branch_entries AS b
|
||||
LEFT JOIN session_entries AS e ON e.session_id = b.session_id AND e.id = b.entry_id
|
||||
WHERE b.session_id = ? AND b.branch_id = ? AND b.entry_seq BETWEEN ? AND ?
|
||||
)
|
||||
SELECT
|
||||
COUNT(*) AS row_count,
|
||||
COALESCE(SUM(
|
||||
stored_entry_id IS NULL OR
|
||||
(? = 0 AND previous_entry_id IS NULL AND parent_id IS NOT NULL) OR
|
||||
(previous_entry_id IS NOT NULL AND parent_id IS NOT previous_entry_id)
|
||||
), 0) AS invalid_count,
|
||||
COALESCE(MAX(entry_seq = ? AND entry_id = ?), 0) AS contains_leaf
|
||||
FROM path`,
|
||||
)
|
||||
.get<{ row_count: number; invalid_count: number; contains_leaf: number }>(
|
||||
sessionId,
|
||||
branch.branchId,
|
||||
startSeq,
|
||||
branch.leafSeq,
|
||||
startSeq,
|
||||
branch.leafSeq,
|
||||
leafId,
|
||||
);
|
||||
return result?.row_count !== 0 && result?.invalid_count === 0 && result.contains_leaf === 1;
|
||||
}
|
||||
|
||||
export async function readNewestCachedStopSeq(
|
||||
db: SqliteDatabase,
|
||||
sessionId: string,
|
||||
branch: CachedBranch,
|
||||
stopAtType: SessionEntryRow["type"] | undefined,
|
||||
stopAtId: string | undefined,
|
||||
): Promise<number | undefined> {
|
||||
const predicates: string[] = [];
|
||||
const params: unknown[] = [sessionId, branch.branchId, branch.leafSeq];
|
||||
if (stopAtType !== undefined) {
|
||||
predicates.push("e.type = ?");
|
||||
params.push(stopAtType);
|
||||
}
|
||||
if (stopAtId !== undefined) {
|
||||
predicates.push("b.entry_id = ?");
|
||||
params.push(stopAtId);
|
||||
}
|
||||
if (predicates.length === 0) return undefined;
|
||||
const row = await db
|
||||
.prepare(
|
||||
`SELECT MAX(b.entry_seq) AS entry_seq
|
||||
FROM branch_entries AS b
|
||||
JOIN session_entries AS e ON e.session_id = b.session_id AND e.id = b.entry_id
|
||||
WHERE b.session_id = ? AND b.branch_id = ? AND b.entry_seq <= ?
|
||||
AND (${predicates.join(" OR ")})`,
|
||||
)
|
||||
.get<{ entry_seq: number | null }>(...params);
|
||||
return row?.entry_seq ?? undefined;
|
||||
}
|
||||
|
||||
export async function readCachedBranchRows(
|
||||
db: SqliteDatabase,
|
||||
sessionId: string,
|
||||
|
|
@ -40,6 +118,51 @@ export async function readCachedBranchRows(
|
|||
.all<SessionEntryRow>(sessionId, branch.branchId, startSeq, branch.leafSeq);
|
||||
}
|
||||
|
||||
export async function queryCachedBranchRows(
|
||||
db: SqliteDatabase,
|
||||
sessionId: string,
|
||||
branch: CachedBranch,
|
||||
query: CachedBranchQuery,
|
||||
): Promise<SessionEntryRow[]> {
|
||||
const oldestFirst = query.order === "oldestFirst";
|
||||
const boundaryParams: unknown[] = [sessionId, branch.branchId, branch.leafSeq];
|
||||
const stopPredicates: string[] = [];
|
||||
if (query.stopAtType !== undefined) {
|
||||
stopPredicates.push("stop_entry.type = ?");
|
||||
boundaryParams.push(query.stopAtType);
|
||||
}
|
||||
if (query.stopAtId !== undefined) {
|
||||
stopPredicates.push("stop.entry_id = ?");
|
||||
boundaryParams.push(query.stopAtId);
|
||||
}
|
||||
|
||||
const boundary = stopPredicates.length
|
||||
? `WITH boundary AS (
|
||||
SELECT ${oldestFirst ? "MIN" : "MAX"}(stop.entry_seq) AS entry_seq
|
||||
FROM branch_entries AS stop
|
||||
JOIN session_entries AS stop_entry
|
||||
ON stop_entry.session_id = stop.session_id AND stop_entry.id = stop.entry_id
|
||||
WHERE stop.session_id = ? AND stop.branch_id = ? AND stop.entry_seq <= ?
|
||||
AND (${stopPredicates.join(" OR ")})
|
||||
)`
|
||||
: "";
|
||||
const range = stopPredicates.length
|
||||
? `AND b.entry_seq ${oldestFirst ? "<=" : ">="} COALESCE(
|
||||
(SELECT entry_seq FROM boundary), ${oldestFirst ? branch.leafSeq : 0}
|
||||
)`
|
||||
: "";
|
||||
const sql = `${boundary}
|
||||
SELECT e.session_id, e.id, e.entry_seq, e.parent_id, e.type, e.timestamp, e.payload
|
||||
FROM branch_entries AS b
|
||||
JOIN session_entries AS e ON e.session_id = b.session_id AND e.id = b.entry_id
|
||||
WHERE b.session_id = ? AND b.branch_id = ? AND b.entry_seq <= ?
|
||||
${range}
|
||||
ORDER BY b.entry_seq ${oldestFirst ? "ASC" : "DESC"}`;
|
||||
|
||||
const params = [...(stopPredicates.length === 0 ? [] : boundaryParams), sessionId, branch.branchId, branch.leafSeq];
|
||||
return db.prepare(sql).all<SessionEntryRow>(...params);
|
||||
}
|
||||
|
||||
export async function readCachedEntryRowsByType(
|
||||
db: SqliteDatabase,
|
||||
sessionId: string,
|
||||
|
|
|
|||
|
|
@ -8,10 +8,14 @@ import { SessionError, toError } from "@earendil-works/pi-agent-core";
|
|||
import type { SqliteDatabase, SqliteSessionMetadata } from "../types.ts";
|
||||
import {
|
||||
appendEntryToBranchCache,
|
||||
type CachedBranch,
|
||||
isCachedBranchValid,
|
||||
queryCachedBranchRows,
|
||||
readCachedBranch,
|
||||
readCachedBranchRows,
|
||||
readCachedEntryRowsByType,
|
||||
readCachedEntrySeq,
|
||||
readNewestCachedStopSeq,
|
||||
rebuildCachedBranch,
|
||||
} from "./branch-cache.ts";
|
||||
import { decodeEntry, encodeEntry, type SessionEntryRow } from "./session-entries.ts";
|
||||
|
|
@ -81,32 +85,87 @@ export class SqliteSessionConnection implements SessionReader<SqliteSessionMetad
|
|||
throw new RangeError("Session branch query limit must be a positive integer");
|
||||
}
|
||||
if (query.start === null) return [];
|
||||
const fullPath = await this.readPathToRoot(query.start);
|
||||
const path: SessionTreeEntry[] = [];
|
||||
for (let index = fullPath.length - 1; index >= 0; index--) {
|
||||
const entry = fullPath[index]!;
|
||||
path.push(entry);
|
||||
if (entry.id === query.stopAtId || entry.type === query.stopAtType) break;
|
||||
const startId = query.start;
|
||||
let cached = await readCachedBranch(this.db, this.metadata.id, startId);
|
||||
const validationStartSeq =
|
||||
cached && query.order !== "oldestFirst"
|
||||
? await readNewestCachedStopSeq(this.db, this.metadata.id, cached, query.stopAtType, query.stopAtId)
|
||||
: undefined;
|
||||
if (!cached || !(await isCachedBranchValid(this.db, this.metadata.id, cached, startId, validationStartSeq))) {
|
||||
if (query.order !== "oldestFirst" && (query.stopAtId !== undefined || query.stopAtType !== undefined)) {
|
||||
return this.findEntriesOnCanonicalBranch({ ...query, start: startId });
|
||||
}
|
||||
cached = await this.repairBranchCacheForQuery(startId, cached?.branchId);
|
||||
}
|
||||
if (query.order === "oldestFirst") path.reverse();
|
||||
const entries = path.filter(
|
||||
const decoded = decodeEntryRows(await queryCachedBranchRows(this.db, this.metadata.id, cached, query));
|
||||
const filtered = decoded.filter(
|
||||
(entry) =>
|
||||
(query.type === undefined || entry.type === query.type) &&
|
||||
(query.customType === undefined || (entry.type === "custom" && entry.customType === query.customType)),
|
||||
);
|
||||
return query.limit === undefined ? entries : entries.slice(0, query.limit);
|
||||
const entries = query.limit === undefined ? filtered : filtered.slice(0, query.limit);
|
||||
for (const entry of entries) this.byId.set(entry.id, entry);
|
||||
return entries;
|
||||
}
|
||||
|
||||
private async readPathToRoot(leafId: string): Promise<SessionTreeEntry[]> {
|
||||
const cached = await readCachedBranch(this.db, this.metadata.id, leafId);
|
||||
if (cached) {
|
||||
const entries = decodeEntryRows(await readCachedBranchRows(this.db, this.metadata.id, cached, 0));
|
||||
if (this.isValidCachedPath(entries, leafId, null)) {
|
||||
for (const entry of entries) this.byId.set(entry.id, entry);
|
||||
return entries;
|
||||
private async findEntriesOnCanonicalBranch(
|
||||
query: SessionBranchQuery & { start: string },
|
||||
): Promise<SessionTreeEntry[]> {
|
||||
const rows: SessionEntryRow[] = [];
|
||||
const visited = new Set<string>();
|
||||
let currentId: string | null = query.start;
|
||||
while (currentId !== null) {
|
||||
if (visited.has(currentId)) throw invalidSession(`cycle in parent chain at entry ${currentId}`);
|
||||
visited.add(currentId);
|
||||
const row: SessionEntryRow | undefined = await this.db
|
||||
.prepare(
|
||||
"SELECT session_id, id, entry_seq, parent_id, type, timestamp, payload FROM session_entries WHERE session_id = ? AND id = ?",
|
||||
)
|
||||
.get<SessionEntryRow>(this.metadata.id, currentId);
|
||||
if (!row) {
|
||||
if (currentId === query.start) throw new SessionError("not_found", `Entry ${query.start} not found`);
|
||||
throw invalidSession(`Entry ${currentId} not found`);
|
||||
}
|
||||
rows.push(row);
|
||||
if (row.id === query.stopAtId || row.type === query.stopAtType) break;
|
||||
currentId = row.parent_id;
|
||||
}
|
||||
return this.repairBranchCache(leafId, cached?.branchId);
|
||||
const entries = decodeEntryRows(rows).filter(
|
||||
(entry) =>
|
||||
(query.type === undefined || entry.type === query.type) &&
|
||||
(query.customType === undefined || (entry.type === "custom" && entry.customType === query.customType)),
|
||||
);
|
||||
const limited = query.limit === undefined ? entries : entries.slice(0, query.limit);
|
||||
for (const entry of limited) this.byId.set(entry.id, entry);
|
||||
return limited;
|
||||
}
|
||||
|
||||
private async repairBranchCacheForQuery(leafId: string, branchIdToReplace?: string): Promise<CachedBranch> {
|
||||
const visited = new Set<string>();
|
||||
let currentId: string | null = leafId;
|
||||
while (currentId !== null) {
|
||||
if (visited.has(currentId)) throw invalidSession(`cycle in parent chain at entry ${currentId}`);
|
||||
visited.add(currentId);
|
||||
const row: { parent_id: string | null } | undefined = await this.db
|
||||
.prepare("SELECT parent_id FROM session_entries WHERE session_id = ? AND id = ?")
|
||||
.get<{ parent_id: string | null }>(this.metadata.id, currentId);
|
||||
if (!row) {
|
||||
if (currentId === leafId) throw new SessionError("not_found", `Entry ${leafId} not found`);
|
||||
throw invalidSession(`Entry ${currentId} not found`);
|
||||
}
|
||||
currentId = row.parent_id;
|
||||
}
|
||||
try {
|
||||
await rebuildCachedBranch(this.db, this.metadata.id, leafId, branchIdToReplace);
|
||||
} catch (error) {
|
||||
if (error instanceof SessionError) throw error;
|
||||
throw new SessionError("storage", `Failed to rebuild SQLite branch cache at entry ${leafId}`, toError(error));
|
||||
}
|
||||
const cached = await readCachedBranch(this.db, this.metadata.id, leafId);
|
||||
if (!cached || !(await isCachedBranchValid(this.db, this.metadata.id, cached, leafId))) {
|
||||
throw invalidSession(`branch cache repair did not produce a valid path to entry ${leafId}`);
|
||||
}
|
||||
return cached;
|
||||
}
|
||||
|
||||
async readPathToRootOrCompaction(leafId: string | null): Promise<SessionTreeEntry[]> {
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue