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:
mesalogo
@@ -1,4 +1,5 @@
|
||||
import { createHash, randomUUID } from 'node:crypto'
|
||||
import { remoteImageToolMcpName } from '../../shared/remote-image-tool-node'
|
||||
import { remoteQuestionSchema } from '../../shared/remote-question-contracts'
|
||||
import type { AgentQuestionAnswer } from '../../shared/contracts'
|
||||
import { isDeepStrictEqual } from 'node:util'
|
||||
@@ -191,6 +192,7 @@ type SessionRecord = {
|
||||
sessionId: string
|
||||
context: ChannelContext
|
||||
modelBridge?: RemoteModelBridgeSession
|
||||
imageToolUrl?: string
|
||||
sendHistoryWithNextPrompt: boolean
|
||||
ownedPromptAttached?: boolean
|
||||
}
|
||||
@@ -1183,7 +1185,9 @@ export class AcpRemoteRuntime implements AgentRuntime {
|
||||
raw = await this.awaitOperation(
|
||||
context,
|
||||
'启动远端请求',
|
||||
context.channel.preparePrompt(preparation),
|
||||
request.imageToolBinding
|
||||
? context.channel.preparePrompt(preparation, request.imageToolBinding, signal)
|
||||
: context.channel.preparePrompt(preparation),
|
||||
signal
|
||||
)
|
||||
} catch (error) {
|
||||
@@ -1596,7 +1600,7 @@ export class AcpRemoteRuntime implements AgentRuntime {
|
||||
state.agent.loadSession({
|
||||
sessionId: binding.acpSessionId,
|
||||
cwd: this.options.workspacePath,
|
||||
mcpServers: []
|
||||
mcpServers: this.imageMcpServers(binding.bindingId, prepared.acceptance.imageToolUrl)
|
||||
}),
|
||||
signal
|
||||
)
|
||||
@@ -1607,7 +1611,7 @@ export class AcpRemoteRuntime implements AgentRuntime {
|
||||
state.agent.resumeSession({
|
||||
sessionId: binding.acpSessionId,
|
||||
cwd: this.options.workspacePath,
|
||||
mcpServers: []
|
||||
mcpServers: this.imageMcpServers(binding.bindingId, prepared.acceptance.imageToolUrl)
|
||||
}),
|
||||
signal
|
||||
)
|
||||
@@ -1619,6 +1623,7 @@ export class AcpRemoteRuntime implements AgentRuntime {
|
||||
sessionId: binding.acpSessionId,
|
||||
context,
|
||||
modelBridge: prepared.modelBridge,
|
||||
imageToolUrl: prepared.acceptance.imageToolUrl,
|
||||
sendHistoryWithNextPrompt: false
|
||||
}
|
||||
pendingModelBridge = undefined
|
||||
@@ -1668,7 +1673,7 @@ export class AcpRemoteRuntime implements AgentRuntime {
|
||||
'创建会话',
|
||||
state.agent.newSession({
|
||||
cwd: this.options.workspacePath,
|
||||
mcpServers: []
|
||||
mcpServers: this.imageMcpServers(binding.bindingId, prepared.acceptance.imageToolUrl)
|
||||
}),
|
||||
signal
|
||||
)
|
||||
@@ -1697,6 +1702,7 @@ export class AcpRemoteRuntime implements AgentRuntime {
|
||||
sessionId: response.sessionId,
|
||||
context,
|
||||
modelBridge: prepared.modelBridge,
|
||||
imageToolUrl: prepared.acceptance.imageToolUrl,
|
||||
sendHistoryWithNextPrompt: true
|
||||
}
|
||||
pendingModelBridge = undefined
|
||||
@@ -1725,6 +1731,10 @@ export class AcpRemoteRuntime implements AgentRuntime {
|
||||
}
|
||||
}
|
||||
|
||||
private imageMcpServers(bindingId: string, url?: string) {
|
||||
return url ? [{ type: 'http' as const, name: remoteImageToolMcpName(bindingId), url, headers: [] }] : []
|
||||
}
|
||||
|
||||
private async getSession(
|
||||
conversationId: string,
|
||||
request: AgentExecutionRequest,
|
||||
@@ -1782,6 +1792,13 @@ export class AcpRemoteRuntime implements AgentRuntime {
|
||||
current.modelBridge = prepared.modelBridge
|
||||
if (this.options.modelProfile !== undefined) {
|
||||
current.ownedPromptAttached = false
|
||||
} else if (prepared.acceptance.imageToolUrl || current.imageToolUrl) {
|
||||
const state = await this.getState(current.context, signal)
|
||||
await this.awaitOperation(current.context, 'Refresh image tools', state.agent.resumeSession({
|
||||
sessionId: current.sessionId, cwd: this.options.workspacePath,
|
||||
mcpServers: this.imageMcpServers(current.binding.bindingId, prepared.acceptance.imageToolUrl)
|
||||
}), signal)
|
||||
current.imageToolUrl = prepared.acceptance.imageToolUrl
|
||||
}
|
||||
return current
|
||||
} catch (error) {
|
||||
|
||||
@@ -464,6 +464,34 @@ describe('ContinueAgentRuntime', () => {
|
||||
expect(gateway.grantCustomMcp).not.toHaveBeenCalled()
|
||||
})
|
||||
|
||||
it.each([undefined, 'existing-capability'])('cleans up image capabilities after custom MCP preparation fails (existing=%s)', async existingToken => {
|
||||
const gateway = {
|
||||
getEndpoint: () => 'http://127.0.0.1:4567/mcp',
|
||||
bindImageTool: vi.fn(() => existingToken ?? 'image-capability'),
|
||||
grantCustomMcp: vi.fn(() => 'custom-capability'),
|
||||
prepareCustomMcpTools: vi.fn(async () => { throw new Error('Custom MCP unavailable') }),
|
||||
revoke: vi.fn()
|
||||
}
|
||||
const context = { requestId: randomUUID(), conversationId: 'conversation-1', messageId: 'message-1', workMode: 'execute' as const }
|
||||
const runtime = new ContinueAgentRuntime({
|
||||
binaryPath: '', configPath: 'C:\\safe config\\continue.yaml', defaultWorkspace: process.cwd(),
|
||||
hostCacheRoot: 'C:\\safe\\continue-host',
|
||||
knowledgeGateway: gateway as unknown as KnowledgeMcpGateway,
|
||||
mcpServers: [{ id: randomUUID(), name: 'Custom', description: '', enabled: true, allowDynamicTools: false,
|
||||
assignments: ['continue'], secretConfigured: false, transport: 'http', url: 'https://mcp.test' }],
|
||||
createHostAdapter: () => ({ getPreparedHost: mocks.prepareHost, run: mocks.runHost, dispose: mocks.disposeHost })
|
||||
})
|
||||
try {
|
||||
await expect((async () => {
|
||||
for await (const event of runtime.run({ ...context, prompt: 'Draw an image', knowledgeCapabilityToken: existingToken,
|
||||
imageToolBinding: { context, describe: async () => 'Images', call: vi.fn() }
|
||||
}, new AbortController().signal)) void event
|
||||
})()).rejects.toThrow('Custom MCP unavailable')
|
||||
expect(gateway.revoke.mock.calls).toEqual(existingToken ? [['custom-capability']] : [['custom-capability'], ['image-capability']])
|
||||
expect(mocks.runHost).not.toHaveBeenCalled()
|
||||
} finally { await runtime.dispose() }
|
||||
})
|
||||
|
||||
it('adds assigned Skill instructions to the Continue prompt', async () => {
|
||||
let hostOptions: ContinueHostAdapterOptions | undefined
|
||||
const runtime = new ContinueAgentRuntime({
|
||||
|
||||
@@ -455,48 +455,46 @@ export class ContinueAgentRuntime implements AgentRuntime {
|
||||
}
|
||||
|
||||
const execute = request.workMode === 'execute'
|
||||
const imageCapabilityToken = request.imageToolBinding && execute
|
||||
? this.options.knowledgeGateway?.bindImageTool(request.imageToolBinding, signal, request.knowledgeCapabilityToken)
|
||||
: undefined
|
||||
const knowledgeEndpoint = this.options.knowledgeGateway?.getEndpoint()
|
||||
const knowledgeCapability =
|
||||
request.knowledgeCapabilityToken && knowledgeEndpoint
|
||||
(imageCapabilityToken ?? request.knowledgeCapabilityToken) && knowledgeEndpoint
|
||||
? {
|
||||
endpoint: knowledgeEndpoint,
|
||||
token: request.knowledgeCapabilityToken
|
||||
token: (imageCapabilityToken ?? request.knowledgeCapabilityToken)!
|
||||
}
|
||||
: undefined
|
||||
let customMcpCapability:
|
||||
| { endpoint: string; token: string }
|
||||
| undefined
|
||||
if (
|
||||
execute &&
|
||||
knowledgeEndpoint &&
|
||||
this.options.mcpServers?.length
|
||||
) {
|
||||
const token = this.options.knowledgeGateway?.grantCustomMcp(
|
||||
request.requestId,
|
||||
this.options.mcpServers,
|
||||
signal
|
||||
)
|
||||
if (token) {
|
||||
customMcpCapability = {
|
||||
endpoint: knowledgeEndpoint,
|
||||
token
|
||||
}
|
||||
try {
|
||||
await this.options.knowledgeGateway?.prepareCustomMcpTools(
|
||||
token,
|
||||
signal
|
||||
)
|
||||
} catch (error) {
|
||||
this.options.knowledgeGateway?.revoke(token)
|
||||
throw error
|
||||
}
|
||||
}
|
||||
}
|
||||
let result: ContinueHostRunResult
|
||||
const emittedTools = new Map<string, ContinueHostTool>()
|
||||
let emittedText = ''
|
||||
const requestQuestionIds = new Set<string>()
|
||||
try {
|
||||
if (
|
||||
execute &&
|
||||
knowledgeEndpoint &&
|
||||
this.options.mcpServers?.length
|
||||
) {
|
||||
const token = this.options.knowledgeGateway?.grantCustomMcp(
|
||||
request.requestId,
|
||||
this.options.mcpServers,
|
||||
signal
|
||||
)
|
||||
if (token) {
|
||||
customMcpCapability = {
|
||||
endpoint: knowledgeEndpoint,
|
||||
token
|
||||
}
|
||||
await this.options.knowledgeGateway?.prepareCustomMcpTools(
|
||||
token,
|
||||
signal
|
||||
)
|
||||
}
|
||||
}
|
||||
const host = this.getHostAdapter(
|
||||
binaryPath,
|
||||
execute || knowledgeCapability ? 'agent' : 'chat'
|
||||
@@ -643,6 +641,7 @@ export class ContinueAgentRuntime implements AgentRuntime {
|
||||
customMcpCapability.token
|
||||
)
|
||||
}
|
||||
if (imageCapabilityToken && imageCapabilityToken !== request.knowledgeCapabilityToken) this.options.knowledgeGateway?.revoke(imageCapabilityToken)
|
||||
}
|
||||
if (!result.text) {
|
||||
throw new Error('Continue CLI 未返回内容')
|
||||
|
||||
@@ -41,6 +41,64 @@ import { LocalRuntimeRegistry } from './local-runtime-registry'
|
||||
import { createAgentRuntime } from './create-runtime'
|
||||
import { defaultRuntimeSettings } from '../../shared/contracts'
|
||||
import type { ResolvedRuntimeSettings } from '../runtime-settings-store'
|
||||
import { AssistantDatabase } from '../assistant/assistant-database'
|
||||
import { ImageGenerationService } from './image-generation-service'
|
||||
|
||||
vi.mock('electron', () => ({ nativeImage: {} }))
|
||||
|
||||
it('generates and edits durable images through the controlled Harness Main proxy and refreshes it in Ask', async () => {
|
||||
const root = await realpath(await mkdtemp(join(tmpdir(), 'goodbuddy-dsh-images-')))
|
||||
const database = new AssistantDatabase(join(root, 'assistant.sqlite'))
|
||||
database.initialize(root)
|
||||
const conversationId = crypto.randomUUID()
|
||||
const profileId = crypto.randomUUID()
|
||||
const fetcher = vi.fn<typeof fetch>(async () => Response.json({ data: [{ b64_json: 'iVBORw0KGgo=' }] }))
|
||||
const service = new ImageGenerationService({ database, fetcher, getSettings: async () => ({ defaultImageModelProfileId: profileId, modelProfiles: [{
|
||||
id: profileId, name: 'Image fixture', modelName: 'fixture', protocol: 'openai-images-generations',
|
||||
allowConversationInvocation: true, baseUrl: 'https://image.test/v1', authentication: 'none'
|
||||
}] }) })
|
||||
let input: Record<string, unknown> = { intent: 'create', prompt: 'Blue circle' }
|
||||
const inProcess = createInProcessLaunch(root, {
|
||||
stream(options) {
|
||||
const prompt = latestUserText(options)
|
||||
const tool = options.tools?.find(tool => tool.name === 'generate_image')
|
||||
if (prompt === 'ask-image') {
|
||||
expect(tool).toBeUndefined()
|
||||
return textResponse('Ask is read-only')
|
||||
}
|
||||
expect(tool).toBeDefined()
|
||||
expect(tool?.description).toContain(profileId)
|
||||
if (!toolResultText(options, prompt)) return toolCall(prompt, 'generate_image', input)
|
||||
expect(toolResultText(options, prompt)).toContain('completed')
|
||||
return textResponse('IMAGE_SAVED')
|
||||
}
|
||||
})
|
||||
const runtime = new DeepSeekHarnessRuntime({ defaultWorkspace: root, baseUrl: 'https://chat.test/v1', model: 'fixture',
|
||||
launch: inProcess.launch, credentialRefs: { [CREDENTIAL_REF]: 'fixture' }, toolProvider: new ModelToolProvider(root),
|
||||
initializationTimeoutMs: 20_000, promptTimeoutMs: 20_000, shutdownTimeoutMs: 5_000 })
|
||||
try {
|
||||
for (const intent of ['create', 'edit'] as const) {
|
||||
const messageId = crypto.randomUUID()
|
||||
const requestId = crypto.randomUUID()
|
||||
database.saveLocalConversations([{ header: { id: conversationId, title: 'Images', updatedAt: Date.now() },
|
||||
messages: [{ id: messageId, role: 'assistant', content: '', createdAt: Date.now(), state: 'complete' }] }])
|
||||
const events = await collect(runtime.run({ requestId, conversationId, prompt: `${intent}-${requestId}`, workMode: 'execute',
|
||||
imageToolBinding: service.bind({ conversationId, messageId, requestId, workMode: 'execute' }) }, new AbortController().signal, async () => 'once'))
|
||||
expect(events.at(-1)?.type).toBe('done')
|
||||
const operation = database.getConversation(conversationId).messages.find(message => message.id === messageId)!.imageOperations![0]!
|
||||
expect(operation.state).toBe('completed')
|
||||
input = { intent: 'edit', prompt: 'Turn red', sourceArtifactIds: operation.artifactIds }
|
||||
}
|
||||
await collect(runtime.run({ requestId: crypto.randomUUID(), conversationId, prompt: 'ask-image', workMode: 'ask' }, new AbortController().signal))
|
||||
expect(fetcher).toHaveBeenCalledTimes(2)
|
||||
expect(String(fetcher.mock.calls[1]![0])).toMatch(/images\/edits$/u)
|
||||
} finally {
|
||||
await runtime.dispose()
|
||||
await service.dispose()
|
||||
database.close()
|
||||
await rm(root, { recursive: true, force: true })
|
||||
}
|
||||
}, 60_000)
|
||||
|
||||
const CREDENTIAL_REF = 'GOODBUDDY_HARNESS_MODEL_API_KEY'
|
||||
const SKILL_CALL_ID = 'e2e-skill-call'
|
||||
|
||||
@@ -290,7 +290,7 @@ function isMainWebTool(
|
||||
)
|
||||
}
|
||||
|
||||
function dshCompatibleWebInputSchema(
|
||||
function dshCompatibleMainInputSchema(
|
||||
schema: Record<string, unknown>
|
||||
): Record<string, unknown> {
|
||||
const compatible: Record<string, unknown> = {}
|
||||
@@ -311,7 +311,7 @@ function dshCompatibleWebInputSchema(
|
||||
propertySchema &&
|
||||
typeof propertySchema === 'object' &&
|
||||
!Array.isArray(propertySchema)
|
||||
? dshCompatibleWebInputSchema(
|
||||
? dshCompatibleMainInputSchema(
|
||||
propertySchema as Record<string, unknown>
|
||||
)
|
||||
: propertySchema
|
||||
@@ -325,7 +325,7 @@ function dshCompatibleWebInputSchema(
|
||||
typeof value === 'object' &&
|
||||
!Array.isArray(value)
|
||||
) {
|
||||
compatible.items = dshCompatibleWebInputSchema(
|
||||
compatible.items = dshCompatibleMainInputSchema(
|
||||
value as Record<string, unknown>
|
||||
)
|
||||
continue
|
||||
@@ -335,7 +335,7 @@ function dshCompatibleWebInputSchema(
|
||||
candidate &&
|
||||
typeof candidate === 'object' &&
|
||||
!Array.isArray(candidate)
|
||||
? dshCompatibleWebInputSchema(
|
||||
? dshCompatibleMainInputSchema(
|
||||
candidate as Record<string, unknown>
|
||||
)
|
||||
: candidate
|
||||
@@ -350,8 +350,8 @@ function proxyToolInputSchema(
|
||||
ReturnType<ModelToolProviderLike['listTools']>
|
||||
>[number]
|
||||
): Record<string, unknown> {
|
||||
return isMainWebTool(tool)
|
||||
? dshCompatibleWebInputSchema(tool.inputSchema)
|
||||
return isMainWebTool(tool) || tool.name === 'generate_image'
|
||||
? dshCompatibleMainInputSchema(tool.inputSchema)
|
||||
: tool.inputSchema
|
||||
}
|
||||
|
||||
@@ -368,7 +368,7 @@ function boundedProxyToolCatalog(
|
||||
const catalog = tools.filter(
|
||||
(tool) =>
|
||||
isMainWebTool(tool) ||
|
||||
(workMode === 'execute' && tool.source === 'mcp')
|
||||
(workMode === 'execute' && (tool.source === 'mcp' || tool.name === 'generate_image'))
|
||||
)
|
||||
if (catalog.length > MAX_MCP_PROXY_TOOLS) {
|
||||
throw new Error(
|
||||
@@ -386,7 +386,7 @@ function boundedProxyToolCatalog(
|
||||
names.add(tool.name)
|
||||
const description = tool.description.slice(
|
||||
0,
|
||||
MAX_MCP_TOOL_DESCRIPTION_CHARACTERS
|
||||
tool.name === 'generate_image' ? imageToolDescriptionLimit : MAX_MCP_TOOL_DESCRIPTION_CHARACTERS
|
||||
)
|
||||
const inputSchema = proxyToolInputSchema(tool)
|
||||
let serialized: string
|
||||
@@ -838,6 +838,7 @@ export class DeepSeekHarnessRuntime implements AgentRuntime {
|
||||
const run = this.activeRuns.get(params.sessionId)
|
||||
if (!run?.toolProvider || run.closed) return { tools: [] }
|
||||
const context = {
|
||||
imageToolBinding: run.request.imageToolBinding,
|
||||
conversationId:
|
||||
run?.request.conversationId ??
|
||||
'deepseek-harness-tool-catalog',
|
||||
@@ -855,7 +856,7 @@ export class DeepSeekHarnessRuntime implements AgentRuntime {
|
||||
)
|
||||
const catalog = boundedProxyToolCatalog(tools, context.workMode)
|
||||
this.proxyToolCatalogs.set(params.sessionId, tools.filter(
|
||||
(tool) => tool.source === 'mcp' || isMainWebTool(tool)
|
||||
(tool) => tool.source === 'mcp' || isMainWebTool(tool) || tool.name === 'generate_image'
|
||||
))
|
||||
return { tools: catalog }
|
||||
}
|
||||
@@ -881,6 +882,8 @@ export class DeepSeekHarnessRuntime implements AgentRuntime {
|
||||
)
|
||||
}
|
||||
const context = {
|
||||
imageToolBinding: run.request.imageToolBinding,
|
||||
...(name === 'generate_image' ? { toolCallId: typeof params.callId === 'string' ? params.callId : crypto.randomUUID() } : {}),
|
||||
conversationId: run.request.conversationId,
|
||||
browserTabId: run.request.browserTabId,
|
||||
workMode:
|
||||
@@ -894,7 +897,7 @@ export class DeepSeekHarnessRuntime implements AgentRuntime {
|
||||
const tool = tools.find(
|
||||
(candidate) =>
|
||||
candidate.name === name &&
|
||||
(candidate.source === 'mcp' ||
|
||||
(candidate.source === 'mcp' || candidate.name === 'generate_image' ||
|
||||
isMainWebTool(candidate))
|
||||
)
|
||||
if (!tool) {
|
||||
@@ -1754,3 +1757,4 @@ export class DeepSeekHarnessRuntime implements AgentRuntime {
|
||||
).catch(() => undefined)
|
||||
}
|
||||
}
|
||||
import { imageToolDescriptionLimit } from '../../shared/image-generation-contracts'
|
||||
|
||||
@@ -221,7 +221,7 @@ function parseProxyToolCatalog(
|
||||
!/^[a-zA-Z0-9_-]{1,64}$/u.test(tool.name) ||
|
||||
names.has(tool.name) ||
|
||||
typeof tool.description !== 'string' ||
|
||||
tool.description.length > 1_000 ||
|
||||
tool.description.length > (tool.name === 'generate_image' ? imageToolDescriptionLimit : 1_000) ||
|
||||
!tool.inputSchema ||
|
||||
typeof tool.inputSchema !== 'object' ||
|
||||
Array.isArray(tool.inputSchema)
|
||||
@@ -885,6 +885,7 @@ export class GoodBuddyHarnessControlPlane {
|
||||
{
|
||||
sessionId,
|
||||
name: tool.name,
|
||||
callId: randomUUID(),
|
||||
arguments: argumentsValue
|
||||
}
|
||||
)
|
||||
@@ -1450,3 +1451,4 @@ export class GoodBuddyHarnessControlPlane {
|
||||
await this.disposing
|
||||
}
|
||||
}
|
||||
import { imageToolDescriptionLimit } from '../../shared/image-generation-contracts'
|
||||
|
||||
@@ -0,0 +1,231 @@
|
||||
import { randomUUID } from 'node:crypto'
|
||||
import { afterEach, describe, expect, it, vi } from 'vitest'
|
||||
import { Client } from '@modelcontextprotocol/sdk/client/index.js'
|
||||
import { StreamableHTTPClientTransport } from '@modelcontextprotocol/sdk/client/streamableHttp.js'
|
||||
import { AssistantDatabase } from '../assistant/assistant-database'
|
||||
import { ImageGenerationService, type ImageServiceSettings, type ImageGenerationServiceOptions } from './image-generation-service'
|
||||
import { imageOperationSchema, type ImageOperation } from '../../shared/image-generation-contracts'
|
||||
import { KnowledgeMcpGateway } from './knowledge-mcp-gateway'
|
||||
import { ModelToolProvider } from './model-tool-provider'
|
||||
|
||||
vi.mock('electron', () => ({ nativeImage: { createFromBuffer: () => ({ isEmpty: () => false, toPNG: () => Buffer.from(png, 'base64') }) } }))
|
||||
const png = 'iVBORw0KGgo='
|
||||
const cleanups: (() => void | Promise<void>)[] = []
|
||||
afterEach(async () => { for (const cleanup of cleanups.splice(0).reverse()) await cleanup() })
|
||||
|
||||
function setup(fetcher = vi.fn<typeof fetch>(async () => Response.json({ data: [{ b64_json: png }] })), observers: Pick<ImageGenerationServiceOptions, 'onOperation' | 'onUsage' | 'onError'> = {}) {
|
||||
const database = new AssistantDatabase(':memory:')
|
||||
database.initialize(process.cwd())
|
||||
cleanups.push(() => database.close())
|
||||
const context = { conversationId: randomUUID(), messageId: randomUUID(), requestId: randomUUID(), workMode: 'execute' as const }
|
||||
const message = { id: context.messageId, role: 'assistant' as const, content: '', createdAt: Date.now(), state: 'complete' as const }
|
||||
const header = { id: context.conversationId, title: 'Images', updatedAt: Date.now() }
|
||||
database.saveLocalConversations([{ header, messages: [message] }])
|
||||
const profile = { id: randomUUID(), name: 'Image A', modelName: 'image-a', protocol: 'openai-images-generations',
|
||||
allowConversationInvocation: true, baseUrl: 'https://provider.test/v1', authentication: 'api-key' as const, apiKey: 'test-secret' }
|
||||
const settings: ImageServiceSettings = { modelProfiles: [profile], defaultImageModelProfileId: profile.id }
|
||||
const events: ImageOperation[] = []
|
||||
const service = new ImageGenerationService({ database, getSettings: async () => settings, fetcher,
|
||||
onOperation: operation => { imageOperationSchema.parse(operation); events.push(operation); observers.onOperation?.(operation) }, onUsage: observers.onUsage, onError: observers.onError })
|
||||
service.initialize()
|
||||
cleanups.push(() => service.dispose())
|
||||
return { database, service, context, fetcher, settings, profile, events, header, message }
|
||||
}
|
||||
|
||||
describe('Main conversation image service', () => {
|
||||
it('does not let a destroyed Renderer or usage observer prevent submission or overwrite completion', async () => {
|
||||
const onError = vi.fn()
|
||||
const h = setup(vi.fn<typeof fetch>(async () => Response.json({ data: [{ b64_json: png }], usage: { input_tokens: 2, output_tokens: 3 } })), {
|
||||
onOperation: () => { throw new Error('Renderer is destroyed') },
|
||||
onUsage: () => { throw new Error('Usage observer failed') }, onError
|
||||
})
|
||||
const operation = await h.service.bind(h.context).call({ intent: 'create', prompt: 'Blue' }, 'observer')
|
||||
expect(operation.state).toBe('completed')
|
||||
expect(h.service.getOperation(h.context.conversationId, operation.id).state).toBe('completed')
|
||||
expect(h.fetcher).toHaveBeenCalledOnce()
|
||||
expect(onError).toHaveBeenCalled()
|
||||
})
|
||||
|
||||
it('returns the known persistence failure instead of a stale running state when the database stops accepting writes', async () => {
|
||||
const onError = vi.fn()
|
||||
const h = setup(undefined, { onError })
|
||||
h.fetcher.mockImplementation(async () => {
|
||||
vi.spyOn(h.database, 'saveConversationImageOperation').mockImplementation(() => { throw new Error('Disk full') })
|
||||
return Response.json({ data: [{ b64_json: png }] })
|
||||
})
|
||||
const operation = await h.service.bind(h.context).call({ intent: 'create', prompt: 'Blue' }, 'disk-full')
|
||||
expect(operation).toMatchObject({ state: 'failed', error: 'Disk full', artifactIds: [] })
|
||||
expect(onError).toHaveBeenCalledWith(expect.objectContaining({ message: 'Disk full' }))
|
||||
expect(h.fetcher).toHaveBeenCalledOnce()
|
||||
})
|
||||
it('saves before publishing, preserves metadata against stale renderer saves, and edits durable history', async () => {
|
||||
const h = setup()
|
||||
const first = await h.service.bind(h.context).call({ intent: 'create', prompt: 'Blue circle' }, 'first')
|
||||
expect(first.state).toBe('completed')
|
||||
expect(h.database.getArtifact(first.artifactIds[0]!).content).toContain(png)
|
||||
h.database.saveLocalConversations([{ header: h.header, messages: [h.message] }])
|
||||
expect(h.service.getOperation(h.context.conversationId, first.id)).toEqual(first)
|
||||
expect(h.database.getConversation(h.context.conversationId).messages[0]!.artifactIds).toEqual(first.artifactIds)
|
||||
h.settings.defaultImageModelProfileId = randomUUID()
|
||||
const second = await h.service.bind({ ...h.context, requestId: randomUUID() }).call({ intent: 'edit', prompt: 'Turn red', sourceArtifactIds: first.artifactIds }, 'edit')
|
||||
expect(second.modelProfileId).toBe(first.modelProfileId)
|
||||
expect(String(h.fetcher.mock.calls[1]![0])).toMatch(/images\/edits$/u)
|
||||
expect((h.fetcher.mock.calls[1]![1]!.body as FormData).get('image')).toBeInstanceOf(Blob)
|
||||
expect(h.events.map(event => event.state)).toEqual(['running', 'saving', 'completed', 'running', 'saving', 'completed'])
|
||||
})
|
||||
|
||||
it('persists upload bytes and rejects cross-conversation or missing source references', async () => {
|
||||
const h = setup()
|
||||
const ids = h.service.persistUploads(h.context, [{ name: 'upload.png', mediaType: 'image/png', data: png }])
|
||||
h.database.saveLocalConversations([{ header: h.header, messages: [h.message] }])
|
||||
expect(h.database.getConversation(h.context.conversationId).messages[0]!.imageSourceArtifactIds).toEqual(ids)
|
||||
const binding = h.service.bind(h.context)
|
||||
expect((await binding.call({ intent: 'edit', prompt: 'Blue', sourceArtifactIds: ids }, 'upload')).state).toBe('completed')
|
||||
await expect(binding.call({ intent: 'edit', prompt: 'Blue', sourceArtifactIds: [randomUUID()] }, 'foreign')).rejects.toThrow('not available')
|
||||
await expect(binding.call({ intent: 'edit', prompt: 'Blue' }, 'missing')).rejects.toThrow('requires source')
|
||||
expect(h.fetcher).toHaveBeenCalledTimes(1)
|
||||
})
|
||||
|
||||
it('rejects absent, disabled, deleted defaults without silent fallback and rechecks current mode', async () => {
|
||||
const h = setup()
|
||||
let mode: 'ask' | 'execute' = 'execute'
|
||||
const binding = h.service.bind(h.context, () => mode)
|
||||
expect(await binding.describe()).toContain(h.profile.id)
|
||||
h.profile.allowConversationInvocation = false
|
||||
await expect(binding.call({ intent: 'create', prompt: 'Blue' }, 'disabled')).rejects.toThrow('unavailable')
|
||||
h.profile.allowConversationInvocation = true
|
||||
h.settings.defaultImageModelProfileId = randomUUID()
|
||||
await expect(binding.call({ intent: 'create', prompt: 'Blue' }, 'deleted')).rejects.toThrow('unavailable')
|
||||
h.settings.defaultImageModelProfileId = null
|
||||
await expect(binding.call({ intent: 'create', prompt: 'Blue' }, 'absent')).rejects.toThrow('default')
|
||||
mode = 'ask'
|
||||
await expect(binding.call({ intent: 'create', prompt: 'Blue', modelProfileId: h.profile.id }, 'ask')).rejects.toThrow('Execute')
|
||||
expect(h.fetcher).not.toHaveBeenCalled()
|
||||
})
|
||||
|
||||
it('continues after chat abort and pins the accepted profile while a replacement binding starts', async () => {
|
||||
let resolve!: (response: Response) => void
|
||||
const h = setup(vi.fn<typeof fetch>(() => new Promise(done => { resolve = done })))
|
||||
const controller = new AbortController()
|
||||
const pending = h.service.bind(h.context).call({ intent: 'create', prompt: 'Blue' }, 'call', controller.signal)
|
||||
await vi.waitFor(() => expect(h.fetcher).toHaveBeenCalledOnce())
|
||||
const providerSignal = h.fetcher.mock.calls[0]![1]!.signal!
|
||||
controller.abort()
|
||||
expect((await pending).state).toBe('running')
|
||||
expect(providerSignal.aborted).toBe(false)
|
||||
h.profile.modelName = 'replacement'
|
||||
h.profile.allowConversationInvocation = false
|
||||
resolve(Response.json({ data: [{ b64_json: png }] }))
|
||||
await vi.waitFor(() => expect(h.events.at(-1)?.state).toBe('completed'))
|
||||
expect(h.events.at(-1)).toMatchObject({ modelName: 'image-a', modelProfileName: 'Image A' })
|
||||
expect(JSON.parse(h.fetcher.mock.calls[0]![1]!.body as string).model).toBe('image-a')
|
||||
await expect(h.service.bind({ ...h.context, requestId: randomUUID() }).call({ intent: 'create', prompt: 'Other' }, 'other')).rejects.toThrow('unavailable')
|
||||
})
|
||||
|
||||
it('deduplicates one invocation, creates a new operation on manual regenerate, and retains late cancellation output', async () => {
|
||||
let resolve!: (response: Response) => void
|
||||
const h = setup(vi.fn<typeof fetch>(() => new Promise(done => { resolve = done })))
|
||||
const binding = h.service.bind(h.context)
|
||||
const input = { intent: 'create', prompt: 'Blue' }
|
||||
const one = binding.call(input, 'same')
|
||||
const two = binding.call(input, 'same')
|
||||
await vi.waitFor(() => expect(h.events).toHaveLength(1))
|
||||
h.service.cancel(h.context.conversationId, h.events[0]!.id)
|
||||
resolve(Response.json({ data: [{ b64_json: png }] }))
|
||||
const completed = await one
|
||||
expect(await two).toEqual(completed)
|
||||
expect(completed).toMatchObject({ state: 'completed', cancellationRequested: true })
|
||||
expect(h.fetcher).toHaveBeenCalledOnce()
|
||||
h.fetcher.mockImplementation(async () => Response.json({ data: [{ b64_json: png }] }))
|
||||
const regenerated = await h.service.regenerate(h.context, completed.id)
|
||||
expect(regenerated.id).not.toBe(completed.id)
|
||||
expect(h.fetcher).toHaveBeenCalledTimes(2)
|
||||
})
|
||||
|
||||
it('never falls back to a paid generation after an edit rejection and redacts provider errors', async () => {
|
||||
const h = setup(vi.fn<typeof fetch>(async () => Response.json({ error: { message: 'Image editing not supported Authorization: Bearer private-key' } }, { status: 405 })))
|
||||
const ids = h.service.persistUploads(h.context, [{ name: 'source', mediaType: 'image/png', data: png }])
|
||||
const result = await h.service.bind(h.context).call({ intent: 'edit', prompt: 'Blue', sourceArtifactIds: ids }, 'edit')
|
||||
expect(result.state).toBe('failed')
|
||||
expect(result.error).toContain('405')
|
||||
expect(result.error).not.toContain('private-key')
|
||||
expect(h.fetcher).toHaveBeenCalledOnce()
|
||||
})
|
||||
|
||||
it('does not resurrect a deleted conversation or retain its late artifact', async () => {
|
||||
let resolve!: (response: Response) => void
|
||||
const h = setup(vi.fn<typeof fetch>(() => new Promise(done => { resolve = done })))
|
||||
const pending = h.service.bind(h.context).call({ intent: 'create', prompt: 'Blue' }, 'deleted')
|
||||
await vi.waitFor(() => expect(h.fetcher).toHaveBeenCalledOnce())
|
||||
h.database.deleteLocalConversation(h.context.conversationId)
|
||||
resolve(Response.json({ data: [{ b64_json: png }] }))
|
||||
await pending
|
||||
expect(() => h.database.getConversation(h.context.conversationId)).toThrow()
|
||||
expect(h.database.listArtifacts()).toHaveLength(0)
|
||||
})
|
||||
|
||||
it('marks unfinished operations unconfirmed on initialization without issuing requests', () => {
|
||||
const h = setup()
|
||||
h.database.saveConversationImageOperation({ id: randomUUID(), conversationId: h.context.conversationId, messageId: h.context.messageId,
|
||||
requestId: h.context.requestId, callId: 'old', modelProfileId: h.profile.id, modelName: h.profile.name,
|
||||
input: { intent: 'create', prompt: 'Blue', sourceArtifactIds: [] }, state: 'running', artifactIds: [], createdAt: 1, updatedAt: 1 })
|
||||
h.service.initialize()
|
||||
expect(h.database.getConversation(h.context.conversationId).messages[0]!.imageOperations![0]!.state).toBe('unconfirmed')
|
||||
expect(h.fetcher).not.toHaveBeenCalled()
|
||||
})
|
||||
|
||||
it('reports provider-side uncertainty after network loss and local cancellation without retrying', async () => {
|
||||
const h = setup(vi.fn<typeof fetch>(async () => { throw new TypeError('fetch failed') }))
|
||||
const uncertain = await h.service.bind(h.context).call({ intent: 'create', prompt: 'Blue' }, 'network')
|
||||
expect(uncertain.state).toBe('unconfirmed')
|
||||
h.fetcher.mockImplementation((_url, init) => new Promise((_resolve, reject) => {
|
||||
init!.signal!.addEventListener('abort', () => reject(init!.signal!.reason), { once: true })
|
||||
}))
|
||||
const pending = h.service.bind(h.context).call({ intent: 'create', prompt: 'Blue' }, 'cancel')
|
||||
await vi.waitFor(() => expect(h.fetcher).toHaveBeenCalledTimes(2))
|
||||
h.service.cancel(h.context.conversationId, h.events.at(-1)!.id)
|
||||
expect(await pending).toMatchObject({ state: 'stopped', cancellationRequested: true })
|
||||
expect(h.fetcher).toHaveBeenCalledTimes(2)
|
||||
})
|
||||
|
||||
it('reuses an accepted invocation across bindings and does not skip a disabled historical model', async () => {
|
||||
const h = setup()
|
||||
const created = await h.service.bind(h.context).call({ intent: 'create', prompt: 'Blue' }, 'same')
|
||||
expect((await h.service.bind(h.context).call({ intent: 'create', prompt: 'Blue' }, 'same')).id).toBe(created.id)
|
||||
const other = { ...h.profile, id: randomUUID(), name: 'Other', modelName: 'other' }
|
||||
h.settings.modelProfiles = [h.profile, other]
|
||||
h.settings.defaultImageModelProfileId = other.id
|
||||
h.profile.allowConversationInvocation = false
|
||||
await expect(h.service.bind(h.context).call({ intent: 'edit', prompt: 'Red', sourceArtifactIds: created.artifactIds }, 'blocked-edit')).rejects.toThrow('unavailable')
|
||||
const edited = await h.service.bind(h.context).call({ intent: 'edit', prompt: 'Red', sourceArtifactIds: created.artifactIds, modelProfileId: other.id }, 'explicit-edit')
|
||||
expect(edited.modelProfileId).toBe(other.id)
|
||||
expect(h.fetcher).toHaveBeenCalledTimes(2)
|
||||
})
|
||||
|
||||
it('dispatches the native Model tool through its scoped binding and denies Ask', async () => {
|
||||
const h = setup()
|
||||
const provider = new ModelToolProvider(process.cwd())
|
||||
cleanups.push(() => provider.dispose())
|
||||
const context = { conversationId: h.context.conversationId, workMode: 'execute' as const, imageToolBinding: h.service.bind(h.context), toolCallId: 'native' }
|
||||
const result = await provider.callTool('generate_image', { intent: 'create', prompt: 'Blue' }, new AbortController().signal, context)
|
||||
expect(JSON.stringify(result)).toContain('completed')
|
||||
expect(JSON.stringify(result)).not.toContain(png)
|
||||
await expect(provider.callTool('generate_image', { intent: 'create', prompt: 'Blue' }, new AbortController().signal, { ...context, workMode: 'ask' })).rejects.toThrow()
|
||||
})
|
||||
|
||||
it('round trips image-only MCP discovery/call and does not abort accepted work on capability revocation', async () => {
|
||||
const h = setup()
|
||||
const gateway = new KnowledgeMcpGateway({} as never)
|
||||
await gateway.start()
|
||||
cleanups.push(() => gateway.dispose())
|
||||
const token = gateway.bindImageTool(h.service.bind(h.context), new AbortController().signal)
|
||||
const client = new Client({ name: 'image-test', version: '1' })
|
||||
cleanups.push(() => client.close())
|
||||
await client.connect(new StreamableHTTPClientTransport(new URL(gateway.getEndpoint()!), { requestInit: { headers: { Authorization: `Bearer ${token}` } } }))
|
||||
expect((await client.listTools()).tools.map(tool => tool.name)).toEqual(['generate_image'])
|
||||
const result = await client.callTool({ name: 'generate_image', arguments: { intent: 'create', prompt: 'Blue' } })
|
||||
expect(JSON.stringify(result)).toContain('completed')
|
||||
expect(h.fetcher).toHaveBeenCalledOnce()
|
||||
gateway.revoke(token)
|
||||
})
|
||||
})
|
||||
@@ -0,0 +1,209 @@
|
||||
import { randomUUID } from 'node:crypto'
|
||||
import {
|
||||
imageToolDescription, imageToolInputSchema,
|
||||
type ImageOperation, type ImageRequestContext, type ImageToolInput
|
||||
} from '../../shared/image-generation-contracts'
|
||||
import type { AssistantDatabase } from '../assistant/assistant-database'
|
||||
import { ModelAgentRuntime, type ModelRuntimeOptions } from './model-runtime'
|
||||
import { withImageConversationContext } from './image-conversation-context'
|
||||
import type { AgentImage, RuntimeEvent } from './runtime'
|
||||
import { redactSensitiveText, safeToolErrorDetail } from './approval-summary'
|
||||
import type { ImageToolBinding } from './image-tool-binding'
|
||||
export type { ImageToolBinding } from './image-tool-binding'
|
||||
|
||||
export type ConversationImageProfile = {
|
||||
id: string
|
||||
name: string
|
||||
protocol: string
|
||||
modelName: string
|
||||
allowConversationInvocation?: boolean
|
||||
} & Pick<ModelRuntimeOptions, 'baseUrl' | 'apiKey' | 'authentication' | 'requestHeaders' | 'requestBody' | 'imageGenerationQuality'>
|
||||
|
||||
export type ImageServiceSettings = {
|
||||
modelProfiles: readonly ConversationImageProfile[]
|
||||
defaultImageModelProfileId?: string | null
|
||||
}
|
||||
export type ImageGenerationServiceOptions = {
|
||||
database: Pick<AssistantDatabase, 'getConversation' | 'getArtifact' | 'saveConversationImageOperation' | 'saveConversationImageSources' | 'markUnfinishedImageOperationsUnconfirmed'>
|
||||
getSettings(): Promise<ImageServiceSettings>
|
||||
onOperation?(operation: ImageOperation): void
|
||||
onUsage?(event: Extract<RuntimeEvent, { type: 'model-usage' }>): void
|
||||
/** Diagnostic sink for observer or terminal-persistence failures; must not throw. */
|
||||
onError?(error: Error): void
|
||||
fetcher?: typeof fetch
|
||||
requestTimeoutMs?: number
|
||||
toolWaitMs?: number
|
||||
}
|
||||
|
||||
/** Main owns requests; request-scoped bindings own only admission and waiting. */
|
||||
export class ImageGenerationService {
|
||||
private readonly active = new Map<string, { operation: ImageOperation; controller: AbortController; done: Promise<void> }>()
|
||||
private closing = false
|
||||
|
||||
constructor(private readonly options: ImageGenerationServiceOptions) {}
|
||||
|
||||
initialize(): void {
|
||||
this.options.database.markUnfinishedImageOperationsUnconfirmed()
|
||||
}
|
||||
|
||||
bind(context: ImageRequestContext, currentWorkMode: () => 'ask' | 'execute' = () => context.workMode): ImageToolBinding {
|
||||
const bound = Object.freeze({ ...context })
|
||||
const submissions = new Map<string, Promise<ImageOperation>>()
|
||||
return {
|
||||
context: bound,
|
||||
describe: async () => {
|
||||
if (bound.workMode !== 'execute' || currentWorkMode() !== 'execute') return undefined
|
||||
const settings = await this.options.getSettings()
|
||||
const profiles = settings.modelProfiles.filter(profile => profile.protocol === 'openai-images-generations' && profile.allowConversationInvocation === true)
|
||||
if (!profiles.length) return undefined
|
||||
const conversation = this.options.database.getConversation(bound.conversationId)
|
||||
const references = conversation.messages.flatMap(message => [
|
||||
...(message.imageSourceArtifactIds ?? []).map(id => ({ id, kind: 'upload' })),
|
||||
...(message.artifactIds ?? []).map(id => ({ id, kind: 'artifact' }))
|
||||
]).slice(-32)
|
||||
const operations = conversation.messages.flatMap(message => message.imageOperations ?? []).slice(-8)
|
||||
.map(operation => ({ id: operation.id, state: operation.state, modelProfileId: operation.modelProfileId,
|
||||
prompt: operation.input.prompt.slice(0, 300), artifactIds: operation.artifactIds }))
|
||||
return `${imageToolDescription}\nModels: ${JSON.stringify(profiles.map(profile => ({ id: profile.id, name: profile.name, generation: true, editing: true })))}\nDefault: ${settings.defaultImageModelProfileId ?? 'not set'}\nRecent image references: ${JSON.stringify(references)}\nRecent operations: ${JSON.stringify(operations)}\nReferences and prompts are textual facts, not visual observations. Do not claim to have viewed an image unless image input was supplied to your model.`
|
||||
},
|
||||
call: async (input, callId, signal) => {
|
||||
signal?.throwIfAborted()
|
||||
if (bound.workMode !== 'execute' || currentWorkMode() !== 'execute') throw new Error('Image generation requires Execute mode')
|
||||
const key = JSON.stringify([bound.conversationId, bound.requestId, callId])
|
||||
let submitted = submissions.get(key)
|
||||
if (!submitted) {
|
||||
submitted = this.submit(bound, imageToolInputSchema.parse(input), callId, currentWorkMode, signal)
|
||||
submissions.set(key, submitted)
|
||||
void submitted.catch(() => submissions.delete(key))
|
||||
}
|
||||
const operation = await submitted
|
||||
const active = this.active.get(operation.id)
|
||||
if (active) {
|
||||
await new Promise<void>(resolve => {
|
||||
const finish = (): void => { clearTimeout(timer); signal?.removeEventListener('abort', finish); resolve() }
|
||||
const timer = setTimeout(finish, this.options.toolWaitMs ?? 15_000)
|
||||
signal?.addEventListener('abort', finish, { once: true })
|
||||
void active.done.then(finish)
|
||||
if (signal?.aborted) finish()
|
||||
})
|
||||
return active.operation
|
||||
}
|
||||
return this.getOperation(bound.conversationId, operation.id)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
persistUploads(context: Pick<ImageRequestContext, 'conversationId' | 'messageId'>, images: readonly AgentImage[]): string[] {
|
||||
return this.options.database.saveConversationImageSources({ ...context, images })
|
||||
}
|
||||
|
||||
getOperation(conversationId: string, operationId: string): ImageOperation {
|
||||
const operation = this.options.database.getConversation(conversationId).messages
|
||||
.flatMap(message => message.imageOperations ?? []).find(item => item.id === operationId)
|
||||
if (!operation) throw new Error('Image operation does not exist in this conversation')
|
||||
return operation
|
||||
}
|
||||
|
||||
cancel(conversationId: string, operationId: string): ImageOperation {
|
||||
const previous = this.getOperation(conversationId, operationId)
|
||||
const active = this.active.get(operationId)
|
||||
if (!active) return previous
|
||||
active.operation = this.save({ ...active.operation, state: 'cancelling', cancellationRequested: true })
|
||||
active.controller.abort(new Error('Stopped waiting; the image provider may still generate the image'))
|
||||
return active.operation
|
||||
}
|
||||
|
||||
regenerate(context: ImageRequestContext, operationId: string, currentWorkMode?: () => 'ask' | 'execute'): Promise<ImageOperation> {
|
||||
const previous = this.getOperation(context.conversationId, operationId)
|
||||
return this.bind(context, currentWorkMode).call({ ...previous.input, modelProfileId: previous.modelProfileId }, randomUUID())
|
||||
}
|
||||
|
||||
cancelConversation(conversationId: string): void {
|
||||
for (const active of this.active.values()) {
|
||||
if (active.operation.conversationId === conversationId) this.cancel(conversationId, active.operation.id)
|
||||
}
|
||||
}
|
||||
|
||||
async dispose(): Promise<void> {
|
||||
this.closing = true
|
||||
for (const active of this.active.values()) active.controller.abort(new Error('Application is closing; image result is unconfirmed'))
|
||||
await Promise.all([...this.active.values()].map(active => active.done))
|
||||
}
|
||||
|
||||
private save(operation: ImageOperation, image?: { mimeType: 'image/png' | 'image/jpeg' | 'image/webp'; data: string }): ImageOperation {
|
||||
const saved = this.options.database.saveConversationImageOperation({ ...operation, updatedAt: Date.now() }, image)
|
||||
try {
|
||||
this.options.onOperation?.(saved)
|
||||
} catch (error) {
|
||||
// A closed Renderer must not stop an accepted provider request or change its outcome.
|
||||
this.report(error)
|
||||
}
|
||||
return saved
|
||||
}
|
||||
|
||||
private report(error: unknown): void {
|
||||
this.options.onError?.(new Error(redactSensitiveText(safeToolErrorDetail(error) ?? 'Image operation notification or persistence failed')))
|
||||
}
|
||||
|
||||
private async submit(context: ImageRequestContext, input: ImageToolInput, callId: string, currentWorkMode: () => 'ask' | 'execute', signal?: AbortSignal): Promise<ImageOperation> {
|
||||
if (this.closing) throw new Error('Image service is closing')
|
||||
const settings = await this.options.getSettings()
|
||||
const conversation = this.options.database.getConversation(context.conversationId)
|
||||
const existing = conversation.messages.flatMap(message => message.imageOperations ?? [])
|
||||
.find(operation => operation.requestId === context.requestId && operation.callId === callId)
|
||||
if (existing) return existing
|
||||
if (!conversation.messages.some(message => message.id === context.messageId)) throw new Error('Originating image message must be persisted before submission')
|
||||
const allowed = new Set(conversation.messages.flatMap(message => [
|
||||
...(message.artifactIds ?? []), ...(message.imageSourceArtifactIds ?? []),
|
||||
...(message.imageOperations?.flatMap(operation => operation.artifactIds) ?? [])
|
||||
]))
|
||||
if (input.sourceArtifactIds.some(id => !allowed.has(id))) throw new Error('Source image is not available in this conversation')
|
||||
const origins = new Set(conversation.messages.flatMap(message => message.imageOperations ?? [])
|
||||
.filter(operation => operation.artifactIds.some(id => input.sourceArtifactIds.includes(id)))
|
||||
.map(operation => operation.modelProfileId))
|
||||
if (!input.modelProfileId && origins.size > 1) throw new Error('Source images use different models; specify the image model')
|
||||
const profileId = input.modelProfileId ?? (input.intent === 'edit' ? [...origins][0] : undefined) ?? settings.defaultImageModelProfileId
|
||||
if (!profileId) throw new Error('Specify an image model or set the default image model')
|
||||
const profile = settings.modelProfiles.find(item => item.id === profileId)
|
||||
if (!profile || !profile.allowConversationInvocation || profile.protocol !== 'openai-images-generations') throw new Error('Selected image model is unavailable; select another model')
|
||||
if (profile.authentication === 'api-key' && !profile.apiKey) throw new Error('Selected image model has no API key configured')
|
||||
const connection = structuredClone(profile)
|
||||
const request = withImageConversationContext({ requestId: context.requestId, conversationId: context.conversationId, prompt: input.prompt, imageContextArtifactIds: input.sourceArtifactIds }, id => this.options.database.getArtifact(id))
|
||||
if (input.intent === 'edit' && (request.imageContextNotice || request.images?.length !== input.sourceArtifactIds.length)) throw new Error('Source image is missing or unsupported; select the image again')
|
||||
signal?.throwIfAborted()
|
||||
if (this.closing) throw new Error('Image service is closing')
|
||||
if (currentWorkMode() !== 'execute') throw new Error('Image generation requires Execute mode')
|
||||
const operation = this.save({ id: randomUUID(), conversationId: context.conversationId, messageId: context.messageId, requestId: context.requestId,
|
||||
callId, modelProfileId: connection.id, modelName: connection.modelName, modelProfileName: connection.name,
|
||||
input: { ...input, quality: input.quality ?? connection.imageGenerationQuality ?? 'auto' },
|
||||
state: 'running', artifactIds: [], createdAt: Date.now(), updatedAt: Date.now() })
|
||||
const runtime = new ModelAgentRuntime({ ...connection, model: connection.modelName, protocol: 'openai-images-generations',
|
||||
imageGenerationQuality: input.quality ?? connection.imageGenerationQuality, fetcher: this.options.fetcher, requestTimeoutMs: this.options.requestTimeoutMs })
|
||||
const active = { operation, controller: new AbortController(), done: Promise.resolve() }
|
||||
this.active.set(operation.id, active)
|
||||
active.done = (async () => {
|
||||
try {
|
||||
for await (const event of runtime.generateImage(request, active.controller.signal)) {
|
||||
if (event.type === 'model-usage') {
|
||||
try { this.options.onUsage?.(event) } catch (error) { this.report(error) }
|
||||
}
|
||||
if (event.type === 'generated-image') {
|
||||
active.operation = this.save({ ...active.operation, state: 'saving' })
|
||||
active.operation = this.save(active.operation, event)
|
||||
}
|
||||
}
|
||||
} catch (error) {
|
||||
const state = this.closing ? 'unconfirmed' : active.controller.signal.aborted ? 'stopped' : /timeout|timed out|fetch failed|network|超时/iu.test(String(error)) ? 'unconfirmed' : 'failed'
|
||||
const detail = safeToolErrorDetail(error) ?? 'Image request failed'
|
||||
const suffix = detail.match(/（HTTP [^）]+）$/u)?.[0] ?? ''
|
||||
const publicDetail = redactSensitiveText(suffix ? detail.slice(0, -suffix.length) : detail) + suffix
|
||||
active.operation = { ...active.operation, state, error: publicDetail.slice(0, 2_000), updatedAt: Date.now() }
|
||||
try { active.operation = this.save(active.operation) } catch (storageError) { this.report(storageError) }
|
||||
} finally {
|
||||
await runtime.dispose()
|
||||
this.active.delete(operation.id)
|
||||
}
|
||||
})()
|
||||
return operation
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,95 @@
|
||||
// @vitest-environment node
|
||||
import { randomUUID } from 'node:crypto'
|
||||
import { existsSync } from 'node:fs'
|
||||
import { mkdtemp, rm } from 'node:fs/promises'
|
||||
import { createServer } from 'node:http'
|
||||
import { tmpdir } from 'node:os'
|
||||
import { join } from 'node:path'
|
||||
import { expect, it, vi } from 'vitest'
|
||||
import { AssistantDatabase } from '../assistant/assistant-database'
|
||||
import { ImageGenerationService } from './image-generation-service'
|
||||
import { KnowledgeMcpGateway } from './knowledge-mcp-gateway'
|
||||
import { ModelAgentRuntime } from './model-runtime'
|
||||
import { OpenCodeRuntime } from './opencode-runtime'
|
||||
import { ContinueAgentRuntime } from './continue-runtime'
|
||||
|
||||
vi.mock('electron', () => ({ nativeImage: {} }))
|
||||
const openCodeBinary = join(process.cwd(), '.runtime-resources', process.arch, process.platform === 'win32' ? 'opencode.exe' : 'opencode')
|
||||
const continueBinary = join(process.cwd(), 'node_modules', '@continuedev', 'cli', 'dist', 'cn.js')
|
||||
|
||||
it.for(['model', 'opencode', 'continue'] as const)('generates and edits through the real local %s tool loop', { timeout: 120_000 }, async (name, { skip }) => {
|
||||
if ((name === 'opencode' && !existsSync(openCodeBinary)) || (name === 'continue' && !existsSync(continueBinary))) return skip()
|
||||
const root = await mkdtemp(join(tmpdir(), 'goodbuddy-image-runtime-'))
|
||||
const database = new AssistantDatabase(join(root, 'assistant.sqlite'))
|
||||
database.initialize(root)
|
||||
const conversationId = randomUUID()
|
||||
const imageProfileId = randomUUID()
|
||||
let toolInput: Record<string, unknown> = { intent: 'create', prompt: 'Blue circle' }
|
||||
const textRequests: unknown[] = []
|
||||
const model = createServer((request, response) => {
|
||||
void (async () => {
|
||||
let content = ''
|
||||
for await (const chunk of request) content += String(chunk)
|
||||
const body = JSON.parse(content) as { messages?: { role: string; content: unknown }[]; tools?: { function: { name: string } }[] }
|
||||
textRequests.push(body)
|
||||
let lastUser = -1
|
||||
body.messages?.forEach((message, index) => { if (message.role === 'user') lastUser = index })
|
||||
const results = body.messages?.slice(lastUser + 1).filter(message => message.role === 'tool') ?? []
|
||||
const tool = body.tools?.find(tool => tool.function.name.endsWith('generate_image'))
|
||||
if (!tool) throw new Error('Image tool was not exposed')
|
||||
const call = results.length === 0
|
||||
const delta = call ? { role: 'assistant', tool_calls: [{ index: 0, id: randomUUID(), type: 'function',
|
||||
function: { name: tool.function.name, arguments: JSON.stringify(toolInput) } }] } : { role: 'assistant', content: 'IMAGE_SAVED' }
|
||||
response.writeHead(200, { 'content-type': 'text/event-stream' })
|
||||
for (const [value, finish] of [[delta, null], [{}, call ? 'tool_calls' : 'stop']]) {
|
||||
response.write(`data: ${JSON.stringify({ id: randomUUID(), object: 'chat.completion.chunk', created: 1, model: 'fixture',
|
||||
choices: [{ index: 0, delta: value, finish_reason: finish }] })}\n\n`)
|
||||
}
|
||||
response.end('data: [DONE]\n\n')
|
||||
})().catch(error => { response.writeHead(500).end(String(error)) })
|
||||
})
|
||||
await new Promise<void>(resolve => model.listen(0, '127.0.0.1', resolve))
|
||||
const port = (model.address() as { port: number }).port
|
||||
const fetcher = vi.fn<typeof fetch>(async () => Response.json({ data: [{ b64_json: 'iVBORw0KGgo=' }] }))
|
||||
const service = new ImageGenerationService({ database, fetcher, getSettings: async () => ({ defaultImageModelProfileId: imageProfileId, modelProfiles: [{
|
||||
id: imageProfileId, name: 'Image fixture', modelName: 'fixture', protocol: 'openai-images-generations',
|
||||
allowConversationInvocation: true, authentication: 'none', baseUrl: 'https://image.test/v1'
|
||||
}] }) })
|
||||
const gateway = new KnowledgeMcpGateway({} as never)
|
||||
await gateway.start()
|
||||
const profile = { id: randomUUID(), name: 'Chat fixture', modelName: 'fixture', protocol: 'openai-chat-completions' as const,
|
||||
authentication: 'none' as const, baseUrl: `http://127.0.0.1:${port}/v1` }
|
||||
const common = { defaultWorkspace: root, knowledgeGateway: gateway, modelProfile: profile }
|
||||
const runtime = name === 'model'
|
||||
? new ModelAgentRuntime({ ...profile, model: profile.modelName, defaultWorkspace: root })
|
||||
: name === 'opencode'
|
||||
? new OpenCodeRuntime({ ...common, embedded: true, binaryPath: openCodeBinary, configPath: '' })
|
||||
: new ContinueAgentRuntime({ ...common, binaryPath: '', bundledBinaryPath: continueBinary, configPath: '', hostCacheRoot: join(root, 'continue') })
|
||||
try {
|
||||
for (const intent of ['create', 'edit'] as const) {
|
||||
const messageId = randomUUID()
|
||||
const requestId = randomUUID()
|
||||
database.saveLocalConversations([{ header: { id: conversationId, title: 'Image integration', updatedAt: Date.now() },
|
||||
messages: [{ id: messageId, role: 'assistant', content: '', createdAt: Date.now(), state: 'complete' }] }])
|
||||
const binding = service.bind({ conversationId, messageId, requestId, workMode: 'execute' })
|
||||
const events = []
|
||||
for await (const event of runtime.run({ requestId, conversationId, prompt: `${intent} image ${requestId}`, workMode: 'execute', imageToolBinding: binding }, new AbortController().signal, async () => 'once')) events.push(event)
|
||||
expect(events.some(event => event.type === 'done')).toBe(true)
|
||||
const operation = database.getConversation(conversationId).messages.find(message => message.id === messageId)!.imageOperations?.[0]
|
||||
expect(operation?.state).toBe('completed')
|
||||
expect(database.getArtifact(operation!.artifactIds[0]!).content).toContain('iVBORw0KGgo=')
|
||||
toolInput = { intent: 'edit', prompt: 'Turn red', sourceArtifactIds: operation!.artifactIds }
|
||||
}
|
||||
expect(fetcher).toHaveBeenCalledTimes(2)
|
||||
expect(String(fetcher.mock.calls[1]![0])).toMatch(/images\/edits$/u)
|
||||
expect(textRequests.length).toBeGreaterThanOrEqual(4)
|
||||
} finally {
|
||||
await runtime.dispose()
|
||||
await gateway.dispose()
|
||||
await service.dispose()
|
||||
model.closeAllConnections()
|
||||
await new Promise<void>(resolve => model.close(() => resolve()))
|
||||
database.close()
|
||||
await rm(root, { recursive: true, force: true })
|
||||
}
|
||||
})
|
||||
@@ -0,0 +1,16 @@
|
||||
import { z } from 'zod'
|
||||
import { imageToolInputSchema, imageToolName, type ImageOperation, type ImageRequestContext } from '../../shared/image-generation-contracts'
|
||||
|
||||
export type ImageToolBinding = {
|
||||
context: Readonly<ImageRequestContext>
|
||||
describe(): Promise<string | undefined>
|
||||
call(input: unknown, callId: string, waitSignal?: AbortSignal): Promise<ImageOperation>
|
||||
}
|
||||
|
||||
export async function imageToolDefinition(binding: ImageToolBinding | undefined) {
|
||||
const description = await binding?.describe()
|
||||
if (!description) return undefined
|
||||
const inputSchema = z.toJSONSchema(imageToolInputSchema, { target: 'draft-7', io: 'input' })
|
||||
Reflect.deleteProperty(inputSchema, '$schema')
|
||||
return { name: imageToolName, displayName: 'Generate or edit image', description, inputSchema, source: 'builtin' as const }
|
||||
}
|
||||
@@ -102,6 +102,7 @@ export {
|
||||
}
|
||||
|
||||
export type GoodBuddyBuiltinToolName =
|
||||
| 'generate_image'
|
||||
| ScopedDataToolName
|
||||
| BrowserToolName
|
||||
|
||||
@@ -184,6 +185,7 @@ export type KnowledgeLibraryListItem = {
|
||||
}
|
||||
|
||||
type Capability = {
|
||||
imageToolBinding?: ImageToolBinding
|
||||
requestId: string
|
||||
libraryIds: readonly string[]
|
||||
magicNotesAccess: MagicNotesCapabilityAccess
|
||||
@@ -530,6 +532,20 @@ export class KnowledgeMcpGateway {
|
||||
}
|
||||
}
|
||||
|
||||
bindImageTool(binding: ImageToolBinding, signal: AbortSignal, existingToken?: string): string {
|
||||
signal.throwIfAborted()
|
||||
if (existingToken) {
|
||||
const capability = this.getCapability(existingToken)
|
||||
if (capability.requestId !== binding.context.requestId) throw new Error('Image capability request does not match')
|
||||
capability.imageToolBinding = binding
|
||||
return existingToken
|
||||
}
|
||||
return this.storeCapability({
|
||||
requestId: binding.context.requestId, libraryIds: [], magicNotesAccess: 'none', configAccess: 'none',
|
||||
customMcpServers: [], signal, imageToolBinding: binding
|
||||
})
|
||||
}
|
||||
|
||||
grantCustomMcp(
|
||||
requestId: string,
|
||||
servers: readonly ResolvedMcpServer[],
|
||||
@@ -760,6 +776,7 @@ export class KnowledgeMcpGateway {
|
||||
getAvailableToolNames(token: string): GoodBuddyBuiltinToolName[] {
|
||||
const capability = this.getCapability(token)
|
||||
return [
|
||||
...(capability.imageToolBinding ? ['generate_image' as const] : []),
|
||||
...(capability.libraryIds.length > 0
|
||||
? knowledgeToolNames
|
||||
: []),
|
||||
@@ -1570,6 +1587,7 @@ export class KnowledgeMcpGateway {
|
||||
}
|
||||
)
|
||||
const capability = this.getCapability(token)
|
||||
const imageTool = await imageToolDefinition(capability.imageToolBinding)
|
||||
const browserTools = capability.browserConversationId && capability.browserTabId
|
||||
? new BrowserModelTools({
|
||||
service: this.browserService!,
|
||||
@@ -1596,6 +1614,7 @@ export class KnowledgeMcpGateway {
|
||||
session.listedTools = true
|
||||
return {
|
||||
tools: [
|
||||
...(imageTool ? [{ name: imageTool.name, description: imageTool.description, inputSchema: imageTool.inputSchema as Tool['inputSchema'], annotations: { readOnlyHint: false, destructiveHint: false } }] : []),
|
||||
...scopedTools,
|
||||
...browserDefinitions,
|
||||
...[...customBindings.values()].map(
|
||||
@@ -1610,6 +1629,12 @@ export class KnowledgeMcpGateway {
|
||||
async (call, extra) => {
|
||||
const name = call.params.name
|
||||
const input = call.params.arguments ?? {}
|
||||
if (name === 'generate_image') {
|
||||
const capability = this.getCapability(token)
|
||||
if (!capability.imageToolBinding) throw new Error('Image capability is unavailable')
|
||||
const operation = await capability.imageToolBinding.call(input, `${session.registryKey}:${extra.requestId}`, AbortSignal.any([extra.signal, capability.signal]))
|
||||
return { content: [{ type: 'text' as const, text: JSON.stringify(operation) }] }
|
||||
}
|
||||
if (availableTools.has(name as ScopedDataToolName)) {
|
||||
const definition = scopedDataToolByName.get(
|
||||
name as ScopedDataToolName
|
||||
@@ -1868,3 +1893,4 @@ export class KnowledgeMcpGateway {
|
||||
}
|
||||
}
|
||||
}
|
||||
import { imageToolDefinition, type ImageToolBinding } from './image-tool-binding'
|
||||
|
||||
@@ -2681,9 +2681,15 @@ export class ModelAgentRuntime implements AgentRuntime {
|
||||
return history
|
||||
}
|
||||
|
||||
/** Strict image service entry point; direct-image conversations retain their fallback policy. */
|
||||
generateImage(request: AgentExecutionRequest, signal: AbortSignal): AsyncGenerator<RuntimeEvent, void, void> {
|
||||
return this.runImageGeneration(request, signal, false)
|
||||
}
|
||||
|
||||
private async *runImageGeneration(
|
||||
request: AgentExecutionRequest,
|
||||
signal: AbortSignal
|
||||
signal: AbortSignal,
|
||||
allowEditingFallback = true
|
||||
): AsyncGenerator<RuntimeEvent, void, void> {
|
||||
const images = request.images ?? []
|
||||
const prompt = imageConversationPrompt(request)
|
||||
@@ -2744,7 +2750,7 @@ export class ModelAgentRuntime implements AgentRuntime {
|
||||
)
|
||||
const providerMessage = getErrorMessage(errorPayload)
|
||||
if (
|
||||
images.length > 0 &&
|
||||
allowEditingFallback && images.length > 0 &&
|
||||
isImageEditingUnavailable(response.status, providerMessage)
|
||||
) {
|
||||
signal.throwIfAborted()
|
||||
@@ -3417,6 +3423,7 @@ export class ModelAgentRuntime implements AgentRuntime {
|
||||
wakeNestedEvents = undefined
|
||||
}
|
||||
const toolContext: ModelToolCallContext = {
|
||||
imageToolBinding: request.imageToolBinding,
|
||||
conversationId: request.conversationId,
|
||||
browserTabId: request.browserTabId,
|
||||
browserConversationId: request.browserConversationId,
|
||||
@@ -3743,6 +3750,7 @@ export class ModelAgentRuntime implements AgentRuntime {
|
||||
let roundContextBytes = 0
|
||||
const roundCallIds = new Set<string>()
|
||||
for (const call of response.toolCalls) {
|
||||
toolContext.toolCallId = call.id
|
||||
signal.throwIfAborted()
|
||||
if (roundCallIds.has(call.id)) {
|
||||
throw new Error('模型重复使用了工具调用 ID')
|
||||
|
||||
@@ -290,6 +290,8 @@ export type ModelToolResult = {
|
||||
}
|
||||
|
||||
export type ModelToolCallContext = {
|
||||
imageToolBinding?: ImageToolBinding
|
||||
toolCallId?: string
|
||||
conversationId: string
|
||||
browserTabId?: BrowserTabId
|
||||
browserConversationId?: string
|
||||
@@ -1290,6 +1292,7 @@ export class ModelToolProvider implements ModelToolProviderLike {
|
||||
): Promise<ModelToolDefinition[]> {
|
||||
signal.throwIfAborted()
|
||||
const scopedTools = this.getScopedTools(context)
|
||||
const imageTool = context.workMode === 'execute' ? await imageToolDefinition(context.imageToolBinding) : undefined
|
||||
const webTools = this.webSearchEnabled
|
||||
? this.getWebSearchDefinitions()
|
||||
: []
|
||||
@@ -1312,6 +1315,7 @@ export class ModelToolProvider implements ModelToolProviderLike {
|
||||
const browserTools = this.getBrowserTools(context)
|
||||
return [
|
||||
...workspaceTools,
|
||||
...(imageTool ? [imageTool] : []),
|
||||
...processTools,
|
||||
...subagentTools,
|
||||
...outputTools,
|
||||
@@ -1439,6 +1443,10 @@ export class ModelToolProvider implements ModelToolProviderLike {
|
||||
): Promise<ModelToolResult> {
|
||||
signal.throwIfAborted()
|
||||
assertToolAuthorizedForWorkMode(name, context)
|
||||
if (name === 'generate_image') {
|
||||
if (!context.imageToolBinding || !context.toolCallId) throw new Error('Image tool request binding is unavailable')
|
||||
return createTextToolResult(JSON.stringify(await context.imageToolBinding.call(argumentsValue, context.toolCallId, signal)))
|
||||
}
|
||||
if (name === 'knowledge_list') {
|
||||
if (
|
||||
!this.knowledgeGateway ||
|
||||
@@ -2020,3 +2028,4 @@ function emitDirectModelSubagentUsage(
|
||||
)
|
||||
})
|
||||
}
|
||||
import { imageToolDefinition, type ImageToolBinding } from './image-tool-binding'
|
||||
|
||||
@@ -2161,9 +2161,13 @@ export class OpenCodeRuntime implements AgentRuntime {
|
||||
let customMcpToken: string | undefined;
|
||||
const attemptedMcpNames: string[] = [];
|
||||
const subscriptionController = new AbortController();
|
||||
const imageCapabilityToken = request.imageToolBinding && request.workMode === 'execute'
|
||||
? this.options.knowledgeGateway?.bindImageTool(request.imageToolBinding, signal, request.knowledgeCapabilityToken)
|
||||
: undefined;
|
||||
const scopedCapabilityToken = imageCapabilityToken ?? request.knowledgeCapabilityToken;
|
||||
try {
|
||||
if (
|
||||
request.knowledgeCapabilityToken &&
|
||||
scopedCapabilityToken &&
|
||||
this.usesEmbeddedPermissionMediation() &&
|
||||
this.options.knowledgeGateway?.getEndpoint()
|
||||
) {
|
||||
@@ -2186,7 +2190,7 @@ export class OpenCodeRuntime implements AgentRuntime {
|
||||
url: knowledgeEndpoint,
|
||||
enabled: true,
|
||||
headers: {
|
||||
Authorization: `Bearer ${request.knowledgeCapabilityToken}`,
|
||||
Authorization: `Bearer ${scopedCapabilityToken}`,
|
||||
},
|
||||
oauth: false,
|
||||
},
|
||||
@@ -2212,7 +2216,7 @@ export class OpenCodeRuntime implements AgentRuntime {
|
||||
// experimental/tool/ids. Its model tool namespace is deterministic:
|
||||
// "<MCP server name>_<declared tool name>".
|
||||
knowledgeToolIds = this.options.knowledgeGateway
|
||||
.getAvailableToolNames(request.knowledgeCapabilityToken)
|
||||
.getAvailableToolNames(scopedCapabilityToken)
|
||||
.map((toolName) => `${knowledgeMcpName}_${toolName}`);
|
||||
}
|
||||
if (
|
||||
@@ -2996,6 +3000,7 @@ export class OpenCodeRuntime implements AgentRuntime {
|
||||
}
|
||||
} finally {
|
||||
subscriptionController.abort();
|
||||
if (imageCapabilityToken && imageCapabilityToken !== request.knowledgeCapabilityToken) this.options.knowledgeGateway?.revoke(imageCapabilityToken);
|
||||
for (const name of attemptedMcpNames) {
|
||||
await this.mutateMcp(new AbortController().signal, () => {
|
||||
// Queue wait must not consume the disconnect timeout.
|
||||
|
||||
@@ -18,6 +18,7 @@ import type {
|
||||
remoteSemanticTranscriptPageResultSchema,
|
||||
} from '../../shared/remote-agent-contracts'
|
||||
import type { z } from 'zod'
|
||||
import type { ImageToolBinding } from './image-tool-binding'
|
||||
import type { RemoteQuestionResponse } from '../../shared/remote-question-contracts'
|
||||
import type {
|
||||
ModelBridgePolicy
|
||||
@@ -111,7 +112,9 @@ export interface RemoteRuntimeChannel {
|
||||
): Promise<RuntimeSessionBindingCursors>
|
||||
/** Accepts the prompt and starts its Ask or Execute Runtime process. */
|
||||
preparePrompt(
|
||||
preparation: RemotePromptOperationPreparation
|
||||
preparation: RemotePromptOperationPreparation,
|
||||
imageToolBinding?: ImageToolBinding,
|
||||
waitSignal?: AbortSignal
|
||||
): Promise<RemotePromptOperationAcceptance>
|
||||
/** Starts an Agent-owned ACP prompt exactly once for its stable operation. */
|
||||
startOwnedPrompt?(
|
||||
|
||||
@@ -146,6 +146,7 @@ export type AgentExecutionRequest = Omit<AgentRequest, 'workMode'> & {
|
||||
trustedInstructions?: string
|
||||
/** Main-process-only request-scoped authorization for built-in data tools. */
|
||||
knowledgeCapabilityToken?: string
|
||||
imageToolBinding?: ImageToolBinding
|
||||
/** Main-process-only browser tab fixed for this request. */
|
||||
browserTabId?: BrowserTabId
|
||||
/**
|
||||
@@ -160,3 +161,4 @@ export type AgentExecutionRequest = Omit<AgentRequest, 'workMode'> & {
|
||||
/** Main-only durable native subagent state for partial recovery updates. */
|
||||
remoteRecoveredSubagents?: readonly RemoteRecoveredSubagent[]
|
||||
}
|
||||
import type { ImageToolBinding } from './image-tool-binding'
|
||||
|
||||
@@ -15,8 +15,43 @@ import {
|
||||
ASSISTANT_DATABASE_SCHEMA_VERSION
|
||||
} from './assistant-database'
|
||||
import { agentRuntimeSelectionKey } from '../../shared/runtime-selection-contracts'
|
||||
import type { ImageOperation } from '../../shared/image-generation-contracts'
|
||||
|
||||
const temporaryDirectories: string[] = []
|
||||
|
||||
it('atomically completes conversation images and retains operations and upload references across saves and reopening', async () => {
|
||||
const database = await createDatabase()
|
||||
const conversationId = randomUUID()
|
||||
const messageId = randomUUID()
|
||||
const header = { id: conversationId, title: 'Image operation', updatedAt: Date.now() }
|
||||
const message = { id: messageId, role: 'assistant' as const, content: '', state: 'complete' as const, createdAt: Date.now() }
|
||||
database.saveLocalConversations([{ header, messages: [message] }])
|
||||
const operation: ImageOperation = { id: randomUUID(), conversationId, messageId, requestId: randomUUID(), callId: 'image-call',
|
||||
modelProfileId: randomUUID(), modelName: 'Image model', input: { intent: 'create', prompt: 'Blue', sourceArtifactIds: [] },
|
||||
state: 'running', createdAt: Date.now(), updatedAt: Date.now(), artifactIds: [] }
|
||||
database.saveConversationImageOperation(operation)
|
||||
const sources = database.saveConversationImageSources({ conversationId, messageId,
|
||||
images: [{ name: 'source.png', mediaType: 'image/png', data: 'iVBORw0KGgo=' }] })
|
||||
const completed = database.saveConversationImageOperation(operation, { mimeType: 'image/png', data: 'iVBORw0KGgo=' })
|
||||
expect(database.saveConversationImageOperation(operation, { mimeType: 'image/png', data: 'iVBORw0KGgo=' })).toEqual(completed)
|
||||
expect(database.listArtifacts()).toHaveLength(2)
|
||||
database.saveLocalConversations([{ header, messages: [message] }])
|
||||
database.close()
|
||||
database.initialize(process.cwd())
|
||||
const restored = database.getConversation(conversationId).messages[0]!
|
||||
expect(restored.imageOperations).toEqual([completed])
|
||||
expect(restored.artifactIds).toEqual(completed.artifactIds)
|
||||
expect(restored.imageSourceArtifactIds).toEqual(sources)
|
||||
expect(database.getArtifact(sources[0]!).content).toContain('iVBORw0KGgo=')
|
||||
const failed = { ...operation, id: randomUUID() }
|
||||
database.saveConversationImageOperation(failed)
|
||||
expect(() => database.saveConversationImageOperation(failed, { mimeType: 'image/png', data: 'a'.repeat(5 * 1024 * 1024) })).toThrow('5MB')
|
||||
expect(database.listArtifacts()).toHaveLength(2)
|
||||
expect(database.getConversation(conversationId).messages[0]!.imageOperations?.find(item => item.id === failed.id)?.state).toBe('running')
|
||||
database.deleteLocalConversation(conversationId)
|
||||
expect(() => database.saveConversationImageOperation(failed, { mimeType: 'image/png', data: 'iVBORw0KGgo=' })).toThrow('no longer exists')
|
||||
database.close()
|
||||
})
|
||||
const channelDefaultProfileId =
|
||||
'00000000-0000-4000-8000-000000000001'
|
||||
|
||||
|
||||
@@ -279,6 +279,8 @@ type MessageMetadata = {
|
||||
sourceReferences?: ConversationSnapshot['messages'][number]['sourceReferences']
|
||||
knowledgeRetrieval?: ConversationSnapshot['messages'][number]['knowledgeRetrieval']
|
||||
artifactIds?: string[]
|
||||
imageOperations?: ImageOperation[]
|
||||
imageSourceArtifactIds?: string[]
|
||||
imageContextNotice?: ConversationSnapshot['messages'][number]['imageContextNotice']
|
||||
task?: ConversationSnapshot['messages'][number]['task']
|
||||
attachments?: ConversationSnapshot['messages'][number]['attachments']
|
||||
@@ -1177,6 +1179,8 @@ function toConversationSnapshot(
|
||||
sourceReferences: metadata.sourceReferences,
|
||||
knowledgeRetrieval: metadata.knowledgeRetrieval,
|
||||
artifactIds: metadata.artifactIds,
|
||||
imageOperations: metadata.imageOperations,
|
||||
imageSourceArtifactIds: metadata.imageSourceArtifactIds,
|
||||
imageContextNotice: metadata.imageContextNotice,
|
||||
task: metadata.task,
|
||||
attachments: metadata.attachments,
|
||||
@@ -1204,6 +1208,8 @@ function serializeConversationMessageMetadata(
|
||||
sourceReferences: message.sourceReferences,
|
||||
knowledgeRetrieval: message.knowledgeRetrieval,
|
||||
artifactIds: message.artifactIds,
|
||||
imageOperations: message.imageOperations,
|
||||
imageSourceArtifactIds: message.imageSourceArtifactIds,
|
||||
imageContextNotice: message.imageContextNotice,
|
||||
task: message.task,
|
||||
attachments: message.attachments,
|
||||
@@ -2862,7 +2868,7 @@ export class AssistantDatabase {
|
||||
WHERE id = ? AND channel IS NULL`
|
||||
)
|
||||
const findMessage = database.prepare(
|
||||
`SELECT conversation_id, role,
|
||||
`SELECT conversation_id, role, metadata_json,
|
||||
EXISTS(
|
||||
SELECT 1 FROM tasks
|
||||
WHERE tasks.id = messages.request_id
|
||||
@@ -2941,6 +2947,7 @@ export class AssistantDatabase {
|
||||
| {
|
||||
conversation_id: string
|
||||
role: MessageRow['role']
|
||||
metadata_json: string
|
||||
remote_recoverable: number
|
||||
}
|
||||
| undefined
|
||||
@@ -2957,7 +2964,20 @@ export class AssistantDatabase {
|
||||
updateMessage.run(
|
||||
message.content,
|
||||
message.state,
|
||||
serializeConversationMessageMetadata(message),
|
||||
serializeConversationMessageMetadata({
|
||||
...message,
|
||||
...(() => {
|
||||
const stored = JSON.parse(existingMessage.metadata_json) as MessageMetadata
|
||||
return {
|
||||
imageOperations: stored.imageOperations,
|
||||
imageSourceArtifactIds: stored.imageSourceArtifactIds,
|
||||
artifactIds: stored.imageOperations?.some(operation => operation.artifactIds.length) ? [...new Set([
|
||||
...(message.artifactIds ?? []),
|
||||
...(stored.imageOperations?.flatMap(operation => operation.artifactIds) ?? [])
|
||||
])].slice(-8) : message.artifactIds
|
||||
}
|
||||
})()
|
||||
}),
|
||||
message.id
|
||||
)
|
||||
continue
|
||||
@@ -5705,6 +5725,101 @@ export class AssistantDatabase {
|
||||
})
|
||||
}
|
||||
|
||||
saveConversationImageOperation(
|
||||
operation: ImageOperation,
|
||||
image?: { mimeType: 'image/png' | 'image/jpeg' | 'image/webp'; data: string }
|
||||
): ImageOperation {
|
||||
const database = this.requireDatabase()
|
||||
database.exec('BEGIN IMMEDIATE')
|
||||
try {
|
||||
const row = database.prepare(`SELECT m.metadata_json, c.project_id FROM messages m
|
||||
JOIN conversations c ON c.id = m.conversation_id
|
||||
WHERE m.id = ? AND m.conversation_id = ?`).get(operation.messageId, operation.conversationId) as
|
||||
{ metadata_json: string; project_id: string | null } | undefined
|
||||
if (!row) throw new Error('Image operation conversation or message no longer exists')
|
||||
const metadata = JSON.parse(row.metadata_json) as MessageMetadata
|
||||
const previous = metadata.imageOperations?.find(item => item.id === operation.id)
|
||||
if (previous?.state === 'completed') {
|
||||
database.exec('COMMIT')
|
||||
return previous
|
||||
}
|
||||
let saved = operation
|
||||
if (image) {
|
||||
const artifact = this.createImageArtifact({
|
||||
projectId: row.project_id ?? undefined,
|
||||
title: operation.input.prompt.slice(0, 120),
|
||||
mimeType: image.mimeType,
|
||||
base64: image.data
|
||||
})
|
||||
saved = { ...operation, state: 'completed', artifactIds: [artifact.id] }
|
||||
metadata.artifactIds = [...new Set([...(metadata.artifactIds ?? []), artifact.id])].slice(-8)
|
||||
}
|
||||
metadata.imageOperations = [...(metadata.imageOperations ?? []).filter(item => item.id !== saved.id), saved]
|
||||
database.prepare('UPDATE messages SET metadata_json = ? WHERE id = ? AND conversation_id = ?')
|
||||
.run(JSON.stringify(metadata), saved.messageId, saved.conversationId)
|
||||
database.prepare('UPDATE conversations SET updated_at = MAX(updated_at, ?) WHERE id = ?')
|
||||
.run(new Date(saved.updatedAt).toISOString(), saved.conversationId)
|
||||
database.exec('COMMIT')
|
||||
return saved
|
||||
} catch (error) {
|
||||
database.exec('ROLLBACK')
|
||||
throw error
|
||||
}
|
||||
}
|
||||
|
||||
saveConversationImageSources(input: {
|
||||
conversationId: string
|
||||
messageId: string
|
||||
images: readonly { name: string; mediaType: 'image/png' | 'image/jpeg'; data: string }[]
|
||||
}): string[] {
|
||||
if (input.images.length > 8) throw new Error('At most eight source images are supported')
|
||||
const database = this.requireDatabase()
|
||||
database.exec('BEGIN IMMEDIATE')
|
||||
try {
|
||||
const row = database.prepare(`SELECT m.metadata_json, c.project_id FROM messages m
|
||||
JOIN conversations c ON c.id = m.conversation_id
|
||||
WHERE m.id = ? AND m.conversation_id = ?`).get(input.messageId, input.conversationId) as
|
||||
{ metadata_json: string; project_id: string | null } | undefined
|
||||
if (!row) throw new Error('Image source conversation or message no longer exists')
|
||||
const metadata = JSON.parse(row.metadata_json) as MessageMetadata
|
||||
if (!metadata.imageSourceArtifactIds?.length) {
|
||||
metadata.imageSourceArtifactIds = input.images.map(image => this.createImageArtifact({
|
||||
projectId: row.project_id ?? undefined, title: image.name,
|
||||
mimeType: image.mediaType, base64: image.data
|
||||
}).id)
|
||||
database.prepare('UPDATE messages SET metadata_json = ? WHERE id = ?').run(JSON.stringify(metadata), input.messageId)
|
||||
}
|
||||
database.exec('COMMIT')
|
||||
return metadata.imageSourceArtifactIds
|
||||
} catch (error) {
|
||||
database.exec('ROLLBACK')
|
||||
throw error
|
||||
}
|
||||
}
|
||||
|
||||
markUnfinishedImageOperationsUnconfirmed(): void {
|
||||
const database = this.requireDatabase()
|
||||
database.exec('BEGIN IMMEDIATE')
|
||||
try {
|
||||
const rows = database.prepare("SELECT id, metadata_json FROM messages WHERE json_type(metadata_json, '$.imageOperations') = 'array'")
|
||||
.all() as { id: string; metadata_json: string }[]
|
||||
for (const row of rows) {
|
||||
const metadata = JSON.parse(row.metadata_json) as MessageMetadata
|
||||
let changed = false
|
||||
metadata.imageOperations = metadata.imageOperations?.map(operation => {
|
||||
if (!['running', 'saving', 'cancelling'].includes(operation.state)) return operation
|
||||
changed = true
|
||||
return { ...operation, state: 'unconfirmed', updatedAt: Date.now() }
|
||||
})
|
||||
if (changed) database.prepare('UPDATE messages SET metadata_json = ? WHERE id = ?').run(JSON.stringify(metadata), row.id)
|
||||
}
|
||||
database.exec('COMMIT')
|
||||
} catch (error) {
|
||||
database.exec('ROLLBACK')
|
||||
throw error
|
||||
}
|
||||
}
|
||||
|
||||
createImageArtifact(input: {
|
||||
projectId?: string
|
||||
taskId?: string
|
||||
@@ -9764,3 +9879,4 @@ export class AssistantDatabase {
|
||||
return this.database
|
||||
}
|
||||
}
|
||||
import type { ImageOperation } from '../../shared/image-generation-contracts'
|
||||
|
||||
@@ -137,6 +137,30 @@ afterEach(async () => {
|
||||
})
|
||||
|
||||
describe('CapabilityService', () => {
|
||||
it('derives image assignments from saved profiles without persisting or exposing connection data', async () => {
|
||||
const image = { id: '00000000-0000-4000-8000-000000000001', name: 'Image', protocol: 'openai-images-generations', allowConversationInvocation: true, apiKey: 'must-not-leak' }
|
||||
let profiles = [image]
|
||||
const getSavedModelProfiles = vi.fn(async () => profiles)
|
||||
const { service, filePath, builtinRoot, importedRoot } = await createService({ getSavedModelProfiles })
|
||||
const all = ['model', 'opencode', 'continue', 'deepseek-harness']
|
||||
expect((await service.getSnapshot()).imageGeneration).toEqual({ modelProfiles: [{ id: image.id, name: image.name }], assignments: all })
|
||||
profiles = [image, { ...image, id: '00000000-0000-4000-8000-000000000002' }]
|
||||
expect((await service.getSnapshot()).imageGeneration?.modelProfiles).toHaveLength(2)
|
||||
profiles = [{ ...image, allowConversationInvocation: false }, profiles[1]!]
|
||||
expect((await service.getSnapshot()).imageGeneration).toMatchObject({ assignments: all, modelProfiles: [{ id: profiles[1]!.id }] })
|
||||
profiles = [{ ...image, allowConversationInvocation: false }, { ...image, protocol: 'openai-chat-completions' }]
|
||||
expect((await service.getSnapshot()).imageGeneration).toEqual({ modelProfiles: [], assignments: [] })
|
||||
profiles = []
|
||||
expect((await service.getSnapshot()).imageGeneration?.assignments).toEqual([])
|
||||
profiles = [image]
|
||||
const reloaded = new CapabilityService(filePath, builtinRoot, importedRoot, cipher, { getSavedModelProfiles })
|
||||
expect((await reloaded.getSnapshot()).imageGeneration?.assignments).toEqual(all)
|
||||
await service.setBuiltinMcpServerEnabled('knowledge-base', false)
|
||||
expect(await readFile(filePath, 'utf8')).not.toContain('imageGeneration')
|
||||
expect((await service.getSnapshot()).imageGeneration?.assignments).toEqual(all)
|
||||
expect((await new CapabilityService(filePath, builtinRoot, importedRoot, cipher).getSnapshot()).imageGeneration).toEqual({ modelProfiles: [], assignments: [] })
|
||||
})
|
||||
|
||||
it('memoizes concurrent loads and retries after a failed load', async () => {
|
||||
const initialStore = {
|
||||
load: vi.fn(async () => undefined),
|
||||
|
||||
@@ -273,6 +273,12 @@ export type SkillImportInspection = {
|
||||
}
|
||||
|
||||
export type CapabilityServiceOptions = Readonly<{
|
||||
getSavedModelProfiles?: () => Promise<ReadonlyArray<{
|
||||
id: string
|
||||
name: string
|
||||
protocol: string
|
||||
allowConversationInvocation?: boolean
|
||||
}>>
|
||||
platform?: NodeJS.Platform
|
||||
architecture?: string
|
||||
electronTarget?: boolean
|
||||
@@ -702,6 +708,7 @@ async function parseSkillZip(archivePath: string): Promise<ParsedSkillZip> {
|
||||
}
|
||||
|
||||
export class CapabilityService {
|
||||
private readonly getSavedModelProfiles: NonNullable<CapabilityServiceOptions['getSavedModelProfiles']>
|
||||
private state?: StoredCapabilities
|
||||
private loadPromise?: Promise<StoredCapabilities>
|
||||
private warnings: SettingsWarning[] = []
|
||||
@@ -721,6 +728,8 @@ export class CapabilityService {
|
||||
private readonly cipher: CapabilityCipher,
|
||||
options: CapabilityServiceOptions = {}
|
||||
) {
|
||||
this.getSavedModelProfiles =
|
||||
options.getSavedModelProfiles ?? (async () => [])
|
||||
this.platform = options.platform ?? process.platform
|
||||
this.architecture = options.architecture ?? process.arch
|
||||
this.electronTarget =
|
||||
@@ -971,12 +980,25 @@ export class CapabilityService {
|
||||
}
|
||||
|
||||
async getSnapshot(): Promise<CapabilitySnapshot> {
|
||||
const [state, catalog, browserProfileState] = await Promise.all([
|
||||
const [state, catalog, browserProfileState, savedModelProfiles] = await Promise.all([
|
||||
this.load(),
|
||||
this.getSkillCatalog(),
|
||||
this.browserProfiles.getSnapshot()
|
||||
this.browserProfiles.getSnapshot(),
|
||||
this.getSavedModelProfiles()
|
||||
])
|
||||
const imageProfiles = savedModelProfiles
|
||||
.filter((profile) =>
|
||||
profile.protocol === 'openai-images-generations' &&
|
||||
profile.allowConversationInvocation === true
|
||||
)
|
||||
.map(({ id, name }) => ({ id, name }))
|
||||
return {
|
||||
imageGeneration: {
|
||||
modelProfiles: imageProfiles,
|
||||
assignments: imageProfiles.length > 0
|
||||
? ['model', 'opencode', 'continue', 'deepseek-harness']
|
||||
: []
|
||||
},
|
||||
skills: catalog
|
||||
.map((skill) => ({
|
||||
...skill,
|
||||
|
||||
+25
-2
@@ -31,6 +31,7 @@ import {
|
||||
type SelectedRuntimeTarget
|
||||
} from './agent/runtime-selection'
|
||||
import { CapabilityService } from './capabilities/capability-service'
|
||||
import { ImageGenerationService } from './agent/image-generation-service'
|
||||
import { ContextManager } from './context-manager'
|
||||
import {
|
||||
registerIpcHandlers,
|
||||
@@ -211,6 +212,7 @@ const localRuntimeRegistry = new LocalRuntimeRegistry()
|
||||
let knowledgeService: KnowledgeService | undefined
|
||||
let knowledgeGateway: KnowledgeMcpGateway | undefined
|
||||
let assistantDatabase: AssistantDatabase | undefined
|
||||
let imageGenerationService: ImageGenerationService | undefined
|
||||
let browserService: BrowserService | undefined
|
||||
let globalTlsPolicy: GlobalTlsPolicy | undefined
|
||||
let feedbackService: FeedbackService | undefined
|
||||
@@ -750,7 +752,8 @@ if (hasSingleInstanceLock) {
|
||||
? join(process.resourcesPath, 'skills')
|
||||
: join(app.getAppPath(), 'resources', 'skills'),
|
||||
join(app.getPath('userData'), 'skills', 'imported'),
|
||||
secureCipher
|
||||
secureCipher,
|
||||
{ getSavedModelProfiles: () => settingsStore.getSavedModelProfiles() }
|
||||
)
|
||||
const channelSettingsStore = new ChannelSettingsStore(
|
||||
join(app.getPath('userData'), 'channel-settings.json'),
|
||||
@@ -910,6 +913,23 @@ if (hasSingleInstanceLock) {
|
||||
}
|
||||
)
|
||||
assistantDatabase = startupAssistantDatabase
|
||||
imageGenerationService = new ImageGenerationService({
|
||||
database: startupAssistantDatabase,
|
||||
getSettings: () => settingsStore.getResolvedSettings(),
|
||||
onError: (error) => {
|
||||
void desktopDiagnostics.recordFailure({ component: 'desktop', stage: 'image-generation', code: 'image-generation.observer.failed', error }).catch(() => undefined)
|
||||
},
|
||||
onOperation: (operation) => {
|
||||
if (mainWindow && !mainWindow.isDestroyed() && !mainWindow.webContents.isDestroyed()) {
|
||||
mainWindow.webContents.send(ipcChannels.imageOperationChanged, operation)
|
||||
}
|
||||
},
|
||||
onUsage: (event) => startupAssistantDatabase.upsertModelUsageCall({
|
||||
requestId: event.requestId, callId: event.callId, runtime: event.runtime,
|
||||
provider: event.provider, model: event.model, input: event.inputTokens,
|
||||
output: event.outputTokens, cacheRead: event.cacheReadTokens, cacheWrite: event.cacheWriteTokens
|
||||
})
|
||||
})
|
||||
directModelSubagentScheduler = new SubagentScheduler({
|
||||
concurrency: 3,
|
||||
queueLimit: 20,
|
||||
@@ -1130,6 +1150,7 @@ if (hasSingleInstanceLock) {
|
||||
createConfiguredRuntime(initialResolvedSettings),
|
||||
initializeAssistant: () => {
|
||||
startupAssistantDatabase.initialize(defaultWorkspace)
|
||||
imageGenerationService!.initialize()
|
||||
startupAssistantDatabase.ensureChannelProjects(
|
||||
defaultWorkspace,
|
||||
initialRuntimeSettings.defaultModelProfileId
|
||||
@@ -1426,7 +1447,8 @@ if (hasSingleInstanceLock) {
|
||||
agentPackageManager,
|
||||
startupRemoteAgentServices.connectionManager,
|
||||
terminalSessionManager,
|
||||
startupLocalToolEnvironmentService
|
||||
startupLocalToolEnvironmentService,
|
||||
imageGenerationService
|
||||
)
|
||||
removeFeedbackIpcHandler = registerFeedbackIpcHandler(
|
||||
mainWindow,
|
||||
@@ -1522,6 +1544,7 @@ app.on('before-quit', (event) => {
|
||||
() => feedbackService?.dispose(),
|
||||
() => removeFeedbackIpcHandler?.(),
|
||||
() => dshExtensionInstaller?.dispose(),
|
||||
() => imageGenerationService?.dispose(),
|
||||
() => removeIpcHandlers?.()
|
||||
],
|
||||
[() => stopRuntimeReconfiguration?.()],
|
||||
|
||||
+89
-4
@@ -28,6 +28,8 @@ import type {
|
||||
} from '../shared/ssh-host-contracts'
|
||||
import { defaultKnowledgeOntologySettings } from '../shared/knowledge-ontology'
|
||||
import { AssistantDatabase } from './assistant/assistant-database'
|
||||
import { ImageGenerationService } from './agent/image-generation-service'
|
||||
import type { AgentExecutionRequest } from './agent/runtime'
|
||||
import { KnowledgeService } from './knowledge/knowledge-service'
|
||||
import { BrowserNavigationStoppedError } from './browser/browser-service'
|
||||
import { RemotePromptCancelledError } from './agent/acp-remote-runtime'
|
||||
@@ -1310,6 +1312,7 @@ describe('registerIpcHandlers model download source routing', () => {
|
||||
})
|
||||
|
||||
vi.mock('electron', () => ({
|
||||
nativeImage: { createFromBuffer: (buffer: Buffer) => ({ isEmpty: () => false, toPNG: () => buffer }) },
|
||||
app: {
|
||||
getName: vi.fn(() => 'GoodBuddy'),
|
||||
getVersion: vi.fn(() => '0.1.0')
|
||||
@@ -5058,7 +5061,9 @@ describe('registerIpcHandlers agent terminal state', () => {
|
||||
magicNotesEnabled = false,
|
||||
goodbuddyConfigService?: Record<string, unknown>,
|
||||
capabilityServiceOverride?: Record<string, unknown>,
|
||||
browserControl?: Record<string, unknown>
|
||||
browserControl?: Record<string, unknown>,
|
||||
imageService?: ImageGenerationService,
|
||||
imageDatabase?: AssistantDatabase
|
||||
) {
|
||||
const assistantDatabase = {
|
||||
createTask: vi.fn(),
|
||||
@@ -5229,7 +5234,7 @@ describe('registerIpcHandlers agent terminal state', () => {
|
||||
})
|
||||
)
|
||||
const onRuntimeSettingsChanged = vi.fn(async () => {})
|
||||
const dispose = registerIpcHandlers(
|
||||
const args: Parameters<typeof registerIpcHandlers> = [
|
||||
window as never,
|
||||
runtime as never,
|
||||
'CommandOrControl+Shift+Space',
|
||||
@@ -5256,7 +5261,7 @@ describe('registerIpcHandlers agent terminal state', () => {
|
||||
(knowledgeServiceOverride ?? {
|
||||
database: { listKnowledgeBases: vi.fn(() => []) }
|
||||
}) as never,
|
||||
assistantDatabase as never,
|
||||
(imageDatabase ?? assistantDatabase) as never,
|
||||
approvalBroker as never,
|
||||
{} as never,
|
||||
onRuntimeSettingsChanged,
|
||||
@@ -5279,7 +5284,9 @@ describe('registerIpcHandlers agent terminal state', () => {
|
||||
undefined,
|
||||
undefined,
|
||||
goodbuddyConfigService as never
|
||||
)
|
||||
]
|
||||
args[43] = imageService
|
||||
const dispose = registerIpcHandlers(...args)
|
||||
return {
|
||||
approvalBroker,
|
||||
assistantDatabase,
|
||||
@@ -5370,6 +5377,84 @@ describe('registerIpcHandlers agent terminal state', () => {
|
||||
knowledgeLibraryIds: []
|
||||
})
|
||||
|
||||
it('binds the production image service before text dispatch and routes regenerate, Ask rejection, and deletion through IPC', async () => {
|
||||
const database = new AssistantDatabase(':memory:')
|
||||
database.initialize(process.cwd())
|
||||
const conversationId = crypto.randomUUID()
|
||||
const userId = crypto.randomUUID()
|
||||
const messageId = crypto.randomUUID()
|
||||
const header = { id: conversationId, title: 'Image IPC', updatedAt: Date.now(), workMode: 'execute' as const }
|
||||
database.saveLocalConversations([{ header, messages: [] }])
|
||||
const profile = { id: crypto.randomUUID(), name: 'Image model', modelName: 'image-test', protocol: 'openai-images-generations',
|
||||
baseUrl: 'https://image.test/v1', authentication: 'none' as const, allowConversationInvocation: true }
|
||||
const fetcher = vi.fn<typeof fetch>(async () => Response.json({ data: [{ b64_json: 'iVBORw0KGgo=' }] }))
|
||||
const service = new ImageGenerationService({ database, getSettings: async () => ({ modelProfiles: [profile], defaultImageModelProfileId: profile.id }), fetcher })
|
||||
service.initialize()
|
||||
const run = vi.fn(async function* (request: AgentExecutionRequest) {
|
||||
expect(database.getConversation(conversationId).messages.map(message => message.id)).toEqual([userId, messageId])
|
||||
expect(request.imageToolBinding?.context.messageId).toBe(messageId)
|
||||
expect(await request.imageToolBinding?.describe()).toContain(profile.id)
|
||||
expect((await request.imageToolBinding!.call({ intent: 'create', prompt: 'A blue circle' }, 'image-call')).state).toBe('completed')
|
||||
yield { type: 'done' as const, requestId: request.requestId }
|
||||
})
|
||||
const harness = createHarness({ runtimeId: 'model', capability: 'text', supportsToolExecution: true, run },
|
||||
undefined, 'always', undefined, false, undefined, undefined, undefined, false, undefined, undefined, undefined, service, database)
|
||||
const event = trustedEvent(harness.webContents)
|
||||
try {
|
||||
const requestId = crypto.randomUUID()
|
||||
await harness.handler!(event, { requestId, conversationId, prompt: 'Draw a circle', workMode: 'execute',
|
||||
currentUserMessageId: userId, currentAssistantMessageId: messageId })
|
||||
await vi.waitFor(() => expect(database.getTask(requestId).status).toBe('completed'))
|
||||
const original = database.getConversation(conversationId).messages[1]!.imageOperations![0]!
|
||||
expect(database.getArtifact(original.artifactIds[0]!).content).toContain('iVBORw0KGgo=')
|
||||
const regenerate = electronMocks.handlers.get(ipcChannels.imageOperationRegenerate)!
|
||||
const cancel = electronMocks.handlers.get(ipcChannels.imageOperationCancel)!
|
||||
expect(() => cancel(event, { conversationId, operationId: original.id, prompt: 'injected' })).toThrow()
|
||||
const regenerated = await regenerate(event, { conversationId, operationId: original.id }) as typeof original
|
||||
expect(regenerated.id).not.toBe(original.id)
|
||||
expect(regenerated.messageId).toBe(messageId)
|
||||
expect(fetcher).toHaveBeenCalledTimes(2)
|
||||
database.saveLocalConversations([{ header: { ...header, workMode: 'ask' }, messages: [] }])
|
||||
await expect(regenerate(event, { conversationId, operationId: original.id })).rejects.toThrow('Execute')
|
||||
expect(fetcher).toHaveBeenCalledTimes(2)
|
||||
const cancelConversation = vi.spyOn(service, 'cancelConversation')
|
||||
await electronMocks.handlers.get(ipcChannels.conversationsDeleteLocal)!(event, conversationId)
|
||||
expect(cancelConversation).toHaveBeenCalledWith(conversationId)
|
||||
expect(() => database.getConversation(conversationId)).toThrow()
|
||||
} finally {
|
||||
await harness.dispose()
|
||||
await service.dispose()
|
||||
database.close()
|
||||
}
|
||||
})
|
||||
|
||||
it('persists uploads before dispatch while withholding image bytes from a nonvisual text model', async () => {
|
||||
const database = new AssistantDatabase(':memory:')
|
||||
database.initialize(process.cwd())
|
||||
const conversationId = crypto.randomUUID()
|
||||
const userId = crypto.randomUUID()
|
||||
database.saveLocalConversations([{ header: { id: conversationId, title: 'Upload', updatedAt: Date.now() }, messages: [] }])
|
||||
const service = new ImageGenerationService({ database, getSettings: async () => ({ modelProfiles: [] }) })
|
||||
const run = vi.fn(async function* (request: AgentExecutionRequest) {
|
||||
expect(request.images).toBeUndefined()
|
||||
expect(request.trustedInstructions).toContain('cannot view')
|
||||
const sources = database.getConversation(conversationId).messages.find(message => message.id === userId)!.imageSourceArtifactIds!
|
||||
expect(sources).toHaveLength(1)
|
||||
expect(database.getArtifact(sources[0]!).content).toContain('iVBORw0KGgo=')
|
||||
yield { type: 'done' as const, requestId: request.requestId }
|
||||
})
|
||||
const harness = createHarness({ runtimeId: 'model', capability: 'text', supportsToolExecution: true, run },
|
||||
undefined, 'always', undefined, false, undefined, undefined, undefined, false, undefined, undefined, undefined, service, database)
|
||||
harness.contextManager.enrichRequest.mockImplementation(request => ({ ...request, images: [{ name: 'Upload', mediaType: 'image/png', data: 'iVBORw0KGgo=' }] }))
|
||||
try {
|
||||
const requestId = crypto.randomUUID()
|
||||
await harness.handler!(trustedEvent(harness.webContents), { requestId, conversationId, prompt: 'Turn it red', workMode: 'execute',
|
||||
currentUserMessageId: userId, currentAssistantMessageId: crypto.randomUUID() })
|
||||
await vi.waitFor(() => expect(database.getTask(requestId).status).toBe('completed'))
|
||||
expect(run).toHaveBeenCalledOnce()
|
||||
} finally { await harness.dispose(); await service.dispose(); database.close() }
|
||||
})
|
||||
|
||||
it('restores image context at the IPC boundary and persists the quiet notice on its artifact event', async () => {
|
||||
const png = 'iVBORw0KGgo='
|
||||
const artifactId = '00000000-0000-4000-8000-000000000301'
|
||||
|
||||
+84
-2
@@ -382,6 +382,8 @@ import {
|
||||
} from './magic-notes/magic-note-analyzer'
|
||||
import { AgentEventBuffer } from './agent-event-buffer'
|
||||
import { withImageConversationContext } from './agent/image-conversation-context'
|
||||
import type { ImageGenerationService } from './agent/image-generation-service'
|
||||
import { imageOperationTargetSchema } from '../shared/image-operation-ipc'
|
||||
import {
|
||||
ExecutionSpaceResolver,
|
||||
REMOTE_EXECUTION_SPACE_UNAVAILABLE
|
||||
@@ -1334,7 +1336,8 @@ export function registerIpcHandlers(
|
||||
'getHostConnectionState' | 'onHostConnectionStateChange'
|
||||
>,
|
||||
terminalSessionManager?: TerminalSessionManager,
|
||||
localToolEnvironmentService?: LocalToolEnvironmentService
|
||||
localToolEnvironmentService?: LocalToolEnvironmentService,
|
||||
imageGenerationService?: ImageGenerationService
|
||||
): () => Promise<void> {
|
||||
type ActiveRequestLease = {
|
||||
controller: AbortController
|
||||
@@ -4018,6 +4021,34 @@ export function registerIpcHandlers(
|
||||
const attachedRequest = contextManager.enrichRequest(
|
||||
parsedRequest
|
||||
)
|
||||
let uploadedImageSourceIds: string[] = []
|
||||
if (imageGenerationService && parsedRequest.currentUserMessageId && parsedRequest.currentAssistantMessageId) {
|
||||
const conversation = assistantDatabase.getConversation(parsedRequest.conversationId)
|
||||
const now = Date.now()
|
||||
assistantDatabase.saveLocalConversations([{
|
||||
header: {
|
||||
id: conversation.id, projectId: conversation.projectId, title: conversation.title,
|
||||
updatedAt: now, workMode: normalizedWorkMode, runtimeSelection: parsedRequest.runtimeSelection,
|
||||
knowledgeLibraryIds: conversation.knowledgeLibraryIds,
|
||||
knowledgeRetrievalMode: conversation.knowledgeRetrievalMode,
|
||||
contextMetrics: conversation.contextMetrics, contextCompressionState: conversation.contextCompressionState,
|
||||
branch: conversation.branch
|
||||
},
|
||||
messages: [
|
||||
...(conversation.messages.some(message => message.id === parsedRequest.currentUserMessageId) ? [] : [{
|
||||
id: parsedRequest.currentUserMessageId, role: 'user' as const,
|
||||
content: parsedRequest.prompt, createdAt: now, state: 'complete' as const
|
||||
}]),
|
||||
...(conversation.messages.some(message => message.id === parsedRequest.currentAssistantMessageId) ? [] : [{
|
||||
id: parsedRequest.currentAssistantMessageId, role: 'assistant' as const,
|
||||
content: '', createdAt: now, state: 'streaming' as const
|
||||
}])
|
||||
]
|
||||
}])
|
||||
if (attachedRequest.images?.length) {
|
||||
uploadedImageSourceIds = imageGenerationService.persistUploads({ conversationId: parsedRequest.conversationId, messageId: parsedRequest.currentUserMessageId }, attachedRequest.images)
|
||||
}
|
||||
}
|
||||
const enrichedRequest = imageGeneration
|
||||
? withImageConversationContext(
|
||||
attachedRequest,
|
||||
@@ -4045,7 +4076,7 @@ export function registerIpcHandlers(
|
||||
!agentRuntimeSelected
|
||||
? capabilityService.getWebSearchCapabilityStatus?.()
|
||||
: undefined,
|
||||
configAccess !== 'none' && !enrichedRequest.projectId
|
||||
(configAccess !== 'none' && !enrichedRequest.projectId) || (imageGenerationService && !imageGeneration && enrichedRequest.images?.length)
|
||||
? settingsStore.getResolvedSettings()
|
||||
: undefined,
|
||||
selectedRuntimeTarget
|
||||
@@ -4077,6 +4108,15 @@ export function registerIpcHandlers(
|
||||
: undefined
|
||||
: resolvedRuntimeSettings?.workspacePath
|
||||
const controller = new AbortController()
|
||||
const imageToolBinding = !imageGeneration && imageGenerationService && enrichedRequest.currentAssistantMessageId
|
||||
? imageGenerationService.bind({
|
||||
conversationId: enrichedRequest.conversationId,
|
||||
messageId: enrichedRequest.currentAssistantMessageId,
|
||||
requestId: enrichedRequest.requestId,
|
||||
workMode: normalizedWorkMode
|
||||
}, () => normalizeInteractiveWorkMode(assistantDatabase.getConversation(enrichedRequest.conversationId).workMode))
|
||||
: undefined
|
||||
const imageToolAvailable = Boolean(await imageToolBinding?.describe())
|
||||
const scopedCapability = await grantScopedDataCapability({
|
||||
gateway: knowledgeGateway,
|
||||
browserControl: browserControl as BrowserCapabilityControl | undefined,
|
||||
@@ -4101,6 +4141,7 @@ export function registerIpcHandlers(
|
||||
})
|
||||
const knowledgeCapabilityToken = scopedCapability.token
|
||||
const availableTools = [
|
||||
...(imageToolAvailable ? ['generate_image'] : []),
|
||||
...(webSearchEnabled ? ['web_search', 'web_fetch'] : []),
|
||||
...(!agentRuntimeSelected && !imageGeneration
|
||||
? ['workspace_rg', 'workspace_read_text', 'output_read', 'subagent_delegate']
|
||||
@@ -4142,6 +4183,28 @@ export function registerIpcHandlers(
|
||||
if (scopedCapability.browserTabId) {
|
||||
request.browserTabId = scopedCapability.browserTabId
|
||||
}
|
||||
if (imageToolBinding) {
|
||||
request.imageToolBinding = imageToolBinding
|
||||
const imageReferenceIds = [...uploadedImageSourceIds, ...(request.imageContextArtifactIds ?? [])]
|
||||
if (imageReferenceIds.length) {
|
||||
request.trustedInstructions = [request.trustedInstructions,
|
||||
`Image references for this request: ${JSON.stringify(imageReferenceIds)}. These are references, not visual observations; use explicit sourceArtifactIds when editing.`].filter(Boolean).join('\n')
|
||||
}
|
||||
if (request.images?.length && resolvedRuntimeSettings) {
|
||||
const effective = request.runtimeSelection
|
||||
? applyRuntimeSelection(resolvedRuntimeSettings, request.runtimeSelection).settings
|
||||
: resolvedRuntimeSettings
|
||||
const supportsImages = selectedRuntimeTarget === 'model' ? effective.supportsImageInput
|
||||
: selectedRuntimeTarget === 'opencode' ? effective.opencodeModelProfile?.supportsImageInput
|
||||
: selectedRuntimeTarget === 'continue' ? effective.continueModelProfile?.supportsImageInput
|
||||
: effective.deepseekHarnessModelProfile?.supportsImageInput
|
||||
if (supportsImages !== true) {
|
||||
request.images = undefined
|
||||
request.trustedInstructions = [request.trustedInstructions,
|
||||
'The uploaded images were saved as conversation image references for the image tool. Your chat model cannot view them. Do not claim visual observations; use the user description to edit them.'].filter(Boolean).join('\n')
|
||||
}
|
||||
}
|
||||
}
|
||||
const managedSshExecution =
|
||||
agentRuntimeSelected &&
|
||||
configExecutionSpace?.kind === 'ssh'
|
||||
@@ -6762,6 +6825,7 @@ export function registerIpcHandlers(
|
||||
(event, input: unknown) => {
|
||||
assertTrustedSender(event, window)
|
||||
const conversationId = assistantIdSchema.parse(input)
|
||||
imageGenerationService?.cancelConversation(conversationId)
|
||||
const queuedItems =
|
||||
assistantDatabase.listConversationQueueItems(conversationId)
|
||||
for (const item of queuedItems) {
|
||||
@@ -6789,6 +6853,24 @@ export function registerIpcHandlers(
|
||||
}
|
||||
)
|
||||
|
||||
registerHandler(ipcChannels.imageOperationCancel, (event, input: unknown) => {
|
||||
assertTrustedSender(event, window)
|
||||
const target = imageOperationTargetSchema.parse(input)
|
||||
if (!imageGenerationService) throw new Error('Image service is unavailable')
|
||||
return imageGenerationService.cancel(target.conversationId, target.operationId)
|
||||
})
|
||||
registerHandler(ipcChannels.imageOperationRegenerate, (event, input: unknown) => {
|
||||
assertTrustedSender(event, window)
|
||||
const target = imageOperationTargetSchema.parse(input)
|
||||
if (!imageGenerationService) throw new Error('Image service is unavailable')
|
||||
const previous = imageGenerationService.getOperation(target.conversationId, target.operationId)
|
||||
const currentWorkMode = (): 'ask' | 'execute' => normalizeInteractiveWorkMode(assistantDatabase.getConversation(target.conversationId).workMode)
|
||||
return imageGenerationService.regenerate({
|
||||
conversationId: target.conversationId, messageId: previous.messageId,
|
||||
requestId: randomUUID(), workMode: currentWorkMode()
|
||||
}, target.operationId, currentWorkMode)
|
||||
})
|
||||
|
||||
registerHandler(
|
||||
ipcChannels.conversationQueueList,
|
||||
(event, input: unknown) => {
|
||||
|
||||
@@ -0,0 +1,56 @@
|
||||
import type { ImageToolBinding } from '../agent/image-tool-binding'
|
||||
import type { RuntimeProtocolBinaryChannel } from './protocol-remote-runtime-channel'
|
||||
import {
|
||||
decodeRemoteImageToolMessage, encodeRemoteImageToolMessage,
|
||||
remoteImageToolCallSchema, remoteImageToolReplySchema
|
||||
} from '../../shared/remote-image-tool-contracts'
|
||||
|
||||
export class MainImageToolSession {
|
||||
private readonly wait = new AbortController()
|
||||
private readonly unsubscribe: () => void
|
||||
private readonly removeAbort: () => void
|
||||
|
||||
constructor(
|
||||
private readonly channel: RuntimeProtocolBinaryChannel,
|
||||
private readonly binding: ImageToolBinding,
|
||||
signal?: AbortSignal
|
||||
) {
|
||||
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.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 (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)
|
||||
}
|
||||
}
|
||||
@@ -1,4 +1,5 @@
|
||||
import { describe, expect, it, vi } from 'vitest'
|
||||
import type { ImageToolBinding } from '../agent/image-tool-binding'
|
||||
import {
|
||||
AGENT_PROTOCOL_LIMITS,
|
||||
ChannelProtocolError,
|
||||
@@ -489,6 +490,46 @@ async function openChannel(fixture = connectionFixture()) {
|
||||
}
|
||||
|
||||
describe('ProtocolRemoteRuntimeChannel', () => {
|
||||
it('refreshes image binding descriptions per prompt, hides Ask, and closes failed preparations', async () => {
|
||||
const fixture = await openChannel()
|
||||
const describe = vi.fn(async () => 'First image catalog')
|
||||
const binding: ImageToolBinding = {
|
||||
context: { conversationId: 'conversation', messageId: 'message', requestId: 'operation-1', workMode: 'execute' },
|
||||
describe, call: vi.fn()
|
||||
}
|
||||
const prepare = (requestId: string, workMode: 'ask' | 'execute' = 'execute') => ({
|
||||
bindingId: 'binding-1', operationId: requestId, requestId, workMode,
|
||||
controllerId: 'controller-1', controllerGeneration: 1, connectionGeneration: fixture.channel.generation,
|
||||
channelEpoch: fixture.channel.channelEpoch, hostId: 'host-1', hostRevision: 1, hostKeyGeneration: 1,
|
||||
workspaceIdentity: 'workspace-1', agentInstallationId: 'agent-1', runtimeId: 'opencode',
|
||||
runtimeBundleDigest, runtimeAdapterDigest: digest('d'), deadlineAt: new Date(Date.now() + 10000).toISOString(),
|
||||
budget: { maximumInputBytes: 1024, maximumOutputBytes: 1024 }, promptSequence: 0
|
||||
})
|
||||
fixture.client.responder = (method, params) => {
|
||||
if (method !== 'runtime/preparePrompt') return defaultResponse(method, params)
|
||||
const input = params as ReturnType<typeof prepare> & { imageTool?: { description: string } }
|
||||
return { bindingId: input.bindingId, operationId: input.operationId, requestId: input.requestId,
|
||||
workMode: input.workMode, deadlineAt: input.deadlineAt, acceptedAt: new Date().toISOString(),
|
||||
...(input.imageTool ? { imageToolUrl: 'http://127.0.0.1:1234/image' } : {}) }
|
||||
}
|
||||
try {
|
||||
const wait = new AbortController()
|
||||
const first = await fixture.channel.preparePrompt(prepare('operation-1'), binding, wait.signal)
|
||||
fixture.channel.setRecoveryBoundary(first.deadlineAt, wait.signal)
|
||||
expect(fixture.client.blobBinaries[0]!.closed).toBe(false)
|
||||
describe.mockResolvedValue('Refreshed image catalog')
|
||||
await fixture.channel.preparePrompt(prepare('operation-2'), binding, wait.signal)
|
||||
expect(fixture.client.blobBinaries[0]!.closed).toBe(true)
|
||||
expect(fixture.client.requests.at(-1)!.params).toMatchObject({ imageTool: { description: 'Refreshed image catalog' } })
|
||||
await fixture.channel.preparePrompt(prepare('operation-3', 'ask'), binding, wait.signal)
|
||||
expect(describe).toHaveBeenCalledTimes(2)
|
||||
expect(fixture.client.blobBinaries[1]!.closed).toBe(true)
|
||||
expect(fixture.client.requests.at(-1)!.params).not.toHaveProperty('imageTool', expect.anything())
|
||||
fixture.client.responder = () => { throw new Error('preparation response lost') }
|
||||
await expect(fixture.channel.preparePrompt(prepare('operation-4'), binding, wait.signal)).rejects.toThrow()
|
||||
expect(fixture.client.blobBinaries[2]!.closed).toBe(true)
|
||||
} finally { await fixture.channel.close() }
|
||||
})
|
||||
it('sends question replies on the authenticated control channel and rejects foreign bindings', async () => {
|
||||
const fixture = await openChannel()
|
||||
const response = { bindingId: 'binding-1', operationId: 'operation-1', questionId: 'question-1', answers: [] }
|
||||
|
||||
@@ -1,4 +1,6 @@
|
||||
import { z } from 'zod'
|
||||
import type { ImageToolBinding } from '../agent/image-tool-binding'
|
||||
import { MainImageToolSession } from './main-image-tool-session'
|
||||
import { remoteQuestionResponseSchema, type RemoteQuestionResponse } from '../../shared/remote-question-contracts'
|
||||
import {
|
||||
AGENT_PROTOCOL_LIMITS,
|
||||
@@ -584,38 +586,51 @@ export class ProtocolRemoteRuntimeChannel
|
||||
}
|
||||
|
||||
async preparePrompt(
|
||||
preparation: z.input<typeof remotePromptOperationPreparationSchema>
|
||||
preparation: z.input<typeof remotePromptOperationPreparationSchema>,
|
||||
imageToolBinding?: ImageToolBinding,
|
||||
waitSignal?: AbortSignal
|
||||
): Promise<
|
||||
z.infer<typeof remotePromptOperationAcceptanceSchema>
|
||||
> {
|
||||
const parsed = remotePromptOperationPreparationSchema.parse(
|
||||
preparation
|
||||
)
|
||||
this.#assertPreparationIdentity(parsed)
|
||||
const raw = await this.#request(
|
||||
'runtime/preparePrompt',
|
||||
parsed
|
||||
)
|
||||
const acceptance = parseRemoteResult(
|
||||
remotePromptOperationAcceptanceSchema,
|
||||
raw,
|
||||
'Remote Runtime prompt acceptance is invalid'
|
||||
)
|
||||
this.#imageTool?.close()
|
||||
this.#imageTool = undefined
|
||||
const description = preparation.workMode === 'execute' ? await imageToolBinding?.describe() : undefined
|
||||
waitSignal?.throwIfAborted()
|
||||
let imageTool: z.infer<typeof remotePromptOperationPreparationSchema>['imageTool']
|
||||
if (description && imageToolBinding) {
|
||||
this.#assertCurrent()
|
||||
const binary = this.#state.client.allocateBinaryChannel?.({ kind: 'blob' })
|
||||
if (!binary) throw new Error('Remote image tool transport is unavailable')
|
||||
imageTool = { channelId: binary.channelId, channelEpoch: binary.channelEpoch, description }
|
||||
this.#imageTool = new MainImageToolSession(binary, imageToolBinding, waitSignal)
|
||||
}
|
||||
try {
|
||||
assertRemotePromptAcceptanceMatchesPreparation(
|
||||
parsed,
|
||||
acceptance
|
||||
const parsed = remotePromptOperationPreparationSchema.parse({ ...preparation, imageTool })
|
||||
this.#assertPreparationIdentity(parsed)
|
||||
const raw = await this.#request('runtime/preparePrompt', parsed)
|
||||
const acceptance = parseRemoteResult(
|
||||
remotePromptOperationAcceptanceSchema, raw,
|
||||
'Remote Runtime prompt acceptance is invalid'
|
||||
)
|
||||
} catch {
|
||||
throw new ProtocolRemoteRuntimeChannelError(
|
||||
'Remote Runtime prompt acceptance identity is invalid',
|
||||
'binding-mismatch'
|
||||
)
|
||||
try {
|
||||
assertRemotePromptAcceptanceMatchesPreparation(parsed, acceptance)
|
||||
if (Boolean(imageTool) !== Boolean(acceptance.imageToolUrl)) throw new Error('Image tool registration mismatch')
|
||||
} catch {
|
||||
throw new ProtocolRemoteRuntimeChannelError(
|
||||
'Remote Runtime prompt acceptance identity is invalid', 'binding-mismatch'
|
||||
)
|
||||
}
|
||||
this.#assertCurrent()
|
||||
return acceptance
|
||||
} catch (error) {
|
||||
this.#imageTool?.close()
|
||||
this.#imageTool = undefined
|
||||
throw error
|
||||
}
|
||||
this.#assertCurrent()
|
||||
return acceptance
|
||||
}
|
||||
|
||||
#imageTool?: MainImageToolSession
|
||||
|
||||
async startOwnedPrompt(
|
||||
request: z.input<typeof remoteOwnedPromptStartRequestSchema>
|
||||
): Promise<z.infer<typeof remoteOwnedPromptStartResultSchema>> {
|
||||
@@ -754,7 +769,8 @@ export class ProtocolRemoteRuntimeChannel
|
||||
'protocol'
|
||||
)
|
||||
}
|
||||
this.clearRecoveryBoundary()
|
||||
this.#removeRecoveryAbort?.()
|
||||
this.#removeRecoveryAbort = undefined
|
||||
this.#recoveryDeadlineAt = deadline
|
||||
this.#recoverySignal = signal
|
||||
const abort = (): void => {
|
||||
@@ -774,6 +790,8 @@ export class ProtocolRemoteRuntimeChannel
|
||||
}
|
||||
|
||||
clearRecoveryBoundary(): void {
|
||||
this.#imageTool?.close()
|
||||
this.#imageTool = undefined
|
||||
this.#removeRecoveryAbort?.()
|
||||
this.#removeRecoveryAbort = undefined
|
||||
this.#recoverySignal = undefined
|
||||
@@ -1628,6 +1646,7 @@ export class ProtocolRemoteRuntimeChannel
|
||||
this.#state.capabilityGeneration
|
||||
)
|
||||
this.#closing = true
|
||||
this.#imageTool?.close()
|
||||
this.clearRecoveryBoundary()
|
||||
this.#lifetime.abort(
|
||||
new DOMException('Remote Runtime channel closed', 'AbortError')
|
||||
|
||||
@@ -80,6 +80,57 @@ afterEach(async () => {
|
||||
})
|
||||
|
||||
describe('RuntimeSettingsStore', () => {
|
||||
it('round-trips conversation image settings and retains unavailable defaults without resolving credentials', async () => {
|
||||
const { store, filePath } = await createStore()
|
||||
const initial = await store.getPublicSettings()
|
||||
expect(initial.modelProfiles[0]?.allowConversationInvocation).toBe(false)
|
||||
expect(initial.defaultImageModelProfileId).toBeUndefined()
|
||||
const imageId = '00000000-0000-4000-8000-000000000099'
|
||||
const text = {
|
||||
id: initial.modelProfiles[0]!.id, name: 'Text',
|
||||
baseUrl: 'https://example.com', modelName: 'text',
|
||||
protocol: 'openai-chat-completions' as const,
|
||||
authentication: 'api-key' as const,
|
||||
imageGenerationQuality: 'auto' as const,
|
||||
apiKey: { action: 'keep' as const }
|
||||
}
|
||||
const image = {
|
||||
...text, id: imageId, name: 'Image model',
|
||||
protocol: 'openai-images-generations' as const,
|
||||
allowConversationInvocation: true,
|
||||
apiKey: { action: 'replace' as const, value: 'image-secret' }
|
||||
}
|
||||
const input = settings({ modelProfiles: [text, image], defaultModelProfileId: text.id, defaultImageModelProfileId: imageId })
|
||||
await store.update(runtimeSettingsInputSchema.parse(input))
|
||||
const reloaded = new RuntimeSettingsStore(filePath, cipher, {})
|
||||
expect((await reloaded.getPublicSettings()).modelProfiles[1]).toMatchObject({ allowConversationInvocation: true, apiKeyConfigured: true })
|
||||
expect((await reloaded.getResolvedSettings()).defaultImageModelProfileId).toBe(imageId)
|
||||
expect((await reloaded.getResolvedSettings()).modelProfiles[1]).toMatchObject({ allowConversationInvocation: true, apiKey: 'image-secret' })
|
||||
expect((await reloaded.getPublicSettings()).configured?.modelProfiles[1]?.allowConversationInvocation).toBe(true)
|
||||
const decrypt = vi.fn(() => { throw new Error('Credential unavailable') })
|
||||
const credentialFree = new RuntimeSettingsStore(filePath, { ...cipher, decrypt }, {})
|
||||
expect(await credentialFree.getSavedModelProfiles()).toEqual([
|
||||
{ id: text.id, name: text.name, protocol: text.protocol, allowConversationInvocation: false },
|
||||
{ id: imageId, name: image.name, protocol: image.protocol, allowConversationInvocation: true }
|
||||
])
|
||||
expect(decrypt).not.toHaveBeenCalled()
|
||||
await reloaded.update(settings())
|
||||
expect((await reloaded.getSavedModelProfiles())[1]?.allowConversationInvocation).toBe(true)
|
||||
expect((await reloaded.getPublicSettings()).defaultImageModelProfileId).toBe(imageId)
|
||||
const disabled = await reloaded.update({ ...input, modelProfiles: [text, { ...image, allowConversationInvocation: false, apiKey: { action: 'keep' } }] })
|
||||
expect(disabled.defaultImageModelProfileId).toBe(imageId)
|
||||
expect(disabled.modelProfiles[1]).toMatchObject({ allowConversationInvocation: false, apiKeyConfigured: true })
|
||||
const removed = await reloaded.update({ ...input, modelProfiles: [text] })
|
||||
expect(removed.defaultImageModelProfileId).toBe(imageId)
|
||||
expect((await new RuntimeSettingsStore(filePath, cipher, {}).getPublicSettings()).defaultImageModelProfileId).toBe(imageId)
|
||||
expect((await reloaded.update({ ...input, modelProfiles: [text], defaultImageModelProfileId: null })).defaultImageModelProfileId).toBeNull()
|
||||
})
|
||||
|
||||
it('rejects malformed conversation image settings', () => {
|
||||
expect(runtimeSettingsInputSchema.safeParse(settings({ defaultImageModelProfileId: 'not-a-uuid' })).success).toBe(false)
|
||||
expect(runtimeSettingsInputSchema.safeParse({ ...settings(), modelProfiles: [{ id: '00000000-0000-4000-8000-000000000001', name: 'Image', baseUrl: 'https://example.com', modelName: 'image', protocol: 'openai-images-generations', authentication: 'none', imageGenerationQuality: 'auto', apiKey: { action: 'keep' }, allowConversationInvocation: 'true' }] }).success).toBe(false)
|
||||
})
|
||||
|
||||
it('defaults to automatic Execute tool authorization and preserves an explicitly saved deny policy', async () => {
|
||||
const { store, filePath } = await createStore()
|
||||
await expect(store.getPublicSettings()).resolves.toMatchObject({ toolApproval: 'always' })
|
||||
|
||||
@@ -292,12 +292,14 @@ const storedSettingsSchema = version20StoredSettingsSchema
|
||||
modelProfiles: z
|
||||
.array(
|
||||
currentStoredModelProfileSchema.extend({
|
||||
allowConversationInvocation: z.boolean().optional(),
|
||||
requestHeaders: modelRequestHeadersSchema,
|
||||
requestBody: modelRequestBodySchema
|
||||
})
|
||||
)
|
||||
.min(1)
|
||||
.max(20)
|
||||
.max(20),
|
||||
defaultImageModelProfileId: z.string().uuid().nullable().optional()
|
||||
})
|
||||
|
||||
type StoredSettings = z.infer<typeof storedSettingsSchema>
|
||||
@@ -424,6 +426,7 @@ export type ResolvedRuntimeSettings = {
|
||||
apiKey?: string
|
||||
modelProfiles: ResolvedModelProfile[]
|
||||
defaultModelProfileId: string
|
||||
defaultImageModelProfileId?: string | null
|
||||
opencodeModelProfile?: ResolvedModelProfile
|
||||
continueModelProfile?: ResolvedModelProfile
|
||||
deepseekHarnessModelProfile?: ResolvedModelProfile
|
||||
@@ -458,6 +461,7 @@ export type RuntimePolicySettings = Pick<
|
||||
>
|
||||
|
||||
export type ResolvedModelProfile = {
|
||||
allowConversationInvocation?: boolean
|
||||
id: string
|
||||
name: string
|
||||
baseUrl: string
|
||||
@@ -1599,6 +1603,7 @@ export class RuntimeSettingsStore {
|
||||
protocol: effective.protocol,
|
||||
authentication: effective.authentication,
|
||||
supportsImageInput: effective.supportsImageInput,
|
||||
allowConversationInvocation: profile.allowConversationInvocation ?? false,
|
||||
contextWindowTokens: effective.contextWindowTokens,
|
||||
maximumOutputTokens: effective.maximumOutputTokens,
|
||||
imageGenerationQuality: effective.imageGenerationQuality,
|
||||
@@ -1619,6 +1624,7 @@ export class RuntimeSettingsStore {
|
||||
protocol: profile.protocol,
|
||||
authentication: profile.authentication,
|
||||
supportsImageInput: profile.supportsImageInput,
|
||||
allowConversationInvocation: profile.allowConversationInvocation ?? false,
|
||||
contextWindowTokens: profile.contextWindowTokens,
|
||||
maximumOutputTokens: profile.maximumOutputTokens,
|
||||
imageGenerationQuality: profile.imageGenerationQuality,
|
||||
@@ -1704,6 +1710,7 @@ export class RuntimeSettingsStore {
|
||||
protocol: resolved.protocol,
|
||||
authentication: resolved.authentication,
|
||||
supportsImageInput: resolved.supportsImageInput,
|
||||
allowConversationInvocation: profile.allowConversationInvocation ?? false,
|
||||
contextWindowTokens: resolved.contextWindowTokens,
|
||||
maximumOutputTokens: resolved.maximumOutputTokens,
|
||||
imageGenerationQuality:
|
||||
@@ -1728,6 +1735,7 @@ export class RuntimeSettingsStore {
|
||||
protocol: profile.protocol,
|
||||
authentication: profile.authentication,
|
||||
supportsImageInput: profile.supportsImageInput,
|
||||
allowConversationInvocation: profile.allowConversationInvocation ?? false,
|
||||
contextWindowTokens: profile.contextWindowTokens,
|
||||
maximumOutputTokens: profile.maximumOutputTokens,
|
||||
imageGenerationQuality:
|
||||
@@ -1860,6 +1868,7 @@ export class RuntimeSettingsStore {
|
||||
credentialSource: effective.credentialSource,
|
||||
modelProfiles,
|
||||
defaultModelProfileId: settings.defaultModelProfileId,
|
||||
defaultImageModelProfileId: settings.defaultImageModelProfileId,
|
||||
opencodeModelSource: agent.opencodeBaseUrl
|
||||
? { kind: 'platform' }
|
||||
: settings.opencodeModelSource,
|
||||
@@ -1891,6 +1900,18 @@ export class RuntimeSettingsStore {
|
||||
return this.toPublicSettings(await this.load())
|
||||
}
|
||||
|
||||
async getSavedModelProfiles(): Promise<Array<Pick<
|
||||
RuntimeSettings['modelProfiles'][number],
|
||||
'id' | 'name' | 'protocol' | 'allowConversationInvocation'
|
||||
>>> {
|
||||
return (await this.load()).modelProfiles.map((profile) => ({
|
||||
id: profile.id,
|
||||
name: profile.name,
|
||||
protocol: profile.protocol,
|
||||
allowConversationInvocation: profile.allowConversationInvocation ?? false
|
||||
}))
|
||||
}
|
||||
|
||||
captureRollback(): Promise<RuntimeSettingsRollback> {
|
||||
let result: RuntimeSettingsRollback | undefined
|
||||
const operation = this.updateQueue.then(async () => {
|
||||
@@ -1990,6 +2011,7 @@ export class RuntimeSettingsStore {
|
||||
modelProfiles,
|
||||
defaultModelProfileId: settings.defaultModelProfileId,
|
||||
opencodeModelProfile,
|
||||
defaultImageModelProfileId: settings.defaultImageModelProfileId,
|
||||
continueModelProfile,
|
||||
deepseekHarnessModelProfile,
|
||||
...agent,
|
||||
@@ -2128,6 +2150,7 @@ export class RuntimeSettingsStore {
|
||||
protocol: input.modelProtocol,
|
||||
authentication: input.modelAuthentication,
|
||||
supportsImageInput: profile.supportsImageInput,
|
||||
allowConversationInvocation: profile.allowConversationInvocation ?? false,
|
||||
contextWindowTokens: profile.contextWindowTokens,
|
||||
maximumOutputTokens: profile.maximumOutputTokens,
|
||||
imageGenerationQuality: input.imageGenerationQuality,
|
||||
@@ -2143,6 +2166,7 @@ export class RuntimeSettingsStore {
|
||||
protocol: profile.protocol,
|
||||
authentication: profile.authentication,
|
||||
supportsImageInput: profile.supportsImageInput,
|
||||
allowConversationInvocation: profile.allowConversationInvocation ?? false,
|
||||
contextWindowTokens: profile.contextWindowTokens,
|
||||
maximumOutputTokens: profile.maximumOutputTokens,
|
||||
imageGenerationQuality: profile.imageGenerationQuality,
|
||||
@@ -2197,6 +2221,9 @@ export class RuntimeSettingsStore {
|
||||
protocol: profile.protocol,
|
||||
authentication: profile.authentication,
|
||||
supportsImageInput: profile.supportsImageInput ?? false,
|
||||
allowConversationInvocation:
|
||||
profile.allowConversationInvocation ??
|
||||
existing?.allowConversationInvocation ?? false,
|
||||
contextWindowTokens: profile.contextWindowTokens,
|
||||
maximumOutputTokens: profile.maximumOutputTokens,
|
||||
imageGenerationQuality: profile.imageGenerationQuality,
|
||||
@@ -2491,6 +2518,9 @@ export class RuntimeSettingsStore {
|
||||
modelProfiles,
|
||||
defaultModelProfileId,
|
||||
opencodeModelSource,
|
||||
defaultImageModelProfileId: input.defaultImageModelProfileId === undefined
|
||||
? current.defaultImageModelProfileId
|
||||
: input.defaultImageModelProfileId,
|
||||
continueModelSource,
|
||||
deepseekHarnessModelSource:
|
||||
requestedDeepSeekHarnessSource,
|
||||
|
||||
Reference in New Issue
Block a user