xiaozuoassistant 0.2.7 → 0.2.8

package/dist/server/core/memories/short-term.js

@@ -0,0 +1,512 @@
+import fs from 'fs-extra';
+import path from 'path';
+import { v4 as uuidv4 } from 'uuid';
+const SESSIONS_DIR = path.resolve(process.cwd(), 'sessions');
+const INDEX_FILE = path.join(SESSIONS_DIR, 'index.json');
+function defaultWorkspacePath() {
+    return path.resolve(process.cwd(), 'workspace');
+}
+function normalizeWorkspace(input) {
+    const workspace = (input ?? '').trim();
+    if (!workspace)
+        return defaultWorkspacePath();
+    return path.resolve(workspace);
+}
+export class ShortTermMemory {
+    constructor() {
+        this.instanceId = Math.random().toString(36).substring(7);
+        // 主存：所有 Session Meta 数据
+        this.sessions = new Map();
+        this.initPromise = null;
+        // 持久化队列，确保串行写入 index.json
+        this.saveQueue = Promise.resolve();
+        // Session级写入队列，避免messages.json并发冲突
+        this.sessionQueues = new Map();
+        console.log(`[ShortTermMemory] Instance created: ${this.instanceId}`);
+        fs.ensureDirSync(SESSIONS_DIR);
+    }
+    static getInstance() {
+        if (!ShortTermMemory.instance) {
+            ShortTermMemory.instance = new ShortTermMemory();
+        }
+        return ShortTermMemory.instance;
+    }
+    async ensureReady() {
+        if (!this.initPromise) {
+            this.initPromise = this.init();
+        }
+        await this.initPromise;
+    }
+    async init() {
+        try {
+            await fs.ensureDir(SESSIONS_DIR);
+            const indexExists = await fs.pathExists(INDEX_FILE);
+            if (indexExists) {
+                await this.loadSessionsFromDisk();
+            }
+            else {
+                await this.migrateLegacyIfNeeded();
+                // 迁移后再次尝试加载，如果还没有则初始化为空
+                if (await fs.pathExists(INDEX_FILE)) {
+                    await this.loadSessionsFromDisk();
+                }
+                else {
+                    this.sessions.clear();
+                    await this.persistIndex();
+                }
+            }
+            console.log(`[ShortTermMemory] Initialized with ${this.sessions.size} sessions.`);
+        }
+        catch (e) {
+            console.error('[ShortTermMemory] Initialization failed:', e);
+            // 即使失败，sessions 也是空的 Map，不会崩溃，但数据可能丢失
+        }
+    }
+    async loadSessionsFromDisk() {
+        try {
+            // Invalidate cache by reading file directly
+            const content = await fs.readFile(INDEX_FILE, 'utf-8');
+            const data = JSON.parse(content);
+            console.log(`[ShortTermMemory:${this.instanceId}] Loaded ${Array.isArray(data) ? data.length : 0} sessions from disk.`);
+            // MERGE instead of replace to prevent wiping out pending changes in current process
+            if (Array.isArray(data)) {
+                for (const meta of data) {
+                    if (meta && typeof meta.id === 'string') {
+                        // Only update if not exists or if disk version is newer (optional, but for now just trusting disk for existence)
+                        // But wait, if we have a pending write, our memory version is newer.
+                        // However, we are reloading because we are MISSING a session.
+                        // So we should add missing ones.
+                        if (!this.sessions.has(meta.id)) {
+                            this.sessions.set(meta.id, meta);
+                        }
+                        else {
+                            // If it exists, should we update? 
+                            // If we update, we might overwrite in-memory changes that haven't been flushed.
+                            // Safest is to only ADD missing ones.
+                            // But if another process updated a session, we want that update.
+                            // Complex conflict resolution needed? 
+                            // For now, let's assume if we are reloading, it's because we need data.
+                            // Let's just add missing ones.
+                            // NO, wait. If we don't update existing ones, we miss updates from other processes.
+                            // But if we overwrite, we lose our own updates.
+                            // Ideally, we only reload when we *know* we are missing something or on init.
+                            // Compromise: Update everything. Why? Because `persistIndex` queues writes.
+                            // If we have a pending write, it will overwrite disk later anyway.
+                            // So momentarily reverting to disk state is fine, as long as our write queue eventually flushes OUR state.
+                            this.sessions.set(meta.id, meta);
+                        }
+                    }
+                }
+            }
+        }
+        catch (e) {
+            console.error(`[ShortTermMemory:${this.instanceId}] Failed to load index.json:`, e);
+        }
+    }
+    async persistIndex() {
+        // 将当前所有 sessions 转为数组并排序
+        const metas = Array.from(this.sessions.values());
+        metas.sort((a, b) => (b.updatedAt ?? 0) - (a.updatedAt ?? 0));
+        // 加入写入队列
+        this.saveQueue = this.saveQueue.then(async () => {
+            try {
+                await fs.writeJson(INDEX_FILE, metas, { spaces: 2 });
+            }
+            catch (e) {
+                console.error('[ShortTermMemory] Failed to persist index.json:', e);
+            }
+        });
+        return this.saveQueue;
+    }
+    getSessionDir(sessionId) {
+        return path.join(SESSIONS_DIR, sessionId);
+    }
+    getMessagesFile(sessionId) {
+        return path.join(this.getSessionDir(sessionId), 'messages.json');
+    }
+    getContextFile(sessionId) {
+        return path.join(this.getSessionDir(sessionId), 'context.txt');
+    }
+    getRunsFile(sessionId) {
+        return path.join(this.getSessionDir(sessionId), 'runs.json');
+    }
+    async migrateLegacyIfNeeded() {
+        const files = await fs.readdir(SESSIONS_DIR);
+        const legacyFiles = files.filter(f => f.endsWith('.json') && f !== 'index.json');
+        if (legacyFiles.length === 0)
+            return;
+        console.log('[ShortTermMemory] Migrating legacy sessions...');
+        const metas = [];
+        for (const file of legacyFiles) {
+            const filePath = path.join(SESSIONS_DIR, file);
+            try {
+                const legacy = await fs.readJson(filePath);
+                if (!legacy || typeof legacy.id !== 'string')
+                    continue;
+                const id = legacy.id;
+                const createdAt = typeof legacy.createdAt === 'number' ? legacy.createdAt : Date.now();
+                const updatedAt = typeof legacy.updatedAt === 'number' ? legacy.updatedAt : Date.now();
+                const messages = Array.isArray(legacy.messages) ? legacy.messages : [];
+                const meta = {
+                    id,
+                    createdAt,
+                    updatedAt,
+                    lastActiveAt: updatedAt,
+                    workspace: defaultWorkspacePath()
+                };
+                const sessionDir = this.getSessionDir(id);
+                await fs.ensureDir(sessionDir);
+                await fs.writeJson(this.getMessagesFile(id), messages, { spaces: 2 });
+                metas.push(meta);
+            }
+            catch (e) {
+                console.error(`[ShortTermMemory] Legacy migration failed for ${file}:`, e);
+            }
+        }
+        if (metas.length > 0) {
+            metas.sort((a, b) => (b.updatedAt ?? 0) - (a.updatedAt ?? 0));
+            for (const m of metas) {
+                this.sessions.set(m.id, m);
+            }
+            await this.persistIndex();
+        }
+    }
+    async createSession(input) {
+        await this.ensureReady();
+        const id = input?.id || uuidv4();
+        if (this.sessions.has(id)) {
+            // If session already exists, return it (idempotent for external channels)
+            const session = await this.getSession(id);
+            if (session)
+                return session;
+        }
+        console.log(`[ShortTermMemory:${this.instanceId}] Creating session: ${id}`);
+        const now = Date.now();
+        const meta = {
+            id,
+            alias: (input?.alias ?? '').trim() || undefined,
+            workspace: normalizeWorkspace(input?.workspace),
+            createdAt: now,
+            updatedAt: now,
+            lastActiveAt: now,
+            projectId: input?.projectId,
+            projectIds: input?.projectIds,
+            notebookId: input?.notebookId,
+            lastArchivedAt: now,
+        };
+        await fs.ensureDir(meta.workspace);
+        const sessionDir = this.getSessionDir(id);
+        await fs.ensureDir(sessionDir);
+        await fs.writeJson(this.getMessagesFile(id), [], { spaces: 2 });
+        // IMPORTANT: Reload from disk BEFORE adding new session to avoid overwriting other processes' changes
+        await this.loadSessionsFromDisk();
+        // 更新内存
+        this.sessions.set(id, meta);
+        // 异步持久化
+        await this.persistIndex();
+        console.log(`[ShortTermMemory:${this.instanceId}] Session created and persisted: ${id}`);
+        return { meta, messages: [] };
+    }
+    async listSessions() {
+        await this.ensureReady();
+        const metas = Array.from(this.sessions.values());
+        metas.sort((a, b) => (b.updatedAt ?? 0) - (a.updatedAt ?? 0));
+        return metas;
+    }
+    async getSessionMeta(id) {
+        await this.ensureReady();
+        let meta = this.sessions.get(id) || null;
+        // 如果内存中没有，尝试重新加载磁盘（可能是其他进程创建的）
+        if (!meta) {
+            console.log(`[ShortTermMemory:${this.instanceId}] Cache miss for ${id}, reloading from disk...`);
+            await this.loadSessionsFromDisk();
+            meta = this.sessions.get(id) || null;
+        }
+        if (!meta) {
+            console.log(`[ShortTermMemory:${this.instanceId}] Session not found in memory after reload: ${id}. Total sessions: ${this.sessions.size}`);
+            // Log all available IDs for debugging
+            // console.log(`[ShortTermMemory:${this.instanceId}] Available IDs: ${Array.from(this.sessions.keys()).join(', ')}`);
+        }
+        return meta;
+    }
+    async updateSessionMeta(id, patch) {
+        await this.ensureReady();
+        const current = this.sessions.get(id);
+        if (!current)
+            throw new Error('Session not found');
+        const next = {
+            ...current,
+            alias: patch.alias === null ? undefined : (patch.alias !== undefined ? (String(patch.alias).trim() || undefined) : current.alias),
+            workspace: patch.workspace !== undefined ? normalizeWorkspace(String(patch.workspace)) : current.workspace,
+            projectId: patch.projectId === null ? undefined : (patch.projectId !== undefined ? patch.projectId : current.projectId),
+            projectIds: patch.projectIds === null ? undefined : (patch.projectIds !== undefined ? patch.projectIds : current.projectIds),
+            notebookId: patch.notebookId === null ? undefined : (patch.notebookId !== undefined ? patch.notebookId : current.notebookId),
+            lastArchivedAt: patch.lastArchivedAt !== undefined ? patch.lastArchivedAt : current.lastArchivedAt,
+            updatedAt: Date.now()
+        };
+        await fs.ensureDir(next.workspace);
+        if (patch.context !== undefined) {
+            const ctxFile = this.getContextFile(id);
+            const val = String(patch.context ?? '').trim();
+            if (!val) {
+                try {
+                    await fs.remove(ctxFile);
+                }
+                catch { /* ignore */ }
+                delete next.context;
+            }
+            else {
+                await fs.ensureDir(this.getSessionDir(id));
+                await fs.writeFile(ctxFile, val, 'utf-8');
+                next.context = val;
+            }
+        }
+        // 更新内存并持久化
+        this.sessions.set(id, next);
+        await this.persistIndex();
+        return next;
+    }
+    async persistSession(id) {
+        await this.ensureReady();
+        const current = this.sessions.get(id);
+        if (!current)
+            throw new Error('Session not found');
+        const now = Date.now();
+        let context = '';
+        try {
+            const ctxFile = this.getContextFile(id);
+            if (await fs.pathExists(ctxFile)) {
+                context = String(await fs.readFile(ctxFile, 'utf-8') || '').trim();
+            }
+        }
+        catch {
+            // ignore
+        }
+        const next = { ...current, context: context || undefined, updatedAt: now, lastActiveAt: now };
+        this.sessions.set(id, next);
+        await this.persistIndex();
+        await fs.ensureDir(this.getSessionDir(id));
+        await fs.ensureFile(this.getMessagesFile(id));
+        await fs.ensureFile(this.getRunsFile(id));
+        return next;
+    }
+    async getSession(id) {
+        const meta = await this.getSessionMeta(id);
+        if (!meta)
+            return null;
+        const msgFile = this.getMessagesFile(id);
+        let messages = [];
+        try {
+            const data = await fs.readJson(msgFile);
+            messages = Array.isArray(data) ? data : [];
+        }
+        catch {
+            // 如果消息文件不存在或读取失败，尝试重新创建
+            await fs.ensureDir(this.getSessionDir(id));
+            await fs.writeJson(msgFile, [], { spaces: 2 });
+            messages = [];
+        }
+        let context = '';
+        try {
+            const ctxFile = this.getContextFile(id);
+            if (await fs.pathExists(ctxFile)) {
+                context = String(await fs.readFile(ctxFile, 'utf-8') || '').trim();
+            }
+        }
+        catch {
+            // ignore
+        }
+        return { meta: { ...meta, context: context || undefined }, messages };
+    }
+    async createRun(sessionId, userContent, runId) {
+        await this.ensureReady();
+        // Check if session exists, try reload if not
+        if (!this.sessions.has(sessionId)) {
+            console.log(`[ShortTermMemory:${this.instanceId}] Session ${sessionId} not found for createRun, reloading...`);
+            await this.loadSessionsFromDisk();
+        }
+        if (!this.sessions.has(sessionId))
+            throw new Error('Session not found');
+        const now = Date.now();
+        const id = runId || uuidv4();
+        const run = {
+            id,
+            sessionId,
+            status: 'queued',
+            userContent,
+            createdAt: now,
+            updatedAt: now
+        };
+        const file = this.getRunsFile(sessionId);
+        let runs = [];
+        try {
+            const existing = await fs.readJson(file);
+            runs = Array.isArray(existing) ? existing : [];
+        }
+        catch {
+            await fs.ensureDir(this.getSessionDir(sessionId));
+        }
+        runs.push(run);
+        await fs.writeJson(file, runs, { spaces: 2 });
+        return run;
+    }
+    async updateRun(sessionId, runId, patch) {
+        await this.ensureReady();
+        const file = this.getRunsFile(sessionId);
+        let runs = [];
+        try {
+            const existing = await fs.readJson(file);
+            runs = Array.isArray(existing) ? existing : [];
+        }
+        catch {
+            runs = [];
+        }
+        const idx = runs.findIndex(r => r.id === runId);
+        if (idx < 0)
+            throw new Error('Run not found');
+        const next = {
+            ...runs[idx],
+            ...patch,
+            status: patch.status || runs[idx].status,
+            updatedAt: patch.updatedAt ?? Date.now()
+        };
+        runs[idx] = next;
+        await fs.writeJson(file, runs, { spaces: 2 });
+        return next;
+    }
+    async listRuns(sessionId) {
+        await this.ensureReady();
+        const file = this.getRunsFile(sessionId);
+        try {
+            const existing = await fs.readJson(file);
+            const runs = Array.isArray(existing) ? existing : [];
+            runs.sort((a, b) => (b.createdAt ?? 0) - (a.createdAt ?? 0));
+            return runs;
+        }
+        catch {
+            return [];
+        }
+    }
+    async getMessages(sessionId) {
+        // 复用 getSession 逻辑，它包含了文件读取和错误处理
+        const session = await this.getSession(sessionId);
+        return session ? session.messages : [];
+    }
+    async addMessage(sessionId, message) {
+        await this.ensureReady();
+        let sessionMeta = this.sessions.get(sessionId);
+        // 如果内存中没有，尝试重载
+        if (!sessionMeta) {
+            console.log(`[ShortTermMemory:${this.instanceId}] Session ${sessionId} not found in memory for addMessage, reloading...`);
+            await this.loadSessionsFromDisk();
+            sessionMeta = this.sessions.get(sessionId);
+        }
+        // 如果仍然没有，尝试直接从 index.json 读取一次，绕过缓存逻辑
+        if (!sessionMeta) {
+            try {
+                const content = await fs.readFile(INDEX_FILE, 'utf-8');
+                const data = JSON.parse(content);
+                const found = Array.isArray(data) ? data.find((m) => m.id === sessionId) : null;
+                if (found) {
+                    console.log(`[ShortTermMemory:${this.instanceId}] Session ${sessionId} found in fresh disk read, recovering...`);
+                    this.sessions.set(found.id, found);
+                    sessionMeta = found;
+                }
+            }
+            catch (e) {
+                console.error(`[ShortTermMemory:${this.instanceId}] Emergency disk read failed:`, e);
+            }
+        }
+        if (!sessionMeta) {
+            console.error(`[ShortTermMemory:${this.instanceId}] Session ${sessionId} not found even after emergency reload.`);
+            throw new Error(`Session ${sessionId} not found`);
+        }
+        const msg = { ...message };
+        if (!msg.timestamp)
+            msg.timestamp = Date.now();
+        // 使用队列确保对 messages.json 的串行写入
+        const queue = this.sessionQueues.get(sessionId) || Promise.resolve();
+        const nextTask = queue.then(async () => {
+            // 更新 messages.json
+            const msgFile = this.getMessagesFile(sessionId);
+            let messages = [];
+            try {
+                const data = await fs.readJson(msgFile);
+                messages = Array.isArray(data) ? data : [];
+            }
+            catch {
+                messages = [];
+            }
+            messages.push(msg);
+            await fs.writeJson(msgFile, messages, { spaces: 2 });
+        }).catch(e => {
+            console.error(`[ShortTermMemory] Failed to write message for session ${sessionId}:`, e);
+        });
+        this.sessionQueues.set(sessionId, nextTask);
+        // 等待本次写入完成
+        await nextTask;
+        // 清理已完成的队列（可选，防止 Map 无限增长，但需小心并发）
+        // 简单策略：不清理，或者定期清理。由于 Session 数量有限，暂时保留。
+        // 更新 Session Meta
+        const nextMeta = {
+            ...sessionMeta,
+            updatedAt: Date.now(),
+            lastActiveAt: Date.now()
+        };
+        this.sessions.set(sessionId, nextMeta);
+        await this.persistIndex();
+    }
+    async clearMessages(sessionId) {
+        await this.ensureReady();
+        const sessionMeta = this.sessions.get(sessionId);
+        if (!sessionMeta)
+            throw new Error(`Session ${sessionId} not found`);
+        await fs.writeJson(this.getMessagesFile(sessionId), [], { spaces: 2 });
+        // Update timestamp
+        const nextMeta = {
+            ...sessionMeta,
+            updatedAt: Date.now(),
+            lastActiveAt: Date.now()
+        };
+        this.sessions.set(sessionId, nextMeta);
+        await this.persistIndex();
+    }
+    async deleteSession(id) {
+        await this.ensureReady();
+        if (!this.sessions.has(id))
+            throw new Error('Session not found');
+        this.sessions.delete(id);
+        await this.persistIndex();
+        try {
+            await fs.remove(this.getSessionDir(id));
+        }
+        catch (e) {
+            console.error(`[ShortTermMemory] Failed to delete session dir ${id}:`, e);
+        }
+    }
+    async cleanupOldSessions(maxAgeDays) {
+        await this.ensureReady();
+        const now = Date.now();
+        const msPerDay = 24 * 60 * 60 * 1000;
+        const toDelete = [];
+        for (const meta of this.sessions.values()) {
+            const updatedAt = typeof meta.updatedAt === 'number' ? meta.updatedAt : 0;
+            const ageDays = (now - updatedAt) / msPerDay;
+            if (ageDays > maxAgeDays)
+                toDelete.push(meta.id);
+        }
+        for (const id of toDelete) {
+            this.sessions.delete(id);
+            try {
+                await fs.remove(this.getSessionDir(id));
+            }
+            catch (e) {
+                console.error(`[ShortTermMemory] Failed to delete session ${id}:`, e);
+            }
+        }
+        if (toDelete.length > 0) {
+            await this.persistIndex();
+        }
+        return toDelete.length;
+    }
+}