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

This commit is contained in:
NightWatcher314
2026-05-27 11:16:17 +08:00
committed by GitHub
parent 166b711fec
commit de5dc97988
9 changed files with 314 additions and 32 deletions
+73 -1
View File
@@ -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<any, any>('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<any, any>('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();
+1
View File
@@ -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.
`))();
/**
+57 -6
View File
@@ -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<ReturnType<typeof registerGeneratedImage> | 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<ReturnType<typeof buildHapiMcpBridge>>['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<string, unknown>) => {
let codexEventQueue: Promise<void> | null = null;
const handleCodexEvent = async (msg: Record<string, unknown>): Promise<void> => {
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));
});
}
}
});
+32 -2
View File
@@ -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<void> {
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<void> {
}
);
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<any, any>(
'display_image',
{
description: 'Display a local image file inline in the current HAPI chat session',
title: 'Display Image',
inputSchema: displayImageInputSchema,
},
async (args: Record<string, unknown>) => {
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);
+1
View File
@@ -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.
`);
/**
@@ -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()
})
})
+78 -21
View File
@@ -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
}
+1 -2
View File
@@ -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
}
+1
View File
@@ -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.
`);
/**