diff --git a/packages/runtime-host/src/__tests__/storage-maintenance.test.ts b/packages/runtime-host/src/__tests__/storage-maintenance.test.ts index a9f60545251..fd6f240fbc4 100644 --- a/packages/runtime-host/src/__tests__/storage-maintenance.test.ts +++ b/packages/runtime-host/src/__tests__/storage-maintenance.test.ts @@ -50,6 +50,7 @@ test('maintenance waits for start, yields between bounded batches, and stops on collections += 1; return { deletedBlobs: 64, deletedBytes: 1024, hasMore: true }; }, + reclaimFreePages: async () => ({ reclaimedPages: 0, hasMore: false }), }, onError: assert.fail, }); @@ -95,6 +96,7 @@ test('failed lanes back off independently, retry, and reset after success', asyn if (attempts < 3) throw new Error('disk failure'); return { deletedBlobs: 0, deletedBytes: 0, hasMore: false }; }, + reclaimFreePages: async () => ({ reclaimedPages: 0, hasMore: false }), }, onError: (name) => { errors.push(name); @@ -224,3 +226,41 @@ test('the retention lane sweeps a second apart while work remains and every quar assert.equal(sweeps, 4); await maintenance.close(); }); + +test('context-offload page reclamation lane runs with bounded batches', async (t) => { + t.mock.timers.enable({ apis: ['setTimeout'] }); + let reclaimCalls = 0; + const maintenance = new HostStorageMaintenance({ + artifacts: { + reclaimUpgradeResidue: async () => ({ + nextAfter: null, + processedPaths: 0, + failedPaths: 0, + }), + }, + contextOffload: { + collectGarbage: async () => ({ deletedBlobs: 0, deletedBytes: 0, hasMore: false }), + reclaimFreePages: async (input) => { + assert.equal(input.maxPages, 64); + reclaimCalls += 1; + return { + reclaimedPages: reclaimCalls === 1 ? 32 : 0, + hasMore: reclaimCalls === 1, + }; + }, + }, + onError: assert.fail, + }); + maintenance.start(); + t.mock.timers.tick(100); + await settle(); + assert.equal(reclaimCalls, 1); + t.mock.timers.tick(100); + await settle(); + assert.equal(reclaimCalls, 2); + // Once a batch reports no more work the lane idles instead of retrying every 100 ms. + t.mock.timers.tick(100); + await settle(); + assert.equal(reclaimCalls, 2); + await maintenance.close(); +}); diff --git a/packages/runtime-host/src/server/storage-maintenance.ts b/packages/runtime-host/src/server/storage-maintenance.ts index e08489ba327..5770fb46deb 100644 --- a/packages/runtime-host/src/server/storage-maintenance.ts +++ b/packages/runtime-host/src/server/storage-maintenance.ts @@ -24,6 +24,7 @@ const ACTIVE_DELAY_MS = 100; const IDLE_DELAY_MS = 60_000; const MAX_BATCH_ITEMS = 64; const MAX_BATCH_BYTES = 16 * 1024 * 1024; +const MAX_VACUUM_PAGES = 64; // Archive retention deletes user data, so it moves slower than reclamation: // one bounded batch a second while work remains, otherwise a check every // quarter hour. @@ -49,7 +50,7 @@ export class HostStorageMaintenance { constructor(input: { artifacts: Pick; - contextOffload?: Pick; + contextOffload?: Pick; /** The opt-in archived-task retention sweep; true while candidates remain. */ retention?: { sweep(): Promise }; onError: (name: string, error: unknown) => void; @@ -79,7 +80,7 @@ export class HostStorageMaintenance { }, ]; const context = input.contextOffload; - if (context) + if (context) { this.#lanes.push({ name: 'context garbage collection', failures: 0, @@ -94,6 +95,17 @@ export class HostStorageMaintenance { }) ).hasMore, }); + this.#lanes.push({ + name: 'context-offload page reclamation', + failures: 0, + activeDelay: ACTIVE_DELAY_MS, + idleDelay: IDLE_DELAY_MS, + run: async () => { + const result = await context.reclaimFreePages({ maxPages: MAX_VACUUM_PAGES }); + return result.hasMore; + }, + }); + } const retention = input.retention; if (retention) this.#lanes.push({ 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 1eb41126c90..0ced27f540e 100644 --- a/packages/storage/src/__tests__/sqlite-context-offload-store.test.ts +++ b/packages/storage/src/__tests__/sqlite-context-offload-store.test.ts @@ -1079,6 +1079,72 @@ test('lifecycle queries use Session and garbage eligibility indexes', async (t) assert.match(JSON.stringify(fileDeletionPlan), /context_file_deletions_pending/u); }); +test('reclaimFreePages returns garbage-collected pages to the file in bounded batches', async (t) => { + const fixture = await createFixture(t); + const inlineBytes = 1_000_000; + const blob = new Uint8Array(inlineBytes).fill(1); + for (let index = 0; index < 4; index += 1) { + const stored = await fixture.store.put({ + sessionId: 'session-1', + owner: { kind: 'tool_result_archive', ownerId: `tool-${index}` }, + bytes: blob, + mediaType: 'application/octet-stream', + }); + assert.equal(stored.ok, true); + if (!stored.ok) return; + await fixture.store.releaseReference({ sessionId: 'session-1', refId: stored.record.refId }); + } + await fixture.store.collectGarbage({ + olderThan: 1_001, + maxBlobs: 64, + maxBytes: inlineBytes * 4, + }); + // Checkpoint through a second connection so page counts and file size + // describe the main database file, independent of WAL timing. + const inspect = async () => { + const db = new DatabaseSync(fixture.path); + try { + db.exec('PRAGMA wal_checkpoint(TRUNCATE)'); + const read = (pragma: string) => + Object.values(db.prepare(`PRAGMA ${pragma}`).get() as Record)[0]!; + return { + freePages: read('freelist_count'), + pages: read('page_count'), + fileBytes: (await stat(fixture.path)).size, + }; + } finally { + db.close(); + } + }; + const before = await inspect(); + assert.ok(before.freePages > 64, 'the fixture frees more pages than one batch'); + + const first = await fixture.store.reclaimFreePages({ maxPages: 64 }); + assert.deepEqual(first, { reclaimedPages: 64, hasMore: true }); + // Before any external checkpoint: the store's own PASSIVE checkpoint has + // already moved this batch's shrink into the main database file. + assert.ok( + (await stat(fixture.path)).size < before.fileBytes, + 'one batch shrinks the main file without an external checkpoint', + ); + assert.equal((await inspect()).freePages, before.freePages - 64); + + let reclaimed = first.reclaimedPages; + for (let batch = first; batch.hasMore; reclaimed += batch.reclaimedPages) { + batch = await fixture.store.reclaimFreePages({ maxPages: 64 }); + assert.ok(batch.reclaimedPages > 0 && batch.reclaimedPages <= 64); + } + const after = await inspect(); + assert.equal(after.freePages, 0); + assert.equal(reclaimed, before.freePages); + assert.ok(after.pages <= before.pages - before.freePages); + assert.ok(after.fileBytes < before.fileBytes, 'the main database file shrinks'); + assert.deepEqual(await fixture.store.reclaimFreePages({ maxPages: 64 }), { + reclaimedPages: 0, + hasMore: false, + }); +}); + function putInput(sessionId: string, ownerId: string, bytes: Uint8Array) { return { sessionId, diff --git a/packages/storage/src/context-offload-store.ts b/packages/storage/src/context-offload-store.ts index 1aff861e803..6fdce744f4e 100644 --- a/packages/storage/src/context-offload-store.ts +++ b/packages/storage/src/context-offload-store.ts @@ -60,6 +60,10 @@ export interface InteractiveContextOffloadWriter extends Omit; close(): Promise; } @@ -238,6 +242,10 @@ function createWriterFacade( const accepted = Object.freeze({ ...input }); return run(() => store.collectGarbage(accepted)); }, + reclaimFreePages: (input) => { + const accepted = Object.freeze({ ...input }); + return run(() => store.reclaimFreePages(accepted)); + }, usage: (sessionId) => run(() => store.usage(sessionId)), close: () => { if (closeTask) return closeTask; diff --git a/packages/storage/src/sqlite-context-offload-store.ts b/packages/storage/src/sqlite-context-offload-store.ts index edf92b96fe6..1eb96c88be4 100644 --- a/packages/storage/src/sqlite-context-offload-store.ts +++ b/packages/storage/src/sqlite-context-offload-store.ts @@ -507,6 +507,44 @@ export class SqliteContextOffloadStore implements ContextOffloadStore { }); } + /** + * Returns up to `maxPages` free pages to the filesystem with one bounded + * `incremental_vacuum`. The schema guarantees `auto_vacuum = INCREMENTAL`. + */ + async reclaimFreePages(input: { + readonly maxPages: number; + }): Promise<{ readonly reclaimedPages: number; readonly hasMore: boolean }> { + assertPositiveSafeInteger(input.maxPages, 'Context page reclamation limit'); + this.#assertOpen(); + // Probe under a read transaction so an idle tick never takes the write lock. + if (this.#readTransaction(() => this.#freelistPages()) === 0) { + return { reclaimedPages: 0, hasMore: false }; + } + const result = this.#writeTransaction(() => { + const before = this.#freelistPages(); + // incremental_vacuum(0) would free every page, so an emptied list stops here. + if (before === 0) return { reclaimedPages: 0, hasMore: false }; + this.#database.exec(`PRAGMA incremental_vacuum(${Math.min(input.maxPages, before)})`); + const after = this.#freelistPages(); + // A batch that frees nothing reports no more work instead of spinning the lane. + return { reclaimedPages: before - after, hasMore: after > 0 && after < before }; + }); + // PASSIVE never waits on readers; it moves the shrink into the main file. + if (result.reclaimedPages > 0) this.#database.exec('PRAGMA wal_checkpoint(PASSIVE)'); + return result; + } + + #freelistPages(): number { + const row = this.#database.prepare('PRAGMA freelist_count').get() as + | { freelist_count?: unknown } + | undefined; + const pages = row?.freelist_count; + if (typeof pages !== 'number' || !Number.isSafeInteger(pages) || pages < 0) { + throw new Error('Invalid SQLite freelist_count'); + } + return pages; + } + async usage(sessionId?: string): Promise { if (sessionId !== undefined) assertBoundedIdentity(sessionId, 'Session id'); this.#assertOpen();