切换到使用pg数据库

src/results/store.ts

@@ -1,20 +1,19 @@
 import { randomUUID } from "node:crypto";
-import { join } from "node:path";
+import { join, resolve } from "node:path";
+
+import { type QueryResultRow } from "pg";
 
 import { config } from "../config.js";
+import { AgentDatabase, getAgentDatabase } from "../db/index.js";
 import { logger } from "../logger.js";
 import {
   atomicWriteJson,
   ensureDirectory,
-  getFileStat,
-  listJsonFiles,
   readJsonFile,
   removeFileIfExists,
-  toProjectKey,
 } from "../utils/fileStore.js";
 
 export const RESULT_REF_PATTERN = /^res-[a-f0-9-]{8,64}$/;
-const RESULT_REF_FILE_PATTERN = /^(res-[a-f0-9-]{8,64})(?:\.json)?$/;
 
 export const RESULT_REFERENCE_KIND = {
   dynamicHttpResult: "dynamic-http-result",
@@ -44,7 +43,6 @@ export type ResultPreview = {
 export type ResultReferenceRecord = {
   resultRef: string;
   actorKey: string;
-  clientSessionId: string;
   createdAt: string;
   data: unknown;
   kind: ResultReferenceKind;
@@ -56,11 +54,12 @@ export type ResultReferenceRecord = {
   sizeBytes: number;
   source: ResultReferenceSource;
   traceId: string;
+  payloadPath?: string;
+  objectKey?: string;
 };
 
 export type StoreResultInput = {
   actorKey: string;
-  clientSessionId: string;
   data: unknown;
   kind: ResultReferenceKind;
   projectId?: string;
@@ -69,11 +68,13 @@ export type StoreResultInput = {
   sessionId: string;
   source: ResultReferenceSource;
   traceId: string;
+  payloadPath?: string;
+  objectKey?: string;
 };
 
 export type RetrievalContext = {
   actorKey: string;
-  clientSessionId?: string;
+  sessionId?: string;
   projectId?: string;
 };
 
@@ -84,18 +85,37 @@ export type ResultReferencePeek = {
   storedAt: string;
 };
 
-type PartialRecord = Partial<ResultReferenceRecord> & { data?: unknown };
+type ResultReferenceRow = QueryResultRow & {
+  result_ref: string;
+  actor_key: string;
+  session_id: string;
+  created_at: Date | string;
+  kind: ResultReferenceKind;
+  preview: ResultPreview;
+  project_id: string | null;
+  project_key: string;
+  schema_version: number;
+  size_bytes: number;
+  source: ResultReferenceSource;
+  trace_id: string;
+  payload_path: string | null;
+  object_key: string | null;
+};
 
 export class ResultReferenceStore {
   private cleanupTimer: NodeJS.Timeout | null = null;
+  private readonly managedPayloadDir: string;
 
   constructor(
+    private readonly db: AgentDatabase = getAgentDatabase(),
     private readonly baseDir = config.RESULT_REF_STORAGE_DIR,
     private readonly ttlMs = config.RESULT_REF_TTL_HOURS * 60 * 60 * 1000,
-  ) {}
+  ) {
+    this.managedPayloadDir = join(this.baseDir, "payloads");
+  }
 
   async initialize() {
-    await ensureDirectory(this.baseDir);
+    await Promise.all([this.db.initialize(), ensureDirectory(this.managedPayloadDir)]);
   }
 
   startCleanupLoop() {
@@ -119,24 +139,77 @@ export class ResultReferenceStore {
 
   async store(input: StoreResultInput) {
     const resultRef = `res-${randomUUID().slice(0, 16)}`;
-    const record: ResultReferenceRecord = {
+    const createdAt = new Date().toISOString();
+    const payloadPath =
+      input.payloadPath ??
+      (input.objectKey
+        ? undefined
+        : this.managedPayloadPath(resultRef));
+
+    if (!payloadPath && !input.objectKey) {
+      throw new Error("result ref requires payloadPath or objectKey");
+    }
+    if (!input.payloadPath && payloadPath) {
+      await atomicWriteJson(payloadPath, wrapPayload(input.data, createdAt, input.projectId));
+    }
+
+    const preview = buildPreview(input.data);
+    const sizeBytes = estimateBytes(input.data);
+    await this.db.query(
+      `
+        INSERT INTO ${this.db.table("result_refs")} (
+          result_ref,
+          actor_key,
+          session_id,
+          created_at,
+          kind,
+          preview,
+          project_id,
+          project_key,
+          schema_version,
+          size_bytes,
+          source,
+          trace_id,
+          payload_path,
+          object_key
+        )
+        VALUES ($1, $2, $3, $4, $5, $6::jsonb, $7, $8, $9, $10, $11, $12, $13, $14)
+      `,
+      [
+        resultRef,
+        input.actorKey,
+        input.sessionId,
+        createdAt,
+        input.kind,
+        JSON.stringify(preview),
+        input.projectId ?? null,
+        input.projectKey,
+        input.schemaVersion,
+        sizeBytes,
+        input.source,
+        input.traceId,
+        payloadPath ?? null,
+        input.objectKey ?? null,
+      ],
+    );
+
+    return {
       resultRef,
       actorKey: input.actorKey,
-      clientSessionId: input.clientSessionId,
-      createdAt: new Date().toISOString(),
+      createdAt,
       data: input.data,
       kind: input.kind,
-      preview: buildPreview(input.data),
+      preview,
       projectId: input.projectId,
       projectKey: input.projectKey,
       schemaVersion: input.schemaVersion,
       sessionId: input.sessionId,
-      sizeBytes: estimateBytes(input.data),
+      sizeBytes,
       source: input.source,
       traceId: input.traceId,
-    };
-    await atomicWriteJson(this.filePath(resultRef), record);
-    return record;
+      payloadPath,
+      objectKey: input.objectKey,
+    } satisfies ResultReferenceRecord;
   }
 
   async getAuthorizedRecord(resultRef: string, context: RetrievalContext) {
@@ -145,26 +218,49 @@ export class ResultReferenceStore {
       return null;
     }
 
-    const rawRecord = await readJsonFile<unknown>(this.filePath(normalizedResultRef));
-    const record =
-      normalizeResultReferenceRecord(rawRecord) ??
-      normalizeLegacyRenderReferenceRecord(rawRecord, normalizedResultRef, context);
-    if (!record) {
+    const result = await this.db.query<ResultReferenceRow>(
+      `
+        SELECT *
+        FROM ${this.db.table("result_refs")}
+        WHERE result_ref = $1
+        LIMIT 1
+      `,
+      [normalizedResultRef],
+    );
+    const row = result.rows[0];
+    if (!row) {
       return null;
     }
-    if (record.actorKey !== context.actorKey) {
+    if (row.actor_key !== context.actorKey) {
       return null;
     }
-    if ((record.projectId ?? "") !== (context.projectId ?? "")) {
+    if ((row.project_id ?? "") !== (context.projectId ?? "")) {
       return null;
     }
-    if (
-      context.clientSessionId &&
-      record.clientSessionId !== context.clientSessionId
-    ) {
+    if (context.sessionId && row.session_id !== context.sessionId) {
       return null;
     }
-    return record;
+    const data = await this.readPayload(row);
+    if (data === null) {
+      return null;
+    }
+    return {
+      resultRef: row.result_ref,
+      actorKey: row.actor_key,
+      createdAt: toIsoString(row.created_at),
+      data,
+      kind: row.kind,
+      preview: row.preview,
+      projectId: row.project_id ?? undefined,
+      projectKey: row.project_key,
+      schemaVersion: row.schema_version,
+      sessionId: row.session_id,
+      sizeBytes: row.size_bytes,
+      source: row.source,
+      traceId: row.trace_id,
+      payloadPath: row.payload_path ?? undefined,
+      objectKey: row.object_key ?? undefined,
+    } satisfies ResultReferenceRecord;
   }
 
   async peekAuthorized(
@@ -184,196 +280,82 @@ export class ResultReferenceStore {
   }
 
   async listBySession(sessionId: string) {
-    const files = await listJsonFiles(this.baseDir);
-    const records = await Promise.all(
-      files.map(async (filePath) =>
-        normalizeResultReferenceRecord(await readJsonFile<unknown>(filePath)),
-      ),
+    const result = await this.db.query<ResultReferenceRow>(
+      `
+        SELECT *
+        FROM ${this.db.table("result_refs")}
+        WHERE session_id = $1
+        ORDER BY created_at DESC
+      `,
+      [sessionId],
     );
-    return records
-      .filter((record): record is ResultReferenceRecord => Boolean(record))
-      .filter((record) => record.sessionId === sessionId)
-      .sort((left, right) => right.createdAt.localeCompare(left.createdAt));
+    const records = await Promise.all(
+      result.rows.map(async (row) => {
+        const data = await this.readPayload(row);
+        if (data === null) {
+          return null;
+        }
+        return {
+          resultRef: row.result_ref,
+          actorKey: row.actor_key,
+          createdAt: toIsoString(row.created_at),
+          data,
+          kind: row.kind,
+          preview: row.preview,
+          projectId: row.project_id ?? undefined,
+          projectKey: row.project_key,
+          schemaVersion: row.schema_version,
+          sessionId: row.session_id,
+          sizeBytes: row.size_bytes,
+          source: row.source,
+          traceId: row.trace_id,
+          payloadPath: row.payload_path ?? undefined,
+          objectKey: row.object_key ?? undefined,
+        } satisfies ResultReferenceRecord;
+      }),
+    );
+    return records.filter(Boolean) as ResultReferenceRecord[];
   }
 
   async cleanupExpired() {
-    const files = await listJsonFiles(this.baseDir);
-    const now = Date.now();
-    for (const filePath of files) {
-      const stats = await getFileStat(filePath);
-      if (!stats) {
-        continue;
-      }
-      if (now - stats.mtimeMs > this.ttlMs) {
-        await removeFileIfExists(filePath);
-      }
-    }
+    const result = await this.db.query<ResultReferenceRow>(
+      `
+        DELETE FROM ${this.db.table("result_refs")}
+        WHERE created_at < NOW() - ($1 * INTERVAL '1 millisecond')
+        RETURNING *
+      `,
+      [this.ttlMs],
+    );
+    await Promise.all(
+      result.rows.map(async (row) => {
+        if (row.payload_path && isManagedPayloadPath(row.payload_path, this.managedPayloadDir)) {
+          await removeFileIfExists(row.payload_path);
+        }
+      }),
+    );
   }
 
-  private filePath(resultRef: string) {
-    return join(this.baseDir, `${resultRef}.json`);
+  private managedPayloadPath(resultRef: string) {
+    return join(this.managedPayloadDir, `${resultRef}.json`);
+  }
+
+  private async readPayload(row: ResultReferenceRow) {
+    if (row.payload_path) {
+      return unwrapStoredPayload(await readJsonFile<unknown>(row.payload_path));
+    }
+    if (row.object_key) {
+      logger.warn({ resultRef: row.result_ref, objectKey: row.object_key }, "object storage payloads are not implemented");
+      return null;
+    }
+    return null;
   }
 }
 
-export const normalizeResultReferenceRecord = (
-  value: unknown,
-): ResultReferenceRecord | null => {
-  if (!isRecord(value)) {
-    return null;
-  }
-
-  const partial = value as PartialRecord;
-  if (
-    !isValidResultRef(partial.resultRef) ||
-    typeof partial.actorKey !== "string" ||
-    typeof partial.clientSessionId !== "string" ||
-    typeof partial.createdAt !== "string" ||
-    !("data" in partial) ||
-    !isResultPreview(partial.preview) ||
-    typeof partial.projectKey !== "string" ||
-    typeof partial.sessionId !== "string" ||
-    typeof partial.sizeBytes !== "number" ||
-    !Number.isFinite(partial.sizeBytes) ||
-    typeof partial.traceId !== "string"
-  ) {
-    return null;
-  }
-
-  const kind = normalizeResultReferenceKind(partial.kind);
-  const source = normalizeResultReferenceSource(partial.source);
-  const schemaVersion =
-    typeof partial.schemaVersion === "number" &&
-    Number.isInteger(partial.schemaVersion) &&
-    partial.schemaVersion > 0
-      ? partial.schemaVersion
-      : 1;
-
-  if (!kind || !source) {
-    return null;
-  }
-
-  if (
-    partial.projectId !== undefined &&
-    typeof partial.projectId !== "string"
-  ) {
-    return null;
-  }
-
-  return {
-    resultRef: partial.resultRef,
-    actorKey: partial.actorKey,
-    clientSessionId: partial.clientSessionId,
-    createdAt: partial.createdAt,
-    data: partial.data,
-    kind,
-    preview: partial.preview,
-    projectId: partial.projectId,
-    projectKey: partial.projectKey,
-    schemaVersion,
-    sessionId: partial.sessionId,
-    sizeBytes: partial.sizeBytes,
-    source,
-    traceId: partial.traceId,
-  };
-};
-
-const normalizeResultReferenceKind = (
-  value: unknown,
-): ResultReferenceKind | null => {
-  if (value === undefined) {
-    return RESULT_REFERENCE_KIND.dynamicHttpResult;
-  }
-  return Object.values(RESULT_REFERENCE_KIND).includes(
-    value as ResultReferenceKind,
-  )
-    ? (value as ResultReferenceKind)
-    : null;
-};
-
-const normalizeResultReferenceSource = (
-  value: unknown,
-): ResultReferenceSource | null => {
-  if (value === undefined) {
-    return RESULT_REFERENCE_SOURCE.legacy;
-  }
-  return Object.values(RESULT_REFERENCE_SOURCE).includes(
-    value as ResultReferenceSource,
-  )
-    ? (value as ResultReferenceSource)
-    : null;
-};
-
-const isValidResultRef = (value: unknown): value is string =>
-  typeof value === "string" && RESULT_REF_PATTERN.test(value);
-
 const normalizeResultRef = (value: string) => {
-  const match = value.trim().match(RESULT_REF_FILE_PATTERN);
-  return match?.[1] ?? null;
+  const normalized = value.trim();
+  return RESULT_REF_PATTERN.test(normalized) ? normalized : null;
 };
 
-const normalizeLegacyRenderReferenceRecord = (
-  value: unknown,
-  resultRef: string,
-  context: RetrievalContext,
-): ResultReferenceRecord | null => {
-  const data = extractLegacyRenderPayload(value);
-  if (!data) {
-    return null;
-  }
-
-  const root = isRecord(value) ? value : {};
-  const metadata = isRecord(root.metadata) ? root.metadata : {};
-  const projectId = firstNonEmptyString(root.projectId, metadata.projectId);
-  const createdAt =
-    firstNonEmptyString(root.createdAt, metadata.createdAt) ?? new Date().toISOString();
-
-  return {
-    resultRef,
-    actorKey: context.actorKey,
-    clientSessionId: context.clientSessionId ?? "",
-    createdAt,
-    data,
-    kind: RESULT_REFERENCE_KIND.renderJunctionsPayload,
-    preview: buildPreview(data),
-    projectId,
-    projectKey: toProjectKey(projectId),
-    schemaVersion: 1,
-    sessionId: context.clientSessionId ?? resultRef,
-    sizeBytes: estimateBytes(data),
-    source: RESULT_REFERENCE_SOURCE.legacy,
-    traceId: "legacy-render-ref",
-  };
-};
-
-const extractLegacyRenderPayload = (value: unknown) => {
-  if (!isRecord(value)) {
-    return null;
-  }
-  const candidate = isRecord(value.data) ? value.data : value;
-  if (!isRecord(candidate.node_area_map)) {
-    return null;
-  }
-  return candidate;
-};
-
-const firstNonEmptyString = (...values: unknown[]) => {
-  for (const value of values) {
-    if (typeof value === "string" && value.trim().length > 0) {
-      return value.trim();
-    }
-  }
-  return undefined;
-};
-
-const isResultPreview = (value: unknown): value is ResultPreview =>
-  isRecord(value) &&
-  typeof value.count === "number" &&
-  Number.isFinite(value.count) &&
-  Array.isArray(value.fields) &&
-  value.fields.every((field) => typeof field === "string") &&
-  typeof value.summary === "string" &&
-  "sample" in value;
-
 const estimateBytes = (data: unknown) => Buffer.byteLength(JSON.stringify(data));
 
 const buildPreview = (data: unknown): ResultPreview => {
@@ -414,5 +396,26 @@ const buildPreview = (data: unknown): ResultPreview => {
   };
 };
 
+const wrapPayload = (data: unknown, createdAt: string, projectId?: string) => ({
+  metadata: {
+    createdAt,
+    ...(projectId ? { projectId } : {}),
+  },
+  data,
+});
+
 const isRecord = (value: unknown): value is Record<string, unknown> =>
   typeof value === "object" && value !== null && !Array.isArray(value);
+
+const unwrapStoredPayload = (value: unknown) => {
+  if (!isRecord(value)) {
+    return value;
+  }
+  return "data" in value ? value.data : value;
+};
+
+const toIsoString = (value: Date | string) =>
+  value instanceof Date ? value.toISOString() : new Date(value).toISOString();
+
+const isManagedPayloadPath = (payloadPath: string, managedPayloadDir: string) =>
+  resolve(payloadPath).startsWith(resolve(managedPayloadDir));