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.
84 lines
4.1 KiB
TypeScript
84 lines
4.1 KiB
TypeScript
import type { ImageToolBinding } from '../agent/image-tool-binding'
|
|
import type { RuntimeProtocolBinaryChannel } from './protocol-remote-runtime-channel'
|
|
import {
|
|
decodeRemoteImageToolMessage, encodeRemoteImageToolMessage,
|
|
remoteImageSaveChunkBytes, remoteImageToolCallSchema, remoteImageToolReplySchema
|
|
} from '../../shared/remote-image-tool-contracts'
|
|
|
|
export class MainImageToolSession {
|
|
private readonly wait = new AbortController()
|
|
private readonly unsubscribe: () => void
|
|
private readonly removeAbort: () => void
|
|
/** Encoded bytes of the image currently being streamed to the Agent; one save at a time per prompt. */
|
|
private saveCache?: { key: string; bytes: Buffer }
|
|
|
|
constructor(
|
|
private readonly channel: RuntimeProtocolBinaryChannel,
|
|
private readonly binding: ImageToolBinding | undefined,
|
|
signal?: AbortSignal,
|
|
private readonly storyGraph?: import('../agent/knowledge-mcp-gateway').StoryGraphRemoteBinding
|
|
) {
|
|
this.unsubscribe = channel.onClose(() => this.close())
|
|
const abort = (): void => this.close()
|
|
signal?.addEventListener('abort', abort, { once: true })
|
|
this.removeAbort = () => signal?.removeEventListener('abort', abort)
|
|
if (signal?.aborted) this.close()
|
|
void this.receive().catch(() => this.close())
|
|
}
|
|
|
|
close(): void {
|
|
if (this.wait.signal.aborted) return
|
|
// This aborts binding.call's wait, never the Main-owned image operation.
|
|
this.wait.abort()
|
|
this.saveCache = undefined
|
|
this.removeAbort()
|
|
this.unsubscribe()
|
|
this.channel.close()
|
|
}
|
|
|
|
private async receive(): Promise<void> {
|
|
while (!this.wait.signal.aborted) {
|
|
const frame = await this.channel.receive(this.wait.signal)
|
|
const call = remoteImageToolCallSchema.parse(decodeRemoteImageToolMessage(frame.payload))
|
|
await frame.consume()
|
|
void this.dispatch(call).catch(() => this.close())
|
|
}
|
|
}
|
|
|
|
private async dispatch(call: ReturnType<typeof remoteImageToolCallSchema.parse>): Promise<void> {
|
|
let reply: ReturnType<typeof remoteImageToolReplySchema.parse>
|
|
try {
|
|
this.wait.signal.throwIfAborted()
|
|
if ('name' in call && call.name === 'save_image_read') {
|
|
if (!this.binding?.readForSave || this.binding.context.workMode !== 'execute') throw new Error('Saving images is unavailable in Ask mode')
|
|
const { artifactId, mimeType, offset } = call.input
|
|
const cacheKey = `${artifactId}:${mimeType}`
|
|
let bytes = this.saveCache?.key === cacheKey ? this.saveCache.bytes : undefined
|
|
if (!bytes || offset === 0) {
|
|
bytes = await this.binding.readForSave(artifactId, mimeType)
|
|
this.saveCache = { key: cacheKey, bytes }
|
|
}
|
|
if (offset >= bytes.byteLength) throw new Error('Image chunk offset is out of range')
|
|
const end = Math.min(offset + remoteImageSaveChunkBytes, bytes.byteLength)
|
|
if (end === bytes.byteLength) this.saveCache = undefined
|
|
reply = remoteImageToolReplySchema.parse({ callId: call.callId, imageChunk: {
|
|
mimeType, totalBytes: bytes.byteLength, offset, data: bytes.subarray(offset, end).toString('base64')
|
|
} })
|
|
} else if ('name' in call) {
|
|
if (!this.storyGraph) throw new Error('Story Graph capability unavailable')
|
|
const storyGraphResult = call.name === 'story_graph_list'
|
|
? { available: await this.storyGraph.available() }
|
|
: await this.storyGraph.call(call.name, call.input, this.wait.signal)
|
|
reply = remoteImageToolReplySchema.parse({ callId: call.callId, storyGraphResult })
|
|
} else {
|
|
if (!this.binding || this.binding.context.workMode !== 'execute') throw new Error('Image tools are unavailable in Ask mode')
|
|
const result = await this.binding.call(call.input, call.callId, this.wait.signal)
|
|
reply = remoteImageToolReplySchema.parse({ callId: call.callId, result })
|
|
}
|
|
} catch (error) {
|
|
reply = { callId: call.callId, error: (error instanceof Error ? error.message : 'Image tool failed').slice(0, 4_000) }
|
|
}
|
|
if (!this.wait.signal.aborted) await this.channel.send(encodeRemoteImageToolMessage(reply), this.wait.signal)
|
|
}
|
|
}
|