From de5dc979883a360199568f39dbe1ac250a7beaf4 Mon Sep 17 00:00:00 2001 From: NightWatcher314 Date: Wed, 27 May 2026 11:16:17 +0800 Subject: [PATCH] feat(cli): add image display MCP tool (#700) --- cli/src/claude/utils/startHappyServer.ts | 74 +++++++++++++- cli/src/claude/utils/systemPrompt.ts | 1 + cli/src/codex/codexRemoteLauncher.ts | 63 ++++++++++-- cli/src/codex/happyMcpStdioBridge.ts | 34 ++++++- cli/src/codex/utils/systemPrompt.ts | 1 + .../modules/common/generatedImages.test.ts | 70 +++++++++++++ cli/src/modules/common/generatedImages.ts | 99 +++++++++++++++---- cli/src/modules/common/handlers/files.ts | 3 +- cli/src/opencode/utils/systemPrompt.ts | 1 + 9 files changed, 314 insertions(+), 32 deletions(-) create mode 100644 cli/src/modules/common/generatedImages.test.ts diff --git a/cli/src/claude/utils/startHappyServer.ts b/cli/src/claude/utils/startHappyServer.ts index dac93576..df0b0fd4 100644 --- a/cli/src/claude/utils/startHappyServer.ts +++ b/cli/src/claude/utils/startHappyServer.ts @@ -5,12 +5,14 @@ import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js"; import { createServer } from "node:http"; +import { lstat, readFile } from "node:fs/promises"; import { StreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/streamableHttp.js"; import { AddressInfo } from "node:net"; import { z } from "zod"; import { logger } from "@/ui/logger"; import { ApiSessionClient } from "@/api/apiSession"; import { randomUUID } from "node:crypto"; +import { detectImageMimeType, registerGeneratedImage } from "@/modules/common/generatedImages"; type StartHappyServerOptions = { emitTitleSummary?: boolean; @@ -52,6 +54,11 @@ export async function startHappyServer(client: ApiSessionClient, options: StartH title: z.string().describe('The new title for the chat session'), }); + const displayImageInputSchema: z.ZodTypeAny = z.object({ + path: z.string().describe('Local filesystem path of the image to display to the user'), + title: z.string().optional().describe('Optional display title or filename for the image'), + }); + mcp.registerTool('change_title', { description: 'Change the title of the current chat session', title: 'Change Chat Title', @@ -83,6 +90,71 @@ export async function startHappyServer(client: ApiSessionClient, options: StartH } }); + + mcp.registerTool('display_image', { + description: 'Display a local image file inline in the current HAPI chat session', + title: 'Display Image', + inputSchema: displayImageInputSchema, + }, async (args: { path: string; title?: string }) => { + logger.debug('[hapiMCP] Display image:', args.path); + + try { + const info = await lstat(args.path); + if (!info.isFile()) { + throw new Error('Path is not a regular file'); + } + + const maxImageBytes = 25 * 1024 * 1024; + if (info.size > maxImageBytes) { + throw new Error('Image is too large to display inline'); + } + + const bytes = await readFile(args.path); + const mimeType = detectImageMimeType(bytes); + if (!mimeType) { + throw new Error('Unsupported image content'); + } + + const image = registerGeneratedImage({ + id: randomUUID(), + path: args.path, + fileName: args.title, + mimeType, + bytes + }); + + client.sendAgentMessage({ + type: 'generated-image', + imageId: image.id, + fileName: image.fileName, + mimeType: image.mimeType, + id: randomUUID() + }); + + return { + content: [ + { + type: 'text' as const, + text: `Displayed image: ${image.fileName}`, + }, + ], + isError: false, + }; + } catch (error) { + const message = error instanceof Error ? error.message : String(error); + logger.debug('[hapiMCP] Failed to display image:', message); + return { + content: [ + { + type: 'text' as const, + text: `Failed to display image: ${message}`, + }, + ], + isError: true, + }; + } + }); + const transport = new StreamableHTTPServerTransport({ // NOTE: Returning session id here will result in claude // sdk spawn to fail with `Invalid Request: Server already initialized` @@ -114,7 +186,7 @@ export async function startHappyServer(client: ApiSessionClient, options: StartH return { url: baseUrl.toString(), - toolNames: ['change_title'], + toolNames: ['change_title', 'display_image'], stop: () => { logger.debug('[hapiMCP] Stopping server'); mcp.close(); diff --git a/cli/src/claude/utils/systemPrompt.ts b/cli/src/claude/utils/systemPrompt.ts index 502fff43..238a5c9a 100644 --- a/cli/src/claude/utils/systemPrompt.ts +++ b/cli/src/claude/utils/systemPrompt.ts @@ -6,6 +6,7 @@ import { shouldIncludeCoAuthoredBy } from "./claudeSettings"; */ const BASE_SYSTEM_PROMPT = (() => trimIdent(` ALWAYS when you start a new chat - you must call a tool "mcp__hapi__change_title" to set a chat title. When you think chat title is not relevant anymore - call the tool again to change it. When chat name is too generic and you have a change to make it more specific - call the tool again to change it. This title is needed to easily find the chat in the future. Help human. + When you create or find a local image file that the user should see, call the tool "mcp__hapi__display_image" with the image path so HAPI can show it inline. `))(); /** diff --git a/cli/src/codex/codexRemoteLauncher.ts b/cli/src/codex/codexRemoteLauncher.ts index fb69993d..2cb80fe1 100644 --- a/cli/src/codex/codexRemoteLauncher.ts +++ b/cli/src/codex/codexRemoteLauncher.ts @@ -1,5 +1,6 @@ import React from 'react'; import { randomUUID } from 'node:crypto'; +import { lstat, readFile } from 'node:fs/promises'; import { CodexAppServerClient } from './codexAppServerClient'; import { CodexPermissionHandler } from './utils/permissionHandler'; @@ -13,7 +14,7 @@ import type { CodexSession } from './session'; import type { EnhancedMode } from './loop'; import { hasCodexCliOverrides } from './utils/codexCliOverrides'; import { AppServerEventConverter } from './utils/appServerEventConverter'; -import { registerGeneratedImage } from '@/modules/common/generatedImages'; +import { detectImageMimeType, registerGeneratedImage } from '@/modules/common/generatedImages'; import { registerAppServerPermissionHandlers } from './utils/appServerPermissionAdapter'; import { buildThreadStartParams, buildTurnStartParams } from './utils/appServerConfig'; import type { ThreadGoal, ThreadGoalStatus } from './appServerTypes'; @@ -25,6 +26,35 @@ import { type RemoteLauncherExitReason } from '@/modules/common/remote/RemoteLauncherBase'; + +async function registerGeneratedImageFromPath(args: { id: string; path: string; fileName?: string | null }): Promise | null> { + try { + const info = await lstat(args.path); + if (!info.isFile()) { + throw new Error('Path is not a regular file'); + } + const maxImageBytes = 25 * 1024 * 1024; + if (info.size > maxImageBytes) { + throw new Error('Image is too large to display inline'); + } + const bytes = await readFile(args.path); + const mimeType = detectImageMimeType(bytes); + if (!mimeType) { + throw new Error('Unsupported image content'); + } + return registerGeneratedImage({ + id: args.id, + path: args.path, + fileName: args.fileName, + mimeType, + bytes + }); + } catch (error) { + logger.debug('[CodexRemoteLauncher] Failed to register generated image:', error instanceof Error ? error.message : String(error)); + return null; + } +} + type HappyServer = Awaited>['server']; type QueuedMessage = { message: string; mode: EnhancedMode; isolate: boolean; hash: string }; type ChildAgentRuntime = { @@ -1924,7 +1954,9 @@ class CodexRemoteLauncher extends RemoteLauncherBase { return true; }; - const handleCodexEvent = (msg: Record) => { + let codexEventQueue: Promise | null = null; + + const handleCodexEvent = async (msg: Record): Promise => { const msgType = asString(msg.type); if (!msgType) return; const eventTurnId = asString(msg.turn_id ?? msg.turnId); @@ -2183,12 +2215,13 @@ class CodexRemoteLauncher extends RemoteLauncherBase { const imageId = randomUUID(); const savedPath = asString(msg.saved_path ?? msg.savedPath); if (savedPath) { - const image = registerGeneratedImage({ + const image = await registerGeneratedImageFromPath({ id: imageId, path: savedPath, - fileName: asString(msg.file_name ?? msg.fileName), - mimeType: asString(msg.mime_type ?? msg.mimeType) + fileName: asString(msg.file_name ?? msg.fileName) }); + if (!image) return; + messageBuffer.addMessage(`Generated image: ${image.fileName}`, 'assistant'); session.sendAgentMessage({ type: 'generated-image', @@ -2467,7 +2500,25 @@ class CodexRemoteLauncher extends RemoteLauncherBase { const events = appServerEventConverter.handleNotification(method, params); for (const event of events) { const eventRecord = asRecord(event) ?? { type: undefined }; - handleCodexEvent(eventRecord); + const msgType = asString(eventRecord.type); + const hasGeneratedImagePath = msgType === 'generated_image' && Boolean(asString(eventRecord.saved_path ?? eventRecord.savedPath)); + + if (codexEventQueue || hasGeneratedImagePath) { + const previousQueue = codexEventQueue ?? Promise.resolve(); + const nextQueue = previousQueue + .then(() => handleCodexEvent(eventRecord)) + .catch((error) => logger.debug('[Codex] Failed to handle app-server event:', error instanceof Error ? error.message : String(error))); + const queued = nextQueue.finally(() => { + if (codexEventQueue === queued) { + codexEventQueue = null; + } + }); + codexEventQueue = queued; + } else { + void handleCodexEvent(eventRecord).catch((error) => { + logger.debug('[Codex] Failed to handle app-server event:', error instanceof Error ? error.message : String(error)); + }); + } } }); diff --git a/cli/src/codex/happyMcpStdioBridge.ts b/cli/src/codex/happyMcpStdioBridge.ts index 8ef0b829..7c30617c 100644 --- a/cli/src/codex/happyMcpStdioBridge.ts +++ b/cli/src/codex/happyMcpStdioBridge.ts @@ -1,7 +1,7 @@ /** * HAPI MCP STDIO Bridge * - * Minimal STDIO MCP server exposing a single tool `change_title`. + * Minimal STDIO MCP server exposing HAPI tools such as `change_title` and `display_image`. * On invocation it forwards the tool call to an existing HAPI HTTP MCP server * using the StreamableHTTPClientTransport. * @@ -64,7 +64,7 @@ export async function runHappyMcpStdioBridge(argv: string[]): Promise { version: '1.0.0', }); - // Register the single tool and forward to HTTP MCP + // Register tools and forward to HTTP MCP const changeTitleInputSchema: z.ZodTypeAny = z.object({ title: z.string().describe('The new title for the chat session'), }); @@ -93,6 +93,36 @@ export async function runHappyMcpStdioBridge(argv: string[]): Promise { } ); + + + const displayImageInputSchema: z.ZodTypeAny = z.object({ + path: z.string().describe('Local filesystem path of the image to display to the user'), + title: z.string().optional().describe('Optional display title or filename for the image'), + }); + + server.registerTool( + 'display_image', + { + description: 'Display a local image file inline in the current HAPI chat session', + title: 'Display Image', + inputSchema: displayImageInputSchema, + }, + async (args: Record) => { + try { + const client = await ensureHttpClient(); + const response = await client.callTool({ name: 'display_image', arguments: args }); + return response as any; + } catch (error) { + return { + content: [ + { type: 'text' as const, text: `Failed to display image: ${error instanceof Error ? error.message : String(error)}` }, + ], + isError: true, + }; + } + } + ); + // Start STDIO transport const stdio = new StdioServerTransport(); await server.connect(stdio); diff --git a/cli/src/codex/utils/systemPrompt.ts b/cli/src/codex/utils/systemPrompt.ts index c8be6620..f5a17141 100644 --- a/cli/src/codex/utils/systemPrompt.ts +++ b/cli/src/codex/utils/systemPrompt.ts @@ -17,6 +17,7 @@ export const TITLE_INSTRUCTION = trimIdent(` Prefer calling functions.hapi__change_title. If that exact tool name is unavailable, call an equivalent alias such as hapi__change_title, mcp__hapi__change_title, or hapi_change_title. If the task focus changes significantly later, call the title tool again with a better title. + When you create or find a local image file that the user should see, call functions.hapi__display_image with the image path. If that exact tool name is unavailable, use an equivalent alias such as hapi__display_image, mcp__hapi__display_image, or hapi_display_image. `); /** diff --git a/cli/src/modules/common/generatedImages.test.ts b/cli/src/modules/common/generatedImages.test.ts new file mode 100644 index 00000000..471539e4 --- /dev/null +++ b/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() + }) + +}) diff --git a/cli/src/modules/common/generatedImages.ts b/cli/src/modules/common/generatedImages.ts index 39c45f20..fdf6cb94 100644 --- a/cli/src/modules/common/generatedImages.ts +++ b/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 = { - '.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() +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 } diff --git a/cli/src/modules/common/handlers/files.ts b/cli/src/modules/common/handlers/files.ts index 460d2b7d..0e3d28ba 100644 --- a/cli/src/modules/common/handlers/files.ts +++ b/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 } diff --git a/cli/src/opencode/utils/systemPrompt.ts b/cli/src/opencode/utils/systemPrompt.ts index 97dc0084..fef60a16 100644 --- a/cli/src/opencode/utils/systemPrompt.ts +++ b/cli/src/opencode/utils/systemPrompt.ts @@ -12,6 +12,7 @@ import { trimIdent } from '@/utils/trimIdent'; */ export const TITLE_INSTRUCTION = trimIdent(` ALWAYS when you start a new chat - you must call the tool "hapi_change_title" to set a chat title. When you think chat title is not relevant anymore - call the tool again to change it. When chat name is too generic and you have a chance to make it more specific - call the tool again to change it. This title is needed to easily find the chat in the future. Help human. + When you create or find a local image file that the user should see, call the tool "hapi_display_image" with the image path so HAPI can show it inline. `); /**