feat(cli): add image display MCP tool (#700)

cli/src/modules/common/generatedImages.test.ts

@@ -0,0 +1,70 @@
+import { describe, expect, it } from 'vitest'
+import { clearGeneratedImages, detectImageMimeType, getGeneratedImage, registerGeneratedImage } from './generatedImages'
+
+describe('generatedImages', () => {
+    it('detects supported image MIME types from file bytes', () => {
+        expect(detectImageMimeType(Buffer.from([0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a]))).toBe('image/png')
+        expect(detectImageMimeType(Buffer.from([0xff, 0xd8, 0xff, 0xdb]))).toBe('image/jpeg')
+        expect(detectImageMimeType(Buffer.from('GIF89a'))).toBe('image/gif')
+        expect(detectImageMimeType(Buffer.from('RIFFxxxxWEBP'))).toBe('image/webp')
+        expect(detectImageMimeType(Buffer.from([0x00, 0x00, 0x00, 0x1c, 0x66, 0x74, 0x79, 0x70, 0x61, 0x76, 0x69, 0x66]))).toBe('image/avif')
+    })
+
+    it('rejects non-image bytes even if the path has an image extension', () => {
+        expect(detectImageMimeType(Buffer.from('not really a png'))).toBeNull()
+    })
+
+    it('stores only validated MIME type supplied by the server', () => {
+        const image = registerGeneratedImage({
+            id: 'test-image',
+            path: '/tmp/example.png',
+            mimeType: 'image/png',
+            bytes: Buffer.from('original image bytes')
+        })
+
+        expect(image.mimeType).toBe('image/png')
+        clearGeneratedImages()
+    })
+
+    it('snapshots image bytes at registration time', () => {
+        const source = Buffer.from('original image bytes')
+        const image = registerGeneratedImage({
+            id: 'snapshot-image',
+            path: '/tmp/example.png',
+            mimeType: 'image/png',
+            bytes: source
+        })
+        source.fill(0)
+
+        expect(image.content.toString()).toBe('original image bytes')
+        expect(getGeneratedImage('snapshot-image')?.content.toString()).toBe('original image bytes')
+        clearGeneratedImages()
+    })
+
+    it('rejects oversized image snapshots', () => {
+        expect(() => registerGeneratedImage({
+            id: 'too-large-image',
+            path: '/tmp/large.png',
+            mimeType: 'image/png',
+            bytes: new Uint8Array(25 * 1024 * 1024 + 1)
+        })).toThrow('Image is too large to display inline')
+        clearGeneratedImages()
+    })
+
+    it('evicts oldest image snapshots when the count limit is exceeded', () => {
+        for (let i = 0; i < 101; i += 1) {
+            registerGeneratedImage({
+                id: `image-${i}`,
+                path: `/tmp/image-${i}.png`,
+                mimeType: 'image/png',
+                bytes: Buffer.from(`image-${i}`)
+            })
+        }
+
+        expect(getGeneratedImage('image-0')).toBeNull()
+        expect(getGeneratedImage('image-1')).not.toBeNull()
+        expect(getGeneratedImage('image-100')).not.toBeNull()
+        clearGeneratedImages()
+    })
+
+})

cli/src/modules/common/generatedImages.ts

@@ -1,50 +1,107 @@
-import { basename, extname } from 'path'
+import { basename } from 'path'
 
 export type GeneratedImageMetadata = {
     id: string
-    path: string
     fileName: string
+    content: Buffer
     mimeType: string
     createdAt: number
 }
 
-const IMAGE_MIME_BY_EXTENSION: Record<string, string> = {
-    '.apng': 'image/apng',
-    '.avif': 'image/avif',
-    '.bmp': 'image/bmp',
-    '.gif': 'image/gif',
-    '.ico': 'image/x-icon',
-    '.jpeg': 'image/jpeg',
-    '.jpg': 'image/jpeg',
-    '.png': 'image/png',
-    '.svg': 'image/svg+xml',
-    '.tif': 'image/tiff',
-    '.tiff': 'image/tiff',
-    '.webp': 'image/webp'
-}
+const MAX_GENERATED_IMAGE_BYTES = 25 * 1024 * 1024
+const MAX_GENERATED_IMAGE_TOTAL_BYTES = 100 * 1024 * 1024
+const MAX_GENERATED_IMAGE_COUNT = 100
 
 const generatedImages = new Map<string, GeneratedImageMetadata>()
+let generatedImageBytes = 0
 
-export function resolveGeneratedImageMimeType(path: string): string {
-    return IMAGE_MIME_BY_EXTENSION[extname(path).toLowerCase()] ?? 'application/octet-stream'
+export function detectImageMimeType(bytes: Uint8Array): string | null {
+    if (bytes.length >= 8
+        && bytes[0] === 0x89
+        && bytes[1] === 0x50
+        && bytes[2] === 0x4e
+        && bytes[3] === 0x47
+        && bytes[4] === 0x0d
+        && bytes[5] === 0x0a
+        && bytes[6] === 0x1a
+        && bytes[7] === 0x0a) {
+        return 'image/png'
+    }
+
+    if (bytes.length >= 3 && bytes[0] === 0xff && bytes[1] === 0xd8 && bytes[2] === 0xff) {
+        return 'image/jpeg'
+    }
+
+    if (bytes.length >= 6) {
+        const header = ascii(bytes, 0, 6)
+        if (header === 'GIF87a' || header === 'GIF89a') {
+            return 'image/gif'
+        }
+    }
+
+    if (bytes.length >= 12 && ascii(bytes, 0, 4) === 'RIFF' && ascii(bytes, 8, 12) === 'WEBP') {
+        return 'image/webp'
+    }
+
+    if (bytes.length >= 12
+        && bytes[0] === 0x00
+        && bytes[1] === 0x00
+        && bytes[2] === 0x00
+        && ascii(bytes, 4, 8) === 'ftyp'
+        && (ascii(bytes, 8, 12) === 'avif' || ascii(bytes, 8, 12) === 'avis')) {
+        return 'image/avif'
+    }
+
+    return null
 }
 
-export function registerGeneratedImage(args: { id: string; path: string; mimeType?: string | null; fileName?: string | null }): GeneratedImageMetadata {
+function ascii(bytes: Uint8Array, start: number, end: number): string {
+    return String.fromCharCode(...bytes.subarray(start, end))
+}
+
+export function registerGeneratedImage(args: { id: string; path: string; mimeType: string; bytes: Uint8Array; fileName?: string | null }): GeneratedImageMetadata {
+    const content = Buffer.from(args.bytes)
+    if (content.byteLength > MAX_GENERATED_IMAGE_BYTES) {
+        throw new Error('Image is too large to display inline')
+    }
+
+    const previous = generatedImages.get(args.id)
+    if (previous) {
+        generatedImageBytes -= previous.content.byteLength
+    }
+
     const metadata: GeneratedImageMetadata = {
         id: args.id,
-        path: args.path,
         fileName: args.fileName || basename(args.path) || `${args.id}.png`,
-        mimeType: args.mimeType || resolveGeneratedImageMimeType(args.path),
+        content,
+        mimeType: args.mimeType,
         createdAt: Date.now()
     }
     generatedImages.set(args.id, metadata)
+    generatedImageBytes += content.byteLength
+
+    evictOldGeneratedImages()
+
     return metadata
 }
 
+function evictOldGeneratedImages(): void {
+    while (generatedImages.size > MAX_GENERATED_IMAGE_COUNT || generatedImageBytes > MAX_GENERATED_IMAGE_TOTAL_BYTES) {
+        const oldestId = generatedImages.keys().next().value
+        if (!oldestId) break
+        const oldest = generatedImages.get(oldestId)
+        if (oldest) {
+            generatedImageBytes -= oldest.content.byteLength
+        }
+        generatedImages.delete(oldestId)
+    }
+}
+
 export function getGeneratedImage(id: string): GeneratedImageMetadata | null {
     return generatedImages.get(id) ?? null
 }
 
 export function clearGeneratedImages(): void {
     generatedImages.clear()
+    generatedImageBytes = 0
 }

cli/src/modules/common/handlers/files.ts

@@ -62,10 +62,9 @@ export function registerFileHandlers(rpcHandlerManager: RpcHandlerManager, worki
         }
 
         try {
-            const buffer = await readFile(image.path)
             return {
                 success: true,
-                content: buffer.toString('base64'),
+                content: image.content.toString('base64'),
                 mimeType: image.mimeType,
                 fileName: image.fileName
             }