ai-chatbot/lib/storage/image-byte-cache.ts

91 lines
4.2 KiB
TypeScript

import "server-only";
import { createHash, randomUUID } from "node:crypto";
import { mkdir, readFile, writeFile, rename, readdir, stat, statfs, unlink } from "node:fs/promises";
import { join } from "node:path";
import { tmpdir } from "node:os";
export type ImageBytes = { buffer: Buffer; contentType: string };
const directory = process.env.CANVAS_IMAGE_CACHE_DIR || join(tmpdir(), "canvas-image-bytes-v1");
const ttl = 60 * 60 * 1000;
const maxBytes = 512 * 1024 * 1024;
const state = globalThis as typeof globalThis & {
__imageByteFlights?: Map<string, Promise<ImageBytes>>;
__imageCacheCleanupAt?: number;
};
const flights = (state.__imageByteFlights ??= new Map());
function key(value: string) { return createHash("sha256").update(value).digest("hex"); }
async function cleanup() {
if ((state.__imageCacheCleanupAt || 0) > Date.now() - 60_000) return;
state.__imageCacheCleanupAt = Date.now();
const files = await Promise.all((await readdir(directory)).map(async (name) => {
const path = join(directory, name);
return { path, ...(await stat(path)) };
}));
let total = files.reduce((sum, file) => sum + file.size, 0);
for (const file of files.sort((a, b) => a.mtimeMs - b.mtimeMs)) {
if (file.mtimeMs < Date.now() - ttl || total > maxBytes) {
await unlink(file.path).catch(() => undefined);
total -= file.size;
}
}
}
export async function cacheImageBytes(identity: string, data: ImageBytes) {
if (data.buffer.length > maxBytes / 4) return;
try {
await mkdir(directory, { recursive: true, mode: 0o700 });
const disk = await statfs(directory);
if (disk.bavail * disk.bsize < data.buffer.length + 1024 * 1024 * 1024) return;
// One atomic file keeps metadata and bytes together, even across process restarts.
const header = Buffer.from(JSON.stringify({ type: data.contentType, expires: Date.now() + ttl }) + "\n");
const path = join(directory, key(identity));
const temp = `${path}.${randomUUID()}.tmp`;
try {
await writeFile(temp, Buffer.concat([header, data.buffer]), { mode: 0o600 });
await rename(temp, path);
} finally { await unlink(temp).catch(() => undefined); }
await cleanup();
} catch { /* A full/unavailable cache must never fail generation or downloads. */ }
}
export async function readImageBytes(identity: string): Promise<ImageBytes | null> {
try {
const bytes = await readFile(join(directory, key(identity)));
const end = bytes.indexOf(10);
if (end < 0 || end > 1024) return null;
const header = JSON.parse(bytes.subarray(0, end).toString());
if (header.expires <= Date.now()) return null;
return { buffer: bytes.subarray(end + 1), contentType: header.type };
} catch { return null; }
}
export async function onceImageBytes(identity: string, load: () => Promise<ImageBytes>): Promise<ImageBytes> {
const existing = flights.get(identity);
if (existing) return existing;
const promise = (async () => {
const cached = await readImageBytes(identity);
if (cached) return cached;
const result = await load();
await cacheImageBytes(identity, result);
return result;
})();
flights.set(identity, promise);
try { return await promise; } finally { if (flights.get(identity) === promise) flights.delete(identity); }
}
export async function fetchImageBytes(url: string, options: { headers?: HeadersInit; timeoutMs?: number; scope?: string } = {}) {
return onceImageBytes(`source:${options.scope || ""}:${url}`, async () => {
if (url.startsWith("data:")) {
const match = url.match(/^data:(image\/[a-z0-9.+-]+);base64,([a-z0-9+/=]+)$/i);
if (!match) throw new Error("Invalid inline generated image");
return { buffer: Buffer.from(match[2], "base64"), contentType: match[1] };
}
const response = await fetch(url, {
headers: options.headers, cache: "no-store", redirect: "follow",
// The deadline covers the body as well as the response headers.
signal: AbortSignal.timeout(options.timeoutMs || 60_000),
});
if (!response.ok) { await response.body?.cancel(); throw new Error(`HTTP ${response.status}`); }
const buffer = Buffer.from(await response.arrayBuffer());
if (!buffer.length) throw new Error("Downloaded image is empty");
return { buffer, contentType: response.headers.get("content-type") || "" };
});
}