feat: add computer control and managed browser

src/main/agent/model-tool-provider.ts

@@ -14,6 +14,7 @@ import {
   resolve
 } from 'node:path'
 import { z } from 'zod'
+import { builtinModelTools } from '../../shared/builtin-model-tools'
 import type { ResolvedMcpServer } from '../capabilities/capability-service'
 import { createMcpTransport } from '../capabilities/mcp-client-transport'
 import {
@@ -23,6 +24,11 @@ import {
   readBoundedUtf8File
 } from '../workspace-file-access'
 import type { RuntimeApprovalRequest } from './runtime'
+import {
+  BrowserModelTools,
+  type BrowserToolService
+} from '../browser/browser-model-tools'
+import { BrowserStaleReferenceError } from '../browser/cdp-browser-driver'
 
 const MAX_MODEL_TOOLS = 100
 const MAX_MCP_SERVERS = 16
@@ -31,6 +37,15 @@ const MAX_TOOL_RESULT_BYTES = 256 * 1024
 const MAX_READ_BYTES = 256 * 1024
 const MAX_WRITE_BYTES = 512 * 1024
 const MCP_TIMEOUT_MS = 30_000
+const MCP_CALL_MAX_TOTAL_TIMEOUT_MS = 5 * 60_000
+const MCP_TASK_CANCEL_TIMEOUT_MS = 5_000
+const MAX_MCP_CONTENT_BLOCKS = 100
+const MAX_MCP_IMAGES = 8
+const [
+  workspaceReadTextTool,
+  workspaceListDirectoryTool,
+  workspaceWriteTextTool
+] = builtinModelTools
 
 const workspacePathSchema = z
   .string()
@@ -66,20 +81,62 @@ export type ModelToolDefinition = {
   inputSchema: Record<string, unknown>
   source: 'builtin' | 'mcp'
   serverName?: string
+  taskSupport?: 'forbidden' | 'optional' | 'required'
+}
+
+export type ModelToolResultPart =
+  | {
+      type: 'text'
+      text: string
+    }
+  | {
+      type: 'image'
+      mimeType: 'image/png' | 'image/jpeg' | 'image/webp'
+      data: string
+    }
+
+export type ModelToolResult = {
+  parts: ModelToolResultPart[]
+  contextBytes: number
+}
+
+export type ModelToolCallContext = {
+  conversationId: string
+  workMode: 'ask' | 'plan' | 'execute'
+}
+
+export class RecoverableModelToolError extends Error {
+  readonly nextAction: string
+
+  constructor(
+    message: string,
+    nextAction: string,
+    options?: ErrorOptions
+  ) {
+    super(message, options)
+    this.name = 'RecoverableModelToolError'
+    this.nextAction = nextAction
+  }
 }
 
 export interface ModelToolProviderLike {
-  listTools(signal: AbortSignal): Promise<ModelToolDefinition[]>
+  listTools(
+    context: ModelToolCallContext,
+    signal: AbortSignal
+  ): Promise<ModelToolDefinition[]>
   getApproval(
     tool: ModelToolDefinition,
     argumentsValue: Record<string, unknown>,
-    argumentSummary: string
+    argumentSummary: string,
+    context: ModelToolCallContext
   ): RuntimeApprovalRequest
   callTool(
     name: string,
     argumentsValue: Record<string, unknown>,
-    signal: AbortSignal
-  ): Promise<string>
+    signal: AbortSignal,
+    context: ModelToolCallContext
+  ): Promise<ModelToolResult>
+  releaseConversation(conversationId: string): Promise<void>
   dispose(): Promise<void>
 }
 
@@ -151,64 +208,177 @@ function createMcpToolName(serverId: string, originalName: string): string {
   return `mcp_${serverHash}_${toolHash}_${readable}`.slice(0, 64)
 }
 
