diff --git a/packages/session-backends/sqlite-node/benchmark/session/session-repo-targets.ts b/packages/session-backends/sqlite-node/benchmark/session/session-repo-targets.ts index 23c0664de..16127d8a8 100644 --- a/packages/session-backends/sqlite-node/benchmark/session/session-repo-targets.ts +++ b/packages/session-backends/sqlite-node/benchmark/session/session-repo-targets.ts @@ -1,8 +1,7 @@ import { mkdtemp, rm } from "node:fs/promises"; import { tmpdir } from "node:os"; import { join } from "node:path"; -import { BACKGROUND_CONTEXT } from "@earendil-works/pi-agent-core"; -import type { Session, SessionCreateOptions, SessionMetadata, SessionRepo } from "@earendil-works/pi-agent-core"; +import { BACKGROUND_CONTEXT, type SessionRepo } from "@earendil-works/pi-agent-core"; import { createNodeSqliteFactory, SqliteSessionRepo } from "../../src/index.ts"; import type { BenchmarkTarget } from "../../../../agent/benchmark/session/benchmark.ts"; @@ -12,57 +11,20 @@ export interface SessionRepoBenchmarkFixture extends AsyncDisposable { const NOW = 1_700_000_000_000; -class TrackingSessionRepo implements SessionRepo { - private readonly sessions = new Set>(); - - constructor(private readonly inner: SessionRepo) {} - - async create(options: SessionCreateOptions | undefined, context = BACKGROUND_CONTEXT): Promise> { - return this.track(await this.inner.create(options, context)); - } - - async open(metadata: TMetadata, context = BACKGROUND_CONTEXT): Promise> { - return this.track(await this.inner.open(metadata, context)); - } - - list(options: undefined, context = BACKGROUND_CONTEXT): Promise { - return this.inner.list(options, context); - } - - delete(metadata: TMetadata, context = BACKGROUND_CONTEXT): Promise { - return this.inner.delete(metadata, context); - } - - async fork(source: TMetadata, options = {}, context = BACKGROUND_CONTEXT): Promise> { - return this.track(await this.inner.fork(source, options, context)); - } - - async closeTrackedSessions(): Promise { - await Promise.allSettled([...this.sessions].map((session) => session.close(BACKGROUND_CONTEXT))); - this.sessions.clear(); - } - - private track(session: Session): Session { - this.sessions.add(session); - return session; - } -} - export const sessionRepoBenchmarkTargets = [ { name: "sqlite", async createFixture() { const directory = await mkdtemp(join(tmpdir(), "pi-sqlite-session-repo-benchmark-")); - const sqliteRepo = new SqliteSessionRepo({ + const repo = new SqliteSessionRepo({ directory, databaseFactory: createNodeSqliteFactory(), now: () => NOW, }); - const repo = new TrackingSessionRepo(sqliteRepo); return { repo, async [Symbol.asyncDispose]() { - await repo.closeTrackedSessions(); + await repo.close(BACKGROUND_CONTEXT); await rm(directory, { recursive: true, force: true }); }, }; diff --git a/packages/session-backends/sqlite-node/src/sqlite/repo.ts b/packages/session-backends/sqlite-node/src/sqlite/repo.ts index 81ab4c699..311ea039c 100644 --- a/packages/session-backends/sqlite-node/src/sqlite/repo.ts +++ b/packages/session-backends/sqlite-node/src/sqlite/repo.ts @@ -143,6 +143,9 @@ export class SqliteSessionRepo { private readonly now: () => number; private readonly pendingIds = new Set(); private readonly openStorages = new Map(); + private readonly openSessions = new Set(); + private closed = false; + private closePromise: Promise | undefined; constructor(options: SqliteSessionRepoOptions) { this.directory = options.directory; @@ -151,6 +154,7 @@ export class SqliteSessionRepo { } async create(options: SqliteSessionCreateOptions | undefined, _context: Context): Promise { + this.assertOpen(); options ??= {}; const createdAt = this.now(); const id = options.id ?? uuidv7(createdAt); @@ -203,6 +207,7 @@ export class SqliteSessionRepo { } async open(metadata: SqliteSessionMetadata, _context: Context): Promise { + this.assertOpen(); this.reserveId(metadata.id); let db: SqliteDatabase | undefined; let lease: WriterLeaseRow | undefined; @@ -243,6 +248,7 @@ export class SqliteSessionRepo { } async list(_options: undefined, _context: Context): Promise { + this.assertOpen(); await mkdir(this.directory, { recursive: true }); const names = await readdir(this.directory); const sessions: SqliteSessionMetadata[] = []; @@ -267,6 +273,7 @@ export class SqliteSessionRepo { } async delete(metadata: SqliteSessionMetadata, _context: Context): Promise { + this.assertOpen(); if (this.pendingIds.has(metadata.id)) throw new Error(`Session is open: ${metadata.id}`); await access(metadata.path); const db = await this.databaseFactory.open(metadata.path); @@ -283,6 +290,7 @@ export class SqliteSessionRepo { } async fork(source: SqliteSessionMetadata, options: ForkOptions, context: Context): Promise { + this.assertOpen(); const createdAt = this.now(); const id = options.id ?? uuidv7(createdAt); this.reserveId(id); @@ -369,6 +377,15 @@ export class SqliteSessionRepo { } } + close(context: Context): Promise { + if (this.closePromise !== undefined) return this.closePromise; + this.closed = true; + this.closePromise = Promise.all([...this.openSessions].map((session) => session.close(context))).then( + () => undefined, + ); + return this.closePromise; + } + private openStorageBackedSession( metadata: SqliteSessionMetadata, db: SqliteDatabase, @@ -383,20 +400,27 @@ export class SqliteSessionRepo { const storage = new SqliteStorage(db, { now: this.now, beforeCommit: renew }); this.openStorages.set(metadata.id, storage); const session = new StorageBackedSession(metadata, storage); - return new SqliteOpenSession(session, { + const openSession = new SqliteOpenSession(session, { renewWriterLease: renew, releaseWriterLease: () => db.transaction(() => releaseWriterLease(db, lease.owner_id, lease.fence)), renewIntervalMs: WRITER_LEASE_RENEW_INTERVAL_MS, onClose: () => { db.close(); this.openStorages.delete(metadata.id); + this.openSessions.delete(openSession); this.pendingIds.delete(metadata.id); }, }); + this.openSessions.add(openSession); + return openSession; } private reserveId(id: string): void { if (this.pendingIds.has(id)) throw new Error(`Session is already open: ${id}`); this.pendingIds.add(id); } + + private assertOpen(): void { + if (this.closed) throw new Error("SqliteSessionRepo is closed"); + } } diff --git a/packages/session-backends/sqlite-node/test/repo.test.ts b/packages/session-backends/sqlite-node/test/repo.test.ts index 4e6c4edc8..190e01442 100644 --- a/packages/session-backends/sqlite-node/test/repo.test.ts +++ b/packages/session-backends/sqlite-node/test/repo.test.ts @@ -212,6 +212,24 @@ describe("SqliteSessionRepo", () => { }); }); + it("closes open sessions through repo close and rejects later operations", async () => { + await withTempDir(async (directory) => { + const repo = new SqliteSessionRepo({ + directory, + databaseFactory: createNodeSqliteFactory(), + now: () => 1, + }); + const session = await repo.create({ id: "session" }, BACKGROUND_CONTEXT); + + await repo.close(BACKGROUND_CONTEXT); + + await expect(session.getStats(BACKGROUND_CONTEXT)).rejects.toThrow("Session is closed"); + await expect(repo.list(undefined, BACKGROUND_CONTEXT)).rejects.toThrow("SqliteSessionRepo is closed"); + await expect(repo.create({ id: "other" }, BACKGROUND_CONTEXT)).rejects.toThrow("SqliteSessionRepo is closed"); + await expect(repo.close(BACKGROUND_CONTEXT)).resolves.toBeUndefined(); + }); + }); + it("opens a session through the version gate and rejects a live external writer lease", async () => { await withTempDir(async (directory) => { const repo = new SqliteSessionRepo({