zn-ai/electron/gateway/session-store.ts

import * as fs from 'fs';
import * as path from 'path';
import { app } from 'electron';
import logManager from '@electron/service/logger';
import { normalizeAgentSessionKey } from '@runtime/lib/models';
import type { RawMessage } from '@runtime/shared/chat-model';

let sessionsFilePath: string | null = null;

function getSessionsFilePath(): string {
  if (!sessionsFilePath) {
    sessionsFilePath = path.join(app.getPath('userData'), 'chat-sessions.json');
  }
  return sessionsFilePath;
}

export interface SessionEntry {
  key: string;
  messages: RawMessage[];
  updatedAt: number;
  activeRun?: {
    runId: string;
    abortController: AbortController;
  };
}

class SessionStore {
  private sessions = new Map<string, SessionEntry>();
  private loaded = false;

  private ensureLoaded(): void {
    if (this.loaded) return;
    this.loaded = true;
    this.loadFromDisk();
  }

  private loadFromDisk(): void {
    try {
      const filePath = getSessionsFilePath();
      if (fs.existsSync(filePath)) {
        const data = JSON.parse(fs.readFileSync(filePath, 'utf-8')) as Record<
          string,
          Omit<SessionEntry, 'activeRun'>
        >;
        let migrated = false;
        for (const [key, entry] of Object.entries(data)) {
          const normalizedKey = normalizeAgentSessionKey(key);
          if (normalizedKey !== key) {
            migrated = true;
          }
          this.sessions.set(normalizedKey, {
            ...entry,
            key: normalizedKey,
            activeRun: undefined,
          });
        }
        if (migrated) {
          this.saveToDisk();
        }
      }
    } catch (e) {
      logManager.error('Failed to load sessions from disk:', e);
    }
  }

  saveToDisk(): void {
    try {
      const filePath = getSessionsFilePath();
      const data: Record<string, Omit<SessionEntry, 'activeRun'>> = {};
      for (const [key, entry] of this.sessions) {
        data[key] = {
          key: entry.key,
          messages: entry.messages,
          updatedAt: entry.updatedAt,
        };
      }
      fs.mkdirSync(path.dirname(filePath), { recursive: true });
      fs.writeFileSync(filePath, JSON.stringify(data, null, 2), 'utf-8');
    } catch (e) {
      logManager.error('Failed to save sessions to disk:', e);
    }
  }

  getOrCreate(key: string): SessionEntry {
    this.ensureLoaded();
    const normalizedKey = normalizeAgentSessionKey(key);
    let session = this.sessions.get(normalizedKey);
    if (!session) {
      session = {
        key: normalizedKey,
        messages: [],
        updatedAt: Date.now(),
      };
      this.sessions.set(normalizedKey, session);
    }
    return session;
  }

  get(key: string): SessionEntry | undefined {
    this.ensureLoaded();
    return this.sessions.get(normalizeAgentSessionKey(key));
  }

  getAllKeys(): string[] {
    this.ensureLoaded();
    return Array.from(this.sessions.keys());
  }

  appendMessage(key: string, message: RawMessage): void {
    const session = this.getOrCreate(key);
    session.messages.push(message);
    session.updatedAt = Date.now();
    this.saveToDisk();
  }

  getMessages(key: string, limit = 50): RawMessage[] {
    const session = this.get(key);
    if (!session) return [];
    return session.messages.slice(-limit);
  }

  setActiveRun(key: string, runId: string, abortController: AbortController): void {
    const session = this.getOrCreate(key);
    session.activeRun = { runId, abortController };
  }

  clearActiveRun(key: string): void {
    const session = this.sessions.get(normalizeAgentSessionKey(key));
    if (session) {
      session.activeRun = undefined;
    }
  }

  getActiveRun(key: string): { runId: string; abortController: AbortController } | undefined {
    return this.sessions.get(normalizeAgentSessionKey(key))?.activeRun;
  }

  deleteSession(key: string): void {
    this.sessions.delete(normalizeAgentSessionKey(key));
    this.saveToDisk();
  }
}

export const sessionStore = new SessionStore();