refactor: unify agent session persistence

2026-06-04 15:02:27 +08:00
parent 04ded0ceb0
commit 0ecb2babf3
22 changed files with 542 additions and 497 deletions
@@ -0,0 +1,149 @@
+import { join } from "node:path";
+
+import { config } from "../config.js";
+import {
+  atomicWriteJson,
+  ensureDirectory,
+  listJsonFiles,
+  readJsonFile,
+  removeFileIfExists,
+  slugify,
+} from "../utils/fileStore.js";
+
+export type SessionStatus = "active" | "archived";
+
+export type SessionRecord = {
+  sessionId: string;
+  actorKey: string;
+  ownerUserId?: string;
+  projectId?: string;
+  projectKey: string;
+  parentSessionId?: string;
+  createdAt: string;
+  updatedAt: string;
+  status: SessionStatus;
+  title?: string;
+};
+
+type SessionMetadataContext = {
+  actorKey: string;
+  userId?: string;
+  projectId?: string;
+  projectKey: string;
+};
+
+type EnsureSessionMetadataInput = SessionMetadataContext & {
+  sessionId: string;
+  parentSessionId?: string;
+};
+
+export class SessionMetadataStore {
+  constructor(private readonly baseDir = config.SESSION_METADATA_STORAGE_DIR) {}
+
+  async initialize() {
+    await ensureDirectory(this.baseDir);
+  }
+
+  async ensure(input: EnsureSessionMetadataInput) {
+    const sessionId = normalizeSessionId(input.sessionId);
+    if (!sessionId) {
+      throw new Error("sessionId is required");
+    }
+    const existing = await readJsonFile<SessionRecord>(
+      this.filePath(sessionId),
+    );
+    if (existing) {
+      return { created: false, record: existing };
+    }
+
+    const now = new Date().toISOString();
+    const record: SessionRecord = {
+      sessionId,
+      actorKey: input.actorKey,
+      ownerUserId: input.userId?.trim(),
+      projectId: input.projectId,
+      projectKey: input.projectKey,
+      parentSessionId: normalizeSessionId(input.parentSessionId),
+      createdAt: now,
+      updatedAt: now,
+      status: "active",
+    };
+    await atomicWriteJson(
+      this.filePath(record.sessionId),
+      record,
+    );
+    return { created: true, record };
+  }
+
+  async get(context: SessionMetadataContext, sessionId: string) {
+    const normalizedSessionId = normalizeSessionId(sessionId);
+    if (!normalizedSessionId) {
+      return null;
+    }
+    return await readJsonFile<SessionRecord>(
+      this.filePath(normalizedSessionId),
+    );
+  }
+
+  async touch(
+    record: SessionRecord,
+    updates: Partial<Pick<SessionRecord, "title" | "status">> = {},
+  ) {
+    const next: SessionRecord = {
+      ...record,
+      ...normalizeSessionUpdates(updates),
+      updatedAt: new Date().toISOString(),
+    };
+    await atomicWriteJson(
+      this.filePath(record.sessionId),
+      next,
+    );
+    return next;
+  }
+
+  async list(context: SessionMetadataContext) {
+    const files = await listJsonFiles(this.baseDir);
+    const records = await Promise.all(
+      files.map((file) => readJsonFile<SessionRecord>(file)),
+    );
+    return records
+      .filter((record): record is SessionRecord => Boolean(record))
+      .filter(
+        (record) =>
+          record.actorKey === context.actorKey &&
+          record.projectKey === context.projectKey,
+      )
+      .sort((left, right) => right.updatedAt.localeCompare(left.updatedAt));
+  }
+
+  async remove(record: SessionRecord) {
+    await removeFileIfExists(
+      this.filePath(record.sessionId),
+    );
+  }
+
+  private filePath(sessionId: string) {
+    return join(this.baseDir, `${slugify(sessionId)}.json`);
+  }
+}
+
+const normalizeSessionId = (value?: string) => {
+  const normalized = value?.trim();
+  return normalized ? normalized.slice(0, 128) : undefined;
+};
+
+const normalizeSessionUpdates = (
+  updates: Partial<Pick<SessionRecord, "title" | "status">>,
+) => {
+  const normalized: Partial<Pick<SessionRecord, "title" | "status">> = {};
+  if (updates.status === "active" || updates.status === "archived") {
+    normalized.status = updates.status;
+  }
+  if (typeof updates.title === "string") {
+    const trimmed = updates.title.trim();
+    if (trimmed) {
+      normalized.title = trimmed.slice(0, 120);
+    }
+  }
+  return normalized;
+};
@@ -0,0 +1,45 @@
+import { join } from "node:path";
+
+import { config } from "../config.js";
+import {
+  atomicWriteJson,
+  ensureDirectory,
+  readJsonFile,
+  removeFileIfExists,
+} from "../utils/fileStore.js";
+
+export type SessionRuntimeContext = {
+  accessToken?: string;
+  actorKey: string;
+  allowLearningWrite?: boolean;
+  clientSessionId: string;
+  learningMode?: "interactive" | "review";
+  projectId?: string;
+  projectKey: string;
+  sessionId: string;
+  traceId: string;
+};
+
+export class SessionRuntimeContextStore {
+  constructor(private readonly baseDir = config.SESSION_RUNTIME_CONTEXT_STORAGE_DIR) {}
+
+  async initialize() {
+    await ensureDirectory(this.baseDir);
+  }
+
+  async write(context: SessionRuntimeContext) {
+    await atomicWriteJson(this.filePath(context.sessionId), context);
+  }
+
+  async read(sessionId: string) {
+    return await readJsonFile<SessionRuntimeContext>(this.filePath(sessionId));
+  }
+
+  async remove(sessionId: string) {
+    await removeFileIfExists(this.filePath(sessionId));
+  }
+
+  private filePath(sessionId: string) {
+    return join(this.baseDir, `${sessionId}.json`);
+  }
+}
@@ -0,0 +1,280 @@
+import { join } from "node:path";
+
+import { config } from "../config.js";
+import {
+  atomicWriteJson,
+  ensureDirectory,
+  listJsonFiles,
+  readJsonFile,
+  toStableId,
+} from "../utils/fileStore.js";
+import { sanitizePersistentDocument } from "../utils/persistencePolicy.js";
+
+export type SessionTurnRecord = {
+  id: string;
+  assistantMessage: string;
+  timestamp: string;
+  toolCallCount: number;
+  userMessage: string;
+};
+
+type SessionTranscriptRecord = {
+  actorKey: string;
+  clientSessionId?: string;
+  projectKey: string;
+  sessionId: string;
+  turns: SessionTurnRecord[];
+  updatedAt: string;
+};
+
+export type SessionSearchHit = {
+  matchedField: "assistant" | "user";
+  score: number;
+  sessionId: string;
+  snippet: string;
+  timestamp: string;
+  turnId: string;
+};
+
+type SessionTranscriptContext = {
+  actorKey: string;
+  clientSessionId?: string;
+  projectKey: string;
+  sessionId: string;
+};
+
+export class SessionTranscriptStore {
+  private readonly writeQueues = new Map<string, Promise<void>>();
+
+  constructor(private readonly baseDir = config.SESSION_TRANSCRIPT_STORAGE_DIR) {}
+
+  async initialize() {
+    await ensureDirectory(this.baseDir);
+  }
+
+  async appendTurn(
+    context: SessionTranscriptContext,
+    turn: {
+      assistantMessage: string;
+      toolCallCount: number;
+      userMessage: string;
+    },
+  ) {
+    const key = this.filePath(context);
+    return this.serializeWrite(key, async () => {
+      const transcript = (await this.readTranscript(context)) ?? {
+        actorKey: context.actorKey,
+        clientSessionId: context.clientSessionId,
+        projectKey: context.projectKey,
+        sessionId: context.sessionId,
+        turns: [],
+        updatedAt: new Date().toISOString(),
+      };
+      const userMessage = sanitizePersistentDocument(turn.userMessage, 4000);
+      const assistantMessage = sanitizePersistentDocument(turn.assistantMessage, 4000);
+      if (!userMessage || !assistantMessage) {
+        return transcript;
+      }
+
+      const timestamp = new Date().toISOString();
+      const record: SessionTurnRecord = {
+        id: toStableId(context.sessionId, timestamp, userMessage, assistantMessage),
+        assistantMessage,
+        timestamp,
+        toolCallCount: Math.max(0, turn.toolCallCount),
+        userMessage,
+      };
+      transcript.clientSessionId = context.clientSessionId ?? transcript.clientSessionId;
+      transcript.sessionId = context.sessionId;
+      transcript.turns.push(record);
+      if (transcript.turns.length > config.SESSION_TRANSCRIPT_MAX_TURNS_PER_SESSION) {
+        transcript.turns = transcript.turns.slice(
+          transcript.turns.length - config.SESSION_TRANSCRIPT_MAX_TURNS_PER_SESSION,
+        );
+      }
+      transcript.updatedAt = timestamp;
+      await atomicWriteJson(key, transcript);
+      return transcript;
+    });
+  }
+
+  async getRecentTurns(
+    context: SessionTranscriptContext,
+    limit: number,
+  ): Promise<SessionTurnRecord[]> {
+    const transcript = await this.readTranscript(context);
+    if (!transcript) {
+      return [];
+    }
+    return transcript.turns.slice(-Math.max(1, limit));
+  }
+
+  async cloneThread(
+    sourceContext: SessionTranscriptContext,
+    targetContext: SessionTranscriptContext,
+    keepMessageCount: number,
+  ) {
+    const sourceTranscript = await this.readTranscript(sourceContext);
+    const timestamp = new Date().toISOString();
+    const nextTranscript: SessionTranscriptRecord = {
+      actorKey: targetContext.actorKey,
+      clientSessionId: targetContext.clientSessionId,
+      projectKey: targetContext.projectKey,
+      sessionId: targetContext.sessionId,
+      turns: projectTurnsForFork(sourceTranscript?.turns ?? [], keepMessageCount),
+      updatedAt: timestamp,
+    };
+    await atomicWriteJson(this.filePath(targetContext), nextTranscript);
+    return nextTranscript;
+  }
+
+  async search(
+    context: Pick<SessionTranscriptContext, "actorKey" | "projectKey">,
+    query: string,
+    maxResults = config.SESSION_SEARCH_MAX_RESULTS,
+  ): Promise<SessionSearchHit[]> {
+    const normalizedQuery = query.trim().toLowerCase().slice(0, config.SESSION_SEARCH_MAX_QUERY_CHARS);
+    if (!normalizedQuery) {
+      return [];
+    }
+    const queryTokens = normalizedQuery.split(/\s+/).filter(Boolean);
+    const hits: SessionSearchHit[] = [];
+    const files = await listJsonFiles(this.baseDir);
+    for (const file of files) {
+      const transcript = await readJsonFile<SessionTranscriptRecord>(file);
+      if (!transcript) {
+        continue;
+      }
+      if (
+        transcript.actorKey !== context.actorKey ||
+        transcript.projectKey !== context.projectKey
+      ) {
+        continue;
+      }
+      for (const turn of transcript.turns) {
+        const candidates: Array<["user" | "assistant", string]> = [
+          ["user", turn.userMessage],
+          ["assistant", turn.assistantMessage],
+        ];
+        for (const [matchedField, text] of candidates) {
+          const score = scoreText(text, normalizedQuery, queryTokens);
+          if (score <= 0) {
+            continue;
+          }
+          hits.push({
+            matchedField,
+            score,
+            sessionId: transcript.sessionId,
+            snippet: buildSnippet(text, normalizedQuery),
+            timestamp: turn.timestamp,
+            turnId: turn.id,
+          });
+        }
+      }
+    }
+    return hits.sort((a, b) => b.score - a.score).slice(0, Math.max(1, maxResults));
+  }
+
+  private async readTranscript(context: SessionTranscriptContext) {
+    const direct = await readJsonFile<SessionTranscriptRecord>(this.filePath(context));
+    if (direct) {
+      return direct;
+    }
+
+    const clientSessionId = context.clientSessionId?.trim();
+    if (!clientSessionId) {
+      return null;
+    }
+
+    const files = await listJsonFiles(this.baseDir);
+    const matches: SessionTranscriptRecord[] = [];
+    for (const file of files) {
+      const transcript = await readJsonFile<SessionTranscriptRecord>(file);
+      if (!transcript) {
+        continue;
+      }
+      if (
+        transcript.actorKey !== context.actorKey ||
+        transcript.projectKey !== context.projectKey ||
+        transcript.clientSessionId !== clientSessionId
+      ) {
+        continue;
+      }
+      matches.push(transcript);
+    }
+
+    if (matches.length === 0) {
+      return null;
+    }
+
+    return matches.sort((left, right) => right.updatedAt.localeCompare(left.updatedAt))[0] ?? null;
+  }
+
+  private filePath(context: SessionTranscriptContext) {
+    return join(
+      this.baseDir,
+      `${context.actorKey}__${context.projectKey}__${context.sessionId}.json`,
+    );
+  }
+
+  private async serializeWrite<T>(key: string, task: () => Promise<T>) {
+    const previous = this.writeQueues.get(key) ?? Promise.resolve();
+    const run = previous.catch(() => undefined).then(task);
+    const next = run.then(
+      () => undefined,
+      () => undefined,
+    );
+    this.writeQueues.set(key, next);
+    try {
+      return await run;
+    } finally {
+      if (this.writeQueues.get(key) === next) {
+        this.writeQueues.delete(key);
+      }
+    }
+  }
+}
+
+const scoreText = (text: string, query: string, queryTokens: string[]) => {
+  const normalized = text.toLowerCase();
+  let score = 0;
+  if (normalized.includes(query)) {
+    score += Math.max(10, query.length);
+  }
+  for (const token of queryTokens) {
+    if (token.length >= 2 && normalized.includes(token)) {
+      score += 1;
+    }
+  }
+  return score;
+};
+
+const buildSnippet = (text: string, query: string) => {
+  const compact = text.replace(/\s+/g, " ").trim();
+  const idx = compact.toLowerCase().indexOf(query);
+  if (idx === -1) {
+    return compact.length > 180 ? `${compact.slice(0, 177)}...` : compact;
+  }
+  const start = Math.max(0, idx - 60);
+  const end = Math.min(compact.length, idx + query.length + 100);
+  const snippet = compact.slice(start, end).trim();
+  const prefix = start > 0 ? "..." : "";
+  const suffix = end < compact.length ? "..." : "";
+  return `${prefix}${snippet}${suffix}`;
+};
+
+const projectTurnsForFork = (
+  turns: SessionTurnRecord[],
+  keepMessageCount: number,
+): SessionTurnRecord[] => {
+  if (keepMessageCount <= 0) {
+    return [];
+  }
+
+  const keepTurnCount = Math.floor(keepMessageCount / 2);
+  if (keepTurnCount <= 0) {
+    return [];
+  }
+
+  return turns.slice(0, keepTurnCount);
+};
@@ -0,0 +1,46 @@
+import { join } from "node:path";
+
+import { config } from "../config.js";
+import {
+  atomicWriteJson,
+  ensureDirectory,
+  readJsonFile,
+  removeFileIfExists,
+  slugify,
+} from "../utils/fileStore.js";
+
+export type SessionUiStateRecord = {
+  sessionId: string;
+  isTitleManuallyEdited?: boolean;
+  messages: unknown[];
+  branchGroups: unknown[];
+};
+
+type SessionUiStateContext = {
+  sessionId: string;
+};
+
+export class SessionUiStateStore {
+  constructor(private readonly baseDir = config.SESSION_UI_STATE_STORAGE_DIR) {}
+
+  async initialize() {
+    await ensureDirectory(this.baseDir);
+  }
+
+  async read(context: SessionUiStateContext) {
+    return await readJsonFile<SessionUiStateRecord>(this.filePath(context));
+  }
+
+  async write(context: SessionUiStateContext, state: SessionUiStateRecord) {
+    await atomicWriteJson(this.filePath(context), state);
+    return state;
+  }
+
+  async remove(context: SessionUiStateContext) {
+    await removeFileIfExists(this.filePath(context));
+  }
+
+  private filePath(context: SessionUiStateContext) {
+    return join(this.baseDir, `${slugify(context.sessionId)}.json`);
+  }
+}