-function getMcpResultText(result: unknown): string {
+function createTextToolResult(text: string): ModelToolResult {
+  const contextBytes = Buffer.byteLength(text)
+  if (contextBytes > MAX_TOOL_RESULT_BYTES) {
+    throw new Error('工具结果超过 256KB 安全限制')
+  }
+  return {
+    parts: [{ type: 'text', text }],
+    contextBytes
+  }
+}
+
+function parseMcpImage(
+  content: Record<string, unknown>
+): Extract<ModelToolResultPart, { type: 'image' }> {
+  const mimeType = content.mimeType
+  if (
+    mimeType !== 'image/png' &&
+    mimeType !== 'image/jpeg' &&
+    mimeType !== 'image/webp'
+  ) {
+    throw new Error('MCP 工具返回了不支持的图片格式')
+  }
+  if (
+    typeof content.data !== 'string' ||
+    content.data.length === 0 ||
+    content.data.length % 4 !== 0 ||
+    !/^(?:[A-Za-z0-9+/]{4})*(?:[A-Za-z0-9+/]{2}==|[A-Za-z0-9+/]{3}=)?$/u.test(
+      content.data
+    )
+  ) {
+    throw new Error('MCP 工具返回了无效的 base64 图片')
+  }
+  const decoded = Buffer.from(content.data, 'base64')
+  if (
+    decoded.length === 0 ||
+    decoded.length > MAX_TOOL_RESULT_BYTES ||
+    decoded.toString('base64') !== content.data
+  ) {
+    throw new Error('MCP 工具返回了无效或过大的 base64 图片')
+  }
+  const signatureMatches =
+    mimeType === 'image/png'
+      ? decoded.length >= 8 &&
+        decoded.subarray(0, 8).equals(
+          Buffer.from([
+            0x89, 0x50, 0x4e, 0x47,
+            0x0d, 0x0a, 0x1a, 0x0a
+          ])
+        )
+      : mimeType === 'image/jpeg'
+        ? decoded.length >= 3 &&
+          decoded[0] === 0xff &&
+          decoded[1] === 0xd8 &&
+          decoded[2] === 0xff
+        : decoded.length >= 12 &&
+          decoded.subarray(0, 4).toString('ascii') === 'RIFF' &&
+          decoded.subarray(8, 12).toString('ascii') === 'WEBP'
+  if (!signatureMatches) {
+    throw new Error('MCP 工具图片的 MIME 类型与文件签名不匹配')
+  }
+  return {
+    type: 'image',
+    mimeType,
+    data: content.data
+  }
+}
+
+function normalizeMcpResult(result: unknown): ModelToolResult {
   if (!result || typeof result !== 'object') {
-    return boundedJson(result, 'MCP 工具结果无法序列化')
+    return createTextToolResult(
+      boundedJson(result, 'MCP 工具结果无法序列化')
+    )
   }
   const record = result as Record<string, unknown>
   if (record.isError === true) {
     throw new Error('MCP Server 报告工具执行失败')
   }
   if ('toolResult' in record) {
-    return boundedJson(record.toolResult, 'MCP 工具结果无法序列化')
+    if (
+      record.toolResult &&
+      typeof record.toolResult === 'object' &&
+      (
+        Array.isArray(
+          (record.toolResult as Record<string, unknown>).content
+        ) ||
+        'structuredContent' in
+          (record.toolResult as Record<string, unknown>) ||
+        'isError' in (record.toolResult as Record<string, unknown>)
+      )
+    ) {
+      return normalizeMcpResult(record.toolResult)
+    }
+    return createTextToolResult(
+      boundedJson(record.toolResult, 'MCP 工具结果无法序列化')
+    )
   }
 
-  const sections: string[] = []
+  const parts: ModelToolResultPart[] = []
   if (
     record.structuredContent &&
     typeof record.structuredContent === 'object'
   ) {
-    sections.push(
-      boundedJson(
+    parts.push({
+      type: 'text',
+      text: boundedJson(
         record.structuredContent,
         'MCP 结构化工具结果无法序列化'
       )
-    )
+    })
   }
   if (Array.isArray(record.content)) {
-    for (const item of record.content.slice(0, 100)) {
+    if (record.content.length > MAX_MCP_CONTENT_BLOCKS) {
+      throw new Error('MCP 工具结果内容块数量超过安全限制')
+    }
+    let imageCount = 0
+    for (const item of record.content) {
       if (!item || typeof item !== 'object') {
         continue
       }
       const content = item as Record<string, unknown>
       if (content.type === 'text' && typeof content.text === 'string') {
-        sections.push(content.text)
+        parts.push({ type: 'text', text: content.text })
       } else if (
         content.type === 'resource' &&
         content.resource &&
         typeof content.resource === 'object' &&
         typeof (content.resource as Record<string, unknown>).text === 'string'
       ) {
-        sections.push(
-          (content.resource as Record<string, unknown>).text as string
-        )
+        parts.push({
+          type: 'text',
+          text: (content.resource as Record<string, unknown>).text as string
+        })
       } else if (content.type === 'resource_link') {
-        sections.push(
-          boundedJson(content, 'MCP 资源链接无法序列化')
-        )
-      } else if (content.type === 'image' || content.type === 'audio') {
-        sections.push(`[${String(content.type)} result omitted]`)
+        parts.push({
+          type: 'text',
+          text: boundedJson(content, 'MCP 资源链接无法序列化')
+        })
+      } else if (content.type === 'image') {
+        imageCount += 1
+        if (imageCount > MAX_MCP_IMAGES) {
+          throw new Error('MCP 工具结果图片数量超过安全限制')
+        }
+        parts.push(parseMcpImage(content))
+      } else if (content.type === 'audio') {
+        parts.push({
+          type: 'text',
+          text: '[audio result unsupported]'
+        })
       }
     }
   }
-  const text = sections.join('\n\n').trim()
-  if (!text) {
-    return '{}'
+  if (parts.length === 0) {
+    return createTextToolResult('{}')
   }
-  if (Buffer.byteLength(text) > MAX_TOOL_RESULT_BYTES) {
+  let contextBytes = 0
+  let decodedImageBytes = 0
+  for (const part of parts) {
+    contextBytes += Buffer.byteLength(
+      part.type === 'text' ? part.text : part.data
+    )
+    if (part.type === 'image') {
+      decodedImageBytes += Buffer.from(part.data, 'base64').length
+    }
+  }
+  if (
+    contextBytes > MAX_TOOL_RESULT_BYTES ||
+    decodedImageBytes > MAX_TOOL_RESULT_BYTES
+  ) {
     throw new Error('工具结果超过 256KB 安全限制')
   }
-  return text
+  return { parts, contextBytes }
 }
 
 export class ModelToolProvider implements ModelToolProviderLike {
@@ -218,9 +388,21 @@ export class ModelToolProvider implements ModelToolProviderLike {
 
   constructor(
     private readonly workspace: string,
-    private readonly mcpServers: ResolvedMcpServer[] = []
+    private readonly mcpServers: ResolvedMcpServer[] = [],
+    private readonly browserService?: BrowserToolService
   ) {}
 
+  private getBrowserTools(
+    context: ModelToolCallContext
+  ): BrowserModelTools | undefined {
+    return this.browserService && context.workMode === 'execute'
+      ? new BrowserModelTools({
+          service: this.browserService,
+          conversationId: context.conversationId
+        })
+      : undefined
+  }
+
   private async getWorkspace(): Promise<string> {
     this.canonicalWorkspace ??= getCanonicalWorkspace(
       this.workspace,
@@ -289,10 +471,9 @@ export class ModelToolProvider implements ModelToolProviderLike {
   private getBuiltinTools(): ModelToolDefinition[] {
     return [
       {
-        name: 'workspace_read_text',
-        displayName: '读取工作区文本',
-        description:
-          '读取当前工作区内一个不超过 256KB 的 UTF-8 文本文件。',
+        name: workspaceReadTextTool.name,
+        displayName: workspaceReadTextTool.displayName,
+        description: workspaceReadTextTool.description,
         inputSchema: {
           type: 'object',
           properties: {
@@ -307,10 +488,9 @@ export class ModelToolProvider implements ModelToolProviderLike {
         source: 'builtin'
       },
       {
-        name: 'workspace_list_directory',
-        displayName: '列出工作区目录',
-        description:
-          '列出当前工作区内目录的直属内容，最多返回 200 项。',
+        name: workspaceListDirectoryTool.name,
+        displayName: workspaceListDirectoryTool.displayName,
+        description: workspaceListDirectoryTool.description,
         inputSchema: {
           type: 'object',
           properties: {
@@ -324,10 +504,9 @@ export class ModelToolProvider implements ModelToolProviderLike {
         source: 'builtin'
       },
       {
-        name: 'workspace_write_text',
-        displayName: '写入工作区文本',
-        description:
-          '在当前工作区内新建或覆盖一个不超过 512KB 的 UTF-8 文本文件；父目录必须已存在。',
+        name: workspaceWriteTextTool.name,
+        displayName: workspaceWriteTextTool.displayName,
+        description: workspaceWriteTextTool.description,
         inputSchema: {
           type: 'object',
           properties: {
@@ -366,7 +545,9 @@ export class ModelToolProvider implements ModelToolProviderLike {
         timeout: MCP_TIMEOUT_MS,
         signal
       })
-      if (result.tools.length > MAX_MODEL_TOOLS - 3) {
+      const builtinToolCount =
+        this.getBuiltinTools().length + (this.browserService ? 7 : 0)
+      if (result.tools.length > MAX_MODEL_TOOLS - builtinToolCount) {
         throw new Error(
           `MCP Server「${server.name}」提供的工具数量超过安全限制`
         )
@@ -386,7 +567,8 @@ export class ModelToolProvider implements ModelToolProviderLike {
             .slice(0, 1_000),
           inputSchema: normalizeToolSchema(tool.inputSchema),
           source: 'mcp',
-          serverName: server.name
+          serverName: server.name,
+          taskSupport: tool.execution?.taskSupport
         }
       }))
       if (
@@ -423,9 +605,11 @@ export class ModelToolProvider implements ModelToolProviderLike {
     )
       .then((connections) => {
         const bindings = new Map<string, McpToolBinding>()
+        const builtinToolCount =
+          this.getBuiltinTools().length + (this.browserService ? 7 : 0)
         for (const connection of connections) {
           for (const binding of connection.tools) {
-            if (bindings.size + 3 >= MAX_MODEL_TOOLS) {
+            if (bindings.size + builtinToolCount >= MAX_MODEL_TOOLS) {
               throw new Error('直连模型工具总数超过 100 个安全限制')
             }
             if (bindings.has(binding.definition.name)) {
@@ -448,11 +632,16 @@ export class ModelToolProvider implements ModelToolProviderLike {
     return this.mcpBindings
   }
 
-  async listTools(signal: AbortSignal): Promise<ModelToolDefinition[]> {
+  async listTools(
+    context: ModelToolCallContext,
+    signal: AbortSignal
+  ): Promise<ModelToolDefinition[]> {
     signal.throwIfAborted()
     const bindings = await this.getMcpBindings(signal)
+    const browserTools = this.getBrowserTools(context)
     return [
       ...this.getBuiltinTools(),
+      ...(browserTools?.listTools() ?? []),
       ...[...bindings.values()].map((binding) => binding.definition)
     ]
   }
@@ -460,8 +649,17 @@ export class ModelToolProvider implements ModelToolProviderLike {
   getApproval(
     tool: ModelToolDefinition,
     argumentsValue: Record<string, unknown>,
-    argumentSummary: string
+    argumentSummary: string,
+    context: ModelToolCallContext
   ): RuntimeApprovalRequest {
+    const browserTools = this.getBrowserTools(context)
+    if (browserTools?.ownsTool(tool.name)) {
+      return browserTools.getApproval(
+        tool,
+        argumentsValue,
+        argumentSummary
+      )
+    }
     const path =
       typeof argumentsValue.path === 'string'
         ? argumentsValue.path.slice(0, 500)
@@ -490,20 +688,38 @@ export class ModelToolProvider implements ModelToolProviderLike {
   async callTool(
     name: string,
     argumentsValue: Record<string, unknown>,
-    signal: AbortSignal
-  ): Promise<string> {
+    signal: AbortSignal,
+    context: ModelToolCallContext
+  ): Promise<ModelToolResult> {
     signal.throwIfAborted()
+    const browserTools = this.getBrowserTools(context)
+    if (browserTools?.ownsTool(name)) {
+      try {
+        return await browserTools.callTool(name, argumentsValue, signal)
+      } catch (error) {
+        if (error instanceof BrowserStaleReferenceError) {
+          throw new RecoverableModelToolError(
+            error.message,
+            '调用 browser_snapshot 获取新快照，然后用新引用重试刚才的操作',
+            { cause: error }
+          )
+        }
+        throw error
+      }
+    }
     if (name === 'workspace_read_text') {
       const input = readInputSchema.parse(argumentsValue)
       const filePath = await this.resolveExistingPath(input.path, 'file')
-      return (
-        await readBoundedUtf8File(
-          filePath,
-          MAX_READ_BYTES,
-          '工作区文本文件超过 256KB 安全限制',
-          '工作区读取目标不是有效 UTF-8 文本'
-        )
-      ).content
+      return createTextToolResult(
+        (
+          await readBoundedUtf8File(
+            filePath,
+            MAX_READ_BYTES,
+            '工作区文本文件超过 256KB 安全限制',
+            '工作区读取目标不是有效 UTF-8 文本'
+          )
+        ).content
+      )
     }
     if (name === 'workspace_list_directory') {
       const input = listInputSchema.parse(argumentsValue)
@@ -515,21 +731,23 @@ export class ModelToolProvider implements ModelToolProviderLike {
         directoryPath,
         200
       )
-      return boundedJson(
-        {
-          entries: listing.entries
-            .sort((left, right) => left.name.localeCompare(right.name))
-            .map((entry) => ({
-              name: entry.name,
-              type: entry.isDirectory()
-                ? 'directory'
-                : entry.isFile()
-                  ? 'file'
-                  : 'other'
-            })),
-          truncated: listing.truncated
-        },
-        '工作区目录结果无法序列化'
+      return createTextToolResult(
+        boundedJson(
+          {
+            entries: listing.entries
+              .sort((left, right) => left.name.localeCompare(right.name))
+              .map((entry) => ({
+                name: entry.name,
+                type: entry.isDirectory()
+                  ? 'directory'
+                  : entry.isFile()
+                    ? 'file'
+                    : 'other'
+              })),
+            truncated: listing.truncated
+          },
+          '工作区目录结果无法序列化'
+        )
       )
     }
     if (name === 'workspace_write_text') {
@@ -551,12 +769,14 @@ export class ModelToolProvider implements ModelToolProviderLike {
         await rm(temporaryPath, { force: true }).catch(() => undefined)
         throw new Error('无法安全写入工作区文件', { cause: error })
       }
-      return boundedJson(
-        {
-          path: input.path,
-          bytesWritten: Buffer.byteLength(input.content)
-        },
-        '工作区写入结果无法序列化'
+      return createTextToolResult(
+        boundedJson(
+          {
+            path: input.path,
+            bytesWritten: Buffer.byteLength(input.content)
+          },
+          '工作区写入结果无法序列化'
+        )
       )
     }
 
@@ -564,18 +784,52 @@ export class ModelToolProvider implements ModelToolProviderLike {
     if (!binding) {
       throw new Error('模型请求了未知工具')
     }
-    const result = await binding.client.callTool(
-      {
-        name: binding.originalName,
-        arguments: argumentsValue
-      },
-      undefined,
-      {
-        timeout: MCP_TIMEOUT_MS,
-        signal
+    const params = {
+      name: binding.originalName,
+      arguments: argumentsValue
+    }
+    const options = {
+      timeout: MCP_TIMEOUT_MS,
+      signal,
+      onprogress: () => undefined,
+      resetTimeoutOnProgress: true,
+      maxTotalTimeout: MCP_CALL_MAX_TOTAL_TIMEOUT_MS
+    }
+    if (binding.definition.taskSupport !== 'required') {
+      return normalizeMcpResult(
+        await binding.client.callTool(params, undefined, options)
+      )
+    }
+
+    let taskId: string | undefined
+    try {
+      for await (const message of binding.client.experimental.tasks.callToolStream(
+        params,
+        undefined,
+        options
+      )) {
+        if (
+          (message.type === 'taskCreated' ||
+            message.type === 'taskStatus') &&
+          typeof message.task.taskId === 'string'
+        ) {
+          taskId = message.task.taskId
+        } else if (message.type === 'result') {
+          return normalizeMcpResult(message.result)
+        } else if (message.type === 'error') {
+          throw message.error
+        }
       }
-    )
-    return getMcpResultText(result)
+      throw new Error('MCP 任务工具未返回最终结果')
+    } catch (error) {
+      if (taskId) {
+        await binding.client.experimental.tasks.cancelTask(taskId, {
+          timeout: MCP_TASK_CANCEL_TIMEOUT_MS,
+          maxTotalTimeout: MCP_TASK_CANCEL_TIMEOUT_MS
+        }).catch(() => undefined)
+      }
+      throw error
+    }
   }
 
   async dispose(): Promise<void> {
@@ -584,4 +838,14 @@ export class ModelToolProvider implements ModelToolProviderLike {
     this.mcpBindings = undefined
     await Promise.allSettled(clients.map((client) => client.close()))
   }
+
+  async releaseConversation(conversationId: string): Promise<void> {
+    if (!this.browserService) {
+      return
+    }
+    await new BrowserModelTools({
+      service: this.browserService,
+      conversationId
+    }).release()
+  }
 }