源码
packages/desktop/src/main/attachments.ts
import { randomUUID } from "node:crypto";
import { copyFile, mkdir, open, readdir, readFile, realpath, rm, stat, writeFile } from "node:fs/promises";
import { basename, extname, isAbsolute, join } from "node:path";
import { contentText } from "@zpi-dev/ai";
import {
type Attachment,
createReadTool,
type FileAttachment,
type ImageAttachment,
imageFormat,
imageLimits,
isImageAttachment,
prepareInlineImage,
readFileBounded,
validateImageSource,
} from "@zpi-dev/coding-agent";
import { attachmentLimitMessage } from "@zpi-dev/ui/errors";
import { atomicJson } from "./storage.ts";
type Saved = Attachment & {
used: boolean;
queued?: boolean;
cached?: boolean;
};
// 与 ZCode 普通附件一致：已知文本预读，二进制只交付路径，避免把乱码塞进模型上下文。
const textExtensions =
/\.(cjs|conf|cpp|cs|css|csv|go|h|hpp|html|ini|java|js|json|jsx|log|md|mjs|py|rs|sh|sql|toml|ts|tsx|txt|xml|yaml|yml)$/i;
const mimeTypes: Record<string, string> = {
pdf: "application/pdf",
json: "application/json",
csv: "text/csv",
md: "text/markdown",
mp4: "video/mp4",
m4v: "video/x-m4v",
webm: "video/webm",
mov: "video/quicktime",
mkv: "video/x-matroska",
avi: "video/x-msvideo",
svg: "image/svg+xml",
mp3: "audio/mpeg",
wav: "audio/wav",
};
export class AttachmentStore {
private root: string;
constructor(root: string) {
this.root = root;
}
private directory(sessionId: string) {
if (!/^[a-zA-Z0-9_-]+$/.test(sessionId)) throw new Error("invalid_input: 无效会话 ID");
return join(this.root, sessionId);
}
private imagePath(sessionId: string, id: string, mimeType = "image/png") {
if (typeof id !== "string" || !/^[a-zA-Z0-9_-]+$/.test(id)) throw new Error("invalid_input: 无效附件 ID");
const extension = {
"image/png": "png",
"image/jpeg": "jpg",
"image/gif": "gif",
"image/webp": "webp",
"image/bmp": "bmp",
}[mimeType];
if (!extension) throw new Error("invalid_input: 图片格式无效");
return join(this.directory(sessionId), `${id}.${extension}`);
}
private filePath(sessionId: string, id: string) {
return this.imagePath(sessionId, id).replace(/\.png$/, ".file");
}
async importFile(sessionId: string, path: string): Promise<Attachment> {
path = await realpath(path);
const info = await stat(path);
if (!info.isFile()) throw new Error("invalid_input: 附件必须是文件");
const name = basename(path);
if (/\.(png|jpe?g|gif|webp|bmp)$/i.test(name)) {
if (info.size > imageLimits.sourceBytes) throw new Error("invalid_input: 图片单张上限为 10 MiB");
return this.import(sessionId, name, await readFileBounded(path, imageLimits.sourceBytes));
}
return this.saveFile(sessionId, name, path, info.size);
}
async importBytes(sessionId: string, name: string, bytes: Uint8Array): Promise<Attachment> {
if (!name || name.length > 250 || !(bytes instanceof Uint8Array) || bytes.length > 20 * 1024 * 1024)
throw new Error("invalid_input: 附件导入参数无效（最多 20 MiB）");
if (imageFormat(bytes)) return this.import(sessionId, name, bytes);
return this.saveFile(sessionId, basename(name), undefined, bytes.length, bytes);
}
private saveFile(
sessionId: string,
name: string,
path: string | undefined,
size: number,
bytes?: Uint8Array,
) {
return this.serial(sessionId, async () => {
const entries = await this.list(sessionId);
if (entries.filter((e) => !e.used && !e.queued).length >= imageLimits.count)
throw new Error(`invalid_input: ${attachmentLimitMessage}`);
const id = randomUUID();
const item: FileAttachment = {
id,
name,
path: path ?? this.filePath(sessionId, id),
size,
mimeType:
mimeTypes[extname(name).slice(1).toLowerCase()] ??
(textExtensions.test(name) ? "text/plain" : "application/octet-stream"),
};
await mkdir(this.directory(sessionId), { recursive: true });
if (bytes) await writeFile(item.path, bytes, { mode: 0o600 });
try {
atomicJson(join(this.directory(sessionId), "index.json"), [
...entries,
{ ...item, used: false, ...(bytes ? { cached: true } : {}) },
]);
} catch (error) {
if (bytes) await rm(item.path, { force: true });
throw error;
}
return item;
});
}
private async list(sessionId: string): Promise<Saved[]> {
try {
return JSON.parse(await readFile(join(this.directory(sessionId), "index.json"), "utf8"));
} catch (e) {
if ((e as NodeJS.ErrnoException).code === "ENOENT") return [];
throw e;
}
}
private operations = new Map<string, Promise<unknown>>();
private serial<T>(id: string, action: () => Promise<T>): Promise<T> {
const result = (this.operations.get(id) ?? Promise.resolve()).catch(() => {}).then(action);
this.operations.set(id, result);
void result
.finally(() => {
if (this.operations.get(id) === result) this.operations.delete(id);
})
.catch(() => {});
return result;
}
async cleanUnsent(
onError: (error: unknown, sessionId: string) => void = () => {},
recoverUsed?: (sessionId: string, candidates: string[]) => string[],
): Promise<void> {
await mkdir(this.root, { recursive: true });
for (const dir of await readdir(this.root, { withFileTypes: true })) {
if (!dir.isDirectory() || !/^[a-zA-Z0-9_-]+$/.test(dir.name)) continue;
try {
const candidates = (await this.list(dir.name)).filter((item) => !item.used).map((item) => item.id);
if (candidates.length && recoverUsed)
await this.markConsumed(dir.name, recoverUsed(dir.name, candidates));
const recovered = await this.list(dir.name);
await this.markQueued(
dir.name,
recovered.filter((item) => item.used && item.queued).map((item) => item.id),
false,
);
for (const item of recovered) if (!item.used && !item.queued) await this.remove(dir.name, item.id);
const keep = new Set(
(await this.list(dir.name)).map((item) =>
basename(
isImageAttachment(item)
? this.imagePath(dir.name, item.id, item.mimeType)
: this.filePath(dir.name, item.id),
),
),
);
for (const file of await readdir(this.directory(dir.name)))
if (/\.(png|jpg|gif|webp|bmp|file)$/.test(file) && !keep.has(file))
await rm(join(this.directory(dir.name), file), { force: true });
} catch (error) {
onError(error, dir.name);
}
}
}
async import(sessionId: string, name: string, bytes: Uint8Array): Promise<ImageAttachment> {
if (
typeof name !== "string" ||
name.length > 250 ||
!name ||
!(bytes instanceof Uint8Array) ||
bytes.length > imageLimits.sourceBytes
)
throw new Error("invalid_input: 图片导入参数无效（单张最多 10 MiB）");
return this.serial(sessionId, async () => {
const input = Buffer.from(bytes);
const processed = await validateImageSource(input);
const entries = await this.list(sessionId);
if (entries.filter((e) => !e.used && !e.queued).length >= imageLimits.count)
throw new Error(`invalid_input: ${attachmentLimitMessage}`);
const item: Saved = {
id: randomUUID(),
name,
mimeType: processed.mimeType,
width: processed.width,
height: processed.height,
size: input.length,
used: false,
...(processed.warning ? { warning: processed.warning } : {}),
};
await mkdir(this.directory(sessionId), { recursive: true });
await writeFile(this.imagePath(sessionId, item.id, item.mimeType), input, { mode: 0o600 });
try {
atomicJson(join(this.directory(sessionId), "index.json"), [...entries, item]);
} catch (e) {
await rm(this.imagePath(sessionId, item.id, item.mimeType), { force: true });
throw e;
}
const { used: _, queued: _queued, cached: _cached, ...metadata } = item;
return metadata;
});
}
async read(sessionId: string, id: string) {
return this.readSaved(
sessionId,
(await this.list(sessionId)).find((e) => e.id === id),
);
}
private async readSaved(sessionId: string, item: Saved | undefined) {
const image = await this.readSavedBytes(sessionId, item);
return { metadata: image.metadata, data: image.data.toString("base64") };
}
private async readSavedBytes(sessionId: string, item: Saved | undefined) {
if (!item) throw new Error("not_found: 附件不存在或不属于此会话");
this.imagePath(sessionId, item.id);
const { used: _, queued: _queued, cached: _cached, ...metadata } = item;
if (!isImageAttachment(metadata)) {
if (!isAbsolute(metadata.path) || (item.cached && metadata.path !== this.filePath(sessionId, item.id)))
throw new Error("invalid_input: 附件文件路径无效");
return { metadata, data: Buffer.alloc(0) };
}
const data = await readFileBounded(
this.imagePath(sessionId, item.id, item.mimeType),
imageLimits.sourceBytes,
);
// Old indexes and interrupted writes must not bypass source validation or pixel limits.
const source = await validateImageSource(data);
const { warning: _warning, ...savedMetadata } = metadata;
return {
metadata: {
...savedMetadata,
mimeType: source.mimeType,
width: source.width,
height: source.height,
size: data.length,
...(source.warning ? { warning: source.warning } : {}),
},
data,
};
}
async load(sessionId: string, ids: string[]) {
if (
!Array.isArray(ids) ||
ids.length > imageLimits.count ||
ids.some((id) => typeof id !== "string") ||
new Set(ids).size !== ids.length
)
throw new Error("invalid_input: 附件列表无效（最多 8 个）");
const entries = new Map<string, Saved>();
for (const item of ids.length ? await this.list(sessionId) : [])
if (!entries.has(item.id)) entries.set(item.id, item);
const loaded = await Promise.all(ids.map((id) => this.readSavedBytes(sessionId, entries.get(id))));
const images = loaded.filter((e): e is typeof e & { metadata: ImageAttachment } =>
isImageAttachment(e.metadata),
);
if (images.reduce((sum, e) => sum + e.metadata.width * e.metadata.height, 0) > imageLimits.messagePixels)
throw new Error("invalid_input: 本条消息图片总像素超限");
const inlineImages = await Promise.all(
images.map(async ({ metadata, data }) => {
const prepared = await prepareInlineImage(data, metadata.mimeType);
if (!prepared.ok) throw new Error(`invalid_input: ${prepared.message}`);
return {
image: { type: "image" as const, data: prepared.data, mimeType: prepared.mimeType },
note: prepared.hints.length
? `Attached image ${JSON.stringify(metadata.name)}: ${prepared.hints.join("\n")}`
: "",
};
}),
);
const files: string[] = inlineImages.map((prepared) => prepared.note).filter(Boolean);
for (const { metadata } of loaded) {
if (isImageAttachment(metadata)) continue;
const info = await stat(metadata.path);
if (!info.isFile()) throw new Error("invalid_input: 附件文件不可用");
const header = `Attached ${metadata.mimeType}: ${JSON.stringify(metadata.path)} (${metadata.name})`;
const reference = `${header}\nUse the available file reading tools to inspect this file when needed.`;
if (textExtensions.test(metadata.name)) {
// Limit native-file I/O as well as model output. Large attachments remain available by path.
const file = await open(metadata.path, "r");
let prefix: string;
try {
const buffer = Buffer.alloc(64 * 1024);
const { bytesRead } = await file.read(buffer, 0, buffer.length, 0);
const bytes = buffer.subarray(0, bytesRead);
if (bytes.includes(0)) {
files.push(reference);
continue;
}
try {
prefix = new TextDecoder("utf-8", { fatal: true }).decode(bytes, {
stream: info.size > bytesRead,
});
} catch {
files.push(reference);
continue;
}
} finally {
await file.close();
}
const result = await createReadTool(this.directory(sessionId), {
operations: { access: async () => undefined, readFile: async () => Buffer.from(prefix) },
}).execute("attachment", { path: metadata.path });
files.push(
`${header}\n${contentText(result.content)}${info.size > 64 * 1024 ? "\n[Partial attachment preview. Use read with this path to continue.]" : ""}`,
);
} else files.push(reference);
}
return {
metadata: loaded.map((e) => e.metadata),
images: inlineImages.map((prepared) => prepared.image),
fileContext: files.join("\n\n"),
};
}
private async mark(sessionId: string, ids: string[], flags: Partial<Pick<Saved, "used" | "queued">>) {
if (!ids.length) return;
const selected = new Set(ids);
await this.serial(sessionId, async () => {
const entries = await this.list(sessionId);
atomicJson(
join(this.directory(sessionId), "index.json"),
entries.map((item) => (selected.has(item.id) ? { ...item, ...flags } : item)),
);
});
}
markUsed(sessionId: string, ids: string[]) {
return this.mark(sessionId, ids, { used: true });
}
markQueued(sessionId: string, ids: string[], queued: boolean) {
return this.mark(sessionId, ids, { queued });
}
markConsumed(sessionId: string, ids: string[]) {
return this.mark(sessionId, ids, { used: true, queued: false });
}
async remove(sessionId: string, id: string) {
this.imagePath(sessionId, id);
await this.serial(sessionId, async () => {
const entries = await this.list(sessionId);
const item = entries.find((e) => e.id === id);
if (!item || item.used || item.queued) return;
atomicJson(
join(this.directory(sessionId), "index.json"),
entries.filter((e) => e.id !== id),
);
if (isImageAttachment(item)) await rm(this.imagePath(sessionId, id, item.mimeType), { force: true });
else if (item.cached) await rm(this.filePath(sessionId, id), { force: true });
});
}
async copyTo(sourceId: string, targetId: string, ids: string[]): Promise<Map<string, string>> {
const relocated = new Map<string, string>();
if (!ids.length) return relocated;
const entries = new Map((await this.list(sourceId)).map((item) => [item.id, item]));
await this.serial(targetId, async () => {
const metadata = new Map((await this.list(targetId)).map((item) => [item.id, item]));
await mkdir(this.directory(targetId), { recursive: true });
for (const id of new Set(ids)) {
const image = await this.readSavedBytes(sourceId, entries.get(id));
if (isImageAttachment(image.metadata)) {
await writeFile(this.imagePath(targetId, id, image.metadata.mimeType), image.data, { mode: 0o600 });
metadata.set(id, { ...image.metadata, used: true });
} else if (entries.get(id)?.cached) {
const path = this.filePath(targetId, id);
await copyFile(this.filePath(sourceId, id), path);
metadata.set(id, { ...image.metadata, path, used: true, cached: true });
relocated.set(image.metadata.path, path);
} else metadata.set(id, { ...image.metadata, used: true });
}
atomicJson(join(this.directory(targetId), "index.json"), [...metadata.values()]);
});
return relocated;
}
async deleteSession(id: string) {
await this.operations.get(id)?.catch(() => {});
await rm(this.directory(id), { recursive: true, force: true });
}
}