fix(cua): await screenshot projection during disposal

packages/experimental/cua/README.i18n.yaml

@@ -15,8 +15,8 @@
   en: d152484eb41ac6b4
   zh: 7b30b973a849247d
 /deepseek-ai-dsh-experimental-cua/use-this-package:
-  en: ecfefc8c313586d1
-  zh: 263af55e29462082
+  en: 8f59929834c395dc
+  zh: 5c34c8f3b06ed15f
 /deepseek-ai-dsh-experimental-cua/use-this-package/connect-an-external-execution-environment:
   en: dc2f1cd930a6b6ba
   zh: 12b4f57fa745afbf

packages/experimental/cua/README.md

@@ -49,6 +49,8 @@ The [configuration catalog](../../../docs/config-catalog.md#deepseek-aidsh-exper
 
 Key chords use `+` between modifiers and the key. Use `+` for the plus key, `cmd++` for Command-plus, or the key name `plus`; an unfinished chord such as `cmd+` is rejected.
 
+Agent disposal and plugin unload cancel and await the complete tool call, including image admission and attachment storage. An image write already in progress must settle before disposal returns; a canceled call cannot return a successful screenshot.
+
 ### Connect an external execution environment
 
 To run through persistent PTC, add `--patch packages/experimental/cua/examples/persistent-ptc.cordis.yml` after the macOS overlay. This provider keeps variables for the live Agent and runs without a file sandbox. The caller imports the bridge and supplies the current authorized `tools.cua` binding. This integration snippet assumes an execution host has supplied `tools`; the package does not install a global object or keep variables alive:

packages/experimental/cua/README.zh.md

@@ -49,6 +49,8 @@ pnpm dsh headless --patch packages/experimental/cua/examples/macos.cordis.yml --
 
 组合键用 `+` 分隔修饰键与按键。加号键写作 `+`，Command 加号写作 `cmd++`，也可以使用键名 `plus`；`cmd+` 这类缺少按键的输入会被拒绝。
 
+Agent 销毁和插件卸载会取消并等待完整工具调用，包括图像接入检查和附件存储。已经开始的图像写入必须结束后，销毁才会返回；已取消的调用不会返回成功截图。
+
 ### 接入外部执行环境
 
 使用 persistent PTC 时，在 macOS 配置后追加 `--patch packages/experimental/cua/examples/persistent-ptc.cordis.yml`。该 provider 为存活的 Agent 保留变量，不启用文件沙箱。调用方导入适配模块，并提供当前执行中已授权的 `tools.cua` 绑定。下面的接入代码假设执行宿主已提供 `tools`；本包不安装全局对象，也不负责保留变量：

packages/experimental/cua/src/index.ts

@@ -9,7 +9,7 @@ import { z } from 'zod'
 import { NativeClient } from './native-client.ts'
 import { Request, ScreenshotResult } from './protocol.ts'
 import type {} from '@deepseek-ai/dsh-system-prompt'
-import type {} from '@deepseek-ai/dsh-tools'
+import type { ToolExecution } from '@deepseek-ai/dsh-tools'
 import type {} from '@deepseek-ai/dsh-subprocess'
 import type {} from '@deepseek-ai/dsh-computer-use'
 
@@ -99,8 +99,33 @@ export function apply(ctx: Context, config: Config): void {
   const clients = new Map<Agent, NativeClient>()
   const lifetime = new AbortController()
   const pending = new Map<Promise<unknown>, Agent>()
-  const ownerCleanups = new Map<Agent, { close: () => Promise<void>; release: () => Promise<void> }>()
+  const ownerCleanups = new Map<Agent, { close: () => Promise<void>; release: () => Promise<void>; signal: AbortSignal }>()
   const closedOwners = new WeakSet<Agent>()
+  const projectionExecutions = new WeakMap<Readonly<ToolExecution>, ToolExecution>()
+  function ownerFor(agent: Agent) {
+    if (closedOwners.has(agent) || ctx.agents.get(agent.id) !== agent) throw new Error('dsh-cua agent has been disposed')
+    const existing = ownerCleanups.get(agent)
+    if (existing) return existing
+    const controller = new AbortController()
+    let closing: Promise<void> | undefined
+    const close = (): Promise<void> => {
+      closedOwners.add(agent)
+      controller.abort()
+      return closing ??= Promise.resolve().then(async () => {
+        try {
+          await clients.get(agent)?.dispose()
+        } finally {
+          await Promise.allSettled([...pending].filter(([, owner]) => owner === agent).map(([task]) => task))
+          clients.delete(agent)
+          ownerCleanups.delete(agent)
+        }
+      })
+    }
+    const release = agent.ctx.effect(() => close, 'dsh-cua.agent')
+    const owner = { close, release, signal: controller.signal }
+    ownerCleanups.set(agent, owner)
+    return owner
+  }
   ctx.effect(function* () {
     yield ctx.computerUse.register(ComputerUseProviderName('dsh-cua'))
     yield async () => {
@@ -113,7 +138,7 @@ export function apply(ctx: Context, config: Config): void {
       for (const result of [...results, ...releases]) if (result.status === 'rejected') errors.push(result.reason)
       if (errors.length > 0) throw new AggregateError(errors, 'dsh-cua helper cleanup failed')
     }
-    yield ctx.tools.register(createMcpToolDefinition(ctx, {
+    const definition = createMcpToolDefinition(ctx, {
       name: 'cua', rawName: 'cua',
       description: 'Inspect macOS permissions, select an application and window, read AX state, capture screenshots or dispatch input. Use appHandle from getApp; verify actions by reading fresh state.',
       inputSchema: { ...z.record(z.string(), z.json()).parse(z.toJSONSchema(Request)), type: 'object' },
@@ -122,58 +147,65 @@ export function apply(ctx: Context, config: Config): void {
         const params = z.record(z.string(), z.json()).parse(parsedParams)
         const agent = execution.agent
         if (!agent) throw new Error('dsh-cua requires an agent-owned execution')
-        const signal = AbortSignal.any([execution.signal, lifetime.signal])
+        const signal = execution.signal
         signal.throwIfAborted()
+        const previous = clients.get(agent)
+        if (previous?.disposed) await previous.dispose()
+        signal.throwIfAborted()
+        let client = clients.get(agent)
+        if (!client || client.disposed) {
+          client = new NativeClient(ctx, {
+            helperPath: config.helperPath, cwd: process.cwd(), timeoutMs: config.timeoutMs,
+            maxRequestBytes: config.maxRequestBytes, maxResponseBytes: config.maxResponseBytes, graceMs: config.graceMs,
+            textEventIntervalMs: config.textEventIntervalMs,
+          })
+          clients.set(agent, client)
+        }
+        const result = await client.request(method, method === 'getState'
+          ? { ...params, maxElements: config.maxElements, maxDepth: config.maxDepth }
+          : method === 'drag' ? { ...params, durationMs: config.dragDurationMs, steps: config.dragSteps }
+            : params, signal, execution.callId)
+        signal.throwIfAborted()
+        if (method === 'screenshot') {
+          const { data, mimeType, ...metadata } = ScreenshotResult.parse(result)
+          return { content: [{ type: 'text', text: JSON.stringify(metadata) }, { type: 'image', mimeType, data }], structuredContent: result }
+        }
+        return { content: [{ type: 'text', text: JSON.stringify(result) }], structuredContent: result }
+      },
+    })
+    yield ctx.tools.register({
+      ...definition,
+      async execute(args, execution) {
+        const agent = execution.agent
+        if (!agent) throw new Error('dsh-cua requires an agent-owned execution')
+        execution.signal.throwIfAborted()
+        lifetime.signal.throwIfAborted()
+        const owner = ownerFor(agent)
+        const signal = AbortSignal.any([execution.signal, lifetime.signal, owner.signal])
+        signal.throwIfAborted()
+        const scoped = { ...execution, signal }
+        projectionExecutions.set(execution, scoped)
         const operation = (async () => {
-          const previous = clients.get(agent)
-          if (previous?.disposed) await previous.dispose()
+          const result = await definition.execute(args, scoped)
           signal.throwIfAborted()
-          if (closedOwners.has(agent) || ctx.agents.get(agent.id) !== agent) throw new Error('dsh-cua agent has been disposed')
-          if (!ownerCleanups.has(agent)) {
-            let closing: Promise<void> | undefined
-            const close = (): Promise<void> => {
-              closedOwners.add(agent)
-              return closing ??= Promise.resolve().then(async () => {
-                try {
-                  await clients.get(agent)?.dispose()
-                } finally {
-                  await Promise.allSettled([...pending].filter(([, owner]) => owner === agent).map(([task]) => task))
-                  clients.delete(agent)
-                  ownerCleanups.delete(agent)
-                }
-              })
-            }
-            const release = agent.ctx.effect(() => close, 'dsh-cua.agent')
-            ownerCleanups.set(agent, { close, release })
-          }
-          let client = clients.get(agent)
-          if (!client || client.disposed) {
-            client = new NativeClient(ctx, {
-              helperPath: config.helperPath, cwd: process.cwd(), timeoutMs: config.timeoutMs,
-              maxRequestBytes: config.maxRequestBytes, maxResponseBytes: config.maxResponseBytes, graceMs: config.graceMs,
-              textEventIntervalMs: config.textEventIntervalMs,
-            })
-            clients.set(agent, client)
-          }
-          return client.request(method, method === 'getState'
-            ? { ...params, maxElements: config.maxElements, maxDepth: config.maxDepth }
-            : method === 'drag' ? { ...params, durationMs: config.dragDurationMs, steps: config.dragSteps }
-              : params, signal, execution.callId)
+          return result
         })()
         pending.set(operation, agent)
         try {
-          const result = await operation
-          signal.throwIfAborted()
-          if (method === 'screenshot') {
-            const { data, mimeType, ...metadata } = ScreenshotResult.parse(result)
-            return { content: [{ type: 'text', text: JSON.stringify(metadata) }, { type: 'image', mimeType, data }], structuredContent: result }
-          }
-          return { content: [{ type: 'text', text: JSON.stringify(result) }], structuredContent: result }
+          return await operation
+        } catch (error) {
+          projectionExecutions.delete(execution)
+          throw error
         } finally {
           pending.delete(operation)
         }
       },
-    }))
+      projectContent(execution, result) {
+        const scoped = projectionExecutions.get(execution)
+        projectionExecutions.delete(execution)
+        return scoped === undefined ? undefined : definition.projectContent?.(scoped, result)
+      },
+    })
     yield ctx.systemPrompt.section({ name: 'computer-use:dsh-cua', order: ctx.systemPrompt.getSectionOrder('TOOL_COMPUTER_USE'), text: GUIDANCE })
   }, 'dsh-cua.runtime')
 }

packages/experimental/cua/tests/provider-lifecycle.spec.ts

@@ -31,6 +31,16 @@ const config = {
   graceMs: 10, textEventIntervalMs: 8, maxElements: 64, maxDepth: 8, dragDurationMs: 200, dragSteps: 10,
 } satisfies NativeProvider.Config
 
+class ImageAdapter extends LlmAdapter {
+  override resolveModel(provider: string, model: string): Promise<LlmResolvedModelInfo> {
+    return Promise.resolve({ provider, id: model, name: model, inputModalities: ['text', 'image'] })
+  }
+  stream(): AsyncIterable<StreamChunk> { throw new Error('This direct tool test does not request model generation') }
+}
+
+const imageData = 'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAIAAACQd1PeAAAACXBIWXMAAAPoAAAD6AG1e1JrAAAADElEQVQImWNgZGIGAAAOAAeCcsnOAAAAAElFTkSuQmCC'
+const imageMetadata = { pixelWidth: 1, pixelHeight: 1, bounds: { x: 2, y: 3, width: 1, height: 1 }, coordinateSpace: 'window-points', windowId: 7 }
+
 async function setup(model?: { provider: string; model: string }) {
   native.request.mockReset().mockResolvedValue([])
   native.dispose.mockReset()
@@ -227,19 +237,13 @@ it('supplies configured resource bounds and clears registrations on unload', asy
 })
 
 it('preserves screenshot coordinates and image attachments on every host', async () => {
-  class ImageAdapter extends LlmAdapter {
-    override resolveModel(provider: string, model: string): Promise<LlmResolvedModelInfo> {
-      return Promise.resolve({ provider, id: model, name: model, inputModalities: ['text', 'image'] })
-    }
-    stream(): AsyncIterable<StreamChunk> { throw new Error('This direct tool test does not request model generation') }
-  }
   const f = await setup({ provider: 'image-fixture', model: 'vision' })
   f.ctx.llm.registerAdapter(['image-fixture'], new ImageAdapter())
   const root = await mkdtemp(join(tmpdir(), 'dsh-cua-image-'))
   onTestFinished(async () => { await f.ctx.fiber.dispose(); await rm(root, { recursive: true, force: true }) })
   await f.ctx.plugin(LocalAttachmentStore, { dshHome: root })
-  const data = 'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAIAAACQd1PeAAAACXBIWXMAAAPoAAAD6AG1e1JrAAAADElEQVQImWNgZGIGAAAOAAeCcsnOAAAAAElFTkSuQmCC'
-  const metadata = { pixelWidth: 1, pixelHeight: 1, bounds: { x: 2, y: 3, width: 1, height: 1 }, coordinateSpace: 'window-points', windowId: 7 }
+  const data = imageData
+  const metadata = imageMetadata
   native.request.mockResolvedValue({ ...metadata, mimeType: 'image/png', data })
   const result = await f.execute('screenshot', { appHandle: 'app-one', windowId: 7 })
   expect(result.isError).toBe(false)
@@ -252,3 +256,53 @@ it('preserves screenshot coordinates and image attachments on every host', async
   expect(Buffer.from(stored.data).toString('base64')).toBe(data)
   expect(result.value).toMatchObject({ structuredContent: { ...metadata, mimeType: 'image/png', data } })
 })
+
+it.each([
+  { owner: 'provider', stage: 'admission' },
+  { owner: 'provider', stage: 'storage' },
+  { owner: 'agent', stage: 'admission' },
+  { owner: 'agent', stage: 'storage' },
+])('waits for screenshot $stage and cancels its result during $owner disposal', async ({ owner, stage }) => {
+  const f = await setup({ provider: 'image-fixture', model: 'vision' })
+  f.ctx.llm.registerAdapter(['image-fixture'], new ImageAdapter())
+  const root = await mkdtemp(join(tmpdir(), 'dsh-cua-image-disposal-'))
+  onTestFinished(async () => { try { await f.ctx.fiber.dispose() } finally { await rm(root, { recursive: true, force: true }) } })
+  await f.ctx.plugin(LocalAttachmentStore, { dshHome: root })
+  native.request.mockResolvedValue({ ...imageMetadata, mimeType: 'image/png', data: imageData })
+  const entered: PromiseWithResolvers<void> = Promise.withResolvers()
+  const release: PromiseWithResolvers<void> = Promise.withResolvers()
+  onTestFinished(() => { release.resolve() })
+  let admissionSignal: AbortSignal | undefined
+  const resolveModel = f.ctx.llm.resolveModelInfo.bind(f.ctx.llm)
+  const admission = vi.spyOn(f.ctx.llm, 'resolveModelInfo').mockImplementation(async (provider, model, signal) => {
+    admissionSignal = signal
+    if (stage === 'admission') { entered.resolve(); await release.promise }
+    return resolveModel(provider, model, signal)
+  })
+  const saveImages = f.ctx.attachments.saveImages.bind(f.ctx.attachments)
+  const storage = vi.spyOn(f.ctx.attachments, 'saveImages').mockImplementation(async (images) => {
+    if (stage === 'storage') { entered.resolve(); await release.promise }
+    return saveImages(images)
+  })
+  const pending = f.execute('screenshot', { appHandle: 'app-one', windowId: 7 })
+  let closing: Promise<void> | undefined
+  try {
+    const reached = await Promise.race([entered.promise.then(() => true), pending.then(() => false)])
+    expect(reached).toBe(true)
+    let closed = false
+    closing = (owner === 'provider' ? f.provider.dispose() : f.handle.dispose()).then(() => { closed = true })
+    await setImmediate()
+    expect(admissionSignal?.aborted).toBe(true)
+    expect(closed).toBe(false)
+    release.resolve()
+    expect((await pending).isError).toBe(true)
+    await closing
+    expect(closed).toBe(true)
+    expect(storage).toHaveBeenCalledTimes(stage === 'storage' ? 1 : 0)
+  } finally {
+    release.resolve()
+    await Promise.allSettled([pending, closing])
+    admission.mockRestore()
+    storage.mockRestore()
+  }
+})