From 9b041958027301426152a217cd7cf10c175976a0 Mon Sep 17 00:00:00 2001 From: likun Date: Sat, 29 Aug 2026 18:18:43 +0800 Subject: [PATCH 1/2] refactor(storage): add bounded context offload lifecycle Generated-by: OpenAI Codex --- packages/core/src/context-offload.ts | 41 +++ .../sqlite-context-offload-store.test.ts | 220 +++++++++++- .../src/sqlite-context-offload-schema.ts | 37 +- .../src/sqlite-context-offload-store.ts | 336 +++++++++++++++++- 4 files changed, 624 insertions(+), 10 deletions(-) diff --git a/packages/core/src/context-offload.ts b/packages/core/src/context-offload.ts index e08abad618..165a1db2c3 100644 --- a/packages/core/src/context-offload.ts +++ b/packages/core/src/context-offload.ts @@ -82,6 +82,30 @@ export interface ContextOffloadUsage { readonly physicalBytes: number; } +export type ContextOffloadCopyResult = + | { + readonly ok: true; + readonly copied: readonly { + readonly sourceRefId: string; + readonly targetRefId: string; + }[]; + } + | { + readonly ok: false; + readonly reason: 'not_found' | 'session_quota_exceeded' | 'identity_conflict' | 'unavailable'; + }; + +export interface ContextOffloadRetirementResult { + readonly releasedReferences: number; + readonly releasedLogicalBytes: number; +} + +export interface ContextOffloadGarbageCollectionResult { + readonly deletedBlobs: number; + readonly deletedBytes: number; + readonly hasMore: boolean; +} + /** * Storage contract for capped, whole-object Agent context offload. * @@ -103,8 +127,25 @@ export interface ContextOffloadStore { readonly maxBytes: number; }): Promise; + copyReferences(input: { + readonly sourceSessionId: string; + readonly targetSessionId: string; + readonly references: readonly { + readonly sourceRefId: string; + readonly targetOwner: ContextOffloadOwner; + }[]; + }): Promise; + releaseReference(input: { readonly sessionId: string; readonly refId: string }): Promise; + retireSession(sessionId: string): Promise; + + collectGarbage(input: { + readonly olderThan: number; + readonly maxBlobs: number; + readonly maxBytes: number; + }): Promise; + /** * Session-scoped reference/logical usage when sessionId is supplied. Physical * bytes always describe the workspace because shared bytes have no one owner. diff --git a/packages/storage/src/__tests__/sqlite-context-offload-store.test.ts b/packages/storage/src/__tests__/sqlite-context-offload-store.test.ts index 903954c23e..d5ffd3704a 100644 --- a/packages/storage/src/__tests__/sqlite-context-offload-store.test.ts +++ b/packages/storage/src/__tests__/sqlite-context-offload-store.test.ts @@ -36,7 +36,7 @@ test('creates the dedicated WAL schema with incremental auto-vacuum', async (t) const database = new DatabaseSync(fixture.path); t.after(() => database.close()); - assert.equal(pragmaNumber(database, 'user_version'), 1); + assert.equal(pragmaNumber(database, 'user_version'), 2); assert.equal(pragmaNumber(database, 'auto_vacuum'), 2); assert.equal(pragmaText(database, 'journal_mode'), 'wal'); assert.deepEqual( @@ -48,7 +48,13 @@ test('creates the dedicated WAL schema with incremental auto-vacuum', async (t) ) .all() .map((row) => row.name), - ['context_blobs', 'context_refs', 'context_session_usage', 'context_store_usage'], + [ + 'context_blobs', + 'context_gc_candidates', + 'context_refs', + 'context_session_usage', + 'context_store_usage', + ], ); }); @@ -135,7 +141,7 @@ test('rejects a database schema newer than this authority understands', async (t const path = join(root, CONTEXT_OFFLOAD_DATABASE_NAME); t.after(() => rm(root, { recursive: true, force: true })); const database = new DatabaseSync(path); - database.exec('PRAGMA user_version = 2'); + database.exec('PRAGMA user_version = 3'); database.close(); assert.throws( @@ -147,7 +153,7 @@ test('rejects a database schema newer than this authority understands', async (t workspacePhysicalBytes: 1, }, }), - /schema 2 is newer than supported version 1/u, + /schema 3 is newer than supported version 2/u, ); }); @@ -358,6 +364,212 @@ test('releases only the authorized Session reference without deleting shared byt }); }); +test('copies references atomically without copying physical bytes', async (t) => { + const fixture = await createFixture(t, { + ownerMaxBytes: TEST_OWNER_MAX_BYTES, + sessionLogicalBytes: 12, + workspacePhysicalBytes: 16, + }); + const first = await fixture.store.put( + putInput('source', 'source-1', new TextEncoder().encode('first')), + ); + const second = await fixture.store.put( + putInput('source', 'source-2', new TextEncoder().encode('second')), + ); + assert.equal(first.ok, true); + assert.equal(second.ok, true); + if (!first.ok || !second.ok) return; + + const copyInput = { + sourceSessionId: 'source', + targetSessionId: 'target', + references: [ + { + sourceRefId: first.record.refId, + targetOwner: { kind: 'tool_result_archive' as const, ownerId: 'target-1' }, + }, + { + sourceRefId: second.record.refId, + targetOwner: { kind: 'tool_result_archive' as const, ownerId: 'target-2' }, + }, + ], + }; + const copied = await fixture.store.copyReferences(copyInput); + assert.equal(copied.ok, true); + if (!copied.ok) return; + assert.deepEqual(await fixture.store.copyReferences(copyInput), copied); + assert.deepEqual(await fixture.store.usage('target'), { + references: 2, + logicalBytes: 11, + physicalBytes: 11, + }); + assert.equal( + ( + await fixture.store.read({ + sessionId: 'target', + refId: copied.copied[0]?.targetRefId ?? '', + maxBytes: 16, + }) + ).ok, + true, + ); + + assert.deepEqual( + await fixture.store.copyReferences({ + sourceSessionId: 'source', + targetSessionId: 'target', + references: [ + { + sourceRefId: first.record.refId, + targetOwner: { kind: 'tool_result_archive', ownerId: 'target-over-quota' }, + }, + ], + }), + { ok: false, reason: 'session_quota_exceeded' }, + ); + + assert.deepEqual( + await fixture.store.copyReferences({ + sourceSessionId: 'source', + targetSessionId: 'target', + references: [ + { + sourceRefId: second.record.refId, + targetOwner: { kind: 'tool_result_archive', ownerId: 'new-before-conflict' }, + }, + { + sourceRefId: second.record.refId, + targetOwner: { kind: 'tool_result_archive', ownerId: 'target-1' }, + }, + ], + }), + { ok: false, reason: 'identity_conflict' }, + ); + assert.equal((await fixture.store.usage('target')).references, 2); +}); + +test('retires only one Session and collects shared blobs after the last reference', async (t) => { + const fixture = await createFixture(t); + const bytes = new TextEncoder().encode('shared'); + await fixture.store.put(putInput('session-1', 'owner-1', bytes)); + await fixture.store.put(putInput('session-2', 'owner-2', bytes)); + + assert.deepEqual(await fixture.store.retireSession('session-1'), { + releasedReferences: 1, + releasedLogicalBytes: bytes.byteLength, + }); + assert.deepEqual( + await fixture.store.collectGarbage({ olderThan: 1_001, maxBlobs: 1, maxBytes: 16 }), + { deletedBlobs: 0, deletedBytes: 0, hasMore: false }, + ); + assert.equal((await fixture.store.usage('session-2')).references, 1); + + assert.deepEqual(await fixture.store.retireSession('session-2'), { + releasedReferences: 1, + releasedLogicalBytes: bytes.byteLength, + }); + assert.deepEqual( + await fixture.store.collectGarbage({ olderThan: 1_000, maxBlobs: 1, maxBytes: 16 }), + { deletedBlobs: 0, deletedBytes: 0, hasMore: false }, + ); + assert.deepEqual( + await fixture.store.collectGarbage({ olderThan: 1_001, maxBlobs: 1, maxBytes: 16 }), + { deletedBlobs: 1, deletedBytes: bytes.byteLength, hasMore: false }, + ); + assert.deepEqual(await fixture.store.usage(), { + references: 0, + logicalBytes: 0, + physicalBytes: 0, + }); +}); + +test('garbage collection obeys both batch limits and rolls back failed deletion', async (t) => { + let failGc = false; + const fixture = await createFixture(t, undefined, (point) => { + if (point === 'after_gc_blob_delete' && failGc) throw new Error('injected GC failure'); + }); + const bytes = new TextEncoder().encode('four'); + const first = await fixture.store.put(putInput('session-1', 'owner-1', bytes)); + const second = await fixture.store.put( + putInput('session-1', 'owner-2', new TextEncoder().encode('five')), + ); + assert.equal(first.ok, true); + assert.equal(second.ok, true); + if (!first.ok || !second.ok) return; + await fixture.store.releaseReference({ sessionId: 'session-1', refId: first.record.refId }); + await fixture.store.releaseReference({ sessionId: 'session-1', refId: second.record.refId }); + + assert.deepEqual( + await fixture.store.collectGarbage({ olderThan: 1_001, maxBlobs: 2, maxBytes: 4 }), + { deletedBlobs: 1, deletedBytes: 4, hasMore: true }, + ); + failGc = true; + await assert.rejects( + fixture.store.collectGarbage({ olderThan: 1_001, maxBlobs: 1, maxBytes: 8 }), + /injected GC failure/u, + ); + assert.equal((await fixture.store.usage()).physicalBytes, 4); + failGc = false; + assert.deepEqual( + await fixture.store.collectGarbage({ olderThan: 1_001, maxBlobs: 1, maxBytes: 8 }), + { deletedBlobs: 1, deletedBytes: 4, hasMore: false }, + ); +}); + +test('migrates v1 orphan blobs into the indexed garbage candidate set', async (t) => { + const fixture = await createFixture(t); + const stored = await fixture.store.put( + putInput('session-1', 'owner-1', new TextEncoder().encode('orphan')), + ); + assert.equal(stored.ok, true); + if (!stored.ok) return; + await fixture.store.releaseReference({ sessionId: 'session-1', refId: stored.record.refId }); + fixture.store.close(); + + const database = new DatabaseSync(fixture.path); + database.exec('DROP TABLE context_gc_candidates; PRAGMA user_version = 1'); + database.close(); + const migrated = new SqliteContextOffloadStore(fixture.path, { limits: fixture.limits }); + t.after(() => migrated.close()); + assert.deepEqual(await migrated.collectGarbage({ olderThan: 1_001, maxBlobs: 1, maxBytes: 16 }), { + deletedBlobs: 1, + deletedBytes: 6, + hasMore: false, + }); +}); + +test('lifecycle queries use Session and garbage eligibility indexes', async (t) => { + const fixture = await createFixture(t); + fixture.store.close(); + const database = new DatabaseSync(fixture.path); + t.after(() => database.close()); + + const retirementPlan = database + .prepare( + `EXPLAIN QUERY PLAN + SELECT r.blob_id, b.size_bytes + FROM context_refs r INDEXED BY context_refs_session + JOIN context_blobs b ON b.blob_id = r.blob_id + WHERE r.session_id = ?`, + ) + .all('session-1'); + assert.match(JSON.stringify(retirementPlan), /context_refs_session/u); + + const garbagePlan = database + .prepare( + `EXPLAIN QUERY PLAN + SELECT c.blob_id, b.size_bytes + FROM context_gc_candidates c INDEXED BY context_gc_candidates_eligible + JOIN context_blobs b ON b.blob_id = c.blob_id + WHERE c.unreferenced_at < ? + ORDER BY c.unreferenced_at, c.blob_id + LIMIT ?`, + ) + .all(1_001, 2); + assert.match(JSON.stringify(garbagePlan), /context_gc_candidates_eligible/u); + assert.doesNotMatch(JSON.stringify(garbagePlan), /SCAN b(?:\W|$)/u); +}); + function putInput(sessionId: string, ownerId: string, bytes: Uint8Array) { return { sessionId, diff --git a/packages/storage/src/sqlite-context-offload-schema.ts b/packages/storage/src/sqlite-context-offload-schema.ts index e937111622..5f3495474f 100644 --- a/packages/storage/src/sqlite-context-offload-schema.ts +++ b/packages/storage/src/sqlite-context-offload-schema.ts @@ -19,7 +19,7 @@ import type { DatabaseSync } from 'node:sqlite'; -export const SQLITE_CONTEXT_OFFLOAD_SCHEMA_VERSION = 1; +export const SQLITE_CONTEXT_OFFLOAD_SCHEMA_VERSION = 2; const SQLITE_INITIALIZATION_BUSY_TIMEOUT_MS = 5_000; const SQLITE_INITIALIZATION_RETRY_DELAY_MS = 10; const initializationRetryGate = new Int32Array(new SharedArrayBuffer(4)); @@ -51,6 +51,15 @@ const INITIAL_SCHEMA = ` CREATE INDEX context_refs_blob ON context_refs(blob_id); + CREATE TABLE context_gc_candidates ( + blob_id BLOB PRIMARY KEY + REFERENCES context_blobs(blob_id) ON DELETE CASCADE, + unreferenced_at INTEGER NOT NULL CHECK(unreferenced_at >= 0) + ); + + CREATE INDEX context_gc_candidates_eligible + ON context_gc_candidates(unreferenced_at, blob_id); + CREATE TABLE context_session_usage ( session_id TEXT PRIMARY KEY, reference_count INTEGER NOT NULL CHECK(reference_count >= 0), @@ -67,13 +76,33 @@ const INITIAL_SCHEMA = ` VALUES (1, 0, 0); `; +const SCHEMA_V2_MIGRATION = ` + CREATE TABLE context_gc_candidates ( + blob_id BLOB PRIMARY KEY + REFERENCES context_blobs(blob_id) ON DELETE CASCADE, + unreferenced_at INTEGER NOT NULL CHECK(unreferenced_at >= 0) + ); + + CREATE INDEX context_gc_candidates_eligible + ON context_gc_candidates(unreferenced_at, blob_id); + + INSERT INTO context_gc_candidates(blob_id, unreferenced_at) + SELECT b.blob_id, b.created_at + FROM context_blobs b + WHERE NOT EXISTS ( + SELECT 1 FROM context_refs r WHERE r.blob_id = b.blob_id + ); +`; + const REQUIRED_SCHEMA_OBJECTS = Object.freeze([ ['table', 'context_blobs'], ['table', 'context_refs'], ['table', 'context_session_usage'], ['table', 'context_store_usage'], + ['table', 'context_gc_candidates'], ['index', 'context_refs_session'], ['index', 'context_refs_blob'], + ['index', 'context_gc_candidates_eligible'], ] as const); const REQUIRED_TABLE_COLUMNS = Object.freeze({ @@ -89,11 +118,13 @@ const REQUIRED_TABLE_COLUMNS = Object.freeze({ ], context_session_usage: ['session_id', 'reference_count', 'logical_bytes'], context_store_usage: ['singleton', 'blob_count', 'physical_bytes'], + context_gc_candidates: ['blob_id', 'unreferenced_at'], } as const); const REQUIRED_INDEX_COLUMNS = Object.freeze({ context_refs_session: ['session_id', 'created_at', 'ref_id'], context_refs_blob: ['blob_id'], + context_gc_candidates_eligible: ['unreferenced_at', 'blob_id'], } as const); export function configureSqliteContextOffloadDatabase(db: DatabaseSync): void { @@ -130,6 +161,10 @@ export function migrateSqliteContextOffloadDatabase(db: DatabaseSync): void { db.exec(INITIAL_SCHEMA); db.exec(`PRAGMA user_version = ${SQLITE_CONTEXT_OFFLOAD_SCHEMA_VERSION}`); } + if (current === 1) { + db.exec(SCHEMA_V2_MIGRATION); + db.exec(`PRAGMA user_version = ${SQLITE_CONTEXT_OFFLOAD_SCHEMA_VERSION}`); + } validateSchema(db); db.exec('COMMIT'); } catch (error) { diff --git a/packages/storage/src/sqlite-context-offload-store.ts b/packages/storage/src/sqlite-context-offload-store.ts index d6db8d0be4..1b94763c90 100644 --- a/packages/storage/src/sqlite-context-offload-store.ts +++ b/packages/storage/src/sqlite-context-offload-store.ts @@ -23,11 +23,14 @@ import { createRequire } from 'node:module'; import { dirname } from 'node:path'; import type { DatabaseSync } from 'node:sqlite'; import { + type ContextOffloadCopyResult, + type ContextOffloadGarbageCollectionResult, type ContextOffloadLimits, type ContextOffloadOwner, type ContextOffloadPutResult, type ContextOffloadReadResult, type ContextOffloadRecord, + type ContextOffloadRetirementResult, type ContextOffloadStore, type ContextOffloadUsage, } from '@maka/core/context-offload'; @@ -43,7 +46,10 @@ const require = createRequire(import.meta.url); export const CONTEXT_OFFLOAD_DATABASE_NAME = 'context-offload.sqlite'; -export type SqliteContextOffloadStoreFailpoint = 'after_blob_insert' | 'after_ref_insert'; +export type SqliteContextOffloadStoreFailpoint = + | 'after_blob_insert' + | 'after_ref_insert' + | 'after_gc_blob_delete'; export interface SqliteContextOffloadStoreOptions { readonly limits: ContextOffloadLimits; @@ -79,6 +85,11 @@ interface StoreUsageRow { physical_bytes: unknown; } +interface GarbageCandidateRow { + blob_id: unknown; + size_bytes: unknown; +} + export class SqliteContextOffloadStore implements ContextOffloadStore { readonly #database: DatabaseSync; readonly #limits: ContextOffloadLimits; @@ -171,16 +182,20 @@ export class SqliteContextOffloadStore implements ContextOffloadStore { this.#writeTransaction(() => { const row = this.#database .prepare( - `SELECT r.session_id, b.size_bytes + `SELECT r.session_id, r.blob_id, b.size_bytes FROM context_refs r JOIN context_blobs b ON b.blob_id = r.blob_id WHERE r.ref_id = ?`, ) - .get(input.refId) as { session_id?: unknown; size_bytes?: unknown } | undefined; + .get(input.refId) as + | { session_id?: unknown; blob_id?: unknown; size_bytes?: unknown } + | undefined; if (!row || row.session_id !== input.sessionId) return; if (!isNonNegativeSafeInteger(row.size_bytes)) { throw new Error('Invalid context reference size'); } + const blobId = decodeBlobId(row.blob_id); + if (!blobId) throw new Error('Invalid context reference blob identity'); const deleted = this.#database .prepare('DELETE FROM context_refs WHERE session_id = ? AND ref_id = ?') .run(input.sessionId, input.refId); @@ -199,6 +214,159 @@ export class SqliteContextOffloadStore implements ContextOffloadStore { WHERE session_id = ? AND reference_count = 0 AND logical_bytes = 0`, ) .run(input.sessionId); + this.#markBlobUnreferencedIfEligible(blobId, this.#readNow()); + }); + } + + async copyReferences(input: { + readonly sourceSessionId: string; + readonly targetSessionId: string; + readonly references: readonly { + readonly sourceRefId: string; + readonly targetOwner: ContextOffloadOwner; + }[]; + }): Promise { + assertBoundedIdentity(input.sourceSessionId, 'Source Session id'); + assertBoundedIdentity(input.targetSessionId, 'Target Session id'); + const references = input.references.map((reference) => { + assertBoundedIdentity(reference.sourceRefId, 'Source context reference id'); + assertOwner(reference.targetOwner); + return Object.freeze({ + sourceRefId: reference.sourceRefId, + targetOwner: Object.freeze({ ...reference.targetOwner }), + }); + }); + try { + this.#assertOpen(); + return this.#writeTransaction(() => + this.#copyReferences({ + sourceSessionId: input.sourceSessionId, + targetSessionId: input.targetSessionId, + references, + }), + ); + } catch (error) { + this.#onUnavailable?.(error); + return { ok: false, reason: 'unavailable' }; + } + } + + async retireSession(sessionId: string): Promise { + assertBoundedIdentity(sessionId, 'Session id'); + this.#assertOpen(); + return this.#writeTransaction(() => { + const rows = this.#database + .prepare( + `SELECT r.blob_id, b.size_bytes + FROM context_refs r INDEXED BY context_refs_session + JOIN context_blobs b ON b.blob_id = r.blob_id + WHERE r.session_id = ?`, + ) + .all(sessionId) as unknown as Array<{ blob_id?: unknown; size_bytes?: unknown }>; + let releasedLogicalBytes = 0; + const blobIds = new Map(); + for (const row of rows) { + const blobId = decodeBlobId(row.blob_id); + if (!blobId || !isNonNegativeSafeInteger(row.size_bytes)) { + throw new Error('Invalid retiring context reference'); + } + releasedLogicalBytes = addSafeInteger( + releasedLogicalBytes, + row.size_bytes, + 'Retired context logical bytes', + ); + blobIds.set(Buffer.from(blobId).toString('hex'), blobId); + } + const usage = this.#readSessionUsage(sessionId); + if ( + readNonNegativeInteger(usage.reference_count, 'Session reference count') !== rows.length || + readNonNegativeInteger(usage.logical_bytes, 'Session logical bytes') !== + releasedLogicalBytes + ) { + throw new Error('Context Session usage is inconsistent with retiring references'); + } + const deleted = this.#database + .prepare('DELETE FROM context_refs WHERE session_id = ?') + .run(sessionId); + if (deleted.changes !== rows.length) { + throw new Error('Context Session retirement deleted an unexpected reference count'); + } + this.#database + .prepare('DELETE FROM context_session_usage WHERE session_id = ?') + .run(sessionId); + const unreferencedAt = this.#readNow(); + for (const blobId of blobIds.values()) { + this.#markBlobUnreferencedIfEligible(blobId, unreferencedAt); + } + return { + releasedReferences: rows.length, + releasedLogicalBytes, + }; + }); + } + + async collectGarbage(input: { + readonly olderThan: number; + readonly maxBlobs: number; + readonly maxBytes: number; + }): Promise { + assertNonNegativeSafeInteger(input.olderThan, 'Context garbage watermark'); + assertPositiveSafeInteger(input.maxBlobs, 'Context garbage blob limit'); + assertPositiveSafeInteger(input.maxBytes, 'Context garbage byte limit'); + if (input.maxBlobs === Number.MAX_SAFE_INTEGER) { + throw new Error('Context garbage blob limit is too large'); + } + this.#assertOpen(); + return this.#writeTransaction(() => { + const rows = this.#database + .prepare( + `SELECT c.blob_id, b.size_bytes + FROM context_gc_candidates c INDEXED BY context_gc_candidates_eligible + JOIN context_blobs b ON b.blob_id = c.blob_id + WHERE c.unreferenced_at < ? + ORDER BY c.unreferenced_at, c.blob_id + LIMIT ?`, + ) + .all(input.olderThan, input.maxBlobs + 1) as unknown as GarbageCandidateRow[]; + const selected: Uint8Array[] = []; + let deletedBytes = 0; + for (const row of rows) { + if (selected.length === input.maxBlobs) break; + const blobId = decodeBlobId(row.blob_id); + if (!blobId || !isNonNegativeSafeInteger(row.size_bytes)) { + throw new Error('Invalid context garbage candidate'); + } + if (exceedsLimit(deletedBytes, row.size_bytes, input.maxBytes)) break; + deletedBytes = addSafeInteger(deletedBytes, row.size_bytes, 'Collected context bytes'); + selected.push(blobId); + } + const deleteBlob = this.#database.prepare( + `DELETE FROM context_blobs + WHERE blob_id = ? + AND NOT EXISTS (SELECT 1 FROM context_refs WHERE blob_id = ?)`, + ); + for (const blobId of selected) { + const deleted = deleteBlob.run(blobId, blobId); + if (deleted.changes !== 1) { + throw new Error('Context garbage candidate is still referenced or missing'); + } + this.#failpoint?.('after_gc_blob_delete'); + } + if (selected.length > 0) { + const updated = this.#database + .prepare( + `UPDATE context_store_usage + SET blob_count = blob_count - ?, physical_bytes = physical_bytes - ? + WHERE singleton = 1`, + ) + .run(selected.length, deletedBytes); + if (updated.changes !== 1) throw new Error('Missing context store usage row'); + } + return { + deletedBlobs: selected.length, + deletedBytes, + hasMore: rows.length > selected.length, + }; }); } @@ -276,8 +444,7 @@ export class SqliteContextOffloadStore implements ContextOffloadStore { } } - const createdAt = this.#now(); - assertNonNegativeSafeInteger(createdAt, 'Context creation time'); + const createdAt = this.#readNow(); const refId = this.#idFactory(); assertBoundedIdentity(refId, 'Context reference id'); if (!existingBlob) { @@ -321,6 +488,7 @@ export class SqliteContextOffloadStore implements ContextOffloadStore { logical_bytes = logical_bytes + excluded.logical_bytes`, ) .run(input.sessionId, input.bytes.byteLength); + this.#database.prepare('DELETE FROM context_gc_candidates WHERE blob_id = ?').run(blobIdBytes); this.#failpoint?.('after_ref_insert'); return { ok: true, @@ -336,6 +504,130 @@ export class SqliteContextOffloadStore implements ContextOffloadStore { }; } + #copyReferences(input: { + readonly sourceSessionId: string; + readonly targetSessionId: string; + readonly references: readonly { + readonly sourceRefId: string; + readonly targetOwner: ContextOffloadOwner; + }[]; + }): ContextOffloadCopyResult { + const createdAt = this.#readNow(); + const pendingByOwner = new Map< + string, + { + readonly refId: string; + readonly owner: ContextOffloadOwner; + readonly blobId: string; + readonly sizeBytes: number; + readonly mediaType: string; + } + >(); + const copied: Array<{ sourceRefId: string; targetRefId: string }> = []; + let addedLogicalBytes = 0; + + for (const reference of input.references) { + const sourceRow = this.#database + .prepare( + `SELECT r.ref_id, r.session_id, r.owner_kind, r.owner_id, r.blob_id, + b.size_bytes, r.media_type, r.created_at + FROM context_refs r + JOIN context_blobs b ON b.blob_id = r.blob_id + WHERE r.session_id = ? AND r.ref_id = ?`, + ) + .get(input.sourceSessionId, reference.sourceRefId) as ContextReferenceRow | undefined; + if (!sourceRow) return { ok: false, reason: 'not_found' }; + const source = decodeReferenceRow(sourceRow); + if (!source) throw new Error('Invalid source context reference'); + + const ownerKey = `${reference.targetOwner.kind}\0${reference.targetOwner.ownerId}`; + const pending = pendingByOwner.get(ownerKey); + if (pending) { + if (pending.blobId !== source.blobId) return { ok: false, reason: 'identity_conflict' }; + copied.push({ sourceRefId: reference.sourceRefId, targetRefId: pending.refId }); + continue; + } + + const existing = this.#readReferenceByOwner(input.targetSessionId, reference.targetOwner); + if (existing) { + if (existing.blobId !== source.blobId) { + return { ok: false, reason: 'identity_conflict' }; + } + pendingByOwner.set(ownerKey, { + refId: existing.refId, + owner: reference.targetOwner, + blobId: existing.blobId, + sizeBytes: existing.sizeBytes, + mediaType: existing.mediaType, + }); + copied.push({ sourceRefId: reference.sourceRefId, targetRefId: existing.refId }); + continue; + } + + const refId = this.#idFactory(); + assertBoundedIdentity(refId, 'Context reference id'); + pendingByOwner.set(ownerKey, { + refId, + owner: reference.targetOwner, + blobId: source.blobId, + sizeBytes: source.sizeBytes, + mediaType: source.mediaType, + }); + addedLogicalBytes = addSafeInteger( + addedLogicalBytes, + source.sizeBytes, + 'Copied context logical bytes', + ); + copied.push({ sourceRefId: reference.sourceRefId, targetRefId: refId }); + } + + const targetUsage = this.#readSessionUsage(input.targetSessionId); + const currentLogicalBytes = readNonNegativeInteger( + targetUsage.logical_bytes, + 'Target Session logical bytes', + ); + if (exceedsLimit(currentLogicalBytes, addedLogicalBytes, this.#limits.sessionLogicalBytes)) { + return { ok: false, reason: 'session_quota_exceeded' }; + } + + const newReferences = [...pendingByOwner.values()].filter( + (reference) => !this.#readReferenceByOwner(input.targetSessionId, reference.owner), + ); + const insertReference = this.#database.prepare( + `INSERT INTO context_refs( + ref_id, session_id, owner_kind, owner_id, blob_id, media_type, created_at + ) VALUES (?, ?, ?, ?, ?, ?, ?)`, + ); + const clearCandidate = this.#database.prepare( + 'DELETE FROM context_gc_candidates WHERE blob_id = ?', + ); + for (const reference of newReferences) { + const blobId = Buffer.from(reference.blobId, 'hex'); + insertReference.run( + reference.refId, + input.targetSessionId, + reference.owner.kind, + reference.owner.ownerId, + blobId, + reference.mediaType, + createdAt, + ); + clearCandidate.run(blobId); + } + if (newReferences.length > 0) { + this.#database + .prepare( + `INSERT INTO context_session_usage(session_id, reference_count, logical_bytes) + VALUES (?, ?, ?) + ON CONFLICT(session_id) DO UPDATE SET + reference_count = reference_count + excluded.reference_count, + logical_bytes = logical_bytes + excluded.logical_bytes`, + ) + .run(input.targetSessionId, newReferences.length, addedLogicalBytes); + } + return { ok: true, copied }; + } + #read(input: { readonly sessionId: string; readonly refId: string; @@ -402,6 +694,23 @@ export class SqliteContextOffloadStore implements ContextOffloadStore { return blob !== undefined && blobMatches(blob, blobId, bytes); } + #markBlobUnreferencedIfEligible(blobId: Uint8Array, unreferencedAt: number): void { + this.#database + .prepare( + `INSERT INTO context_gc_candidates(blob_id, unreferenced_at) + SELECT ?, ? + WHERE NOT EXISTS (SELECT 1 FROM context_refs WHERE blob_id = ?) + ON CONFLICT(blob_id) DO NOTHING`, + ) + .run(blobId, unreferencedAt, blobId); + } + + #readNow(): number { + const now = this.#now(); + assertNonNegativeSafeInteger(now, 'Context timestamp'); + return now; + } + #readSessionUsage(sessionId: string): SessionUsageRow { return ( (this.#database @@ -498,6 +807,11 @@ function decodeBytes(value: unknown): Uint8Array | undefined { return value instanceof Uint8Array ? new Uint8Array(value) : undefined; } +function decodeBlobId(value: unknown): Uint8Array | undefined { + const bytes = decodeBytes(value); + return bytes?.byteLength === 32 ? bytes : undefined; +} + function usageFromRows(session: SessionUsageRow, store: StoreUsageRow): ContextOffloadUsage { return { references: readNonNegativeInteger(session.reference_count, 'Context reference count'), @@ -547,6 +861,12 @@ function assertNonNegativeSafeInteger(value: number, label: string): void { } } +function assertPositiveSafeInteger(value: number, label: string): void { + if (!Number.isSafeInteger(value) || value <= 0) { + throw new Error(`${label} must be a positive safe integer`); + } +} + function isNonNegativeSafeInteger(value: unknown): value is number { return typeof value === 'number' && Number.isSafeInteger(value) && value >= 0; } @@ -560,6 +880,12 @@ function exceedsLimit(current: number, added: number, limit: number): boolean { return current > limit - added; } +function addSafeInteger(left: number, right: number, label: string): number { + const result = left + right; + if (!Number.isSafeInteger(result) || result < 0) throw new Error(`Invalid ${label}`); + return result; +} + function loadDatabaseSync(): typeof import('node:sqlite').DatabaseSync { return (require('node:sqlite') as typeof import('node:sqlite')).DatabaseSync; } From 462fb3743a1eb5b74823dd4eefbfafc2386b1c8b Mon Sep 17 00:00:00 2001 From: likun Date: Sat, 29 Aug 2026 18:34:48 +0800 Subject: [PATCH 2/2] fix(storage): reject stalled context garbage batches Generated-by: OpenAI Codex --- .../__tests__/sqlite-context-offload-store.test.ts | 11 ++++++++--- packages/storage/src/sqlite-context-offload-store.ts | 9 ++++++++- 2 files changed, 16 insertions(+), 4 deletions(-) diff --git a/packages/storage/src/__tests__/sqlite-context-offload-store.test.ts b/packages/storage/src/__tests__/sqlite-context-offload-store.test.ts index d5ffd3704a..deb64e40f7 100644 --- a/packages/storage/src/__tests__/sqlite-context-offload-store.test.ts +++ b/packages/storage/src/__tests__/sqlite-context-offload-store.test.ts @@ -491,7 +491,7 @@ test('garbage collection obeys both batch limits and rolls back failed deletion' const bytes = new TextEncoder().encode('four'); const first = await fixture.store.put(putInput('session-1', 'owner-1', bytes)); const second = await fixture.store.put( - putInput('session-1', 'owner-2', new TextEncoder().encode('five')), + putInput('session-1', 'owner-2', new TextEncoder().encode('fives')), ); assert.equal(first.ok, true); assert.equal(second.ok, true); @@ -503,16 +503,21 @@ test('garbage collection obeys both batch limits and rolls back failed deletion' await fixture.store.collectGarbage({ olderThan: 1_001, maxBlobs: 2, maxBytes: 4 }), { deletedBlobs: 1, deletedBytes: 4, hasMore: true }, ); + await assert.rejects( + fixture.store.collectGarbage({ olderThan: 1_001, maxBlobs: 1, maxBytes: 4 }), + /byte limit 4 cannot fit eligible blob of 5 bytes/u, + ); + assert.equal((await fixture.store.usage()).physicalBytes, 5); failGc = true; await assert.rejects( fixture.store.collectGarbage({ olderThan: 1_001, maxBlobs: 1, maxBytes: 8 }), /injected GC failure/u, ); - assert.equal((await fixture.store.usage()).physicalBytes, 4); + assert.equal((await fixture.store.usage()).physicalBytes, 5); failGc = false; assert.deepEqual( await fixture.store.collectGarbage({ olderThan: 1_001, maxBlobs: 1, maxBytes: 8 }), - { deletedBlobs: 1, deletedBytes: 4, hasMore: false }, + { deletedBlobs: 1, deletedBytes: 5, hasMore: false }, ); }); diff --git a/packages/storage/src/sqlite-context-offload-store.ts b/packages/storage/src/sqlite-context-offload-store.ts index 1b94763c90..d886b1e217 100644 --- a/packages/storage/src/sqlite-context-offload-store.ts +++ b/packages/storage/src/sqlite-context-offload-store.ts @@ -336,7 +336,14 @@ export class SqliteContextOffloadStore implements ContextOffloadStore { if (!blobId || !isNonNegativeSafeInteger(row.size_bytes)) { throw new Error('Invalid context garbage candidate'); } - if (exceedsLimit(deletedBytes, row.size_bytes, input.maxBytes)) break; + if (exceedsLimit(deletedBytes, row.size_bytes, input.maxBytes)) { + if (selected.length === 0) { + throw new Error( + `Context garbage byte limit ${input.maxBytes} cannot fit eligible blob of ${row.size_bytes} bytes`, + ); + } + break; + } deletedBytes = addSafeInteger(deletedBytes, row.size_bytes, 'Collected context bytes'); selected.push(blobId); }