import { runAgentStateStartupRecovery, runServerStartup } from '../instrumentation'; import { MemoryAgentStateStore } from './agent-state-memory'; import { ensureAgentStateStoreReady, getAgentStateStore, readAgentDatabaseUrl, recoverAgentStateOnStartup, resetAgentStateStoreForTests, setAgentStateStoreFactoryForTests } from './agent-state-runtime'; import type { AgentStateStore } from './agent-state-store'; import type { ImageShareStateStore } from './share-store'; import assert from 'node:assert/strict'; import { afterEach, describe, it } from 'node:test'; import { setTimeout as delay } from 'node:timers/promises'; afterEach(async () => { setAgentStateStoreFactoryForTests(undefined); await resetAgentStateStoreForTests(); }); describe('agent-state-runtime recovery scheduling', () => { it('creates a memory store for ephemeral deployments', () => { const store = getAgentStateStore({ AGENT_STATE_BACKEND: 'memory' }); assert.ok(store instanceof MemoryAgentStateStore); }); it('closes a cached store before clearing test state', async () => { const store = createFakeStore(); let closeCalls = 0; store.close = async () => { closeCalls += 1; }; setAgentStateStoreFactoryForTests(() => store); getAgentStateStore({ AGENT_STATE_BACKEND: 'sqlite', AGENT_SQLITE_PATH: 'agent.sqlite' }); await resetAgentStateStoreForTests(); assert.equal(closeCalls, 1); }); it('waits for initialization to finish before closing test state', async () => { const store = createFakeStore(); let releaseInitialization: (() => void) | undefined; const initialization = new Promise((resolve) => { releaseInitialization = resolve; }); let closeCalls = 0; store.init = async () => { store.initCalls += 1; await initialization; }; store.close = async () => { assert.equal(store.initCalls, 1); closeCalls += 1; }; setAgentStateStoreFactoryForTests(() => store); getAgentStateStore({ AGENT_STATE_BACKEND: 'sqlite', AGENT_SQLITE_PATH: 'agent.sqlite' }); const reset = resetAgentStateStoreForTests(); await delay(10); assert.equal(closeCalls, 0); releaseInitialization?.(); await reset; assert.equal(closeCalls, 1); }); it('waits for an active recovery before closing test state', async () => { const store = createFakeStore(); let releaseRecovery: (() => void) | undefined; const recovery = new Promise((resolve) => { releaseRecovery = resolve; }); let recoveryFinished = false; let closeCalls = 0; store.recoverExpiredRequests = async () => { store.recoveryCalls += 1; await recovery; recoveryFinished = true; return 0; }; store.close = async () => { assert.equal(recoveryFinished, true); closeCalls += 1; }; setAgentStateStoreFactoryForTests(() => store); const env = { AGENT_STATE_BACKEND: 'sqlite', AGENT_SQLITE_PATH: 'agent.sqlite', AGENT_RECOVERY_INTERVAL_MS: '1000' }; const ready = ensureAgentStateStoreReady(env, new Date('2026-05-12T00:00:00.000Z')); await delay(10); const reset = resetAgentStateStoreForTests(); await delay(10); assert.equal(closeCalls, 0); releaseRecovery?.(); await ready; await reset; assert.equal(closeCalls, 1); }); it('waits for explicit startup recovery before closing test state', async () => { const store = createFakeStore(); let releaseRecovery: (() => void) | undefined; const recovery = new Promise((resolve) => { releaseRecovery = resolve; }); let recoveryFinished = false; let closeCalls = 0; store.recoverExpiredRequests = async () => { store.recoveryCalls += 1; await recovery; recoveryFinished = true; return 1; }; store.close = async () => { assert.equal(recoveryFinished, true); closeCalls += 1; }; setAgentStateStoreFactoryForTests(() => store); const env = { AGENT_STATE_BACKEND: 'sqlite', AGENT_SQLITE_PATH: 'agent.sqlite' }; const startupRecovery = recoverAgentStateOnStartup(env); await delay(10); const reset = resetAgentStateStoreForTests(); await delay(10); assert.equal(closeCalls, 0); releaseRecovery?.(); await startupRecovery; await reset; assert.equal(closeCalls, 1); }); it('rejects a configuration change instead of abandoning an active store', async () => { const store = createFakeStore(); let closeCalls = 0; store.close = async () => { closeCalls += 1; }; setAgentStateStoreFactoryForTests(() => store); getAgentStateStore({ AGENT_STATE_BACKEND: 'sqlite', AGENT_SQLITE_PATH: 'first.sqlite' }); assert.throws( () => getAgentStateStore({ AGENT_STATE_BACKEND: 'sqlite', AGENT_SQLITE_PATH: 'second.sqlite' }), /configuration cannot change/ ); assert.equal(closeCalls, 0); await resetAgentStateStoreForTests(); assert.equal(closeCalls, 1); }); it('throttles request-time recovery checks by interval', async () => { const store = createFakeStore(); setAgentStateStoreFactoryForTests(() => store); const env = { AGENT_STATE_BACKEND: 'sqlite', AGENT_SQLITE_PATH: 'agent.sqlite', AGENT_RECOVERY_INTERVAL_MS: '1000' }; await ensureAgentStateStoreReady(env, new Date('2026-05-12T00:00:00.000Z')); await ensureAgentStateStoreReady(env, new Date('2026-05-12T00:00:00.500Z')); await ensureAgentStateStoreReady(env, new Date('2026-05-12T00:00:01.001Z')); assert.equal(store.recoveryCalls, 2); }); it('runs share cleanup with the request-time recovery cycle', async () => { const store = createFakeStore(); setAgentStateStoreFactoryForTests(() => store); const env = { AGENT_STATE_BACKEND: 'sqlite', AGENT_SQLITE_PATH: 'agent.sqlite', AGENT_RECOVERY_INTERVAL_MS: '1000' }; await ensureAgentStateStoreReady(env, new Date('2026-05-12T00:00:00.000Z')); await ensureAgentStateStoreReady(env, new Date('2026-05-12T00:00:00.500Z')); await ensureAgentStateStoreReady(env, new Date('2026-05-12T00:00:01.001Z')); assert.equal(store.shareCleanupCalls, 2); }); it('always runs explicit startup recovery', async () => { const store = createFakeStore(); setAgentStateStoreFactoryForTests(() => store); await recoverAgentStateOnStartup({ AGENT_STATE_BACKEND: 'sqlite', AGENT_SQLITE_PATH: 'agent.sqlite' }); await recoverAgentStateOnStartup({ AGENT_STATE_BACKEND: 'sqlite', AGENT_SQLITE_PATH: 'agent.sqlite' }); assert.equal(store.recoveryCalls, 2); assert.equal(store.shareCleanupCalls, 2); }); it('allows the next request to retry recovery after a failed recovery attempt', async () => { const store = createFakeStore({ failFirstRecovery: true }); setAgentStateStoreFactoryForTests(() => store); const env = { AGENT_STATE_BACKEND: 'sqlite', AGENT_SQLITE_PATH: 'agent.sqlite', AGENT_RECOVERY_INTERVAL_MS: '1000' }; await assert.rejects( () => ensureAgentStateStoreReady(env, new Date('2026-05-12T00:00:00.000Z')), /recovery failed/ ); await ensureAgentStateStoreReady(env, new Date('2026-05-12T00:00:00.100Z')); assert.equal(store.recoveryCalls, 2); }); it('clears a failed store init so the next request can retry after the environment recovers', async () => { let shouldFailInit = true; const store = createFakeStore({ failInit: () => shouldFailInit }); let closeCalls = 0; store.close = async () => { closeCalls += 1; }; setAgentStateStoreFactoryForTests(() => store); const env = { AGENT_STATE_BACKEND: 'sqlite', AGENT_SQLITE_PATH: 'agent.sqlite' }; await assert.rejects( () => ensureAgentStateStoreReady(env, new Date('2026-05-12T00:00:00.000Z')), /init failed/ ); assert.equal(closeCalls, 1); shouldFailInit = false; await ensureAgentStateStoreReady(env, new Date('2026-05-12T00:00:00.100Z')); assert.equal(store.initCalls, 2); }); }); describe('runAgentStateStartupRecovery', () => { it('fails startup when agent state recovery fails', async () => { const logs: Array<{ level: 'info' | 'error'; message: string }> = []; await assert.rejects( () => runAgentStateStartupRecovery({ recoverAgentStateOnStartup: async () => { throw new Error('startup recovery failed'); }, appLogger: { info(message) { logs.push({ level: 'info', message }); }, error(message) { logs.push({ level: 'error', message }); } } }), /startup recovery failed/ ); assert.deepEqual( logs.map((entry) => entry.level), ['info', 'error'] ); }); it('logs startup recovery completion', async () => { const logs: Array<{ level: 'info' | 'error'; message: string; context?: unknown }> = []; await runAgentStateStartupRecovery({ recoverAgentStateOnStartup: async () => 3, appLogger: { info(message, context) { logs.push({ level: 'info', message, context }); }, error(message, context) { logs.push({ level: 'error', message, context }); } } }); assert.equal(logs.length, 2); assert.equal(logs[0]?.message, '开始执行 Agent 状态启动恢复。'); assert.equal(logs[1]?.message, 'Agent 状态启动恢复完成。'); }); }); describe('runServerStartup', () => { it('starts WebUI cleanup after Agent state recovery without blocking server startup', async () => { const events: string[] = []; let cleanupStarted = false; let releaseCleanup: (() => void) | undefined; const cleanup = new Promise((resolve) => { releaseCleanup = resolve; }); let startupSettled = false; const startup = runServerStartup({ recoverAgentStateOnStartup: async () => { events.push('agent-recovery'); return 0; }, startWebuiImageCleanupScheduler: async () => { cleanupStarted = true; events.push('webui-cleanup-start'); await cleanup; }, appLogger: { info() {}, error() {} } }); void startup.then(() => { startupSettled = true; }); try { await waitFor(() => cleanupStarted); assert.equal(startupSettled, true); assert.deepEqual(events, ['agent-recovery', 'webui-cleanup-start']); } finally { releaseCleanup?.(); await startup; } }); it('logs WebUI cleanup startup failures without rejecting server startup', async () => { const logs: Array<{ level: 'info' | 'error'; message: string; context?: unknown }> = []; let resolveFailureLogged: (() => void) | undefined; const failureLogged = new Promise((resolve) => { resolveFailureLogged = resolve; }); await runServerStartup({ recoverAgentStateOnStartup: async () => 0, startWebuiImageCleanupScheduler: async () => { throw new Error('cleanup startup failed'); }, appLogger: { info(message, context) { logs.push({ level: 'info', message, context }); }, error(message, context) { logs.push({ level: 'error', message, context }); resolveFailureLogged?.(); } } }); await failureLogged; assert.equal(logs.at(-1)?.level, 'error'); assert.equal(logs.at(-1)?.message, 'WebUI 图片自动清理启动失败。'); assert.ok(logs.at(-1)?.context instanceof Error); assert.match((logs.at(-1)?.context as Error).message, /cleanup startup failed/); }); }); async function waitFor(predicate: () => boolean, timeoutMs = 1000): Promise { const startedAt = Date.now(); while (!predicate()) { if (Date.now() - startedAt > timeoutMs) throw new Error('condition not met before timeout'); await delay(5); } } describe('readAgentDatabaseUrl', () => { const DB_PASSWORD_FIXTURE = ['database', 'password'].join(' '); const ENCODED_DB_PASSWORD_FIXTURE = encodeURIComponent(DB_PASSWORD_FIXTURE); const EXPLICIT_DATABASE_URL_FIXTURE = `postgres://gpt_image:${ENCODED_DB_PASSWORD_FIXTURE}@postgres:5432/gpt_image_playground`; it('prefers an explicit AGENT_DATABASE_URL', () => { assert.equal( readAgentDatabaseUrl({ AGENT_DATABASE_URL: EXPLICIT_DATABASE_URL_FIXTURE }), EXPLICIT_DATABASE_URL_FIXTURE ); }); it('falls back to split PostgreSQL fields when AGENT_DATABASE_URL is blank', () => { assert.equal( readAgentDatabaseUrl({ AGENT_DATABASE_URL: ' ', AGENT_DB_HOST: 'postgres', AGENT_DB_PORT: '5432', AGENT_DB_NAME: 'gpt_image_playground', AGENT_DB_USER: 'gpt_image', AGENT_DB_PASSWORD: DB_PASSWORD_FIXTURE }), `postgres://gpt_image:${ENCODED_DB_PASSWORD_FIXTURE}@postgres:5432/gpt_image_playground` ); }); it('builds a PostgreSQL URL from individual environment fields', () => { assert.equal( readAgentDatabaseUrl({ AGENT_DB_HOST: 'postgres', AGENT_DB_PORT: '5432', AGENT_DB_NAME: 'gpt_image_playground', AGENT_DB_USER: 'gpt_image', AGENT_DB_PASSWORD: DB_PASSWORD_FIXTURE }), `postgres://gpt_image:${ENCODED_DB_PASSWORD_FIXTURE}@postgres:5432/gpt_image_playground` ); }); it('escapes split PostgreSQL user, password, and database fields', () => { const databaseName = ['gpt', 'image playground'].join('/'); const databaseUser = ['gpt', 'image'].join('@'); const databaseCredential = ['p', 'ss/word:?#'].join('@'); const url = readAgentDatabaseUrl({ AGENT_DB_HOST: 'postgres', AGENT_DB_PORT: '5432', AGENT_DB_NAME: databaseName, AGENT_DB_USER: databaseUser, AGENT_DB_PASSWORD: databaseCredential }); assert.equal( url, `postgres://${encodeURIComponent(databaseUser)}:${encodeURIComponent(databaseCredential)}@postgres:5432/${encodeURIComponent(databaseName)}` ); }); }); function createFakeStore(options: { failFirstRecovery?: boolean; failInit?: () => boolean } = {}): AgentStateStore & ImageShareStateStore & { recoveryCalls: number; initCalls: number; shareCleanupCalls: number; } { return { initCalls: 0, recoveryCalls: 0, shareCleanupCalls: 0, async init() { this.initCalls += 1; if (options.failInit?.()) { throw new Error('init failed'); } }, async recoverExpiredRequests() { this.recoveryCalls += 1; if (options.failFirstRecovery && this.recoveryCalls === 1) { throw new Error('recovery failed'); } return 0; }, async purgeExpiredRequests() { return 0; }, async beginRequest() { throw new Error('not implemented'); }, async refreshRequestLease() { return false; }, async saveArtifacts() {}, async completeRequest() {}, async failRequest() {}, async getArtifact() { return undefined; }, async getRequest() { return undefined; }, async getRequestByIdempotencyKey() { return undefined; }, async listArtifactsForRequest() { return []; }, async listArtifactFilepaths() { return []; }, async deleteArtifact() { return false; }, async createImageShareRecord() {}, async readImageShareRecord() { return undefined; }, async deleteExpiredImageShareRecords() { this.shareCleanupCalls += 1; return []; }, async listImageShareRecords() { return []; } }; }