422 lines
11 KiB
TypeScript
422 lines
11 KiB
TypeScript
import { randomUUID } from "node:crypto";
|
|
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,
|
|
readJsonFile,
|
|
removeFileIfExists,
|
|
} from "../utils/fileStore.js";
|
|
|
|
export const RESULT_REF_PATTERN = /^res-[a-f0-9-]{8,64}$/;
|
|
|
|
export const RESULT_REFERENCE_KIND = {
|
|
dynamicHttpResult: "dynamic-http-result",
|
|
renderJunctionsPayload: "render-junctions-payload",
|
|
} as const;
|
|
|
|
export const RESULT_REFERENCE_SOURCE = {
|
|
dynamicHttp: "dynamic_http",
|
|
agentGenerated: "agent_generated",
|
|
legacy: "legacy",
|
|
migration: "migration",
|
|
} as const;
|
|
|
|
export type ResultReferenceKind =
|
|
(typeof RESULT_REFERENCE_KIND)[keyof typeof RESULT_REFERENCE_KIND];
|
|
|
|
export type ResultReferenceSource =
|
|
(typeof RESULT_REFERENCE_SOURCE)[keyof typeof RESULT_REFERENCE_SOURCE];
|
|
|
|
export type ResultPreview = {
|
|
count: number;
|
|
fields: string[];
|
|
sample: unknown;
|
|
summary: string;
|
|
};
|
|
|
|
export type ResultReferenceRecord = {
|
|
resultRef: string;
|
|
actorKey: string;
|
|
createdAt: string;
|
|
data: unknown;
|
|
kind: ResultReferenceKind;
|
|
preview: ResultPreview;
|
|
projectId?: string;
|
|
projectKey: string;
|
|
schemaVersion: number;
|
|
sessionId: string;
|
|
sizeBytes: number;
|
|
source: ResultReferenceSource;
|
|
traceId: string;
|
|
payloadPath?: string;
|
|
objectKey?: string;
|
|
};
|
|
|
|
export type StoreResultInput = {
|
|
actorKey: string;
|
|
data: unknown;
|
|
kind: ResultReferenceKind;
|
|
projectId?: string;
|
|
projectKey: string;
|
|
schemaVersion: number;
|
|
sessionId: string;
|
|
source: ResultReferenceSource;
|
|
traceId: string;
|
|
payloadPath?: string;
|
|
objectKey?: string;
|
|
};
|
|
|
|
export type RetrievalContext = {
|
|
actorKey: string;
|
|
sessionId?: string;
|
|
projectId?: string;
|
|
};
|
|
|
|
export type ResultReferencePeek = {
|
|
resultRef: string;
|
|
kind: ResultReferenceKind;
|
|
preview: ResultPreview;
|
|
storedAt: string;
|
|
};
|
|
|
|
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 Promise.all([this.db.initialize(), ensureDirectory(this.managedPayloadDir)]);
|
|
}
|
|
|
|
startCleanupLoop() {
|
|
if (this.cleanupTimer) {
|
|
return;
|
|
}
|
|
this.cleanupTimer = setInterval(() => {
|
|
void this.cleanupExpired().catch((error) => {
|
|
logger.warn({ err: error }, "result ref cleanup failed");
|
|
});
|
|
}, config.RESULT_REF_CLEANUP_INTERVAL_MS);
|
|
this.cleanupTimer.unref?.();
|
|
}
|
|
|
|
stopCleanupLoop() {
|
|
if (this.cleanupTimer) {
|
|
clearInterval(this.cleanupTimer);
|
|
this.cleanupTimer = null;
|
|
}
|
|
}
|
|
|
|
async store(input: StoreResultInput) {
|
|
const resultRef = `res-${randomUUID().slice(0, 16)}`;
|
|
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,
|
|
createdAt,
|
|
data: input.data,
|
|
kind: input.kind,
|
|
preview,
|
|
projectId: input.projectId,
|
|
projectKey: input.projectKey,
|
|
schemaVersion: input.schemaVersion,
|
|
sessionId: input.sessionId,
|
|
sizeBytes,
|
|
source: input.source,
|
|
traceId: input.traceId,
|
|
payloadPath,
|
|
objectKey: input.objectKey,
|
|
} satisfies ResultReferenceRecord;
|
|
}
|
|
|
|
async getAuthorizedRecord(resultRef: string, context: RetrievalContext) {
|
|
const normalizedResultRef = normalizeResultRef(resultRef);
|
|
if (!normalizedResultRef) {
|
|
return null;
|
|
}
|
|
|
|
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 (row.actor_key !== context.actorKey) {
|
|
return null;
|
|
}
|
|
if ((row.project_id ?? "") !== (context.projectId ?? "")) {
|
|
return null;
|
|
}
|
|
if (context.sessionId && row.session_id !== context.sessionId) {
|
|
return null;
|
|
}
|
|
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(
|
|
resultRef: string,
|
|
context: RetrievalContext,
|
|
): Promise<ResultReferencePeek | null> {
|
|
const record = await this.getAuthorizedRecord(resultRef, context);
|
|
if (!record) {
|
|
return null;
|
|
}
|
|
return {
|
|
resultRef: record.resultRef,
|
|
kind: record.kind,
|
|
preview: record.preview,
|
|
storedAt: record.createdAt,
|
|
};
|
|
}
|
|
|
|
async listBySession(sessionId: string) {
|
|
const result = await this.db.query<ResultReferenceRow>(
|
|
`
|
|
SELECT *
|
|
FROM ${this.db.table("result_refs")}
|
|
WHERE session_id = $1
|
|
ORDER BY created_at DESC
|
|
`,
|
|
[sessionId],
|
|
);
|
|
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 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 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;
|
|
}
|
|
}
|
|
|
|
const normalizeResultRef = (value: string) => {
|
|
const normalized = value.trim();
|
|
return RESULT_REF_PATTERN.test(normalized) ? normalized : null;
|
|
};
|
|
|
|
const estimateBytes = (data: unknown) => Buffer.byteLength(JSON.stringify(data));
|
|
|
|
const buildPreview = (data: unknown): ResultPreview => {
|
|
if (Array.isArray(data)) {
|
|
const sample = data.slice(0, config.MAX_PREVIEW_SAMPLE_ITEMS);
|
|
const fields =
|
|
sample.length > 0 && isRecord(sample[0])
|
|
? Object.keys(sample[0]).slice(0, 30)
|
|
: [];
|
|
return {
|
|
count: data.length,
|
|
fields,
|
|
sample,
|
|
summary: `list[${data.length}]`,
|
|
};
|
|
}
|
|
|
|
if (isRecord(data)) {
|
|
const fields = Object.keys(data).slice(0, 30);
|
|
const sample = Object.fromEntries(
|
|
fields
|
|
.slice(0, config.MAX_PREVIEW_SAMPLE_ITEMS)
|
|
.map((field) => [field, data[field]]),
|
|
);
|
|
return {
|
|
count: fields.length,
|
|
fields,
|
|
sample,
|
|
summary: `object<${fields.length} fields>`,
|
|
};
|
|
}
|
|
|
|
return {
|
|
count: 1,
|
|
fields: [],
|
|
sample: String(data).slice(0, 300),
|
|
summary: `scalar<${typeof data}>`,
|
|
};
|
|
};
|
|
|
|
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));
|