import { useCallback, useEffect, useRef, useState } from "react"; import { encodeImage, type EncodeFailure } from "@/lib/imageEncode"; import type { WebUIIngressLimits } from "@/lib/types"; /** Lifecycle stages of one attachment: * * - ``encoding`` — posted to the Worker / read from disk; chip shows a spinner * - ``ready`` — ``dataUrl`` available; safe to submit * - ``error`` — validation / decode failure; chip shows inline error */ export type AttachmentStatus = "encoding" | "ready" | "error"; export type AttachmentKind = "image" | "file"; export interface AttachedAttachment { id: string; kind: AttachmentKind; file: File; /** Optimistic ``blob:`` preview URL; revoked on ``remove`` / ``clear`` / * unmount. */ previewUrl?: string; status: AttachmentStatus; /** Populated when ``status === "ready"``. */ dataUrl?: string; /** Size of the final encoded payload (base64 bytes decoded). */ encodedBytes?: number; /** Whether the Worker re-encoded the image to hit the size budget. */ normalized?: boolean; /** Human-readable validation / encoding error when ``status === "error"``. */ error?: AttachmentError; } export type AttachedImage = AttachedAttachment; export interface RestoredReadyAttachment { dataUrl: string; name?: string; kind?: AttachmentKind; } export type RestoredReadyImage = RestoredReadyAttachment; /** Machine-readable rejection reasons surfaced as inline chip errors. * * Callers localize these via the ``composer.imageRejected.*`` i18n table. */ export type AttachmentError = | "unsupported_type" // server whitelist excludes this MIME | "empty_file" // backend data-URL decoder rejects empty payloads | "too_many_attachments" // per-message cap (4) reached before enqueue | "total_too_large" // decoded attachments exceed the business-policy total | "transport_too_large" // projected JSON frame exceeds the transport guard | "magic_mismatch" // extension lies about the real content | "decode_failed" // Worker couldn't decode / re-encode | "too_large" // even after normalization we exceed the budget | "io"; // file read failed at the browser layer export const MAX_ATTACHMENTS_PER_MESSAGE = 4; export const MAX_ATTACHMENT_BYTES = 6 * 1024 * 1024; export const MAX_TOTAL_ATTACHMENT_BYTES = 24 * 1024 * 1024; /** MIME whitelist — mirrors the server's and the ```` attr. */ const ACCEPTED_IMAGE_MIMES: ReadonlySet = new Set([ "image/png", "image/jpeg", "image/webp", "image/gif", ]); const DOCUMENT_MIME_BY_EXTENSION: ReadonlyMap = new Map([ [".pdf", "application/pdf"], [".docx", "application/vnd.openxmlformats-officedocument.wordprocessingml.document"], [".xlsx", "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"], [".pptx", "application/vnd.openxmlformats-officedocument.presentationml.presentation"], [".txt", "text/plain"], [".md", "text/markdown"], [".csv", "text/csv"], [".json", "application/json"], [".xml", "application/xml"], [".html", "text/html"], [".htm", "text/html"], [".log", "text/plain"], [".yaml", "application/yaml"], [".yml", "application/yaml"], [".toml", "application/toml"], [".ini", "text/plain"], [".cfg", "text/plain"], ]); const ACCEPTED_DOCUMENT_MIMES: ReadonlySet = new Set(DOCUMENT_MIME_BY_EXTENSION.values()); export const ACCEPT_ATTR = [ ...ACCEPTED_IMAGE_MIMES, ...ACCEPTED_DOCUMENT_MIMES, ...DOCUMENT_MIME_BY_EXTENSION.keys(), ].join(","); function extensionOf(name: string): string { const dot = name.lastIndexOf("."); return dot < 0 ? "" : name.slice(dot).toLowerCase(); } function mimeForFile(file: File): string { const byName = DOCUMENT_MIME_BY_EXTENSION.get(extensionOf(file.name)); if (byName) return byName; if (!file.type || file.type === "application/octet-stream") { return "application/octet-stream"; } return file.type; } function projectedDataUrlBytes( file: File, kind: AttachmentKind, maxFileBytes: number, ): number { const prefixBytes = `data:${mimeForFile(file)};base64,`.length; const decodedBytes = kind === "image" ? Math.min(file.size, maxFileBytes) : file.size; return prefixBytes + 4 * Math.ceil(decodedBytes / 3); } function positiveLimit(value: number | null | undefined, fallback: number): number { return typeof value === "number" && Number.isFinite(value) && value > 0 ? Math.floor(value) : fallback; } function attachmentPayloadBudget(limits: WebUIIngressLimits | null | undefined): number | null { const maxFrameBytes = limits?.transport.max_frame_bytes; if (typeof maxFrameBytes !== "number" || !Number.isFinite(maxFrameBytes)) { return null; } return Math.max( 0, Math.floor(maxFrameBytes) - positiveLimit(limits?.message.max_text_bytes, 0) - positiveLimit(limits?.transport.envelope_reserve_bytes, 0), ); } export function acceptedAttachmentKind(file: File): AttachmentKind | null { if (DOCUMENT_MIME_BY_EXTENSION.has(extensionOf(file.name))) return "file"; if (ACCEPTED_IMAGE_MIMES.has(file.type)) return "image"; const mime = mimeForFile(file); if (ACCEPTED_DOCUMENT_MIMES.has(mime)) return "file"; return null; } function dataUrlMime(dataUrl: string): string { const match = /^data:([^;,]+)[;,]/.exec(dataUrl); return match?.[1] || "image/png"; } function kindFromDataUrl(dataUrl: string): AttachmentKind { return dataUrlMime(dataUrl).startsWith("image/") ? "image" : "file"; } function dataUrlToFile(dataUrl: string, name?: string): File { const mime = dataUrlMime(dataUrl); const fallbackName = `image.${mime.split("/")[1] || "png"}`; try { const [, base64 = ""] = dataUrl.split(",", 2); const binary = atob(base64); const bytes = new Uint8Array(binary.length); for (let i = 0; i < binary.length; i += 1) { bytes[i] = binary.charCodeAt(i); } return new File([bytes], name || fallbackName, { type: mime }); } catch { return new File([], name || fallbackName, { type: mime }); } } function uuid(): string { if (typeof crypto !== "undefined" && "randomUUID" in crypto) { return (crypto as Crypto).randomUUID(); } return `img-${Date.now()}-${Math.random().toString(36).slice(2)}`; } function bufferToBase64(buf: ArrayBuffer): string { const bytes = new Uint8Array(buf); let binary = ""; const chunk = 0x8000; for (let i = 0; i < bytes.length; i += chunk) { binary += String.fromCharCode.apply( null, bytes.subarray(i, i + chunk) as unknown as number[], ); } return btoa(binary); } async function encodeFile(file: File, maxFileBytes: number): Promise<{ ok: true; dataUrl: string; bytes: number; } | { ok: false; reason: AttachmentError; }> { if (file.size > maxFileBytes) return { ok: false, reason: "too_large" }; try { const buffer = await file.arrayBuffer(); return { ok: true, dataUrl: `data:${mimeForFile(file)};base64,${bufferToBase64(buffer)}`, bytes: file.size, }; } catch { return { ok: false, reason: "io" }; } } function mapEncodeFailure(reason: EncodeFailure["reason"]): AttachmentError { switch (reason) { case "invalid_mime": case "magic_mismatch": return "magic_mismatch"; case "too_large_after_normalize": return "too_large"; case "io": return "io"; case "decode_failed": default: return "decode_failed"; } } export interface UseAttachedImagesApi { images: AttachedAttachment[]; /** Enqueue new files. Returns the list of rejected files so the caller can * surface inline errors. Files rejected client-side (wrong MIME, limit) are * *not* added to ``images`` — only recoverable read/encoding failures show * up as error chips. */ enqueue: (files: Iterable) => { rejected: Array<{ file: File; reason: AttachmentError }>; }; remove: (id: string) => { nextFocusId: string | null }; /** Revoke every staged blob URL and drop all attachments. Called after a * successful submit — the optimistic bubble holds onto an independent * ``data:`` URL so tearing down blob previews here is safe. */ clear: () => void; /** Restore already-encoded attachments, e.g. a queued composer draft moving * back into the input. These entries are immediately sendable and use image * ``data:`` URLs as stable previews. */ restoreReadyImages: (images: RestoredReadyAttachment[]) => void; /** ``true`` when at least one attachment is still encoding — Send should wait. */ encoding: boolean; /** ``true`` when we've hit ``MAX_ATTACHMENTS_PER_MESSAGE``. */ full: boolean; } interface UseAttachedImagesOptions { ingressLimits?: WebUIIngressLimits | null; } /** Manage the lifecycle of attachments in the Composer. * * Responsibilities in one place: * - validation (MIME whitelist, count cap) * - blob URL creation + revocation * - Worker orchestration * - focus bookkeeping so keyboard delete doesn't strand the user */ export function useAttachedImages({ ingressLimits = null, }: UseAttachedImagesOptions = {}): UseAttachedImagesApi { const [images, setImages] = useState([]); const maxAttachments = positiveLimit( ingressLimits?.attachments.max_count, MAX_ATTACHMENTS_PER_MESSAGE, ); const maxFileBytes = positiveLimit( ingressLimits?.attachments.max_file_bytes, MAX_ATTACHMENT_BYTES, ); const maxTotalBytes = positiveLimit( ingressLimits?.attachments.max_total_bytes, MAX_TOTAL_ATTACHMENT_BYTES, ); // Ref mirror so ``enqueue`` can see the authoritative length when invoked // multiple times in a single tick (rapid file selection, drag of many // files, paste storms). ``state`` is stale for that second + call. const imagesRef = useRef([]); imagesRef.current = images; const setEntry = useCallback((id: string, patch: Partial) => { setImages((prev) => { const next = prev.map((img) => (img.id === id ? { ...img, ...patch } : img)); imagesRef.current = next; return next; }); }, []); const enqueue = useCallback( (files: Iterable) => { const rejected: Array<{ file: File; reason: AttachmentError }> = []; const toAdd: AttachedAttachment[] = []; let slot = maxAttachments - imagesRef.current.length; const payloadBudget = attachmentPayloadBudget(ingressLimits); let projectedWireBytes = imagesRef.current.reduce( (total, image) => total + ( image.dataUrl?.length ?? projectedDataUrlBytes(image.file, image.kind, maxFileBytes) ), 0, ); let projectedDecodedBytes = imagesRef.current.reduce( (total, image) => total + ( image.encodedBytes ?? (image.kind === "image" ? Math.min(image.file.size, maxFileBytes) : image.file.size) ), 0, ); for (const file of files) { const kind = acceptedAttachmentKind(file); if (!kind) { rejected.push({ file, reason: "unsupported_type" }); continue; } if (file.size === 0) { rejected.push({ file, reason: "empty_file" }); continue; } if (kind === "file" && file.size > maxFileBytes) { rejected.push({ file, reason: "too_large" }); continue; } if (slot <= 0) { rejected.push({ file, reason: "too_many_attachments" }); continue; } const nextDecodedBytes = kind === "image" ? Math.min(file.size, maxFileBytes) : file.size; if (projectedDecodedBytes + nextDecodedBytes > maxTotalBytes) { rejected.push({ file, reason: "total_too_large" }); continue; } const nextWireBytes = projectedDataUrlBytes(file, kind, maxFileBytes); if (payloadBudget !== null && projectedWireBytes + nextWireBytes > payloadBudget) { rejected.push({ file, reason: "transport_too_large" }); continue; } slot -= 1; projectedDecodedBytes += nextDecodedBytes; projectedWireBytes += nextWireBytes; toAdd.push({ id: uuid(), kind, file, ...(kind === "image" ? { previewUrl: URL.createObjectURL(file) } : {}), status: "encoding", }); } if (toAdd.length > 0) { const next = [...imagesRef.current, ...toAdd]; imagesRef.current = next; setImages(next); // Fire the Worker after the commit so chips render first (good INP). for (const entry of toAdd) { queueMicrotask(() => { const work = entry.kind === "image" ? encodeImage(entry.file) : encodeFile(entry.file, maxFileBytes); work.then( (result) => { if (result.ok) { setEntry(entry.id, { status: "ready", dataUrl: result.dataUrl, encodedBytes: result.bytes, normalized: "normalized" in result ? result.normalized : false, }); } else { setEntry(entry.id, { status: "error", error: entry.kind === "image" ? mapEncodeFailure(result.reason as EncodeFailure["reason"]) : result.reason as AttachmentError, }); } }, () => { setEntry(entry.id, { status: "error", error: "decode_failed", }); }, ); }); } } return { rejected }; }, [ingressLimits, maxAttachments, maxFileBytes, maxTotalBytes, setEntry], ); const remove = useCallback((id: string) => { let nextFocusId: string | null = null; setImages((prev) => { const idx = prev.findIndex((img) => img.id === id); if (idx === -1) return prev; const target = prev[idx]; if (target.previewUrl) { try { URL.revokeObjectURL(target.previewUrl); } catch { // No-op: previewUrl revocation is best-effort. } } const next = [...prev.slice(0, idx), ...prev.slice(idx + 1)]; imagesRef.current = next; // Prefer moving focus to the chip at the same index, else previous. const candidate = next[idx] ?? next[idx - 1]; nextFocusId = candidate?.id ?? null; return next; }); return { nextFocusId }; }, []); const clear = useCallback(() => { setImages((prev) => { for (const img of prev) { if (img.previewUrl) { try { URL.revokeObjectURL(img.previewUrl); } catch { // revoke is best-effort } } } imagesRef.current = []; return []; }); }, []); const restoreReadyImages = useCallback((restored: RestoredReadyAttachment[]) => { const toRestore = restored .filter((img) => acceptedAttachmentKind(dataUrlToFile(img.dataUrl, img.name))) .slice(0, maxAttachments) .map((img): AttachedAttachment => { const file = dataUrlToFile(img.dataUrl, img.name); const kind = img.kind ?? kindFromDataUrl(img.dataUrl); return { id: uuid(), kind, file, ...(kind === "image" ? { previewUrl: img.dataUrl } : {}), status: "ready", dataUrl: img.dataUrl, encodedBytes: file.size, }; }); setImages((prev) => { for (const img of prev) { if (img.previewUrl) { try { URL.revokeObjectURL(img.previewUrl); } catch { // revoke is best-effort } } } imagesRef.current = toRestore; return toRestore; }); }, [maxAttachments]); // Final safety net: revoke any outstanding blob URLs on unmount. Safe // under StrictMode double-invoke because revoked blob URLs are only // referenced from in-hook chip state, which is rebuilt on remount. useEffect(() => { return () => { for (const img of imagesRef.current) { if (img.previewUrl) { try { URL.revokeObjectURL(img.previewUrl); } catch { // best-effort cleanup on unmount } } } }; }, []); const encoding = images.some((img) => img.status === "encoding"); const full = images.length >= maxAttachments; return { images, enqueue, remove, clear, restoreReadyImages, encoding, full }; }