Spaces:
Running
Running
File size: 6,070 Bytes
c7d34c1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 | import path from 'path';
import { readFileSync } from 'fs';
import {
readAgentRecoveryIntervalMs,
readAgentStateBackend,
readAgentSqlitePath,
type AgentStateBackend
} from './agent-api-contracts';
import { MemoryAgentStateStore } from './agent-state-memory';
import { PostgresAgentStateStore } from './agent-state-postgres';
import { SqliteAgentStateStore } from './agent-state-sqlite';
import type { AgentStateStore } from './agent-state-store';
import { purgeExpiredImageSharesForStore } from './share-store';
type CachedStore = {
backend: AgentStateBackend;
key: string;
store: AgentStateStore;
initPromise: Promise<void>;
lastRecoveryAtMs?: number;
recoveryPromise?: Promise<number>;
};
let cachedStore: CachedStore | undefined;
let storeFactoryForTests: ((backend: AgentStateBackend, key: string, env: Record<string, string | undefined>) => AgentStateStore) | undefined;
function cacheAgentStateStore(backend: AgentStateBackend, key: string, store: AgentStateStore): AgentStateStore {
const initPromise = store.init().catch((error) => {
if (cachedStore?.store === store && cachedStore.initPromise === initPromise) {
cachedStore = undefined;
}
throw error;
});
cachedStore = { backend, key, store, initPromise };
return store;
}
function readEnvValue(env: Record<string, string | undefined>, fieldName: string): string | undefined {
const value = env[fieldName]?.trim();
return value ? value : undefined;
}
function readEnvSecret(env: Record<string, string | undefined>, fieldName: string, fileFieldName: string): string | undefined {
const directValue = readEnvValue(env, fieldName);
if (directValue) return directValue;
const filePath = readEnvValue(env, fileFieldName);
if (!filePath) return undefined;
return readFileSync(filePath, 'utf8').trim() || undefined;
}
export function readAgentDatabaseUrl(env: Record<string, string | undefined> = process.env): string | undefined {
const configuredUrl = readEnvValue(env, 'AGENT_DATABASE_URL');
if (configuredUrl) return configuredUrl;
const password = readEnvSecret(env, 'AGENT_DB_PASSWORD', 'AGENT_DB_PASSWORD_FILE');
if (!password) return undefined;
const host = readEnvValue(env, 'AGENT_DB_HOST') || 'localhost';
const port = readEnvValue(env, 'AGENT_DB_PORT') || '5432';
const database = readEnvValue(env, 'AGENT_DB_NAME') || 'gpt_image_playground';
const user = readEnvValue(env, 'AGENT_DB_USER') || 'gpt_image';
return `postgres://${encodeURIComponent(user)}:${encodeURIComponent(password)}@${host}:${port}/${encodeURIComponent(database)}`;
}
export function resetAgentStateStoreForTests(): void {
cachedStore = undefined;
}
export function setAgentStateStoreFactoryForTests(
factory: ((backend: AgentStateBackend, key: string, env: Record<string, string | undefined>) => AgentStateStore) | undefined
): void {
storeFactoryForTests = factory;
}
export function getAgentStateStore(env: Record<string, string | undefined> = process.env): AgentStateStore {
const backend = readAgentStateBackend(env);
const databaseUrl = backend === 'postgres' ? readAgentDatabaseUrl(env) : undefined;
const key =
backend === 'postgres'
? databaseUrl || ''
: backend === 'memory'
? 'memory'
: path.resolve(/* turbopackIgnore: true */ process.cwd(), readAgentSqlitePath(env));
if (cachedStore && cachedStore.backend === backend && cachedStore.key === key) {
return cachedStore.store;
}
if (storeFactoryForTests) {
return cacheAgentStateStore(backend, key, storeFactoryForTests(backend, key, env));
}
if (backend === 'postgres') {
if (!databaseUrl) {
throw new Error('AGENT_STATE_BACKEND=postgres 时必须设置 AGENT_DATABASE_URL 或 AGENT_DB_PASSWORD。');
}
return cacheAgentStateStore(backend, key, new PostgresAgentStateStore(databaseUrl));
}
if (backend === 'memory') {
return cacheAgentStateStore(backend, key, new MemoryAgentStateStore());
}
return cacheAgentStateStore(backend, key, new SqliteAgentStateStore(key));
}
export async function ensureAgentStateStoreReady(
env: Record<string, string | undefined> = process.env,
now = new Date()
): Promise<AgentStateStore> {
const store = getAgentStateStore(env);
await cachedStore?.initPromise;
await recoverAgentStateIfDue(store, env, now);
return store;
}
export async function recoverAgentStateOnStartup(env: Record<string, string | undefined> = process.env): Promise<number> {
const store = getAgentStateStore(env);
await cachedStore?.initPromise;
const recovered = await store.recoverExpiredRequests();
await store.purgeExpiredRequests();
await purgeExpiredImageSharesForStore(store, new Date(), { purgeOrphanFiles: false });
if (cachedStore) {
cachedStore.lastRecoveryAtMs = Date.now();
}
return recovered;
}
async function recoverAgentStateIfDue(store: AgentStateStore, env: Record<string, string | undefined>, now: Date): Promise<void> {
if (!cachedStore) return;
const nowMs = now.getTime();
const intervalMs = readAgentRecoveryIntervalMs(env);
if (cachedStore.recoveryPromise) {
await cachedStore.recoveryPromise;
return;
}
if (cachedStore.lastRecoveryAtMs !== undefined && nowMs - cachedStore.lastRecoveryAtMs < intervalMs) {
return;
}
cachedStore.recoveryPromise = (async () => {
try {
await store.recoverExpiredRequests(now);
await store.purgeExpiredRequests(now);
await purgeExpiredImageSharesForStore(store, now, { purgeOrphanFiles: false });
if (cachedStore) {
cachedStore.lastRecoveryAtMs = nowMs;
}
return 0;
} finally {
if (cachedStore) {
cachedStore.recoveryPromise = undefined;
}
}
})();
await cachedStore.recoveryPromise;
}
|