85 lines
5.1 KiB
TypeScript
85 lines
5.1 KiB
TypeScript
import { app } from 'electron'
|
||
import { randomUUID } from 'node:crypto'
|
||
import { promises as fs } from 'node:fs'
|
||
import { basename, extname, join, relative, resolve, sep } from 'node:path'
|
||
import type { ChatAttachment } from '@shared/conversation'
|
||
|
||
const MAX_BYTES = 20 * 1024 * 1024
|
||
const MIME: Record<string, string> = { '.png': 'image/png', '.jpg': 'image/jpeg', '.jpeg': 'image/jpeg', '.gif': 'image/gif', '.webp': 'image/webp', '.svg': 'image/svg+xml', '.pdf': 'application/pdf', '.txt': 'text/plain', '.md': 'text/markdown', '.json': 'application/json', '.csv': 'text/csv' }
|
||
function root(): string { return join(app.getPath('userData'), 'attachments') }
|
||
function validId(id: string): boolean { return /^[0-9a-f-]{36}$/.test(id) }
|
||
const committedIds = new Set<string>()
|
||
|
||
export async function stageAttachment(name: string, mimeType: string, bytes: Buffer): Promise<ChatAttachment> {
|
||
if (!bytes.length || bytes.length > MAX_BYTES) throw new Error('附件为空或超过 20 MB')
|
||
const id = randomUUID()
|
||
const safeName = basename(name.replace(/\\/g, '/')).replace(/[<>:"|?*\x00-\x1f]/g, '_') || 'attachment'
|
||
if (['.', '..', 'metadata.json'].includes(safeName.toLowerCase())) throw new Error('附件文件名无效')
|
||
const mime = MIME[extname(safeName).toLowerCase()] ?? mimeType ?? 'application/octet-stream'
|
||
const directory = join(root(), id)
|
||
await fs.mkdir(directory, { recursive: true })
|
||
const path = join(directory, safeName)
|
||
const attachment: ChatAttachment = { id, name: safeName, mimeType: mime, size: bytes.length, path,
|
||
previewUrl: mime.startsWith('image/') ? `data:${mime};base64,${bytes.toString('base64')}` : undefined }
|
||
try {
|
||
await fs.writeFile(path, bytes)
|
||
await fs.writeFile(join(directory, 'metadata.json'), JSON.stringify(attachment))
|
||
} catch (error) { await removeAttachment(id); throw error }
|
||
return attachment
|
||
}
|
||
export async function stageAttachmentPath(path: string): Promise<ChatAttachment> {
|
||
const handle = await fs.open(path, 'r')
|
||
try {
|
||
const stat = await handle.stat()
|
||
if (!stat.isFile() || stat.size > MAX_BYTES) throw new Error('不支持该附件或超过 20 MB')
|
||
const bytes = Buffer.alloc(stat.size + 1)
|
||
let length = 0
|
||
while (length < bytes.length) {
|
||
const result = await handle.read(bytes, length, bytes.length - length, null)
|
||
if (!result.bytesRead) break
|
||
length += result.bytesRead
|
||
}
|
||
if (length > MAX_BYTES) throw new Error('附件超过 20 MB')
|
||
return stageAttachment(basename(path), '', bytes.subarray(0, length))
|
||
} finally { await handle.close() }
|
||
}
|
||
export async function getAttachment(id: string): Promise<ChatAttachment> {
|
||
if (!validId(id)) throw new Error('附件标识无效')
|
||
const attachment = JSON.parse(await fs.readFile(join(root(), id, 'metadata.json'), 'utf8')) as ChatAttachment
|
||
if (typeof attachment.name !== 'string' || basename(attachment.name) !== attachment.name || attachment.name === 'metadata.json') throw new Error('附件元数据无效')
|
||
const path = join(root(), id, attachment.name)
|
||
return { ...attachment, id, path }
|
||
}
|
||
export async function commitAttachments(attachments: ChatAttachment[]): Promise<void> {
|
||
for (const attachment of attachments) committedIds.add(attachment.id)
|
||
for (const attachment of attachments) {
|
||
const stored = await getAttachment(attachment.id)
|
||
await fs.writeFile(join(root(), attachment.id, 'metadata.json'), JSON.stringify({ ...stored, committed: true }))
|
||
}
|
||
}
|
||
export async function removeAttachment(id: string): Promise<void> {
|
||
if (!validId(id)) throw new Error('附件标识无效')
|
||
if (committedIds.has(id)) throw new Error('附件已用于会话，不能从草稿中删除原文件')
|
||
const metadata = await fs.readFile(join(root(), id, 'metadata.json'), 'utf8').catch((error) => {
|
||
if ((error as NodeJS.ErrnoException).code === 'ENOENT') return '{}'
|
||
throw error
|
||
})
|
||
if ((JSON.parse(metadata) as { committed?: boolean }).committed) throw new Error('附件已用于会话，不能从草稿中删除原文件')
|
||
const target = resolve(root(), id)
|
||
const rel = relative(resolve(root()), target)
|
||
if (!rel || rel.startsWith(`..${sep}`) || rel === '..' || rel.includes(sep)) throw new Error('附件目录无效')
|
||
await fs.rm(target, { recursive: true, force: true })
|
||
}
|
||
export async function attachmentPrompt(attachments: ChatAttachment[]): Promise<{ text: string; images: Array<{ type: 'image'; data: string; mimeType: string }> }> {
|
||
const text: string[] = []
|
||
const images: Array<{ type: 'image'; data: string; mimeType: string }> = []
|
||
for (const reference of attachments) {
|
||
const attachment = await getAttachment(reference.id)
|
||
const bytes = await fs.readFile(attachment.path!)
|
||
if (attachment.mimeType.startsWith('image/') && attachment.mimeType !== 'image/svg+xml') images.push({ type: 'image', data: bytes.toString('base64'), mimeType: attachment.mimeType })
|
||
else if (!bytes.includes(0) && bytes.length < 2 * 1024 * 1024) text.push(`附件 ${attachment.name}:\n${bytes.toString('utf8')}`)
|
||
else text.push(`附件 ${attachment.name}（${attachment.mimeType}，${attachment.size} 字节），本地路径：${attachment.path}`)
|
||
}
|
||
return { text: text.join('\n\n'), images }
|
||
}
|