import type { AgentImageResponse } from './agent-api-contracts'; import { discardArtifactFiles, isArtifactFilepathAllowed, moveArtifactFilesForDeletion, restoreArtifactFiles, type MovedFileForDeletion } from './agent-file-utils'; import { addMilliseconds, addSeconds, buildRecoveredResponse, computeRetryAfterSeconds, createRequestId, isoDate, type AgentArtifactRecord, type AgentRequestRecord, type AgentStateStore, type BeginAgentRequestInput, type BeginAgentRequestResult, type CompleteAgentRequestInput, type FailAgentRequestInput } from './agent-state-store'; import type { AgentErrorBody } from './api-error-response'; import type { FeedbackDeleteOptions, FeedbackRecord, FeedbackStateStore, FeedbackSource, FeedbackTarget, FeedbackTargetType, FeedbackValue } from './feedback-store'; import type { ImageShareRecord, ImageShareStateStore } from './share-store'; import crypto from 'crypto'; import { Pool, type PoolClient } from 'pg'; type PostgresRequestRow = { request_id: string; idempotency_key: string; request_hash: string; mode: 'generate' | 'edit'; status: AgentRequestRecord['status']; request_json: unknown; response_json: unknown | null; error_json: unknown | null; locked_until: Date | string | null; created_at: Date | string; updated_at: Date | string; expires_at: Date | string; }; type PostgresArtifactRow = { id: string; request_id: string; filename: string; filepath: string; content_url: string; metadata_url: string; output_format: string; mime_type: string; size_bytes: number | string; width: number | null; height: number | null; model: string; prompt_hash: string; created_at: Date | string; }; type PostgresShareRow = { token: string; source_filename: string; content_filename: string; mime_type: string; size_bytes: number | string; created_at: Date | string; access_code_required: boolean; expires_at: Date | string | null; access_code_salt: string | null; access_code_hash: string | null; }; type PostgresFeedbackRow = { target_type: FeedbackTargetType; target_id: string; value: FeedbackValue; note: string | null; source: FeedbackSource; updated_at: Date | string; }; type PostgresMigrationRow = { id: string; checksum: string | null; }; export class PostgresAgentStateStore implements AgentStateStore, ImageShareStateStore, FeedbackStateStore { private readonly pool: Pool; constructor(connectionString: string) { this.pool = new Pool({ connectionString }); } async init(): Promise { await runPostgresMigrations(this.pool); } async close(): Promise { await this.pool.end(); } async recoverExpiredRequests(now = new Date()): Promise { const client = await this.pool.connect(); const nowIso = isoDate(now); try { await client.query('BEGIN'); const expiredResult = await client.query( "SELECT * FROM agent_requests WHERE status = 'running' AND locked_until IS NOT NULL AND locked_until < $1 FOR UPDATE SKIP LOCKED", [nowIso] ); for (const row of expiredResult.rows as PostgresRequestRow[]) { const record = this.mapRequestRow(row); const artifacts = await this.listArtifactsForRequestInTransaction(client, record.requestId); if (artifacts.length > 0) { await client.query( "UPDATE agent_requests SET status = 'succeeded', response_json = $1, error_json = NULL, locked_until = NULL, updated_at = $2 WHERE request_id = $3", [buildRecoveredResponse(record, artifacts), nowIso, record.requestId] ); } else if (record.errorJson) { await client.query( "UPDATE agent_requests SET status = 'failed', locked_until = NULL, updated_at = $1 WHERE request_id = $2", [nowIso, record.requestId] ); } else { await client.query( "UPDATE agent_requests SET status = 'orphaned', locked_until = NULL, updated_at = $1 WHERE request_id = $2", [nowIso, record.requestId] ); } } if (expiredResult.rowCount && expiredResult.rowCount > 0) { await client.query( 'INSERT INTO agent_recovery_events (id, event_type, details_json, created_at) VALUES ($1, $2, $3, $4)', [crypto.randomUUID(), 'expired_running_requests', { count: expiredResult.rowCount }, nowIso] ); } await client.query('COMMIT'); return expiredResult.rowCount ?? 0; } catch (error) { await client.query('ROLLBACK'); throw error; } finally { client.release(); } } async purgeExpiredRequests(now = new Date()): Promise { const client = await this.pool.connect(); const nowIso = isoDate(now); let transactionStarted = false; let movedFiles: MovedFileForDeletion[] = []; try { await client.query('BEGIN'); transactionStarted = true; const expiredResult = await client.query( "SELECT request_id FROM agent_requests WHERE expires_at < $1 AND status IN ('succeeded', 'failed', 'orphaned') FOR UPDATE SKIP LOCKED", [nowIso] ); const requestIds = expiredResult.rows.map((row: { request_id: string }) => row.request_id); const artifactRows: Array<{ id: string; filepath: string | null }> = requestIds.length > 0 ? ( await client.query('SELECT id, filepath FROM agent_artifacts WHERE request_id = ANY($1)', [ requestIds ]) ).rows : []; const artifactIds = artifactRows.map((row) => row.id); const artifactFilepaths = artifactRows .map((row) => row.filepath) .filter( (filepath): filepath is string => typeof filepath === 'string' && isArtifactFilepathAllowed(filepath) ); movedFiles = await moveArtifactFilesForDeletion([...new Set(artifactFilepaths)]); if (requestIds.length > 0) { if (artifactIds.length > 0) { await client.query( "DELETE FROM result_feedback WHERE target_type = 'agent_artifact' AND target_id = ANY($1)", [artifactIds] ); } await client.query( "DELETE FROM result_feedback WHERE target_type = 'agent_request' AND target_id = ANY($1)", [requestIds] ); await client.query('DELETE FROM agent_artifacts WHERE request_id = ANY($1)', [requestIds]); await client.query('DELETE FROM agent_requests WHERE request_id = ANY($1)', [requestIds]); } await client.query('COMMIT'); transactionStarted = false; await discardArtifactFiles(movedFiles); return requestIds.length; } catch (error) { if (transactionStarted) { await client.query('ROLLBACK').catch(() => {}); } await restoreArtifactFiles(movedFiles); throw error; } finally { client.release(); } } async beginRequest(input: BeginAgentRequestInput): Promise { const client = await this.pool.connect(); try { await client.query('BEGIN'); const result = await this.beginRequestInTransaction(client, input); await client.query('COMMIT'); return result; } catch (error) { await client.query('ROLLBACK'); throw error; } finally { client.release(); } } async refreshRequestLease(input: { requestId: string; leaseMs: number; now?: Date }): Promise { const now = input.now ?? new Date(); const result = await this.pool.query( "UPDATE agent_requests SET locked_until = $1, updated_at = $2 WHERE request_id = $3 AND status IN ('running', 'pending')", [isoDate(addMilliseconds(now, input.leaseMs)), isoDate(now), input.requestId] ); return (result.rowCount ?? 0) > 0; } async completeRequest(input: CompleteAgentRequestInput): Promise { const client = await this.pool.connect(); const nowIso = isoDate(input.now ?? new Date()); try { await client.query('BEGIN'); await this.insertArtifacts(client, input.artifacts); await client.query( "UPDATE agent_requests SET status = 'succeeded', response_json = $1, error_json = NULL, locked_until = NULL, updated_at = $2 WHERE request_id = $3", [input.response, nowIso, input.requestId] ); await client.query('COMMIT'); } catch (error) { await client.query('ROLLBACK'); throw error; } finally { client.release(); } } async saveArtifacts(artifacts: AgentArtifactRecord[]): Promise { const client = await this.pool.connect(); try { await client.query('BEGIN'); await this.insertArtifacts(client, artifacts); await client.query('COMMIT'); } catch (error) { await client.query('ROLLBACK'); throw error; } finally { client.release(); } } async failRequest(input: FailAgentRequestInput): Promise { await this.pool.query( "UPDATE agent_requests SET status = 'failed', response_json = NULL, error_json = $1, locked_until = NULL, updated_at = $2 WHERE request_id = $3", [input.error, isoDate(input.now ?? new Date()), input.requestId] ); } async getRequest(requestId: string): Promise { const result = await this.pool.query('SELECT * FROM agent_requests WHERE request_id = $1', [requestId]); const row = result.rows[0] as PostgresRequestRow | undefined; return row ? this.mapRequestRow(row) : undefined; } async getRequestByIdempotencyKey(idempotencyKey: string): Promise { const result = await this.pool.query('SELECT * FROM agent_requests WHERE idempotency_key = $1', [ idempotencyKey ]); const row = result.rows[0] as PostgresRequestRow | undefined; return row ? this.mapRequestRow(row) : undefined; } async getArtifact(id: string): Promise { const result = await this.pool.query('SELECT * FROM agent_artifacts WHERE id = $1', [id]); const row = result.rows[0] as PostgresArtifactRow | undefined; return row ? this.mapArtifactRow(row) : undefined; } async listArtifactsForRequest(requestId: string): Promise { const result = await this.pool.query( 'SELECT * FROM agent_artifacts WHERE request_id = $1 ORDER BY created_at ASC', [requestId] ); return (result.rows as PostgresArtifactRow[]).map((row) => this.mapArtifactRow(row)); } async listArtifactFilepaths(): Promise { const result = await this.pool.query('SELECT DISTINCT filepath FROM agent_artifacts ORDER BY filepath ASC'); return (result.rows as Array<{ filepath: string }>).map((row) => row.filepath); } async deleteArtifact(id: string): Promise { const client = await this.pool.connect(); try { await client.query('BEGIN'); const result = await client.query('DELETE FROM agent_artifacts WHERE id = $1', [id]); if ((result.rowCount ?? 0) > 0) { await client.query( "DELETE FROM result_feedback WHERE target_type = 'agent_artifact' AND target_id = $1", [id] ); } await client.query('COMMIT'); return (result.rowCount ?? 0) > 0; } catch (error) { await client.query('ROLLBACK').catch(() => {}); throw error; } finally { client.release(); } } async upsertFeedback(record: FeedbackRecord): Promise { await this.pool.query( `INSERT INTO result_feedback (target_type, target_id, value, note, source, updated_at) VALUES ($1, $2, $3, $4, $5, $6) ON CONFLICT (target_type, target_id) DO UPDATE SET value = EXCLUDED.value, note = EXCLUDED.note, source = EXCLUDED.source, updated_at = EXCLUDED.updated_at WHERE result_feedback.updated_at <= EXCLUDED.updated_at`, [record.targetType, record.targetId, record.value, record.note ?? null, record.source, record.updatedAt] ); } async upsertFeedbackBatch(records: FeedbackRecord[]): Promise { if (records.length === 0) return; const client = await this.pool.connect(); try { await client.query('BEGIN'); for (const record of records) { await client.query( `INSERT INTO result_feedback (target_type, target_id, value, note, source, updated_at) VALUES ($1, $2, $3, $4, $5, $6) ON CONFLICT (target_type, target_id) DO UPDATE SET value = EXCLUDED.value, note = EXCLUDED.note, source = EXCLUDED.source, updated_at = EXCLUDED.updated_at WHERE result_feedback.updated_at <= EXCLUDED.updated_at`, [ record.targetType, record.targetId, record.value, record.note ?? null, record.source, record.updatedAt ] ); } await client.query('COMMIT'); } catch (error) { await client.query('ROLLBACK').catch(() => {}); throw error; } finally { client.release(); } } async readFeedback(targetType: FeedbackTargetType, targetId: string): Promise { const result = await this.pool.query( 'SELECT * FROM result_feedback WHERE target_type = $1 AND target_id = $2', [targetType, targetId] ); const row = result.rows[0] as PostgresFeedbackRow | undefined; return row ? this.mapFeedbackRow(row) : undefined; } async listFeedbackByTargets(targets: FeedbackTarget[]): Promise { if (targets.length === 0) return []; const targetTypes = targets.map((target) => target.targetType); const targetIds = targets.map((target) => target.targetId); const result = await this.pool.query( `SELECT feedback.* FROM result_feedback AS feedback JOIN UNNEST($1::text[], $2::text[]) WITH ORDINALITY AS target(target_type, target_id, ord) ON feedback.target_type = target.target_type AND feedback.target_id = target.target_id ORDER BY target.ord`, [targetTypes, targetIds] ); return result.rows.map((row) => this.mapFeedbackRow(row as PostgresFeedbackRow)); } async deleteFeedbackByTargets(targets: FeedbackTarget[], options: FeedbackDeleteOptions = {}): Promise { if (targets.length === 0) return 0; const targetTypes = targets.map((target) => target.targetType); const targetIds = targets.map((target) => target.targetId); const result = await this.pool.query( `DELETE FROM result_feedback AS feedback USING UNNEST($1::text[], $2::text[]) AS target(target_type, target_id) WHERE feedback.target_type = target.target_type AND feedback.target_id = target.target_id AND ($3::timestamptz IS NULL OR feedback.updated_at <= $3)`, [targetTypes, targetIds, options.deletedAt ?? null] ); return result.rowCount ?? 0; } async createImageShareRecord(record: ImageShareRecord): Promise { await this.pool.query( `INSERT INTO image_shares (token, source_filename, content_filename, mime_type, size_bytes, created_at, access_code_required, expires_at, access_code_salt, access_code_hash) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)`, [ record.token, record.sourceFilename, record.contentFilename, record.mimeType, record.sizeBytes, record.createdAt, record.accessCodeRequired, record.expiresAt ?? null, record.accessCodeSalt ?? null, record.accessCodeHash ?? null ] ); } async readImageShareRecord(token: string): Promise { const result = await this.pool.query('SELECT * FROM image_shares WHERE token = $1', [token]); const row = result.rows[0] as PostgresShareRow | undefined; return row ? this.mapShareRow(row) : undefined; } async deleteExpiredImageShareRecords(nowIso: string): Promise { const client = await this.pool.connect(); try { await client.query('BEGIN'); const result = await client.query( 'DELETE FROM image_shares WHERE expires_at IS NOT NULL AND expires_at < $1 RETURNING *', [nowIso] ); await client.query('COMMIT'); return (result.rows as PostgresShareRow[]).map((row) => this.mapShareRow(row)); } catch (error) { await client.query('ROLLBACK').catch(() => {}); throw error; } finally { client.release(); } } async listImageShareRecords(): Promise { const result = await this.pool.query('SELECT * FROM image_shares ORDER BY created_at ASC, token ASC'); return (result.rows as PostgresShareRow[]).map((row) => this.mapShareRow(row)); } private async beginRequestInTransaction( client: PoolClient, input: BeginAgentRequestInput ): Promise { const now = input.now ?? new Date(); const nowIso = isoDate(now); const lockedUntil = isoDate(addMilliseconds(now, input.leaseMs)); const expiresAt = isoDate(addSeconds(now, input.ttlSeconds)); const existingResult = await client.query( 'SELECT * FROM agent_requests WHERE idempotency_key = $1 FOR UPDATE', [input.idempotencyKey] ); const existing = existingResult.rows[0] as PostgresRequestRow | undefined; if (!existing) { const requestId = createRequestId(); const insertResult = await client.query( `INSERT INTO agent_requests (request_id, idempotency_key, request_hash, mode, status, request_json, locked_until, created_at, updated_at, expires_at) VALUES ($1, $2, $3, $4, 'running', $5, $6, $7, $8, $9) ON CONFLICT (idempotency_key) DO NOTHING RETURNING *`, [ requestId, input.idempotencyKey, input.requestHash, input.mode, input.requestJson, lockedUntil, nowIso, nowIso, expiresAt ] ); if (insertResult.rowCount && insertResult.rowCount > 0) { return { type: 'acquired', record: this.mapRequestRow(insertResult.rows[0] as PostgresRequestRow) }; } const conflicted = await client.query( 'SELECT * FROM agent_requests WHERE idempotency_key = $1 FOR UPDATE', [input.idempotencyKey] ); if (!conflicted.rows[0]) { throw new Error('idempotency conflict row disappeared during acquisition'); } return this.beginFromExistingRow( conflicted.rows[0] as PostgresRequestRow, input, now, nowIso, lockedUntil, expiresAt, client ); } return this.beginFromExistingRow(existing, input, now, nowIso, lockedUntil, expiresAt, client); } private async beginFromExistingRow( existing: PostgresRequestRow, input: BeginAgentRequestInput, now: Date, nowIso: string, lockedUntil: string, expiresAt: string, client: PoolClient ): Promise { const record = this.mapRequestRow(existing); if (existing.request_hash !== input.requestHash) { return { type: 'conflict', record }; } if (existing.status === 'succeeded' && existing.response_json) { return { type: 'replay', record, response: existing.response_json as AgentImageResponse }; } if (existing.status === 'failed' && existing.error_json) { return { type: 'failed', record, error: existing.error_json as AgentErrorBody }; } const lockedUntilIso = toIso(existing.locked_until); if ( (existing.status === 'running' || existing.status === 'pending') && lockedUntilIso && lockedUntilIso > nowIso ) { return { type: 'in_progress', record, retryAfterSeconds: computeRetryAfterSeconds(lockedUntilIso, now) }; } await client.query( "UPDATE agent_requests SET status = 'running', locked_until = $1, updated_at = $2, expires_at = $3 WHERE idempotency_key = $4", [lockedUntil, nowIso, expiresAt, input.idempotencyKey] ); const updated = await client.query('SELECT * FROM agent_requests WHERE idempotency_key = $1', [ input.idempotencyKey ]); return { type: 'acquired', record: this.mapRequestRow(updated.rows[0] as PostgresRequestRow) }; } private mapRequestRow(row: PostgresRequestRow): AgentRequestRecord { return { requestId: row.request_id, idempotencyKey: row.idempotency_key, requestHash: row.request_hash, mode: row.mode, status: row.status, requestJson: row.request_json, responseJson: row.response_json ? (row.response_json as AgentImageResponse) : undefined, errorJson: row.error_json ? (row.error_json as AgentErrorBody) : undefined, ...(row.locked_until ? { lockedUntil: toIso(row.locked_until) } : {}), createdAt: toIso(row.created_at), updatedAt: toIso(row.updated_at), expiresAt: toIso(row.expires_at) }; } private mapArtifactRow(row: PostgresArtifactRow): AgentArtifactRecord { return { id: row.id, requestId: row.request_id, filename: row.filename, filepath: row.filepath, contentUrl: row.content_url, metadataUrl: row.metadata_url, outputFormat: row.output_format, mimeType: row.mime_type, sizeBytes: Number(row.size_bytes), width: row.width, height: row.height, model: row.model, promptHash: row.prompt_hash, createdAt: toIso(row.created_at) }; } private mapShareRow(row: PostgresShareRow): ImageShareRecord { return { token: row.token, sourceFilename: row.source_filename, contentFilename: row.content_filename, mimeType: row.mime_type, sizeBytes: Number(row.size_bytes), createdAt: toIso(row.created_at), accessCodeRequired: row.access_code_required, ...(row.expires_at ? { expiresAt: toIso(row.expires_at) } : {}), ...(row.access_code_salt ? { accessCodeSalt: row.access_code_salt } : {}), ...(row.access_code_hash ? { accessCodeHash: row.access_code_hash } : {}) }; } private mapFeedbackRow(row: PostgresFeedbackRow): FeedbackRecord { return { targetType: row.target_type, targetId: row.target_id, value: row.value, source: row.source, updatedAt: toIso(row.updated_at), ...(row.note ? { note: row.note } : {}) }; } private async insertArtifacts(client: PoolClient, artifacts: AgentArtifactRecord[]): Promise { for (const artifact of artifacts) { await this.assertArtifactCanBeInserted(client, artifact); await client.query( `INSERT INTO agent_artifacts (id, request_id, filename, filepath, content_url, metadata_url, output_format, mime_type, size_bytes, width, height, model, prompt_hash, created_at) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14) ON CONFLICT (id) DO NOTHING`, [ artifact.id, artifact.requestId, artifact.filename, artifact.filepath, artifact.contentUrl, artifact.metadataUrl, artifact.outputFormat, artifact.mimeType, artifact.sizeBytes, artifact.width, artifact.height, artifact.model, artifact.promptHash, artifact.createdAt ] ); } } private async assertArtifactCanBeInserted(client: PoolClient, artifact: AgentArtifactRecord): Promise { const result = await client.query('SELECT * FROM agent_artifacts WHERE id = $1', [artifact.id]); const existing = result.rows[0] as PostgresArtifactRow | undefined; if (existing && !sameArtifactRecord(this.mapArtifactRow(existing), artifact)) { throw new Error('artifact metadata conflict'); } } private async listArtifactsForRequestInTransaction( client: PoolClient, requestId: string ): Promise { const result = await client.query( 'SELECT * FROM agent_artifacts WHERE request_id = $1 ORDER BY created_at ASC', [requestId] ); return (result.rows as PostgresArtifactRow[]).map((row) => this.mapArtifactRow(row)); } } function toIso(value: Date | string | null): string { if (!value) return ''; return value instanceof Date ? value.toISOString() : new Date(value).toISOString(); } type PostgresMigration = { id: string; sql: string; }; const POSTGRES_MIGRATION_TABLE_SCHEMA = ` CREATE TABLE IF NOT EXISTS state_schema_migrations ( id TEXT PRIMARY KEY, checksum TEXT NOT NULL, applied_at TIMESTAMPTZ NOT NULL );`; const POSTGRES_MIGRATIONS: PostgresMigration[] = [ { id: '001_agent_state_core', sql: ` CREATE TABLE IF NOT EXISTS agent_requests ( request_id TEXT PRIMARY KEY, idempotency_key TEXT NOT NULL UNIQUE, request_hash TEXT NOT NULL, mode TEXT NOT NULL CHECK (mode IN ('generate', 'edit')), status TEXT NOT NULL CHECK (status IN ('pending', 'running', 'succeeded', 'failed', 'orphaned')), request_json JSONB NOT NULL, response_json JSONB, error_json JSONB, locked_until TIMESTAMPTZ, created_at TIMESTAMPTZ NOT NULL, updated_at TIMESTAMPTZ NOT NULL, expires_at TIMESTAMPTZ NOT NULL ); CREATE INDEX IF NOT EXISTS idx_agent_requests_status_locked_until ON agent_requests(status, locked_until); CREATE INDEX IF NOT EXISTS idx_agent_requests_expires_at ON agent_requests(expires_at); CREATE TABLE IF NOT EXISTS agent_artifacts ( id TEXT PRIMARY KEY, request_id TEXT NOT NULL REFERENCES agent_requests(request_id), filename TEXT NOT NULL UNIQUE, filepath TEXT NOT NULL, content_url TEXT NOT NULL, metadata_url TEXT NOT NULL, output_format TEXT NOT NULL, mime_type TEXT NOT NULL, size_bytes BIGINT NOT NULL, width INTEGER, height INTEGER, model TEXT NOT NULL, prompt_hash TEXT NOT NULL, created_at TIMESTAMPTZ NOT NULL ); CREATE INDEX IF NOT EXISTS idx_agent_artifacts_request_id ON agent_artifacts(request_id); CREATE TABLE IF NOT EXISTS agent_recovery_events ( id TEXT PRIMARY KEY, event_type TEXT NOT NULL, details_json JSONB NOT NULL, created_at TIMESTAMPTZ NOT NULL ); ` }, { id: '002_image_shares', sql: ` CREATE TABLE IF NOT EXISTS image_shares ( token TEXT PRIMARY KEY, source_filename TEXT NOT NULL, content_filename TEXT NOT NULL UNIQUE, mime_type TEXT NOT NULL, size_bytes BIGINT NOT NULL, created_at TIMESTAMPTZ NOT NULL, access_code_required BOOLEAN NOT NULL, expires_at TIMESTAMPTZ, access_code_salt TEXT, access_code_hash TEXT, CHECK ( (access_code_required = FALSE AND access_code_salt IS NULL AND access_code_hash IS NULL) OR (access_code_required = TRUE AND access_code_salt IS NOT NULL AND access_code_hash IS NOT NULL) ) ); CREATE INDEX IF NOT EXISTS idx_image_shares_expires_at ON image_shares(expires_at); ` }, { id: '003_result_feedback', sql: ` CREATE TABLE IF NOT EXISTS result_feedback ( target_type TEXT NOT NULL CHECK (target_type IN ('page_request', 'agent_request', 'agent_artifact')), target_id TEXT NOT NULL, value TEXT NOT NULL CHECK (value IN ('usable', 'needs_revision')), note TEXT, source TEXT NOT NULL CHECK (source IN ('webui', 'agent')), updated_at TIMESTAMPTZ NOT NULL, PRIMARY KEY (target_type, target_id) ); CREATE INDEX IF NOT EXISTS idx_result_feedback_updated_at ON result_feedback(updated_at); ` } ]; async function runPostgresMigrations(pool: Pool): Promise { const client = await pool.connect(); try { await client.query('BEGIN'); await client.query(POSTGRES_MIGRATION_TABLE_SCHEMA); await ensurePostgresMigrationChecksumColumn(client); await client.query('LOCK TABLE state_schema_migrations IN EXCLUSIVE MODE'); const appliedResult = await client.query('SELECT id, checksum FROM state_schema_migrations'); const applied = new Map((appliedResult.rows as PostgresMigrationRow[]).map((row) => [row.id, row.checksum])); for (const migration of POSTGRES_MIGRATIONS) { const checksum = migrationChecksum(migration.sql); if (applied.has(migration.id)) { if (applied.get(migration.id) !== checksum) { throw new Error(`PostgreSQL migration checksum mismatch: ${migration.id}`); } continue; } for (const statement of splitSqlStatements(migration.sql)) { await client.query(statement); } await client.query('INSERT INTO state_schema_migrations (id, checksum, applied_at) VALUES ($1, $2, $3)', [ migration.id, checksum, isoDate(new Date()) ]); } await client.query('COMMIT'); } catch (error) { await client.query('ROLLBACK').catch(() => {}); throw error; } finally { client.release(); } } export const POSTGRES_SCHEMA = [ POSTGRES_MIGRATION_TABLE_SCHEMA, ...POSTGRES_MIGRATIONS.map((migration) => migration.sql) ] .map((sql) => sql.trim()) .join('\n\n'); async function ensurePostgresMigrationChecksumColumn(client: PoolClient): Promise { await client.query('ALTER TABLE state_schema_migrations ADD COLUMN IF NOT EXISTS checksum TEXT'); for (const migration of POSTGRES_MIGRATIONS) { await client.query('UPDATE state_schema_migrations SET checksum = $1 WHERE id = $2 AND checksum IS NULL', [ migrationChecksum(migration.sql), migration.id ]); } await client.query('ALTER TABLE state_schema_migrations ALTER COLUMN checksum SET NOT NULL'); } type SqlStatementSplitState = { statements: string[]; current: string; inSingleQuote: boolean; inDoubleQuote: boolean; inLineComment: boolean; inBlockComment: boolean; singleQuoteBackslashEscapes: boolean; dollarQuoteTag: string | undefined; }; export function splitSqlStatements(sql: string): string[] { const state: SqlStatementSplitState = { statements: [], current: '', inSingleQuote: false, inDoubleQuote: false, inLineComment: false, inBlockComment: false, singleQuoteBackslashEscapes: false, dollarQuoteTag: undefined }; for (let index = 0; index < sql.length; index += 1) { const consumedIndex = consumeSqlCharacter(sql, index, state); if (consumedIndex !== undefined) { index = consumedIndex; } } pushSqlStatement(state); return state.statements; } function consumeSqlCharacter(sql: string, index: number, state: SqlStatementSplitState): number | undefined { const activeSpanIndex = consumeActiveSqlSpan(sql, index, state); if (activeSpanIndex !== undefined) return activeSpanIndex; const startedSpanIndex = startSqlSpan(sql, index, state); if (startedSpanIndex !== undefined) return startedSpanIndex; if (sql[index] === ';') { pushSqlStatement(state); return index; } state.current += sql[index]; return index; } function consumeActiveSqlSpan(sql: string, index: number, state: SqlStatementSplitState): number | undefined { if (state.dollarQuoteTag) return consumeDollarQuote(sql, index, state); if (state.inLineComment) return consumeLineComment(sql, index, state); if (state.inBlockComment) return consumeBlockComment(sql, index, state); if (state.inSingleQuote) return consumeSingleQuote(sql, index, state); if (state.inDoubleQuote) return consumeDoubleQuote(sql, index, state); return undefined; } function startSqlSpan(sql: string, index: number, state: SqlStatementSplitState): number | undefined { const char = sql[index]; const next = sql[index + 1]; if (char === '-' && next === '-') return startTwoCharSpan(state, index, 'inLineComment', char + next); if (char === '/' && next === '*') return startTwoCharSpan(state, index, 'inBlockComment', char + next); if (char === "'") return startSingleQuote(sql, index, state); if (char === '"') return startOneCharSpan(state, index, 'inDoubleQuote', char); return startDollarQuote(sql, index, state); } function consumeDollarQuote(sql: string, index: number, state: SqlStatementSplitState): number { const tag = state.dollarQuoteTag; if (tag && sql.startsWith(tag, index)) { state.current += tag; state.dollarQuoteTag = undefined; return index + tag.length - 1; } state.current += sql[index]; return index; } function consumeLineComment(sql: string, index: number, state: SqlStatementSplitState): number { state.current += sql[index]; if (sql[index] === '\n') state.inLineComment = false; return index; } function consumeBlockComment(sql: string, index: number, state: SqlStatementSplitState): number { state.current += sql[index]; if (sql[index] === '*' && sql[index + 1] === '/') { state.current += sql[index + 1]; state.inBlockComment = false; return index + 1; } return index; } function consumeSingleQuote(sql: string, index: number, state: SqlStatementSplitState): number { state.current += sql[index]; if (state.singleQuoteBackslashEscapes && sql[index] === '\\' && index + 1 < sql.length) { state.current += sql[index + 1]; return index + 1; } if (sql[index] === "'" && sql[index + 1] === "'") { state.current += sql[index + 1]; return index + 1; } if (sql[index] === "'") { state.inSingleQuote = false; state.singleQuoteBackslashEscapes = false; } return index; } function consumeDoubleQuote(sql: string, index: number, state: SqlStatementSplitState): number { state.current += sql[index]; if (sql[index] === '"' && sql[index + 1] === '"') { state.current += sql[index + 1]; return index + 1; } if (sql[index] === '"') state.inDoubleQuote = false; return index; } function startTwoCharSpan( state: SqlStatementSplitState, index: number, key: 'inLineComment' | 'inBlockComment', text: string ): number { state.current += text; state[key] = true; return index + 1; } function startOneCharSpan( state: SqlStatementSplitState, index: number, key: 'inSingleQuote' | 'inDoubleQuote', text: string ): number { state.current += text; state[key] = true; return index; } function startSingleQuote(sql: string, index: number, state: SqlStatementSplitState): number { state.current += sql[index]; state.inSingleQuote = true; state.singleQuoteBackslashEscapes = startsPostgresEscapeString(sql, index); return index; } function startsPostgresEscapeString(sql: string, quoteIndex: number): boolean { const prefix = sql[quoteIndex - 1]; const beforePrefix = sql[quoteIndex - 2]; return (prefix === 'E' || prefix === 'e') && !isSqlIdentifierPart(beforePrefix); } function isSqlIdentifierPart(char: string | undefined): boolean { return Boolean(char && /[A-Za-z0-9_$]/.test(char)); } function startDollarQuote(sql: string, index: number, state: SqlStatementSplitState): number | undefined { const tag = sql[index] === '$' ? sql.slice(index).match(/^\$[A-Za-z_][A-Za-z0-9_]*\$|^\$\$/)?.[0] : undefined; if (!tag) return undefined; state.current += tag; state.dollarQuoteTag = tag; return index + tag.length - 1; } function pushSqlStatement(state: SqlStatementSplitState): void { const statement = state.current.trim(); if (statement) state.statements.push(statement); state.current = ''; } function migrationChecksum(sql: string): string { return crypto.createHash('sha256').update(sql.trim()).digest('hex'); } function sameArtifactRecord(left: AgentArtifactRecord, right: AgentArtifactRecord): boolean { return JSON.stringify(left) === JSON.stringify(right); }