mesalogo
dd29f309f1
Expose save_image across direct model, MCP gateway, DeepSeek Harness and remote Agents (written on the remote host via chunked reads), and show the image MCP tool list in settings.
185 lines
10 KiB
TypeScript
185 lines
10 KiB
TypeScript
import { randomUUID } from 'node:crypto'
|
|
import { mkdir, writeFile } from 'node:fs/promises'
|
|
import { createServer, type Server } from 'node:http'
|
|
import { dirname, isAbsolute } from 'node:path'
|
|
import { Server as McpServer } from '@modelcontextprotocol/sdk/server/index.js'
|
|
import { StreamableHTTPServerTransport } from '@modelcontextprotocol/sdk/server/streamableHttp.js'
|
|
import { CallToolRequestSchema, ListToolsRequestSchema, type CallToolResult } from '@modelcontextprotocol/sdk/types.js'
|
|
import { z } from 'zod'
|
|
import {
|
|
imageSaveMimeTypeForPath, imageSaveToolInputSchema, imageSaveToolName,
|
|
imageToolInputSchema, imageToolName, type ImageSaveResult, type ImageSaveToolInput
|
|
} from '../shared/image-generation-contracts'
|
|
import { storyGraphTools, type StoryGraphToolName } from '../shared/story-graph-tools'
|
|
import {
|
|
decodeRemoteImageToolMessage, encodeRemoteImageToolMessage,
|
|
remoteImageSaveChunkSchema, remoteImageToolMaximumBytes, remoteImageToolReplySchema, type remoteImageToolSchema
|
|
} from '../shared/remote-image-tool-contracts'
|
|
|
|
/** One prompt's local MCP adapter. It never holds provider credentials or image bytes. */
|
|
export class AgentImageToolMcp {
|
|
private readonly calls = new Map<string, { input: string; result: Promise<CallToolResult> }>()
|
|
private readonly pending = new Map<string, {
|
|
resolve: (value: CallToolResult) => void
|
|
reject: (error: Error) => void
|
|
}>()
|
|
private server?: Server
|
|
private closed = false
|
|
url?: string
|
|
|
|
constructor(
|
|
readonly descriptor: z.infer<typeof remoteImageToolSchema>,
|
|
private readonly send: (payload: Uint8Array) => Promise<void>
|
|
) {}
|
|
|
|
async start(): Promise<void> {
|
|
const path = `/${randomUUID()}/mcp`
|
|
const schema = z.toJSONSchema(imageToolInputSchema, { target: 'draft-7', io: 'input' })
|
|
Reflect.deleteProperty(schema, '$schema')
|
|
const saveSchema = z.toJSONSchema(imageSaveToolInputSchema, { target: 'draft-7', io: 'input' })
|
|
Reflect.deleteProperty(saveSchema, '$schema')
|
|
this.server = createServer((request, response) => {
|
|
void (async () => {
|
|
if (this.closed || request.url !== path) { response.writeHead(404).end(); return }
|
|
if (request.method !== 'POST') { response.writeHead(405).end(); return }
|
|
const chunks: Buffer[] = []
|
|
let size = 0
|
|
for await (const chunk of request) {
|
|
const bytes = Buffer.from(chunk)
|
|
size += bytes.length
|
|
if (size > remoteImageToolMaximumBytes) { response.writeHead(413).end(); return }
|
|
chunks.push(bytes)
|
|
}
|
|
const body: unknown = JSON.parse(Buffer.concat(chunks).toString('utf8'))
|
|
const mcp = new McpServer({ name: 'goodbuddy-image', version: '1' }, { capabilities: { tools: {} } })
|
|
mcp.setRequestHandler(ListToolsRequestSchema, async () => {
|
|
const available = this.descriptor.storyGraph ? await this.call({}, 'story_graph_list') : undefined
|
|
const graphEnabled = available?.content.some(part => part.type === 'text' && JSON.parse(part.text).available === true)
|
|
return { tools: [
|
|
...(this.descriptor.description ? [{ name: imageToolName, description: this.descriptor.description, inputSchema: schema as { type: 'object' } }] : []),
|
|
...(this.descriptor.saveDescription ? [{ name: imageSaveToolName, description: this.descriptor.saveDescription, inputSchema: saveSchema as { type: 'object' },
|
|
annotations: { readOnlyHint: false, destructiveHint: true } }] : []),
|
|
...(graphEnabled ? storyGraphTools.map(tool => ({ name: tool.name, description: tool.description,
|
|
inputSchema: z.toJSONSchema(tool.inputSchema, { target: 'draft-7', io: 'input' }) as { type: 'object' },
|
|
annotations: { readOnlyHint: true, destructiveHint: false } })) : [])
|
|
] }
|
|
})
|
|
mcp.setRequestHandler(CallToolRequestSchema, async (call, extra) => {
|
|
const graphTool = storyGraphTools.find(tool => tool.name === call.params.name)
|
|
if (graphTool && this.descriptor.storyGraph) {
|
|
try { return await this.call(call.params.arguments ?? {}, graphTool.name) }
|
|
catch (error) { return { isError: true, content: [{ type: 'text', text: error instanceof Error ? error.message : 'Story Graph read failed' }] } }
|
|
}
|
|
if (call.params.name === imageSaveToolName && this.descriptor.saveDescription) {
|
|
try {
|
|
const saved = await this.saveImage(imageSaveToolInputSchema.parse(call.params.arguments ?? {}))
|
|
return { content: [{ type: 'text', text: JSON.stringify(saved) }] }
|
|
} catch (error) {
|
|
return { isError: true, content: [{ type: 'text', text: error instanceof Error ? error.message : 'Saving image failed' }] }
|
|
}
|
|
}
|
|
if (call.params.name !== imageToolName || !this.descriptor.description) throw new Error('Unknown image tool')
|
|
const input = imageToolInputSchema.parse(call.params.arguments)
|
|
// Repeated delivery of an MCP request returns the same outcome, including unknown delivery.
|
|
const key = `${typeof extra.requestId}:${extra.requestId}`
|
|
let existing = this.calls.get(key)
|
|
if (existing && existing.input !== JSON.stringify(input)) throw new Error('Image tool request ID reused with different input')
|
|
if (!existing) {
|
|
const result = this.call(input).catch((error: unknown) => ({
|
|
isError: true,
|
|
content: [{ type: 'text' as const, text: error instanceof Error ? error.message : 'Image tool delivery failed; do not retry automatically' }]
|
|
}))
|
|
existing = { input: JSON.stringify(input), result }
|
|
this.calls.set(key, existing)
|
|
}
|
|
return await existing.result
|
|
})
|
|
const transport = new StreamableHTTPServerTransport({ sessionIdGenerator: undefined, enableJsonResponse: true })
|
|
response.once('close', () => { void mcp.close().catch(() => undefined) })
|
|
await mcp.connect(transport)
|
|
await transport.handleRequest(request, response, body)
|
|
})().catch(() => {
|
|
if (!response.headersSent) response.writeHead(400)
|
|
response.end()
|
|
})
|
|
})
|
|
await new Promise<void>((resolve, reject) => {
|
|
this.server!.once('error', reject)
|
|
this.server!.listen(0, '127.0.0.1', resolve)
|
|
})
|
|
const address = this.server.address()
|
|
if (!address || typeof address === 'string') throw new Error('Image MCP endpoint did not start')
|
|
this.url = `http://127.0.0.1:${address.port}${path}`
|
|
}
|
|
|
|
onReply(payload: Uint8Array): void {
|
|
const reply = remoteImageToolReplySchema.parse(decodeRemoteImageToolMessage(payload))
|
|
const pending = this.pending.get(reply.callId)
|
|
if (!pending) return
|
|
if (reply.error || (!reply.result && !reply.storyGraphResult && !reply.imageChunk)) pending.reject(new Error(reply.error ?? 'Tool response has no result'))
|
|
else pending.resolve({ content: [{ type: 'text', text: JSON.stringify(reply.imageChunk ?? reply.storyGraphResult ?? reply.result) }] })
|
|
}
|
|
|
|
/** Pulls the encoded image from Desktop in bounded chunks and writes it on this host. */
|
|
private async saveImage(input: ImageSaveToolInput): Promise<ImageSaveResult> {
|
|
const target = input.path
|
|
if (!isAbsolute(target)) throw new Error('path must be an absolute file path')
|
|
const mimeType = imageSaveMimeTypeForPath(target)
|
|
if (!mimeType) throw new Error('path must end with .png, .jpg, .jpeg or .webp')
|
|
const chunks: Buffer[] = []
|
|
let offset = 0
|
|
let totalBytes: number | undefined
|
|
do {
|
|
const result = await this.call({ artifactId: input.artifactId, mimeType, offset }, 'save_image_read')
|
|
const part = result.content[0]
|
|
if (part?.type !== 'text') throw new Error('Image chunk response is invalid')
|
|
const chunk = remoteImageSaveChunkSchema.parse(JSON.parse(part.text))
|
|
if (chunk.offset !== offset || chunk.mimeType !== mimeType || (totalBytes !== undefined && chunk.totalBytes !== totalBytes)) {
|
|
throw new Error('Image chunk response is inconsistent')
|
|
}
|
|
totalBytes = chunk.totalBytes
|
|
const bytes = Buffer.from(chunk.data, 'base64')
|
|
if (bytes.byteLength === 0 || offset + bytes.byteLength > totalBytes) throw new Error('Image chunk response is inconsistent')
|
|
chunks.push(bytes)
|
|
offset += bytes.byteLength
|
|
} while (offset < totalBytes)
|
|
const data = Buffer.concat(chunks)
|
|
await mkdir(dirname(target), { recursive: true })
|
|
try {
|
|
await writeFile(target, data, { flag: input.overwrite ? 'w' : 'wx' })
|
|
} catch (error) {
|
|
if ((error as NodeJS.ErrnoException).code === 'EEXIST') throw new Error('File already exists; choose another path or set overwrite=true', { cause: error })
|
|
throw error
|
|
}
|
|
return { artifactId: input.artifactId, path: target, mimeType, byteSize: data.byteLength }
|
|
}
|
|
|
|
close(): void {
|
|
if (this.closed) return
|
|
this.closed = true
|
|
for (const pending of this.pending.values()) pending.reject(new Error('Desktop image tool disconnected; delivery may be unknown. Do not retry automatically.'))
|
|
this.pending.clear()
|
|
this.server?.closeAllConnections()
|
|
this.server?.close()
|
|
}
|
|
|
|
private async call(input: Record<string, unknown>, name?: StoryGraphToolName | 'story_graph_list' | 'save_image_read'): Promise<CallToolResult> {
|
|
if (this.closed) throw new Error('Image tool prompt has ended')
|
|
const callId = randomUUID()
|
|
let timer: NodeJS.Timeout | undefined
|
|
const result = new Promise<CallToolResult>((resolve, reject) => {
|
|
this.pending.set(callId, { resolve, reject })
|
|
timer = setTimeout(() => reject(new Error('Image tool wait timed out; delivery may be unknown. Do not retry automatically.')), 5 * 60_000)
|
|
timer.unref()
|
|
})
|
|
// Register the wait before sending; send failures must not cause a retry.
|
|
void this.send(encodeRemoteImageToolMessage({ callId, input, ...(name ? { name } : {}) })).catch(() => {
|
|
this.pending.get(callId)?.reject(new Error('Image tool delivery may be unknown. Do not retry automatically.'))
|
|
})
|
|
try { return await result } finally {
|
|
clearTimeout(timer)
|
|
this.pending.delete(callId)
|
|
}
|
|
}
|
|
}
|