From e6a6a673708eedaa527baf8217827708de60e7e4 Mon Sep 17 00:00:00 2001 From: likun Date: Fri, 28 Aug 2026 17:42:10 +0800 Subject: [PATCH] refactor(storage): add SQLite context offload foundation Generated-by: OpenAI Codex --- packages/core/package.json | 1 + packages/core/src/context-offload.ts | 120 ++++ .../sqlite-context-offload-store.test.ts | 367 ++++++++++++ .../src/sqlite-context-offload-schema.ts | 292 +++++++++ .../src/sqlite-context-offload-store.ts | 564 ++++++++++++++++++ 5 files changed, 1344 insertions(+) create mode 100644 packages/core/src/context-offload.ts create mode 100644 packages/storage/src/__tests__/sqlite-context-offload-store.test.ts create mode 100644 packages/storage/src/sqlite-context-offload-schema.ts create mode 100644 packages/storage/src/sqlite-context-offload-store.ts diff --git a/packages/core/package.json b/packages/core/package.json index f0086d4c2a..ed659203e5 100644 --- a/packages/core/package.json +++ b/packages/core/package.json @@ -47,6 +47,7 @@ "./permission-profile-compiler": "./dist/permission-profile-compiler.js", "./user-question": "./dist/user-question.js", "./connections": "./dist/connections.js", + "./context-offload": "./dist/context-offload.js", "./attachments": "./dist/attachments.js", "./artifacts": "./dist/artifacts.js", "./pet": "./dist/pet.js", diff --git a/packages/core/src/context-offload.ts b/packages/core/src/context-offload.ts new file mode 100644 index 0000000000..b9af00981a --- /dev/null +++ b/packages/core/src/context-offload.ts @@ -0,0 +1,120 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +import { MAX_READ_IMAGE_BYTES } from './attachments.js'; + +export interface SessionContextRef { + readonly kind: 'session_context'; + readonly sessionId: string; + readonly refId: string; +} + +export type ContextOffloadOwner = + | { + readonly kind: 'read_image_snapshot'; + readonly ownerId: string; + } + | { + readonly kind: 'tool_result_archive'; + readonly ownerId: string; + }; + +export const CONTEXT_OFFLOAD_OWNER_MAX_BYTES = Object.freeze({ + read_image_snapshot: MAX_READ_IMAGE_BYTES, + tool_result_archive: 4 * 1024 * 1024, +}) satisfies Readonly>; + +export interface ContextOffloadRecord { + readonly refId: string; + readonly sessionId: string; + readonly owner: ContextOffloadOwner; + /** Canonical lowercase SHA-256. */ + readonly blobId: string; + readonly sizeBytes: number; + readonly mediaType: string; + readonly createdAt: number; +} + +export interface ContextOffloadLimits { + /** Logical bytes referenced by one Session, counting shared blobs once per reference. */ + readonly sessionLogicalBytes: number; + /** Physical bytes stored by the workspace, counting each content-addressed blob once. */ + readonly workspacePhysicalBytes: number; +} + +export type ContextOffloadPutResult = + | { readonly ok: true; readonly record: ContextOffloadRecord } + | { + readonly ok: false; + readonly reason: + | 'too_large' + | 'session_quota_exceeded' + | 'workspace_quota_exceeded' + | 'identity_conflict' + | 'unavailable'; + }; + +export type ContextOffloadReadResult = + | { + readonly ok: true; + readonly record: ContextOffloadRecord; + readonly bytes: Uint8Array; + } + | { + readonly ok: false; + readonly reason: 'not_found' | 'session_mismatch' | 'too_large' | 'corrupt' | 'unavailable'; + }; + +export interface ContextOffloadUsage { + readonly references: number; + readonly logicalBytes: number; + readonly physicalBytes: number; +} + +/** + * Storage contract for capped, whole-object Agent context offload. + * + * The asynchronous boundary is intentional even when the first implementation + * uses DatabaseSync, so callers do not depend on the execution substrate. + */ +export interface ContextOffloadStore { + put(input: { + readonly sessionId: string; + readonly owner: ContextOffloadOwner; + readonly bytes: Uint8Array; + readonly mediaType: string; + readonly expectedSha256?: string; + }): Promise; + + read(input: { + readonly sessionId: string; + readonly refId: string; + readonly maxBytes: number; + }): Promise; + + releaseReference(input: { readonly sessionId: string; readonly refId: string }): Promise; + + /** + * Session-scoped reference/logical usage when sessionId is supplied. Physical + * bytes always describe the workspace because shared bytes have no one owner. + */ + usage(sessionId?: string): Promise; + + close(): void; +} diff --git a/packages/storage/src/__tests__/sqlite-context-offload-store.test.ts b/packages/storage/src/__tests__/sqlite-context-offload-store.test.ts new file mode 100644 index 0000000000..07e8302171 --- /dev/null +++ b/packages/storage/src/__tests__/sqlite-context-offload-store.test.ts @@ -0,0 +1,367 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +import assert from 'node:assert/strict'; +import { createHash } from 'node:crypto'; +import { mkdtemp, rm } from 'node:fs/promises'; +import { tmpdir } from 'node:os'; +import { join } from 'node:path'; +import { DatabaseSync } from 'node:sqlite'; +import test, { type TestContext } from 'node:test'; +import { CONTEXT_OFFLOAD_OWNER_MAX_BYTES } from '@maka/core/context-offload'; +import { + CONTEXT_OFFLOAD_DATABASE_NAME, + SqliteContextOffloadStore, +} from '../sqlite-context-offload-store.js'; + +test('creates the dedicated WAL schema with incremental auto-vacuum', async (t) => { + const fixture = await createFixture(t); + fixture.store.close(); + const database = new DatabaseSync(fixture.path); + t.after(() => database.close()); + + assert.equal(pragmaNumber(database, 'user_version'), 1); + assert.equal(pragmaNumber(database, 'auto_vacuum'), 2); + assert.equal(pragmaText(database, 'journal_mode'), 'wal'); + assert.deepEqual( + database + .prepare( + `SELECT name FROM sqlite_schema + WHERE type = 'table' AND name LIKE 'context_%' + ORDER BY name`, + ) + .all() + .map((row) => row.name), + ['context_blobs', 'context_refs', 'context_session_usage', 'context_store_usage'], + ); +}); + +test('atomically persists one idempotent owner identity and verifies reads', async (t) => { + const fixture = await createFixture(t, { + sessionLogicalBytes: 64, + workspacePhysicalBytes: 64, + }); + const bytes = new TextEncoder().encode('snapshot'); + const expectedSha256 = sha256(bytes); + const input = { + sessionId: 'session-1', + owner: { kind: 'read_image_snapshot' as const, ownerId: 'read-call-1' }, + bytes, + mediaType: 'image/png', + expectedSha256, + }; + + const first = await fixture.store.put(input); + const retried = await fixture.store.put(input); + assert.equal(first.ok, true); + assert.deepEqual(retried, first); + if (!first.ok) return; + assert.equal(first.record.blobId, expectedSha256); + assert.deepEqual( + await fixture.store.read({ + sessionId: input.sessionId, + refId: first.record.refId, + maxBytes: bytes.byteLength, + }), + { + ok: true, + record: first.record, + bytes, + }, + ); + assert.deepEqual(await fixture.store.usage('session-1'), { + references: 1, + logicalBytes: bytes.byteLength, + physicalBytes: bytes.byteLength, + }); + + assert.deepEqual( + await fixture.store.put({ ...input, bytes: new TextEncoder().encode('changed') }), + { + ok: false, + reason: 'identity_conflict', + }, + ); + assert.deepEqual(await fixture.store.put({ ...input, expectedSha256: '0'.repeat(64) }), { + ok: false, + reason: 'identity_conflict', + }); +}); + +test('reopens durable records and preserves owner idempotency', async (t) => { + const fixture = await createFixture(t); + const bytes = new TextEncoder().encode('durable'); + const first = await fixture.store.put(putInput('session-1', 'archive-1', bytes)); + assert.equal(first.ok, true); + if (!first.ok) return; + fixture.store.close(); + + const reopened = new SqliteContextOffloadStore(fixture.path, { + limits: fixture.limits, + now: () => 2_000, + idFactory: () => 'unexpected-new-reference', + }); + t.after(() => reopened.close()); + assert.deepEqual(await reopened.put(putInput('session-1', 'archive-1', bytes)), first); + assert.deepEqual( + await reopened.read({ + sessionId: 'session-1', + refId: first.record.refId, + maxBytes: bytes.byteLength, + }), + { ok: true, record: first.record, bytes }, + ); +}); + +test('rejects a database schema newer than this authority understands', async (t) => { + const root = await mkdtemp(join(tmpdir(), 'maka-context-offload-newer-')); + 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.close(); + + assert.throws( + () => + new SqliteContextOffloadStore(path, { + limits: { sessionLogicalBytes: 1, workspacePhysicalBytes: 1 }, + }), + /schema 2 is newer than supported version 1/u, + ); +}); + +test('rejects a current schema missing a query-required index', async (t) => { + const fixture = await createFixture(t); + fixture.store.close(); + const database = new DatabaseSync(fixture.path); + database.exec('DROP INDEX context_refs_session'); + database.close(); + + assert.throws( + () => new SqliteContextOffloadStore(fixture.path, { limits: fixture.limits }), + /missing index context_refs_session/u, + ); +}); + +test('deduplicates physical bytes while quotas count each Session reference logically', async (t) => { + const fixture = await createFixture(t, { + sessionLogicalBytes: 8, + workspacePhysicalBytes: 4, + }); + const bytes = new TextEncoder().encode('same'); + + const first = await fixture.store.put(putInput('session-1', 'owner-1', bytes)); + const crossSession = await fixture.store.put(putInput('session-2', 'owner-2', bytes)); + const secondReference = await fixture.store.put(putInput('session-1', 'owner-3', bytes)); + assert.equal(first.ok, true); + assert.equal(crossSession.ok, true); + assert.equal(secondReference.ok, true); + assert.deepEqual(await fixture.store.usage('session-1'), { + references: 2, + logicalBytes: 8, + physicalBytes: 4, + }); + assert.deepEqual(await fixture.store.usage('session-2'), { + references: 1, + logicalBytes: 4, + physicalBytes: 4, + }); + + assert.deepEqual(await fixture.store.put(putInput('session-1', 'owner-4', bytes)), { + ok: false, + reason: 'session_quota_exceeded', + }); + assert.deepEqual( + await fixture.store.put(putInput('session-2', 'owner-5', new TextEncoder().encode('else'))), + { ok: false, reason: 'workspace_quota_exceeded' }, + ); +}); + +test('fails closed before returning bytes for Session mismatch and size limits', async (t) => { + const fixture = await createFixture(t); + const stored = await fixture.store.put( + putInput('session-1', 'archive-1', new TextEncoder().encode('archive')), + ); + assert.equal(stored.ok, true); + if (!stored.ok) return; + + assert.deepEqual( + await fixture.store.read({ + sessionId: 'session-2', + refId: stored.record.refId, + maxBytes: 100, + }), + { ok: false, reason: 'session_mismatch' }, + ); + assert.deepEqual( + await fixture.store.read({ + sessionId: 'session-1', + refId: stored.record.refId, + maxBytes: 3, + }), + { ok: false, reason: 'too_large' }, + ); + assert.deepEqual( + await fixture.store.read({ + sessionId: 'session-1', + refId: 'missing', + maxBytes: 100, + }), + { ok: false, reason: 'not_found' }, + ); +}); + +test('enforces owner hard caps before commit', async (t) => { + const imageLimit = CONTEXT_OFFLOAD_OWNER_MAX_BYTES.read_image_snapshot; + const archiveLimit = CONTEXT_OFFLOAD_OWNER_MAX_BYTES.tool_result_archive; + const fixture = await createFixture(t, { + sessionLogicalBytes: imageLimit * 2, + workspacePhysicalBytes: imageLimit * 2, + }); + + assert.deepEqual( + await fixture.store.put({ + ...putInput('session-1', 'large-image', new Uint8Array(imageLimit + 1)), + owner: { kind: 'read_image_snapshot', ownerId: 'large-image' }, + }), + { ok: false, reason: 'too_large' }, + ); + assert.deepEqual( + await fixture.store.put({ + ...putInput('session-1', 'large-archive', new Uint8Array(archiveLimit + 1)), + owner: { kind: 'tool_result_archive', ownerId: 'large-archive' }, + }), + { ok: false, reason: 'too_large' }, + ); + assert.deepEqual(await fixture.store.usage(), { + references: 0, + logicalBytes: 0, + physicalBytes: 0, + }); +}); + +test('detects payload corruption instead of returning unverified bytes', async (t) => { + const fixture = await createFixture(t); + const stored = await fixture.store.put( + putInput('session-1', 'archive-1', new TextEncoder().encode('original')), + ); + assert.equal(stored.ok, true); + if (!stored.ok) return; + + const database = new DatabaseSync(fixture.path); + database + .prepare('UPDATE context_blobs SET payload = ?') + .run(new TextEncoder().encode('tampered')); + database.close(); + + assert.deepEqual( + await fixture.store.read({ + sessionId: 'session-1', + refId: stored.record.refId, + maxBytes: 100, + }), + { ok: false, reason: 'corrupt' }, + ); +}); + +test('rolls back blob and reference together when publication fails', async (t) => { + const fixture = await createFixture(t, undefined, (point) => { + if (point === 'after_ref_insert') throw new Error('injected publication failure'); + }); + + assert.deepEqual( + await fixture.store.put( + putInput('session-1', 'archive-1', new TextEncoder().encode('archive')), + ), + { ok: false, reason: 'unavailable' }, + ); + assert.deepEqual(await fixture.store.usage(), { + references: 0, + logicalBytes: 0, + physicalBytes: 0, + }); +}); + +test('releases only the authorized Session reference without deleting shared bytes', async (t) => { + const fixture = await createFixture(t); + const bytes = new TextEncoder().encode('shared'); + const first = await fixture.store.put(putInput('session-1', 'owner-1', bytes)); + const second = await fixture.store.put(putInput('session-2', 'owner-2', bytes)); + assert.equal(first.ok, true); + assert.equal(second.ok, true); + if (!first.ok) return; + + await fixture.store.releaseReference({ sessionId: 'session-2', refId: first.record.refId }); + assert.equal((await fixture.store.usage('session-1')).references, 1); + await fixture.store.releaseReference({ sessionId: 'session-1', refId: first.record.refId }); + await fixture.store.releaseReference({ sessionId: 'session-1', refId: first.record.refId }); + assert.deepEqual(await fixture.store.usage('session-1'), { + references: 0, + logicalBytes: 0, + physicalBytes: bytes.byteLength, + }); +}); + +function putInput(sessionId: string, ownerId: string, bytes: Uint8Array) { + return { + sessionId, + owner: { kind: 'tool_result_archive' as const, ownerId }, + bytes, + mediaType: 'application/json', + }; +} + +async function createFixture( + t: TestContext, + limits = { sessionLogicalBytes: 16 * 1024 * 1024, workspacePhysicalBytes: 32 * 1024 * 1024 }, + failpoint?: ConstructorParameters[1]['failpoint'], +) { + const root = await mkdtemp(join(tmpdir(), 'maka-context-offload-')); + const path = join(root, CONTEXT_OFFLOAD_DATABASE_NAME); + let nextId = 1; + const store = new SqliteContextOffloadStore(path, { + limits, + now: () => 1_000, + idFactory: () => `ref-${nextId++}`, + failpoint, + }); + t.after(async () => { + store.close(); + await rm(root, { recursive: true, force: true }); + }); + return { limits, path, store }; +} + +function sha256(bytes: Uint8Array): string { + return createHash('sha256').update(bytes).digest('hex'); +} + +function pragmaNumber(database: DatabaseSync, name: string): number { + const row = database.prepare(`PRAGMA ${name}`).get() as Record; + const value = row[name]; + if (typeof value !== 'number') throw new Error(`Expected numeric PRAGMA ${name}`); + return value; +} + +function pragmaText(database: DatabaseSync, name: string): string { + const row = database.prepare(`PRAGMA ${name}`).get() as Record; + const value = row[name]; + if (typeof value !== 'string') throw new Error(`Expected text PRAGMA ${name}`); + return value; +} diff --git a/packages/storage/src/sqlite-context-offload-schema.ts b/packages/storage/src/sqlite-context-offload-schema.ts new file mode 100644 index 0000000000..e937111622 --- /dev/null +++ b/packages/storage/src/sqlite-context-offload-schema.ts @@ -0,0 +1,292 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +import type { DatabaseSync } from 'node:sqlite'; + +export const SQLITE_CONTEXT_OFFLOAD_SCHEMA_VERSION = 1; +const SQLITE_INITIALIZATION_BUSY_TIMEOUT_MS = 5_000; +const SQLITE_INITIALIZATION_RETRY_DELAY_MS = 10; +const initializationRetryGate = new Int32Array(new SharedArrayBuffer(4)); + +const INITIAL_SCHEMA = ` + CREATE TABLE context_blobs ( + blob_id BLOB PRIMARY KEY CHECK(length(blob_id) = 32), + payload BLOB NOT NULL, + size_bytes INTEGER NOT NULL CHECK(size_bytes >= 0 AND length(payload) = size_bytes), + created_at INTEGER NOT NULL CHECK(created_at >= 0) + ); + + CREATE TABLE context_refs ( + ref_id TEXT PRIMARY KEY, + session_id TEXT NOT NULL, + owner_kind TEXT NOT NULL CHECK( + owner_kind IN ('read_image_snapshot', 'tool_result_archive') + ), + owner_id TEXT NOT NULL, + blob_id BLOB NOT NULL REFERENCES context_blobs(blob_id) ON DELETE RESTRICT, + media_type TEXT NOT NULL, + created_at INTEGER NOT NULL CHECK(created_at >= 0), + UNIQUE(session_id, owner_kind, owner_id) + ); + + CREATE INDEX context_refs_session + ON context_refs(session_id, created_at, ref_id); + + CREATE INDEX context_refs_blob + ON context_refs(blob_id); + + CREATE TABLE context_session_usage ( + session_id TEXT PRIMARY KEY, + reference_count INTEGER NOT NULL CHECK(reference_count >= 0), + logical_bytes INTEGER NOT NULL CHECK(logical_bytes >= 0) + ); + + CREATE TABLE context_store_usage ( + singleton INTEGER PRIMARY KEY CHECK(singleton = 1), + blob_count INTEGER NOT NULL CHECK(blob_count >= 0), + physical_bytes INTEGER NOT NULL CHECK(physical_bytes >= 0) + ); + + INSERT INTO context_store_usage(singleton, blob_count, physical_bytes) + VALUES (1, 0, 0); +`; + +const REQUIRED_SCHEMA_OBJECTS = Object.freeze([ + ['table', 'context_blobs'], + ['table', 'context_refs'], + ['table', 'context_session_usage'], + ['table', 'context_store_usage'], + ['index', 'context_refs_session'], + ['index', 'context_refs_blob'], +] as const); + +const REQUIRED_TABLE_COLUMNS = Object.freeze({ + context_blobs: ['blob_id', 'payload', 'size_bytes', 'created_at'], + context_refs: [ + 'ref_id', + 'session_id', + 'owner_kind', + 'owner_id', + 'blob_id', + 'media_type', + 'created_at', + ], + context_session_usage: ['session_id', 'reference_count', 'logical_bytes'], + context_store_usage: ['singleton', 'blob_count', 'physical_bytes'], +} as const); + +const REQUIRED_INDEX_COLUMNS = Object.freeze({ + context_refs_session: ['session_id', 'created_at', 'ref_id'], + context_refs_blob: ['blob_id'], +} as const); + +export function configureSqliteContextOffloadDatabase(db: DatabaseSync): void { + db.exec(`PRAGMA busy_timeout = ${SQLITE_INITIALIZATION_BUSY_TIMEOUT_MS}`); + // WAL initialization fixes the database header in a form where changing + // auto_vacuum from NONE is no longer accepted. Configure it first, while a + // brand-new dedicated database still has no application schema objects. + if (readSqliteContextOffloadSchemaVersion(db) === 0 && !hasApplicationSchemaObjects(db)) { + db.exec('PRAGMA auto_vacuum = INCREMENTAL'); + } + ensureWalJournalMode(db); + db.exec('PRAGMA synchronous = FULL'); + db.exec('PRAGMA foreign_keys = ON'); +} + +export function migrateSqliteContextOffloadDatabase(db: DatabaseSync): void { + const observedVersion = readSqliteContextOffloadSchemaVersion(db); + if (observedVersion > SQLITE_CONTEXT_OFFLOAD_SCHEMA_VERSION) { + throw newerSchemaError(observedVersion); + } + if (observedVersion === SQLITE_CONTEXT_OFFLOAD_SCHEMA_VERSION) { + validateSchema(db); + return; + } + + db.exec('BEGIN IMMEDIATE'); + try { + const current = readSqliteContextOffloadSchemaVersion(db); + if (current > SQLITE_CONTEXT_OFFLOAD_SCHEMA_VERSION) throw newerSchemaError(current); + if (current === 0) { + if (hasApplicationSchemaObjects(db)) { + throw new Error('Unversioned context-offload SQLite schema is not supported'); + } + db.exec(INITIAL_SCHEMA); + db.exec(`PRAGMA user_version = ${SQLITE_CONTEXT_OFFLOAD_SCHEMA_VERSION}`); + } + validateSchema(db); + db.exec('COMMIT'); + } catch (error) { + rollback(db); + throw error; + } +} + +export function readSqliteContextOffloadSchemaVersion(db: DatabaseSync): number { + const row = retryWhileSqliteBusy( + () => db.prepare('PRAGMA user_version').get() as { user_version?: unknown } | undefined, + ); + const value = row?.user_version; + if (typeof value !== 'number' || !Number.isSafeInteger(value) || value < 0) { + throw new Error('Invalid context-offload SQLite schema version'); + } + return value; +} + +function validateSchema(db: DatabaseSync): void { + const autoVacuum = db.prepare('PRAGMA auto_vacuum').get() as + | { auto_vacuum?: unknown } + | undefined; + if (autoVacuum?.auto_vacuum !== 2) { + throw new Error('Incomplete context-offload SQLite schema: incremental auto-vacuum required'); + } + const readObject = db.prepare('SELECT type FROM sqlite_schema WHERE name = ?'); + for (const [type, name] of REQUIRED_SCHEMA_OBJECTS) { + const row = readObject.get(name) as { type?: unknown } | undefined; + if (row?.type !== type) { + throw new Error(`Incomplete context-offload SQLite schema: missing ${type} ${name}`); + } + } + for (const [table, requiredColumns] of Object.entries(REQUIRED_TABLE_COLUMNS)) { + const columns = new Set( + (db.prepare(`PRAGMA table_info(${table})`).all() as Array<{ name?: unknown }>).flatMap( + (row) => (typeof row.name === 'string' ? [row.name] : []), + ), + ); + for (const column of requiredColumns) { + if (!columns.has(column)) { + throw new Error( + `Incomplete context-offload SQLite schema: table ${table} is missing column ${column}`, + ); + } + } + } + for (const [index, requiredColumns] of Object.entries(REQUIRED_INDEX_COLUMNS)) { + const columns = ( + db.prepare(`PRAGMA index_info(${index})`).all() as Array<{ + seqno?: unknown; + name?: unknown; + }> + ) + .filter( + (row): row is { seqno: number; name: string } => + typeof row.seqno === 'number' && typeof row.name === 'string', + ) + .sort((left, right) => left.seqno - right.seqno) + .map((row) => row.name); + if (requiredColumns.some((column, position) => columns[position] !== column)) { + throw new Error( + `Incomplete context-offload SQLite schema: index ${index} has incompatible columns`, + ); + } + } + const usage = db + .prepare('SELECT blob_count, physical_bytes FROM context_store_usage WHERE singleton = 1') + .get() as { blob_count?: unknown; physical_bytes?: unknown } | undefined; + if (!isNonNegativeInteger(usage?.blob_count) || !isNonNegativeInteger(usage.physical_bytes)) { + throw new Error('Incomplete context-offload SQLite schema: missing store usage row'); + } +} + +function hasApplicationSchemaObjects(db: DatabaseSync): boolean { + const row = db + .prepare( + `SELECT 1 AS present FROM sqlite_schema + WHERE name NOT LIKE 'sqlite_%' LIMIT 1`, + ) + .get() as { present?: unknown } | undefined; + return row?.present === 1; +} + +function ensureWalJournalMode(db: DatabaseSync): void { + const deadline = Date.now() + SQLITE_INITIALIZATION_BUSY_TIMEOUT_MS; + while (true) { + const journalMode = readJournalMode(db); + if (journalMode === 'wal' || journalMode === 'memory') return; + try { + db.exec('PRAGMA journal_mode = WAL'); + const configuredMode = readJournalMode(db); + if (configuredMode !== 'wal') { + throw new Error( + `Context-offload SQLite requires WAL journal mode, received ${configuredMode}`, + ); + } + return; + } catch (error) { + if (!isSqliteBusy(error) || Date.now() >= deadline) throw error; + Atomics.wait( + initializationRetryGate, + 0, + 0, + Math.min(SQLITE_INITIALIZATION_RETRY_DELAY_MS, Math.max(1, deadline - Date.now())), + ); + } + } +} + +function readJournalMode(db: DatabaseSync): string { + const row = retryWhileSqliteBusy( + () => db.prepare('PRAGMA journal_mode').get() as { journal_mode?: unknown } | undefined, + ); + if (typeof row?.journal_mode !== 'string') { + throw new Error('Invalid context-offload SQLite journal mode'); + } + return row.journal_mode.toLowerCase(); +} + +function retryWhileSqliteBusy(operation: () => T): T { + const deadline = Date.now() + SQLITE_INITIALIZATION_BUSY_TIMEOUT_MS; + while (true) { + try { + return operation(); + } catch (error) { + if (!isSqliteBusy(error) || Date.now() >= deadline) throw error; + Atomics.wait( + initializationRetryGate, + 0, + 0, + Math.min(SQLITE_INITIALIZATION_RETRY_DELAY_MS, Math.max(1, deadline - Date.now())), + ); + } + } +} + +function isSqliteBusy(error: unknown): boolean { + if (!(error instanceof Error)) return false; + const code = 'code' in error ? String(error.code) : ''; + return code === 'SQLITE_BUSY' || /database (?:is )?(?:locked|busy)/i.test(error.message); +} + +function newerSchemaError(version: number): Error { + return new Error( + `Context-offload SQLite schema ${version} is newer than supported version ${SQLITE_CONTEXT_OFFLOAD_SCHEMA_VERSION}`, + ); +} + +function isNonNegativeInteger(value: unknown): value is number { + return typeof value === 'number' && Number.isSafeInteger(value) && value >= 0; +} + +function rollback(db: DatabaseSync): void { + try { + db.exec('ROLLBACK'); + } catch { + // Preserve the migration failure that triggered rollback. + } +} diff --git a/packages/storage/src/sqlite-context-offload-store.ts b/packages/storage/src/sqlite-context-offload-store.ts new file mode 100644 index 0000000000..ff20f8f16f --- /dev/null +++ b/packages/storage/src/sqlite-context-offload-store.ts @@ -0,0 +1,564 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +import { createHash, randomUUID } from 'node:crypto'; +import { mkdirSync } from 'node:fs'; +import { createRequire } from 'node:module'; +import { dirname } from 'node:path'; +import type { DatabaseSync } from 'node:sqlite'; +import { + CONTEXT_OFFLOAD_OWNER_MAX_BYTES, + type ContextOffloadLimits, + type ContextOffloadOwner, + type ContextOffloadPutResult, + type ContextOffloadReadResult, + type ContextOffloadRecord, + type ContextOffloadStore, + type ContextOffloadUsage, +} from '@maka/core/context-offload'; +import { + configureSqliteContextOffloadDatabase, + migrateSqliteContextOffloadDatabase, +} from './sqlite-context-offload-schema.js'; + +const MAX_ID_CODE_POINTS = 512; +const MAX_MEDIA_TYPE_CODE_POINTS = 256; +const SHA256_PATTERN = /^[0-9a-f]{64}$/; +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 interface SqliteContextOffloadStoreOptions { + readonly limits: ContextOffloadLimits; + readonly now?: () => number; + readonly idFactory?: () => string; + readonly failpoint?: (point: SqliteContextOffloadStoreFailpoint) => void; + readonly onUnavailable?: (error: unknown) => void; +} + +interface ContextReferenceRow { + ref_id: unknown; + session_id: unknown; + owner_kind: unknown; + owner_id: unknown; + blob_id: unknown; + size_bytes: unknown; + media_type: unknown; + created_at: unknown; +} + +interface ContextBlobRow { + payload: unknown; + size_bytes: unknown; +} + +interface SessionUsageRow { + reference_count: unknown; + logical_bytes: unknown; +} + +interface StoreUsageRow { + blob_count: unknown; + physical_bytes: unknown; +} + +export class SqliteContextOffloadStore implements ContextOffloadStore { + readonly #database: DatabaseSync; + readonly #limits: ContextOffloadLimits; + readonly #now: () => number; + readonly #idFactory: () => string; + readonly #failpoint?: (point: SqliteContextOffloadStoreFailpoint) => void; + readonly #onUnavailable?: (error: unknown) => void; + #closed = false; + + constructor(path: string, options: SqliteContextOffloadStoreOptions) { + if (!path) throw new Error('Context-offload SQLite path is required'); + this.#limits = validateLimits(options.limits); + this.#now = options.now ?? Date.now; + this.#idFactory = options.idFactory ?? randomUUID; + this.#failpoint = options.failpoint; + this.#onUnavailable = options.onUnavailable; + if (path !== ':memory:') mkdirSync(dirname(path), { recursive: true }); + const Database = loadDatabaseSync(); + this.#database = new Database(path); + try { + configureSqliteContextOffloadDatabase(this.#database); + migrateSqliteContextOffloadDatabase(this.#database); + } catch (error) { + this.#database.close(); + this.#closed = true; + throw error; + } + } + + async put(input: { + readonly sessionId: string; + readonly owner: ContextOffloadOwner; + readonly bytes: Uint8Array; + readonly mediaType: string; + readonly expectedSha256?: string; + }): Promise { + assertBoundedIdentity(input.sessionId, 'Session id'); + assertOwner(input.owner); + assertBoundedText(input.mediaType, MAX_MEDIA_TYPE_CODE_POINTS, 'Context media type'); + if (!(input.bytes instanceof Uint8Array)) { + throw new Error('Context bytes must be a Uint8Array'); + } + if (input.expectedSha256 !== undefined && !SHA256_PATTERN.test(input.expectedSha256)) { + throw new Error('Expected context SHA-256 must be canonical lowercase hexadecimal'); + } + if (input.bytes.byteLength > CONTEXT_OFFLOAD_OWNER_MAX_BYTES[input.owner.kind]) { + return { ok: false, reason: 'too_large' }; + } + + // Snapshot caller-owned bytes before crossing the asynchronous interface. + const bytes = new Uint8Array(input.bytes); + const blobId = createHash('sha256').update(bytes).digest('hex'); + if (input.expectedSha256 !== undefined && input.expectedSha256 !== blobId) { + return { ok: false, reason: 'identity_conflict' }; + } + + try { + this.#assertOpen(); + return this.#writeTransaction(() => this.#put({ ...input, bytes, blobId })); + } catch (error) { + this.#onUnavailable?.(error); + return { ok: false, reason: 'unavailable' }; + } + } + + async read(input: { + readonly sessionId: string; + readonly refId: string; + readonly maxBytes: number; + }): Promise { + assertBoundedIdentity(input.sessionId, 'Session id'); + assertBoundedIdentity(input.refId, 'Context reference id'); + assertNonNegativeSafeInteger(input.maxBytes, 'Context read byte limit'); + try { + this.#assertOpen(); + return this.#readTransaction(() => this.#read(input)); + } catch (error) { + this.#onUnavailable?.(error); + return { ok: false, reason: 'unavailable' }; + } + } + + async releaseReference(input: { + readonly sessionId: string; + readonly refId: string; + }): Promise { + assertBoundedIdentity(input.sessionId, 'Session id'); + assertBoundedIdentity(input.refId, 'Context reference id'); + this.#assertOpen(); + this.#writeTransaction(() => { + const row = this.#database + .prepare( + `SELECT r.session_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; + if (!row || row.session_id !== input.sessionId) return; + if (!isNonNegativeSafeInteger(row.size_bytes)) { + throw new Error('Invalid context reference size'); + } + const deleted = this.#database + .prepare('DELETE FROM context_refs WHERE session_id = ? AND ref_id = ?') + .run(input.sessionId, input.refId); + if (deleted.changes !== 1) return; + this.#database + .prepare( + `UPDATE context_session_usage + SET reference_count = reference_count - 1, + logical_bytes = logical_bytes - ? + WHERE session_id = ?`, + ) + .run(row.size_bytes, input.sessionId); + this.#database + .prepare( + `DELETE FROM context_session_usage + WHERE session_id = ? AND reference_count = 0 AND logical_bytes = 0`, + ) + .run(input.sessionId); + }); + } + + async usage(sessionId?: string): Promise { + if (sessionId !== undefined) assertBoundedIdentity(sessionId, 'Session id'); + this.#assertOpen(); + return this.#readTransaction(() => { + const storeUsage = this.#readStoreUsage(); + if (sessionId === undefined) { + const row = this.#database + .prepare( + `SELECT COALESCE(SUM(reference_count), 0) AS reference_count, + COALESCE(SUM(logical_bytes), 0) AS logical_bytes + FROM context_session_usage`, + ) + .get() as unknown as SessionUsageRow; + return usageFromRows(row, storeUsage); + } + const row = this.#readSessionUsage(sessionId); + return usageFromRows(row, storeUsage); + }); + } + + close(): void { + if (this.#closed) return; + this.#closed = true; + this.#database.close(); + } + + #put(input: { + readonly sessionId: string; + readonly owner: ContextOffloadOwner; + readonly bytes: Uint8Array; + readonly mediaType: string; + readonly blobId: string; + }): ContextOffloadPutResult { + const existingReference = this.#readReferenceByOwner(input.sessionId, input.owner); + if (existingReference) { + if (existingReference.blobId !== input.blobId) { + return { ok: false, reason: 'identity_conflict' }; + } + if (!this.#verifyBlob(input.blobId, input.bytes)) { + throw new Error(`Context blob failed integrity verification: ${input.blobId}`); + } + return { ok: true, record: existingReference }; + } + + const sessionUsage = this.#readSessionUsage(input.sessionId); + const logicalBytes = readNonNegativeInteger( + sessionUsage.logical_bytes, + 'Session logical bytes', + ); + if (exceedsLimit(logicalBytes, input.bytes.byteLength, this.#limits.sessionLogicalBytes)) { + return { ok: false, reason: 'session_quota_exceeded' }; + } + + const blobIdBytes = Buffer.from(input.blobId, 'hex'); + const existingBlob = this.#database + .prepare('SELECT payload, size_bytes FROM context_blobs WHERE blob_id = ?') + .get(blobIdBytes) as ContextBlobRow | undefined; + const storeUsage = this.#readStoreUsage(); + if (existingBlob) { + if (!blobMatches(existingBlob, input.blobId, input.bytes)) { + throw new Error(`Context blob identity is inconsistent: ${input.blobId}`); + } + } else { + const physicalBytes = readNonNegativeInteger( + storeUsage.physical_bytes, + 'Workspace physical bytes', + ); + if ( + exceedsLimit(physicalBytes, input.bytes.byteLength, this.#limits.workspacePhysicalBytes) + ) { + return { ok: false, reason: 'workspace_quota_exceeded' }; + } + } + + const createdAt = this.#now(); + assertNonNegativeSafeInteger(createdAt, 'Context creation time'); + const refId = this.#idFactory(); + assertBoundedIdentity(refId, 'Context reference id'); + if (!existingBlob) { + this.#database + .prepare( + `INSERT INTO context_blobs(blob_id, payload, size_bytes, created_at) + VALUES (?, ?, ?, ?)`, + ) + .run(blobIdBytes, input.bytes, input.bytes.byteLength, createdAt); + this.#database + .prepare( + `UPDATE context_store_usage + SET blob_count = blob_count + 1, + physical_bytes = physical_bytes + ? + WHERE singleton = 1`, + ) + .run(input.bytes.byteLength); + this.#failpoint?.('after_blob_insert'); + } + this.#database + .prepare( + `INSERT INTO context_refs( + ref_id, session_id, owner_kind, owner_id, blob_id, media_type, created_at + ) VALUES (?, ?, ?, ?, ?, ?, ?)`, + ) + .run( + refId, + input.sessionId, + input.owner.kind, + input.owner.ownerId, + blobIdBytes, + input.mediaType, + createdAt, + ); + this.#database + .prepare( + `INSERT INTO context_session_usage(session_id, reference_count, logical_bytes) + VALUES (?, 1, ?) + ON CONFLICT(session_id) DO UPDATE SET + reference_count = reference_count + 1, + logical_bytes = logical_bytes + excluded.logical_bytes`, + ) + .run(input.sessionId, input.bytes.byteLength); + this.#failpoint?.('after_ref_insert'); + return { + ok: true, + record: { + refId, + sessionId: input.sessionId, + owner: { ...input.owner }, + blobId: input.blobId, + sizeBytes: input.bytes.byteLength, + mediaType: input.mediaType, + createdAt, + }, + }; + } + + #read(input: { + readonly sessionId: string; + readonly refId: string; + readonly maxBytes: number; + }): ContextOffloadReadResult { + const row = 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.ref_id = ?`, + ) + .get(input.refId) as ContextReferenceRow | undefined; + if (!row) return { ok: false, reason: 'not_found' }; + const record = decodeReferenceRow(row); + if (!record) return { ok: false, reason: 'corrupt' }; + if (record.sessionId !== input.sessionId) return { ok: false, reason: 'session_mismatch' }; + if ( + record.sizeBytes > input.maxBytes || + record.sizeBytes > CONTEXT_OFFLOAD_OWNER_MAX_BYTES[record.owner.kind] + ) { + return { ok: false, reason: 'too_large' }; + } + const blob = this.#database + .prepare('SELECT payload, size_bytes FROM context_blobs WHERE blob_id = ?') + .get(Buffer.from(record.blobId, 'hex')) as ContextBlobRow | undefined; + if (!blob) return { ok: false, reason: 'corrupt' }; + const bytes = decodeBytes(blob.payload); + if ( + !bytes || + blob.size_bytes !== record.sizeBytes || + bytes.byteLength !== record.sizeBytes || + createHash('sha256').update(bytes).digest('hex') !== record.blobId + ) { + return { ok: false, reason: 'corrupt' }; + } + return { ok: true, record, bytes }; + } + + #readReferenceByOwner( + sessionId: string, + owner: ContextOffloadOwner, + ): ContextOffloadRecord | undefined { + const row = 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.owner_kind = ? AND r.owner_id = ?`, + ) + .get(sessionId, owner.kind, owner.ownerId) as ContextReferenceRow | undefined; + if (!row) return undefined; + const record = decodeReferenceRow(row); + if (!record) throw new Error('Invalid context reference row'); + return record; + } + + #verifyBlob(blobId: string, bytes: Uint8Array): boolean { + const blob = this.#database + .prepare('SELECT payload, size_bytes FROM context_blobs WHERE blob_id = ?') + .get(Buffer.from(blobId, 'hex')) as ContextBlobRow | undefined; + return blob !== undefined && blobMatches(blob, blobId, bytes); + } + + #readSessionUsage(sessionId: string): SessionUsageRow { + return ( + (this.#database + .prepare( + `SELECT reference_count, logical_bytes + FROM context_session_usage WHERE session_id = ?`, + ) + .get(sessionId) as SessionUsageRow | undefined) ?? { + reference_count: 0, + logical_bytes: 0, + } + ); + } + + #readStoreUsage(): StoreUsageRow { + const row = this.#database + .prepare( + `SELECT blob_count, physical_bytes + FROM context_store_usage WHERE singleton = 1`, + ) + .get() as StoreUsageRow | undefined; + if (!row) throw new Error('Missing context store usage row'); + readNonNegativeInteger(row.blob_count, 'Workspace blob count'); + readNonNegativeInteger(row.physical_bytes, 'Workspace physical bytes'); + return row; + } + + #writeTransaction(operation: () => T): T { + this.#database.exec('BEGIN IMMEDIATE'); + try { + const result = operation(); + this.#database.exec('COMMIT'); + return result; + } catch (error) { + rollback(this.#database); + throw error; + } + } + + #readTransaction(operation: () => T): T { + this.#database.exec('BEGIN'); + try { + const result = operation(); + this.#database.exec('COMMIT'); + return result; + } catch (error) { + rollback(this.#database); + throw error; + } + } + + #assertOpen(): void { + if (this.#closed) throw new Error('SQLite Context Offload Store is closed'); + } +} + +function decodeReferenceRow(row: ContextReferenceRow): ContextOffloadRecord | undefined { + if ( + typeof row.ref_id !== 'string' || + typeof row.session_id !== 'string' || + !isOwnerKind(row.owner_kind) || + typeof row.owner_id !== 'string' || + typeof row.media_type !== 'string' || + !isNonNegativeSafeInteger(row.size_bytes) || + !isNonNegativeSafeInteger(row.created_at) + ) { + return undefined; + } + const blobIdBytes = decodeBytes(row.blob_id); + if (!blobIdBytes || blobIdBytes.byteLength !== 32) return undefined; + return { + refId: row.ref_id, + sessionId: row.session_id, + owner: { kind: row.owner_kind, ownerId: row.owner_id }, + blobId: Buffer.from(blobIdBytes).toString('hex'), + sizeBytes: row.size_bytes, + mediaType: row.media_type, + createdAt: row.created_at, + }; +} + +function blobMatches(row: ContextBlobRow, blobId: string, expectedBytes: Uint8Array): boolean { + const storedBytes = decodeBytes(row.payload); + return ( + storedBytes !== undefined && + isNonNegativeSafeInteger(row.size_bytes) && + row.size_bytes === expectedBytes.byteLength && + storedBytes.byteLength === expectedBytes.byteLength && + createHash('sha256').update(storedBytes).digest('hex') === blobId + ); +} + +function decodeBytes(value: unknown): Uint8Array | undefined { + return value instanceof Uint8Array ? new Uint8Array(value) : undefined; +} + +function usageFromRows(session: SessionUsageRow, store: StoreUsageRow): ContextOffloadUsage { + return { + references: readNonNegativeInteger(session.reference_count, 'Context reference count'), + logicalBytes: readNonNegativeInteger(session.logical_bytes, 'Context logical bytes'), + physicalBytes: readNonNegativeInteger(store.physical_bytes, 'Context physical bytes'), + }; +} + +function validateLimits(limits: ContextOffloadLimits): ContextOffloadLimits { + assertNonNegativeSafeInteger(limits.sessionLogicalBytes, 'Session context quota'); + assertNonNegativeSafeInteger(limits.workspacePhysicalBytes, 'Workspace context quota'); + return Object.freeze({ ...limits }); +} + +function assertOwner(owner: ContextOffloadOwner): void { + if (!isOwnerKind(owner.kind)) throw new Error(`Unsupported context owner: ${String(owner.kind)}`); + assertBoundedIdentity(owner.ownerId, 'Context owner id'); +} + +function isOwnerKind(value: unknown): value is ContextOffloadOwner['kind'] { + return value === 'read_image_snapshot' || value === 'tool_result_archive'; +} + +function assertBoundedIdentity(value: string, label: string): void { + assertBoundedText(value, MAX_ID_CODE_POINTS, label); +} + +function assertBoundedText(value: string, maxCodePoints: number, label: string): void { + if (typeof value !== 'string' || value.length === 0 || [...value].length > maxCodePoints) { + throw new Error(`${label} must be a non-empty string of at most ${maxCodePoints} code points`); + } +} + +function assertNonNegativeSafeInteger(value: number, label: string): void { + if (!isNonNegativeSafeInteger(value)) { + throw new Error(`${label} must be a non-negative safe integer`); + } +} + +function isNonNegativeSafeInteger(value: unknown): value is number { + return typeof value === 'number' && Number.isSafeInteger(value) && value >= 0; +} + +function readNonNegativeInteger(value: unknown, label: string): number { + if (!isNonNegativeSafeInteger(value)) throw new Error(`Invalid ${label}`); + return value; +} + +function exceedsLimit(current: number, added: number, limit: number): boolean { + return current > limit - added; +} + +function loadDatabaseSync(): typeof import('node:sqlite').DatabaseSync { + return (require('node:sqlite') as typeof import('node:sqlite')).DatabaseSync; +} + +function rollback(database: DatabaseSync): void { + try { + database.exec('ROLLBACK'); + } catch { + // Preserve the operation failure that triggered rollback. + } +}