切换到使用pg数据库

src/conversations/stateStore.ts

@@ -1,12 +1,6 @@
-import { join } from "node:path";
+import { type QueryResultRow } from "pg";
 
-import { config } from "../config.js";
-import {
-  atomicWriteJson,
-  ensureDirectory,
-  readJsonFile,
-  removeFileIfExists,
-} from "../utils/fileStore.js";
+import { AgentDatabase, getAgentDatabase } from "../db/index.js";
 
 export type ConversationStateRecord = {
   sessionId: string;
@@ -15,27 +9,81 @@ export type ConversationStateRecord = {
   branchGroups: unknown[];
 };
 
+type ConversationStateRow = QueryResultRow & {
+  session_id: string;
+  is_title_manually_edited: boolean;
+  messages: unknown[];
+  branch_groups: unknown[];
+};
+
 export class ConversationStateStore {
-  constructor(private readonly baseDir = config.CONVERSATION_STATE_STORAGE_DIR) {}
+  constructor(private readonly db: AgentDatabase = getAgentDatabase()) {}
 
   async initialize() {
-    await ensureDirectory(this.baseDir);
+    await this.db.initialize();
   }
 
-  async read(sessionScopeKey: string) {
-    return await readJsonFile<ConversationStateRecord>(this.filePath(sessionScopeKey));
+  async read(sessionId: string) {
+    const result = await this.db.query<ConversationStateRow>(
+      `
+        SELECT session_id, is_title_manually_edited, messages, branch_groups
+        FROM ${this.db.table("conversation_states")}
+        WHERE session_id = $1
+        LIMIT 1
+      `,
+      [sessionId],
+    );
+    return mapConversationStateRow(result.rows[0]);
   }
 
-  async write(sessionScopeKey: string, state: ConversationStateRecord) {
-    await atomicWriteJson(this.filePath(sessionScopeKey), state);
-    return state;
+  async write(sessionId: string, state: ConversationStateRecord) {
+    const result = await this.db.query<ConversationStateRow>(
+      `
+        INSERT INTO ${this.db.table("conversation_states")} (
+          session_id,
+          is_title_manually_edited,
+          messages,
+          branch_groups
+        )
+        VALUES ($1, $2, $3::jsonb, $4::jsonb)
+        ON CONFLICT (session_id)
+        DO UPDATE SET
+          is_title_manually_edited = EXCLUDED.is_title_manually_edited,
+          messages = EXCLUDED.messages,
+          branch_groups = EXCLUDED.branch_groups
+        RETURNING session_id, is_title_manually_edited, messages, branch_groups
+      `,
+      [
+        sessionId,
+        state.isTitleManuallyEdited ?? false,
+        JSON.stringify(state.messages),
+        JSON.stringify(state.branchGroups),
+      ],
+    );
+    return mapConversationStateRow(result.rows[0]) ?? state;
   }
 
-  async remove(sessionScopeKey: string) {
-    await removeFileIfExists(this.filePath(sessionScopeKey));
-  }
-
-  private filePath(sessionScopeKey: string) {
-    return join(this.baseDir, `${sessionScopeKey}.json`);
+  async remove(sessionId: string) {
+    await this.db.query(
+      `
+        DELETE FROM ${this.db.table("conversation_states")}
+        WHERE session_id = $1
+      `,
+      [sessionId],
+    );
   }
 }
+
+const mapConversationStateRow = (
+  row?: ConversationStateRow | null,
+): ConversationStateRecord | null => {
+  if (!row) {
+    return null;
+  }
+  return {
+    sessionId: row.session_id,
+    isTitleManuallyEdited: row.is_title_manually_edited,
+    messages: Array.isArray(row.messages) ? row.messages : [],
+    branchGroups: Array.isArray(row.branch_groups) ? row.branch_groups : [],
+  };
+};