134 lines
3.4 KiB
TypeScript
134 lines
3.4 KiB
TypeScript
import * as fs from 'fs';
|
|
import * as path from 'path';
|
|
import { app } from 'electron';
|
|
import logManager from '@electron/service/logger';
|
|
import type { RawMessage } from '@src/pages/home/model/ChatModel';
|
|
|
|
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'>
|
|
>;
|
|
for (const [key, entry] of Object.entries(data)) {
|
|
this.sessions.set(key, {
|
|
...entry,
|
|
activeRun: undefined,
|
|
});
|
|
}
|
|
}
|
|
} 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();
|
|
let session = this.sessions.get(key);
|
|
if (!session) {
|
|
session = {
|
|
key,
|
|
messages: [],
|
|
updatedAt: Date.now(),
|
|
};
|
|
this.sessions.set(key, session);
|
|
}
|
|
return session;
|
|
}
|
|
|
|
get(key: string): SessionEntry | undefined {
|
|
this.ensureLoaded();
|
|
return this.sessions.get(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(key);
|
|
if (session) {
|
|
session.activeRun = undefined;
|
|
}
|
|
}
|
|
|
|
getActiveRun(key: string): { runId: string; abortController: AbortController } | undefined {
|
|
return this.sessions.get(key)?.activeRun;
|
|
}
|
|
|
|
deleteSession(key: string): void {
|
|
this.sessions.delete(key);
|
|
this.saveToDisk();
|
|
}
|
|
}
|
|
|
|
export const sessionStore = new SessionStore();
|