feat: generate and edit conversation images across runtimes
Add model-controlled image invocation and read-only runtime assignments. Route image tools through Main so accepted operations retain their originating conversation when chat runtimes change, with saved references, cancellation, and manual recovery. Keep Node hashing outside browser-shared contracts to restore production builds. Validation includes the production build, typecheck, lint, focused regressions, real Host transport checks, and two successful real image requests. Full UI-driven runtime switching acceptance remains documented as pending. Release note: 开启图片模型的会话调用后,聊天助手可调用图片生成与编辑工具,Runtime 分配自动管理。图片成果保存在原会话,可继续修改,并独立显示生成状态和取消入口。
This commit is contained in:
@@ -87,12 +87,13 @@ export class AgentAcpConnection {
|
||||
sessionId: string,
|
||||
operationId: string,
|
||||
socketPath?: string,
|
||||
workMode?: 'ask' | 'execute'
|
||||
workMode?: 'ask' | 'execute',
|
||||
imageToolName?: string
|
||||
): Promise<void> {
|
||||
if (!this.bridgeOrigin) throw new Error('Shared model bridge is not ready')
|
||||
const response = await fetch(`${this.bridgeOrigin}/session`, {
|
||||
method: 'POST', headers: { 'content-type': 'application/json' },
|
||||
body: JSON.stringify({ sessionId, operationId, ...(socketPath ? { socketPath, workMode } : { release: true }) }),
|
||||
body: JSON.stringify({ sessionId, operationId, ...(socketPath ? { socketPath, workMode, imageToolName } : { release: true }) }),
|
||||
signal: AbortSignal.timeout(10_000)
|
||||
})
|
||||
if (!response.ok) throw new Error(`Shared model bridge route failed (${response.status})`)
|
||||
|
||||
@@ -30,23 +30,27 @@ describe('AgentOwnedAcpPrompt', () => {
|
||||
let client!: Client
|
||||
let finish!: (response: PromptResponse) => void
|
||||
const newSession = vi.fn(async () => ({ sessionId: 'session-1' }))
|
||||
const resumeSession = vi.fn(async () => ({}))
|
||||
let imageUrl: string | undefined
|
||||
const mcpServers = () => imageUrl ? [{ type: 'http' as const, name: 'images', url: imageUrl, headers: [] }] : []
|
||||
const prepareSession = vi.fn(async () => undefined)
|
||||
const completePrompt = vi.fn(async () => undefined)
|
||||
const owner = new AgentOwnedAcpPrompt({
|
||||
bindingId: 'binding-1', controllerId: 'controller-1',
|
||||
workspaceDirectory: '/workspace', process: new MemoryProcess(),
|
||||
transcript, prepareSession, completePrompt,
|
||||
transcript, prepareSession, completePrompt, mcpServers,
|
||||
createConnection: factory => {
|
||||
client = factory()
|
||||
return {
|
||||
initialize: async () => ({ protocolVersion: 1, agentCapabilities: {} }),
|
||||
newSession,
|
||||
newSession, resumeSession,
|
||||
prompt: () => new Promise<PromptResponse>(resolve => { finish = resolve })
|
||||
} as unknown as ClientSideConnection
|
||||
}
|
||||
})
|
||||
try {
|
||||
for (const [index, workMode] of (['execute', 'ask', 'execute'] as const).entries()) {
|
||||
for (const [index, workMode] of (['execute', 'execute', 'ask', 'execute', 'ask'] as const).entries()) {
|
||||
imageUrl = workMode === 'execute' && index !== 1 ? `http://127.0.0.1:1234/prompt-${index}` : undefined
|
||||
const operationId = `operation-${index}`
|
||||
transcript.prepare({
|
||||
bindingId: 'binding-1', operationId, requestId: operationId,
|
||||
@@ -58,6 +62,9 @@ describe('AgentOwnedAcpPrompt', () => {
|
||||
prompt: [{ type: 'text', text: 'Continue with the selected mode' }]
|
||||
}, workMode)
|
||||
expect(prepareSession).toHaveBeenLastCalledWith('session-1', operationId, workMode)
|
||||
if (index !== 2) {
|
||||
expect(index === 0 ? newSession : resumeSession).toHaveBeenLastCalledWith(expect.objectContaining({ mcpServers: mcpServers() }))
|
||||
}
|
||||
const result = await client.requestPermission!({
|
||||
sessionId: 'session-1',
|
||||
toolCall: { toolCallId: `tool-${index}`, kind: 'execute', title: 'Run a test' },
|
||||
@@ -74,7 +81,8 @@ describe('AgentOwnedAcpPrompt', () => {
|
||||
.toMatchObject({ state: 'completed' }))
|
||||
}
|
||||
expect(newSession).toHaveBeenCalledOnce()
|
||||
expect(completePrompt).toHaveBeenCalledTimes(3)
|
||||
expect(resumeSession).toHaveBeenCalledTimes(3)
|
||||
expect(completePrompt).toHaveBeenCalledTimes(5)
|
||||
} finally {
|
||||
owner.close()
|
||||
transcript.close()
|
||||
|
||||
@@ -4,6 +4,7 @@ import {
|
||||
PROTOCOL_VERSION,
|
||||
ndJsonStream,
|
||||
type Client,
|
||||
type McpServer,
|
||||
type PromptResponse,
|
||||
type RequestPermissionRequest,
|
||||
type RequestPermissionResponse,
|
||||
@@ -26,6 +27,7 @@ export type AgentOwnedAcpPromptOptions = {
|
||||
bindingId: string
|
||||
controllerId: string
|
||||
workspaceDirectory: string
|
||||
mcpServers?: () => McpServer[]
|
||||
expectedModel?: string
|
||||
process: RuntimeAcpProcessOwner
|
||||
transport?: AgentAcpConnection
|
||||
@@ -60,6 +62,7 @@ export class AgentOwnedAcpPrompt {
|
||||
#promptSettled?: Promise<void>
|
||||
#closed = false
|
||||
#initialized = false
|
||||
#hasMcpServers = false
|
||||
readonly #questions = new Map<string, {
|
||||
endpoint: string; operationId: string; questionCount: number; notification: SessionNotification
|
||||
}>()
|
||||
@@ -109,6 +112,7 @@ export class AgentOwnedAcpPrompt {
|
||||
)
|
||||
}
|
||||
|
||||
const mcpServers = this.#options.mcpServers?.() ?? []
|
||||
if (!this.#initialized) {
|
||||
const initialization = await this.#whileProcessAlive(
|
||||
this.#transport.initialize({
|
||||
@@ -135,7 +139,7 @@ export class AgentOwnedAcpPrompt {
|
||||
this.#connection.loadSession({
|
||||
sessionId: request.acpSessionId,
|
||||
cwd: this.#options.workspaceDirectory,
|
||||
mcpServers: []
|
||||
mcpServers
|
||||
})
|
||||
)
|
||||
} else if (
|
||||
@@ -145,7 +149,7 @@ export class AgentOwnedAcpPrompt {
|
||||
this.#connection.resumeSession({
|
||||
sessionId: request.acpSessionId,
|
||||
cwd: this.#options.workspaceDirectory,
|
||||
mcpServers: []
|
||||
mcpServers
|
||||
})
|
||||
)
|
||||
} else {
|
||||
@@ -156,7 +160,7 @@ export class AgentOwnedAcpPrompt {
|
||||
const created = await this.#whileProcessAlive(
|
||||
this.#connection.newSession({
|
||||
cwd: this.#options.workspaceDirectory,
|
||||
mcpServers: []
|
||||
mcpServers
|
||||
})
|
||||
)
|
||||
this.#sessionId = created.sessionId
|
||||
@@ -186,7 +190,13 @@ export class AgentOwnedAcpPrompt {
|
||||
request.acpSessionId !== this.#sessionId
|
||||
) {
|
||||
throw new Error('ACP session identity cannot change within a binding')
|
||||
} else if (mcpServers.length > 0 || this.#hasMcpServers) {
|
||||
await this.#whileProcessAlive(this.#connection.resumeSession({
|
||||
sessionId: this.#sessionId!, cwd: this.#options.workspaceDirectory,
|
||||
mcpServers
|
||||
}))
|
||||
}
|
||||
this.#hasMcpServers = mcpServers.length > 0
|
||||
await this.#options.prepareSession?.(this.#sessionId!, request.operationId, workMode)
|
||||
this.#active = {
|
||||
operationId: request.operationId,
|
||||
|
||||
@@ -0,0 +1,146 @@
|
||||
// @vitest-environment node
|
||||
import { randomUUID } from 'node:crypto'
|
||||
import { createServer } from 'node:http'
|
||||
import { Client } from '@modelcontextprotocol/sdk/client/index.js'
|
||||
import { StreamableHTTPClientTransport } from '@modelcontextprotocol/sdk/client/streamableHttp.js'
|
||||
import { afterEach, describe, expect, it, vi } from 'vitest'
|
||||
import { AgentImageToolMcp } from './image-tool-mcp'
|
||||
import { imageOperationSchema } from '../shared/image-generation-contracts'
|
||||
import { AssistantDatabase } from '../main/assistant/assistant-database'
|
||||
import { ImageGenerationService } from '../main/agent/image-generation-service'
|
||||
import { MainImageToolSession } from '../main/remote-agent/main-image-tool-session'
|
||||
import type { RuntimeProtocolBinaryChannel, RuntimeProtocolBinaryFrame } from '../main/remote-agent/protocol-remote-runtime-channel'
|
||||
|
||||
const png = 'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAwMCAO+jRZkAAAAASUVORK5CYII='
|
||||
vi.mock('electron', () => ({ nativeImage: { createFromBuffer: () => ({ isEmpty: () => false, toPNG: () => Buffer.from(png, 'base64') }) } }))
|
||||
const cleanup: (() => void | Promise<void>)[] = []
|
||||
afterEach(async () => { for (const close of cleanup.splice(0).reverse()) await close() })
|
||||
|
||||
async function setup() {
|
||||
const received: { path: string; body: string; authorization?: string }[] = []
|
||||
let release: (() => void) | undefined
|
||||
let hold = false
|
||||
const provider = createServer((req, res) => {
|
||||
void (async () => {
|
||||
let body = ''
|
||||
for await (const chunk of req) body += chunk
|
||||
received.push({ path: req.url!, body, authorization: req.headers.authorization })
|
||||
if (hold) await new Promise<void>(resolve => { release = resolve })
|
||||
res.setHeader('content-type', 'application/json')
|
||||
res.end(JSON.stringify({ data: [{ b64_json: png }] }))
|
||||
})()
|
||||
})
|
||||
await new Promise<void>(resolve => provider.listen(0, '127.0.0.1', resolve))
|
||||
cleanup.push(() => { release?.(); provider.closeAllConnections(); provider.close() })
|
||||
const address = provider.address() as { port: number }
|
||||
const database = new AssistantDatabase(':memory:')
|
||||
database.initialize(process.cwd())
|
||||
cleanup.push(() => database.close())
|
||||
const context = { conversationId: randomUUID(), messageId: randomUUID(), requestId: randomUUID(), workMode: 'execute' as 'ask' | 'execute' }
|
||||
database.saveLocalConversations([{ header: { id: context.conversationId, title: 'Remote images', updatedAt: 1 },
|
||||
messages: [{ id: context.messageId, role: 'assistant', content: '', createdAt: 1, state: 'complete' }] }])
|
||||
const profile = { id: randomUUID(), name: 'Controlled image provider', modelName: 'controlled-image',
|
||||
protocol: 'openai-images-generations', allowConversationInvocation: true,
|
||||
baseUrl: `http://127.0.0.1:${address.port}/v1`, authentication: 'api-key' as const, apiKey: 'controlled-main-only' }
|
||||
const service = new ImageGenerationService({ database, getSettings: async () => ({ modelProfiles: [profile], defaultImageModelProfileId: profile.id }) })
|
||||
service.initialize()
|
||||
cleanup.push(() => service.dispose())
|
||||
const frames: Uint8Array[] = []
|
||||
const connect = async (mode: 'ask' | 'execute' = 'execute') => {
|
||||
const binding = service.bind({ ...context, requestId: randomUUID(), workMode: mode })
|
||||
const queue: RuntimeProtocolBinaryFrame[] = []
|
||||
let receiver: ((frame: RuntimeProtocolBinaryFrame) => void) | undefined
|
||||
let rejectReceive: ((error: unknown) => void) | undefined
|
||||
const listeners = new Set<() => void>()
|
||||
let closed = false
|
||||
const channel: RuntimeProtocolBinaryChannel = {
|
||||
channelId: randomUUID(), channelEpoch: '1',
|
||||
send: async payload => { frames.push(payload); adapter.onReply(payload) },
|
||||
receive: async () => queue.shift() ?? await new Promise((resolve, reject) => { receiver = resolve; rejectReceive = reject }),
|
||||
close: () => { if (closed) return; closed = true; rejectReceive?.(new Error('closed')); for (const fn of listeners) fn() },
|
||||
onClose: fn => { listeners.add(fn); return () => { listeners.delete(fn) } }
|
||||
}
|
||||
const wait = new AbortController()
|
||||
const main = new MainImageToolSession(channel, binding, wait.signal)
|
||||
const adapter = new AgentImageToolMcp({ channelId: channel.channelId, channelEpoch: '1', description: await binding.describe() ?? 'Not available' }, async payload => {
|
||||
if (closed) throw new Error('disconnected')
|
||||
frames.push(payload)
|
||||
const frame = { payload, sequence: '1', consume: async () => {} }
|
||||
if (receiver) { const resolve = receiver; receiver = undefined; resolve(frame) } else queue.push(frame)
|
||||
})
|
||||
await adapter.start()
|
||||
const client = new Client({ name: 'remote-image-test', version: '1' })
|
||||
await client.connect(new StreamableHTTPClientTransport(new URL(adapter.url!)))
|
||||
cleanup.push(async () => { main.close(); adapter.close(); await client.close() })
|
||||
return { client, adapter, wait, main, binding }
|
||||
}
|
||||
return { received, context, profile, service, database, connect, frames,
|
||||
hold: () => { hold = true }, release: () => release?.() }
|
||||
}
|
||||
|
||||
describe('remote image MCP to Main service', () => {
|
||||
it('generates, edits persisted uploads and history through a controlled HTTP provider', async () => {
|
||||
const h = await setup()
|
||||
const remote = await h.connect()
|
||||
expect((await remote.client.listTools()).tools[0]!.description).toContain(h.profile.id)
|
||||
const call = async (args: Record<string, unknown>) => {
|
||||
const result = await remote.client.callTool({ name: 'generate_image', arguments: args })
|
||||
expect(result.isError).not.toBe(true)
|
||||
return imageOperationSchema.parse(JSON.parse((result.content as { text: string }[])[0]!.text))
|
||||
}
|
||||
const first = await call({ intent: 'create', prompt: 'A blue square' })
|
||||
const uploaded = h.service.persistUploads(h.context, [{ name: 'source.png', mediaType: 'image/png', data: png }])
|
||||
const uploadEdit = await call({ intent: 'edit', prompt: 'Make the uploaded square red', sourceArtifactIds: uploaded })
|
||||
const historyEdit = await call({ intent: 'edit', prompt: 'Make the previous result green', sourceArtifactIds: first.artifactIds })
|
||||
expect([first, uploadEdit, historyEdit].map(op => op.state)).toEqual(['completed', 'completed', 'completed'])
|
||||
expect(h.received.map(call => call.path)).toEqual(['/v1/images/generations', '/v1/images/edits', '/v1/images/edits'])
|
||||
expect(h.received.every(call => call.authorization === 'Bearer controlled-main-only')).toBe(true)
|
||||
expect(h.received[1]!.body).toContain('name="image"')
|
||||
expect(h.database.getArtifact(historyEdit.artifactIds[0]!).content).toContain(png)
|
||||
expect(h.database.getConversation(h.context.conversationId).messages[0]!.imageOperations).toHaveLength(3)
|
||||
const wire = Buffer.concat(h.frames).toString()
|
||||
expect(wire).not.toContain(png)
|
||||
expect(wire).not.toContain('controlled-main-only')
|
||||
expect(wire).not.toContain(h.profile.baseUrl)
|
||||
})
|
||||
|
||||
it.each(['chat cancellation', 'disconnect'])('preserves Main operation after %s without replay', async reason => {
|
||||
const h = await setup()
|
||||
h.hold()
|
||||
const remote = await h.connect()
|
||||
const waiting = remote.client.callTool({ name: 'generate_image', arguments: { intent: 'create', prompt: 'Blue' } }).catch(() => undefined)
|
||||
await vi.waitFor(() => expect(h.received).toHaveLength(1))
|
||||
if (reason === 'chat cancellation') remote.wait.abort()
|
||||
else remote.main.close()
|
||||
remote.adapter.close()
|
||||
h.release()
|
||||
await vi.waitFor(() => expect(h.database.getConversation(h.context.conversationId).messages[0]!.imageOperations?.[0]?.state).toBe('completed'))
|
||||
await waiting
|
||||
expect(h.received).toHaveLength(1)
|
||||
})
|
||||
|
||||
it('rejects Ask calls at the desktop binding and refreshes a new prompt description', async () => {
|
||||
const h = await setup()
|
||||
const ask = await h.connect('ask')
|
||||
const denied = await ask.client.callTool({ name: 'generate_image', arguments: { intent: 'create', prompt: 'Blue' } })
|
||||
expect(denied.isError).toBe(true)
|
||||
expect(h.received).toHaveLength(0)
|
||||
h.profile.name = 'Changed next prompt'
|
||||
const execute = await h.connect()
|
||||
expect((await execute.client.listTools()).tools[0]!.description).toContain('Changed next prompt')
|
||||
})
|
||||
|
||||
it('returns unknown delivery without sending a second request for the same MCP call', async () => {
|
||||
const send = vi.fn(async () => { throw new Error('connection lost after send') })
|
||||
const adapter = new AgentImageToolMcp({ channelId: randomUUID(), channelEpoch: '1', description: 'Images' }, send)
|
||||
await adapter.start()
|
||||
cleanup.push(() => adapter.close())
|
||||
const request = { jsonrpc: '2.0', id: 7, method: 'tools/call', params: { name: 'generate_image', arguments: { intent: 'create', prompt: 'Blue' } } }
|
||||
for (let i = 0; i < 2; i++) {
|
||||
const response = await fetch(adapter.url!, { method: 'POST', headers: { 'content-type': 'application/json', accept: 'application/json, text/event-stream' }, body: JSON.stringify(request) })
|
||||
const result = await response.json()
|
||||
expect(JSON.stringify(result)).toContain('unknown')
|
||||
}
|
||||
expect(send).toHaveBeenCalledOnce()
|
||||
})
|
||||
})
|
||||
@@ -0,0 +1,121 @@
|
||||
import { randomUUID } from 'node:crypto'
|
||||
import { createServer, type Server } from 'node:http'
|
||||
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 { imageToolInputSchema, imageToolName } from '../shared/image-generation-contracts'
|
||||
import {
|
||||
decodeRemoteImageToolMessage, encodeRemoteImageToolMessage,
|
||||
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')
|
||||
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 () => ({ tools: [{
|
||||
name: imageToolName, description: this.descriptor.description,
|
||||
inputSchema: schema as { type: 'object' }
|
||||
}] }))
|
||||
mcp.setRequestHandler(CallToolRequestSchema, async (call, extra) => {
|
||||
if (call.params.name !== imageToolName) 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) pending.reject(new Error(reply.error ?? 'Image tool response has no result'))
|
||||
else pending.resolve({ content: [{ type: 'text', text: JSON.stringify(reply.result) }] })
|
||||
}
|
||||
|
||||
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: z.infer<typeof imageToolInputSchema>): 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 })).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)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -133,7 +133,7 @@ const FORWARDED_HEADER_NAMES = [
|
||||
export class ModelBridgeLoopbackProxy {
|
||||
readonly #exchange: ModelBridgeExchange
|
||||
readonly #sharedSessions: boolean
|
||||
readonly #sessionRoutes = new Map<string, { operationId: string; socketPath: string; workMode: 'ask' | 'execute' }>()
|
||||
readonly #sessionRoutes = new Map<string, { operationId: string; socketPath: string; workMode: 'ask' | 'execute'; imageToolName?: string }>()
|
||||
readonly #routeToken: string
|
||||
readonly #maximumConnections: number
|
||||
readonly #requestTimeoutMs: number
|
||||
@@ -317,12 +317,12 @@ export class ModelBridgeLoopbackProxy {
|
||||
const route = this.#sessionRoutes.get(url.searchParams.get('sessionId') ?? '')
|
||||
if (!route) throw new HttpRequestError(409, 'session-inactive')
|
||||
outgoing.setHeader('content-type', 'application/json')
|
||||
outgoing.end(JSON.stringify({ operationId: route.operationId, workMode: route.workMode }))
|
||||
outgoing.end(JSON.stringify({ operationId: route.operationId, workMode: route.workMode, imageToolName: route.imageToolName }))
|
||||
return
|
||||
}
|
||||
if (incoming.method !== 'POST') throw new HttpRequestError(405, 'method-not-allowed')
|
||||
const value = JSON.parse((await readBoundedBody(incoming)).toString('utf8')) as {
|
||||
sessionId?: string; operationId?: string; socketPath?: string; release?: boolean; workMode?: 'ask' | 'execute'
|
||||
sessionId?: string; operationId?: string; socketPath?: string; release?: boolean; workMode?: 'ask' | 'execute'; imageToolName?: string
|
||||
}
|
||||
if (
|
||||
typeof value.sessionId !== 'string' || value.sessionId.length > 128 ||
|
||||
@@ -335,7 +335,8 @@ export class ModelBridgeLoopbackProxy {
|
||||
} else {
|
||||
if (value.workMode !== 'ask' && value.workMode !== 'execute') throw new HttpRequestError(400, 'request-invalid')
|
||||
const socketPath = normalizedAbsolutePath(value.socketPath ?? '', 'Model bridge socket')
|
||||
this.#sessionRoutes.set(value.sessionId, { operationId: value.operationId, socketPath, workMode: value.workMode })
|
||||
if (value.imageToolName !== undefined && !/^goodbuddy_image_[a-f0-9]{24}$/.test(value.imageToolName)) throw new HttpRequestError(400, 'request-invalid')
|
||||
this.#sessionRoutes.set(value.sessionId, { operationId: value.operationId, socketPath, workMode: value.workMode, imageToolName: value.imageToolName })
|
||||
}
|
||||
outgoing.end('{}')
|
||||
return
|
||||
|
||||
@@ -5,6 +5,29 @@ import { describe, expect, it, vi } from 'vitest'
|
||||
import { openCodeSubagentPluginSource } from './opencode-subagent-plugin'
|
||||
|
||||
describe('OpenCode ACP child event plugin', () => {
|
||||
it('hides foreign image tools and removes the current image tool in Ask on a reused Session', async () => {
|
||||
let mode = 'execute'
|
||||
const imageToolName = `goodbuddy_image_${'a'.repeat(24)}`
|
||||
const patch = vi.fn()
|
||||
const factory = runInNewContext(
|
||||
`(${openCodeSubagentPluginSource('http://127.0.0.1:12345/fixture').replace(/^import .*\n/gmu, '').replace('export default ', '')})`,
|
||||
{ process: { stdout: { write: vi.fn() } }, createServer, randomBytes,
|
||||
fetch: async () => ({ ok: true, json: async () => ({ operationId: mode, workMode: mode, imageToolName }) }) }
|
||||
)
|
||||
const hooks = await factory({ client: {
|
||||
_client: { get: vi.fn(), post: vi.fn(), patch },
|
||||
session: { get: async () => ({ data: { id: 'root' } }) }
|
||||
} })
|
||||
try {
|
||||
for (mode of ['execute', 'ask', 'execute']) {
|
||||
await hooks.event({ event: { type: 'session.status', properties: { sessionID: 'root', status: { type: 'idle' } } } })
|
||||
await hooks['chat.message']({ sessionID: 'root' }, { message: { id: mode } })
|
||||
const rules = patch.mock.calls.at(-1)![0].body.permission
|
||||
expect(rules).toContainEqual({ permission: 'goodbuddy_image_*', pattern: '*', action: 'deny' })
|
||||
expect(rules.some((rule: { permission: string; action: string }) => rule.permission === `${imageToolName}_*` && rule.action === 'allow')).toBe(mode === 'execute')
|
||||
}
|
||||
} finally { await hooks.dispose() }
|
||||
})
|
||||
it('attributes child and compaction model rounds to the frozen root operation', async () => {
|
||||
const operations = new Map([['one', 'operation-1'], ['two', 'operation-2']])
|
||||
const fetchRoute = vi.fn(async (url: string) => {
|
||||
@@ -33,7 +56,10 @@ describe('OpenCode ACP child event plugin', () => {
|
||||
try {
|
||||
await hooks['chat.message']({ sessionID: 'one' }, { message: { id: 'first' } })
|
||||
expect(patch).toHaveBeenCalledWith({
|
||||
url: '/session/one', body: { permission: [{ permission: '*', pattern: '*', action: 'allow' }] }, throwOnError: true
|
||||
url: '/session/one', body: { permission: [
|
||||
{ permission: '*', pattern: '*', action: 'allow' },
|
||||
{ permission: 'goodbuddy_image_*', pattern: '*', action: 'deny' }
|
||||
] }, throwOnError: true
|
||||
})
|
||||
expect(await headers('child', 'child-message')).toEqual({
|
||||
'x-goodbuddy-session': 'one', 'x-goodbuddy-operation': 'operation-1'
|
||||
@@ -42,7 +68,8 @@ describe('OpenCode ACP child event plugin', () => {
|
||||
expect(patch).toHaveBeenLastCalledWith({
|
||||
url: '/session/child', body: { permission: [
|
||||
{ permission: '*', pattern: '*', action: 'allow' },
|
||||
{ permission: 'task', pattern: '*', action: 'deny' }
|
||||
{ permission: 'task', pattern: '*', action: 'deny' },
|
||||
{ permission: 'goodbuddy_image_*', pattern: '*', action: 'deny' }
|
||||
] }, throwOnError: true
|
||||
})
|
||||
expect(await headers('two', 'compact')).toEqual({
|
||||
|
||||
@@ -26,7 +26,7 @@ async function plugin(input?: {
|
||||
}) {
|
||||
const tasks = new Map<string, { sessionId: string; callId: string }>()
|
||||
const pending = new Map<string, { sessionId: string; root: string }>()
|
||||
const modelMessages = new Map<string, { sessionId: string; operationId: string; workMode: 'ask' | 'execute' }>()
|
||||
const modelMessages = new Map<string, { sessionId: string; operationId: string; workMode: 'ask' | 'execute'; imageToolName?: string }>()
|
||||
const modelRoute = async (sessionId: string, messageId: string) => {
|
||||
const key = `${sessionId}\0${messageId}`
|
||||
const previous = modelMessages.get(key)
|
||||
@@ -43,8 +43,8 @@ async function plugin(input?: {
|
||||
}
|
||||
const response = await fetch(`${modules.modelBridgeOrigin}/session?sessionId=${encodeURIComponent(root)}`)
|
||||
if (!response.ok) throw new Error('GoodBuddy model operation is no longer active')
|
||||
const { operationId, workMode } = await response.json() as { operationId: string; workMode: 'ask' | 'execute' }
|
||||
const route = { sessionId: root, operationId, workMode }
|
||||
const { operationId, workMode, imageToolName } = await response.json() as { operationId: string; workMode: 'ask' | 'execute'; imageToolName?: string }
|
||||
const route = { sessionId: root, operationId, workMode, imageToolName }
|
||||
modelMessages.set(key, route)
|
||||
return route
|
||||
}
|
||||
@@ -131,6 +131,10 @@ async function plugin(input?: {
|
||||
// Keep its explicit subagent restrictions after the request rules.
|
||||
permission.push(...(child.data.permission ?? []).filter(rule => rule.permission !== '*'))
|
||||
}
|
||||
permission.push({ permission: 'goodbuddy_image_*', pattern: '*', action: 'deny' })
|
||||
if (route.workMode === 'execute' && route.imageToolName) {
|
||||
permission.push({ permission: `${route.imageToolName}_*`, pattern: '*', action: 'allow' })
|
||||
}
|
||||
await input!.client._client.patch({
|
||||
url: `/session/${encodeURIComponent(request.sessionID)}`, body: { permission }, throwOnError: true
|
||||
})
|
||||
|
||||
@@ -1,4 +1,8 @@
|
||||
import { describe, expect, it, vi } from 'vitest'
|
||||
import { Client } from '@modelcontextprotocol/sdk/client/index.js'
|
||||
import { StreamableHTTPClientTransport } from '@modelcontextprotocol/sdk/client/streamableHttp.js'
|
||||
import { decodeRemoteImageToolMessage, encodeRemoteImageToolMessage, remoteImageToolCallSchema } from '../shared/remote-image-tool-contracts'
|
||||
import type { AgentFrame } from '../shared/agent-protocol'
|
||||
import { mkdtemp, rm } from 'node:fs/promises'
|
||||
import { tmpdir } from 'node:os'
|
||||
import { join } from 'node:path'
|
||||
@@ -309,7 +313,7 @@ function harness(input: {
|
||||
throw new Error('broker close failed')
|
||||
}
|
||||
})
|
||||
const blobSink = vi.fn(async () => undefined)
|
||||
const blobSink = vi.fn<(frame: AgentFrame) => Promise<void>>(async () => undefined)
|
||||
let bridgeDispatch: ModelBridgeBrokerDispatch | undefined
|
||||
const launchModelBridges: unknown[] = []
|
||||
const resolved = resolvedBundle()
|
||||
@@ -451,6 +455,49 @@ function harness(input: {
|
||||
}
|
||||
|
||||
describe('RuntimeAcpBackend', () => {
|
||||
it('binds MCP image calls to authenticated blob framing and closes each prompt endpoint', async () => {
|
||||
const fixture = harness({ workMode: 'execute' })
|
||||
const client = new Client({ name: 'image-backend-test', version: '1' })
|
||||
try {
|
||||
await open(fixture)
|
||||
const accepted = await invoke(fixture, 'runtime/preparePrompt', fixture.preparation({
|
||||
imageTool: { channelId: 'image-channel', channelEpoch: '1', description: 'Image catalog' }
|
||||
})) as { imageToolUrl: string }
|
||||
await client.connect(new StreamableHTTPClientTransport(new URL(accepted.imageToolUrl)))
|
||||
expect((await client.listTools()).tools.map(tool => tool.name)).toEqual(['generate_image'])
|
||||
const result = client.callTool({ name: 'generate_image', arguments: { intent: 'create', prompt: 'Blue' } })
|
||||
await vi.waitFor(() => expect(fixture.blobSink).toHaveBeenCalledOnce())
|
||||
const frame = fixture.blobSink.mock.calls[0]![0]
|
||||
expect(frame.header).toMatchObject({ channelId: 'image-channel', channelEpoch: '1', sequence: '1', direction: 'agent-to-main' })
|
||||
expect(fixture.backend.authorizeBlobFrame(frame, fixture.context.controller)).toBe(true)
|
||||
expect(fixture.backend.authorizeBlobFrame(frame, { ...fixture.context.controller, connectionId: 'foreign' })).toBe(false)
|
||||
const call = remoteImageToolCallSchema.parse(decodeRemoteImageToolMessage(frame.payload))
|
||||
const payload = encodeRemoteImageToolMessage({ callId: call.callId, error: 'Controlled test failure' })
|
||||
const reply = { header: { ...frame.header, direction: 'main-to-agent' as const, payloadLength: payload.byteLength }, payload }
|
||||
await expect(fixture.backend.onBlobFrame(reply, { ...fixture.context, channelId: 'image-channel', controller: { ...fixture.context.controller, connectionId: 'foreign' } })).rejects.toThrow('authority')
|
||||
await fixture.backend.onBlobFrame(reply, { ...fixture.context, channelId: 'image-channel' })
|
||||
expect(await result).toMatchObject({ isError: true })
|
||||
await invoke(fixture, 'runtime/completePrompt', { bindingId: 'binding-1', operationId: 'request-1', requestId: 'request-1' })
|
||||
await expect(fetch(accepted.imageToolUrl)).rejects.toThrow()
|
||||
const next = await invoke(fixture, 'runtime/preparePrompt', fixture.preparation({
|
||||
requestId: 'request-2', operationId: 'request-2', promptSequence: 1,
|
||||
imageTool: { channelId: 'image-channel-2', channelEpoch: '2', description: 'Refreshed image catalog' }
|
||||
})) as { imageToolUrl: string }
|
||||
expect(next.imageToolUrl).not.toBe(accepted.imageToolUrl)
|
||||
} finally { await client.close(); await fixture.backend.dispose() }
|
||||
})
|
||||
|
||||
it('rejects image tool preparation in Ask before launching a process or bridge', async () => {
|
||||
const fixture = harness()
|
||||
try {
|
||||
await open(fixture)
|
||||
await expect(invoke(fixture, 'runtime/preparePrompt', fixture.preparation({
|
||||
imageTool: { channelId: 'image-channel', channelEpoch: '1', description: 'Images' }
|
||||
}))).rejects.toThrow('Execute')
|
||||
expect(fixture.launches).toHaveLength(0)
|
||||
expect(fixture.lifecycle).not.toContain('bridge-listen')
|
||||
} finally { await fixture.backend.dispose() }
|
||||
})
|
||||
it.each([false, true])('allows mode changes only on a shared Session (shared=%s)', async (shared) => {
|
||||
const fixture = harness({ agentOwned: true, shareOwnedProcesses: shared, workMode: 'execute' })
|
||||
try {
|
||||
|
||||
@@ -1,4 +1,6 @@
|
||||
import { z } from 'zod'
|
||||
import { AgentImageToolMcp } from './image-tool-mcp'
|
||||
import { remoteImageToolMcpName } from '../shared/remote-image-tool-node'
|
||||
import { remoteQuestionResponseSchema } from '../shared/remote-question-contracts'
|
||||
import { createHash } from 'node:crypto'
|
||||
import {
|
||||
@@ -338,6 +340,7 @@ type BindingState = {
|
||||
ownedPromptStartTimer?: NodeJS.Timeout
|
||||
modelBridgePolicy?: ModelBridgePolicy
|
||||
modelBridgeClient?: ModelBridgeBlobClient
|
||||
imageTool?: AgentImageToolMcp
|
||||
modelBridgeBroker?: ModelBridgeBrokerServer
|
||||
ownedAcp?: AgentOwnedAcpPrompt
|
||||
workspaceDirectory?: string
|
||||
@@ -558,6 +561,18 @@ export class RuntimeAcpBackend {
|
||||
'identity'
|
||||
)
|
||||
}
|
||||
const imageBinding = [...this.#bindings.values()].find(binding =>
|
||||
binding.controllerId === context.controller.controllerId &&
|
||||
binding.controllerGeneration === context.controller.generation &&
|
||||
binding.connectionId === context.controller.connectionId &&
|
||||
binding.activeOperationId !== undefined &&
|
||||
binding.imageTool?.descriptor.channelId === frame.header.channelId &&
|
||||
binding.imageTool.descriptor.channelEpoch === frame.header.channelEpoch &&
|
||||
context.channelId === frame.header.channelId)
|
||||
if (imageBinding) {
|
||||
imageBinding.imageTool!.onReply(frame.payload)
|
||||
return
|
||||
}
|
||||
const matches = [...this.#bindings.values()].filter(
|
||||
(binding) =>
|
||||
binding.controllerId === context.controller.controllerId &&
|
||||
@@ -613,10 +628,10 @@ export class RuntimeAcpBackend {
|
||||
binding.controllerId === controller.controllerId &&
|
||||
binding.controllerGeneration === controller.generation &&
|
||||
binding.connectionId === controller.connectionId &&
|
||||
binding.modelBridgeClient?.binding.channelId ===
|
||||
frame.header.channelId &&
|
||||
binding.modelBridgeClient.binding.channelEpoch ===
|
||||
frame.header.channelEpoch &&
|
||||
((binding.modelBridgeClient?.binding.channelId === frame.header.channelId &&
|
||||
binding.modelBridgeClient.binding.channelEpoch === frame.header.channelEpoch) ||
|
||||
(binding.imageTool?.descriptor.channelId === frame.header.channelId &&
|
||||
binding.imageTool.descriptor.channelEpoch === frame.header.channelEpoch)) &&
|
||||
frame.header.connectionId === binding.connectionId &&
|
||||
frame.header.generation === binding.controllerGeneration
|
||||
)
|
||||
@@ -1140,11 +1155,37 @@ export class RuntimeAcpBackend {
|
||||
}
|
||||
}
|
||||
}
|
||||
if (preparation.imageTool && (preparation.workMode !== 'execute' || !this.#options.blobSink)) {
|
||||
throw new RuntimeAcpBackendError('Image tools require Execute and a desktop transport', 'identity')
|
||||
}
|
||||
if (requestedModelBridgePolicy !== undefined) {
|
||||
await this.#startModelBridge(binding, preparation, workspace)
|
||||
}
|
||||
const acceptedAt = new Date(this.#now()).toISOString()
|
||||
binding.imageTool?.close()
|
||||
binding.imageTool = undefined
|
||||
if (preparation.imageTool) {
|
||||
let sequence = 0n
|
||||
const imageTool = new AgentImageToolMcp(preparation.imageTool, async payload => {
|
||||
if (binding.activeOperationId !== preparation.operationId || binding.transportState === 'detached') {
|
||||
throw new Error('Image tool prompt is no longer connected')
|
||||
}
|
||||
await this.#options.blobSink!({
|
||||
header: {
|
||||
protocolMajor: AGENT_PROTOCOL_VERSION.major, protocolMinor: AGENT_PROTOCOL_VERSION.minor,
|
||||
kind: 'blob', direction: 'agent-to-main',
|
||||
connectionId: binding.connectionId, generation: binding.controllerGeneration,
|
||||
channelId: preparation.imageTool!.channelId, channelEpoch: preparation.imageTool!.channelEpoch,
|
||||
sequence: String(++sequence), payloadLength: payload.byteLength
|
||||
}, payload
|
||||
}, { bindingId: binding.request.bindingId, controllerId: binding.controllerId,
|
||||
controllerGeneration: binding.controllerGeneration })
|
||||
})
|
||||
binding.imageTool = imageTool
|
||||
await imageTool.start()
|
||||
}
|
||||
const acceptance = remotePromptOperationAcceptanceSchema.parse({
|
||||
...(binding.imageTool?.url ? { imageToolUrl: binding.imageTool.url } : {}),
|
||||
bindingId: preparation.bindingId,
|
||||
operationId: preparation.operationId,
|
||||
requestId: preparation.requestId,
|
||||
@@ -1448,6 +1489,9 @@ export class RuntimeAcpBackend {
|
||||
bindingId: binding.request.bindingId,
|
||||
controllerId: binding.controllerId,
|
||||
workspaceDirectory: binding.workspaceDirectory,
|
||||
mcpServers: () => binding.imageTool?.url ? [{
|
||||
type: 'http', name: remoteImageToolMcpName(binding.request.bindingId), url: binding.imageTool.url, headers: []
|
||||
}] : [],
|
||||
...(binding.modelBridgePolicy === undefined
|
||||
? {}
|
||||
: {
|
||||
@@ -1461,7 +1505,8 @@ export class RuntimeAcpBackend {
|
||||
transport: binding.sharedProcess.transport,
|
||||
prepareSession: async (sessionId: string, operationId: string, workMode: 'ask' | 'execute') => {
|
||||
await binding.sharedProcess!.transport.setModelRoute(
|
||||
sessionId, operationId, binding.modelBridgeBroker!.socketPath, workMode
|
||||
sessionId, operationId, binding.modelBridgeBroker!.socketPath, workMode,
|
||||
binding.imageTool ? remoteImageToolMcpName(binding.request.bindingId) : undefined
|
||||
)
|
||||
}
|
||||
} : {}),
|
||||
@@ -1816,6 +1861,8 @@ export class RuntimeAcpBackend {
|
||||
poisonIfActive: boolean,
|
||||
preserveRuntime = false
|
||||
): Promise<void> {
|
||||
binding.imageTool?.close()
|
||||
binding.imageTool = undefined
|
||||
const broker = binding.modelBridgeBroker
|
||||
const client = binding.modelBridgeClient
|
||||
binding.modelBridgeBroker = undefined
|
||||
@@ -2141,6 +2188,8 @@ export class RuntimeAcpBackend {
|
||||
return
|
||||
}
|
||||
binding.transportState = 'detached'
|
||||
binding.imageTool?.close()
|
||||
binding.imageTool = undefined
|
||||
if (binding.modelBridgeClient !== undefined) {
|
||||
await this.#closeModelBridge(binding, false, true).catch(
|
||||
() => undefined
|
||||
|
||||
Reference in New Issue
Block a user