diff --git a/docs/features/experimental-tool-modes/spec.md b/docs/features/experimental-tool-modes/spec.md index 0c89b52b5..a55294c48 100644 --- a/docs/features/experimental-tool-modes/spec.md +++ b/docs/features/experimental-tool-modes/spec.md @@ -40,8 +40,8 @@ Agent Mode 是兼容基线。工具 registry、disabled-tool 配置、权限、 Code Mode 保留当前已启用工具的能力,但不把这些工具逐个发送给模型: - `openai-codex` 路由使用 Codex 风格的 raw JavaScript `exec` 和续跑 `wait`; -- 其他 function-tool 路由只发送 `run_code`,参数固定为 - `{ code: string, description: string }`; +- Function-tool routes expose `run_code` with + `{ code: string, description: string, timeout_ms?: number }`; - 两种入口都进入同一个 `RunCodeRuntimeManager`,不存在 `transport` 或第二套 runtime; - function-tool 路由把当前目录生成成 TypeScript SDK,Codex 路由把嵌套声明写入 `exec` 描述; @@ -188,11 +188,18 @@ IPC allowlist、V8 内存限制、heartbeat 和强制回收的组合。 | READY | 5 秒 | | heartbeat 丢失 | 约 3.5 秒后终止 | | VM 同步执行 slice | 2 秒 | -| cell 总执行时间 | 5 分钟 | +| cell 总执行时间 | Default 5 minutes; optional `timeout_ms` override | | yielded / permission cell lease | 60 秒 | | RSS hard ceiling | 512 MiB | | STOP grace | 500 ms 后 `kill()` | +`timeout_ms` is a positive integer in milliseconds, up to `2147483647`. Function-tool routes pass +it as an optional `run_code` argument; Codex routes use the first-line `exec` pragma. Models omit +it for routine work and request a larger value only when the operation is expected to exceed five +minutes. The deadline covers the cell and its awaited subtools, is not reset by `wait` or permission +continuation, and does not extend a subtool's own timeout. User cancellation remains independent +of the selected duration. Invalid timeout values fail before dispatch. + 成功、异常、取消、超时、进程退出、会话清理和应用退出都进入同一个幂等 cleanup:取消嵌套 调用,清理 timer/listener/active map,发送 `STOP`,超时后强制 `kill()`。失败 cell 不自动 重放,避免重复执行 Shell、文件或 MCP 副作用。 diff --git a/docs/issues/agent-cancellation-code-timeout/spec.md b/docs/issues/agent-cancellation-code-timeout/spec.md new file mode 100644 index 000000000..7e9e498b0 --- /dev/null +++ b/docs/issues/agent-cancellation-code-timeout/spec.md @@ -0,0 +1,92 @@ +# Agent Cancellation and Code Cell Timeouts + +## Context + +[Issue #2221](https://github.com/ThinkInAIXYZ/deepchat/issues/2221) reports a Code Mode run that +stopped progressing and could not be stopped by the user. Its partial trace records a provider +429 followed by a scheduled retry; it does not identify the complete incident sequence. + +Independent reproduction identifies two cancellation boundaries that require protection: + +- A provider iterator can remain pending after the Run signal is aborted. Waiting for its next + event or asynchronous cleanup must not prevent attempt settlement. +- An unresolved browser Promise can leave a YoBrowser CDP command and its activity waiting + indefinitely unless both observe the tool cancellation signal. + +Code cells already have a five-minute execution deadline. That default is useful, but a longer +operation must be able to request a longer cell lifetime without weakening user cancellation. + +## Contract + +### User Cancellation + +- Deliver the Run signal to the provider immediately. The provider-attempt owner also interrupts + its own iterator wait if the provider ignores that signal. +- Request iterator cleanup without waiting indefinitely after cancellation. Observe late + rejections, discard late output, and record the attempt as aborted exactly once. +- Preserve usage observed before cancellation and existing no-replay rules for committed output. +- Browser tools observe cancellation while waiting and check the signal immediately before + dispatch. A canceled readiness wait must not dispatch a CDP command later. +- Cancellation ends browser activity without projecting a late response into the canceled run. + It does not claim to undo browser or external side effects already dispatched. +- Code cell startup and execution both observe cancellation. An abandoned utility process must + be reclaimed even if it becomes ready after cancellation. +- Do not add a whole-run deadline or infer failure merely from a lack of output. + +### Code Cell Execution Limit + +- `run_code` accepts optional `timeout_ms`, an integer duration in milliseconds. Omitting it uses + `300000` milliseconds, or five minutes. `code` and `description` remain the only required fields. +- The valid range is `1` through `2147483647`, the host timer's supported range. Invalid values fail + before a utility process or tool dispatch starts. +- The Code Mode `exec` frontend accepts the same optional field in its first-line pragma, because + both frontends share the same cell runtime. +- The limit covers one cell, including awaited subtools. It starts when the cell is attached and + does not restart on yield, `wait`, or a permission continuation. +- The tool description, parameter description, and SDK prompt all instruct the model to omit the + field for routine work and set a larger value only when the current operation is expected to + exceed five minutes. A cell override does not extend a subtool's own timeout. +- User cancellation always takes precedence over a longer requested timeout. Existing source, + output, concurrency, heartbeat, memory, startup, and abandoned-cell limits remain unchanged. + +## Ownership and Scope + +`DeepChatContextCoordinator` owns provider iterator cancellation and attempt provenance. +`AgentToolManager`, `YoBrowserToolHandler`, and the existing browser command path carry the tool +signal through dispatch and response handling. `ToolService` validates the public timeout input; +`RunCodeRuntimeManager` applies the selected cell deadline. Code Mode definitions and generated +SDK text expose the same contract. + +This change does not modify provider retry budgets, add dependencies, introduce a new task +scheduler, change UI layout, automatically replay canceled work, or add an unlimited timeout mode. + +## Acceptance Criteria + +- Stopping a silent provider settles the attempt without needing another provider event or a + cooperative iterator `return()`; late output does not start a tool or create another outcome. +- A normal long-running operation is not canceled merely because it is silent. +- Stopping a pending CDP call returns control and prevents delayed dispatch after readiness. +- An omitted cell timeout still expires at five minutes; an explicit longer timeout survives that + boundary and expires at the requested duration. +- A cell with an extended timeout still stops promptly when the user cancels. +- Invalid timeout values do not dispatch work, and both model-facing frontends document the + default, omission guidance, and extension semantics. + +## Implementation and Validation + +- [x] Provider and browser waits observe cancellation without requiring downstream cooperation. +- [x] Utility startup observes cancellation; both frontends accept the optional cell timeout. +- [x] The maintained Code Mode contract and model-facing instructions describe the same behavior. +- [x] Regression coverage protects cancellation, late completion, and timeout overrides. +- [x] Formatting, i18n, lint, type checks, and relevant test suites pass. + +Validation uses Node 24.18.0 and pnpm 10.34.5: + +- `pnpm run format`, `pnpm run i18n`, `pnpm run lint`, and `pnpm run typecheck` pass. +- 407 main-process tests pass across 13 relevant suites: Code Mode, ToolService, browser tools, + provider attempts and retries, stream processing, Run lifecycle, and abortable waits. +- All three ChatPage stop-request tests pass, covering duplicate requests, unsuccessful stop + responses, and rejected stop requests. +- Fake-clock cases verify the five-minute default, a ten-minute override, and explicit + cancellation of a provider silent for fifteen minutes and a CDP command silent for ten minutes. + Browser integration cases use mocked Electron boundaries; no live provider credentials are used. diff --git a/src/main/agent/deepchat/loop/contextCoordinator.ts b/src/main/agent/deepchat/loop/contextCoordinator.ts index 325f09bff..5122b059e 100644 --- a/src/main/agent/deepchat/loop/contextCoordinator.ts +++ b/src/main/agent/deepchat/loop/contextCoordinator.ts @@ -23,6 +23,7 @@ import type { MCPToolDefinition } from '@shared/types/core/mcp' import type { ModelConfig } from '@shared/types/provider' import type { DeepChatExecutionContract } from '@shared/types/execution-contract' import { isDeepStrictEqual } from 'node:util' +import { awaitWithAbort } from '@/lib/awaitWithAbort' import type { DeepChatProviderAttemptIdentity, DeepChatProviderAttemptOrigin, @@ -559,30 +560,52 @@ async function* observeProviderAttempt(input: { }): AsyncGenerator { const { observation } = input try { - for await (const event of input.provider.stream(input.streamInput)) { - if (!input.bypassContextBudget && input.isContextOverflowEvent(event)) { - observation.contextOverflowObserved = true - const facts = input.inspectContextOverflow?.(event) - if (facts?.matched) { - observation.contextOverflowFacts = facts - input.onContextOverflowFacts?.(facts) + const stream = input.provider.stream(input.streamInput) + let finished = false + try { + while (true) { + input.streamInput.signal.throwIfAborted() + const next = await awaitWithAbort(stream.next(), input.streamInput.signal) + input.streamInput.signal.throwIfAborted() + if (next.done) { + finished = true + break } - } - if (event.type === 'usage') { - observation.usage = providerAttemptUsageFromEvent(event) - } else if (event.type === 'error') { - observation.errorEvent = event - observation.stopReason = 'error' - } else if (event.type === 'stop') { - observation.stopEvent = event - if (observation.stopReason !== 'error') { - observation.stopReason = event.stop_reason + const event = next.value + if (!input.bypassContextBudget && input.isContextOverflowEvent(event)) { + observation.contextOverflowObserved = true + const facts = input.inspectContextOverflow?.(event) + if (facts?.matched) { + observation.contextOverflowFacts = facts + input.onContextOverflowFacts?.(facts) + } + } + if (event.type === 'usage') { + observation.usage = providerAttemptUsageFromEvent(event) + } else if (event.type === 'error') { + observation.errorEvent = event + observation.stopReason = 'error' + } else if (event.type === 'stop') { + observation.stopEvent = event + if (observation.stopReason !== 'error') { + observation.stopReason = event.stop_reason + } } - } - if (!isProviderControlEvent(event)) { - observation.outputCommitted = true - yield event + if (!isProviderControlEvent(event)) { + observation.outputCommitted = true + yield event + } + } + } finally { + if (!finished) { + // An async generator's return() can wait for its outstanding next() forever. + const closing = stream.return(undefined) + if (input.streamInput.signal.aborted) { + void closing.catch(() => undefined) + } else { + await awaitWithAbort(closing, input.streamInput.signal) + } } } } catch (error) { diff --git a/src/main/desktop/browser/BrowserTab.ts b/src/main/desktop/browser/BrowserTab.ts index be7f35048..ad447da10 100644 --- a/src/main/desktop/browser/BrowserTab.ts +++ b/src/main/desktop/browser/BrowserTab.ts @@ -1,5 +1,6 @@ import { WebContents } from 'electron' import { nanoid } from 'nanoid' +import { awaitWithAbort } from '@/lib/awaitWithAbort' import { BrowserPageStatus, type BrowserPageInfo, @@ -72,10 +73,13 @@ export class BrowserTab { url: string, timeoutMs: number = 30000, beforeDispatch?: () => void, - onDispatched?: () => void + onDispatched?: () => void, + signal?: AbortSignal ): Promise { + signal?.throwIfAborted() this.ensureAvailable() beforeDispatch?.() + signal?.throwIfAborted() this.beginMainFrameNavigation(url) const loadPromise = this.webContents.loadURL(url) @@ -90,13 +94,15 @@ export class BrowserTab { onDispatched?.() try { - await Promise.race([this.waitForInteractiveReady(timeoutMs), loadPromise]) + await Promise.race([this.waitForInteractiveReady(timeoutMs, signal), loadPromise]) + signal?.throwIfAborted() if (!this.interactiveReady) { throw new Error(`Navigation finished before dom-ready for ${url}`) } this.title = this.webContents.getTitle() || url this.updatedAt = Date.now() } catch (error) { + signal?.throwIfAborted() this.markNavigationError(error) throw error } @@ -118,19 +124,24 @@ export class BrowserTab { method: string, params?: Record, beforeDispatch?: () => void, - onDispatched?: () => void + onDispatched?: () => void, + signal?: AbortSignal ): Promise { + signal?.throwIfAborted() if (NAVIGATION_CDP_METHODS.has(method)) { this.ensureAvailable() } else { - await this.ensureInteractiveReadyOrWait(`send CDP command ${method}`) + await this.ensureInteractiveReadyOrWait(`send CDP command ${method}`, undefined, signal) } - const session = await this.ensureSession() + const session = await awaitWithAbort(this.ensureSession(), signal) + signal?.throwIfAborted() beforeDispatch?.() + signal?.throwIfAborted() const responsePromise = session.sendCommand(method, params ?? {}) onDispatched?.() - const response = await responsePromise + const response = await awaitWithAbort(responsePromise, signal) + signal?.throwIfAborted() if (method === 'Page.navigate') { const navigationResponse = response as { @@ -558,8 +569,10 @@ export class BrowserTab { private async ensureInteractiveReadyOrWait( action: string, - timeoutMs: number = INTERACTIVE_READY_WAIT_TIMEOUT_MS + timeoutMs: number = INTERACTIVE_READY_WAIT_TIMEOUT_MS, + signal?: AbortSignal ): Promise { + signal?.throwIfAborted() this.ensureAvailable() if (this.interactiveReady) { @@ -568,7 +581,7 @@ export class BrowserTab { if (this.awaitingMainFrameInteractive || this.status === BrowserPageStatus.Loading) { try { - await this.waitForInteractiveReady(timeoutMs) + await this.waitForInteractiveReady(timeoutMs, signal) } catch (error) { if (!this.isInteractiveReadyTimeoutError(error)) { throw error @@ -579,7 +592,7 @@ export class BrowserTab { return } - if (await this.probeInteractiveReadiness()) { + if (await this.probeInteractiveReadiness(signal)) { return } } @@ -790,7 +803,8 @@ export class BrowserTab { }) } - private async waitForInteractiveReady(timeoutMs: number): Promise { + private async waitForInteractiveReady(timeoutMs: number, signal?: AbortSignal): Promise { + signal?.throwIfAborted() if (this.interactiveReady) { return } @@ -808,6 +822,7 @@ export class BrowserTab { this.webContents.removeListener('dom-ready', onDomReady) this.webContents.removeListener('did-fail-load', onFailLoad as any) this.webContents.removeListener('destroyed', onDestroyed) + signal?.removeEventListener('abort', onAbort) } const onDomReady = () => { @@ -834,6 +849,11 @@ export class BrowserTab { reject(new Error('Page was destroyed before dom-ready')) } + const onAbort = () => { + cleanup() + reject(signal?.reason ?? new DOMException('Aborted', 'AbortError')) + } + timeoutId = setTimeout(() => { cleanup() reject(new Error(`${INTERACTIVE_READY_TIMEOUT_MESSAGE_PREFIX} ${this.url}`)) @@ -842,6 +862,8 @@ export class BrowserTab { this.webContents.once('dom-ready', onDomReady) this.webContents.on('did-fail-load', onFailLoad as any) this.webContents.once('destroyed', onDestroyed) + signal?.addEventListener('abort', onAbort, { once: true }) + if (signal?.aborted) onAbort() }) } @@ -851,12 +873,15 @@ export class BrowserTab { ) } - private async probeInteractiveReadiness(): Promise { + private async probeInteractiveReadiness(signal?: AbortSignal): Promise { try { - const session = await this.ensureSession() - const probe = (await this.cdpManager.evaluateScript( - session, - `(() => { + signal?.throwIfAborted() + const session = await awaitWithAbort(this.ensureSession(), signal) + signal?.throwIfAborted() + const probe = (await awaitWithAbort( + this.cdpManager.evaluateScript( + session, + `(() => { try { return { readyState: document.readyState, @@ -867,7 +892,10 @@ export class BrowserTab { return null } })()` + ), + signal )) as { readyState?: unknown; hasBody?: unknown; href?: unknown } | null + signal?.throwIfAborted() const readyState = typeof probe?.readyState === 'string' ? probe.readyState : '' const hasBody = probe?.hasBody === true @@ -888,6 +916,7 @@ export class BrowserTab { } return true } catch { + signal?.throwIfAborted() return false } } diff --git a/src/main/desktop/browser/YoBrowserPresenter.ts b/src/main/desktop/browser/YoBrowserPresenter.ts index 6ad11393f..93aba72f9 100644 --- a/src/main/desktop/browser/YoBrowserPresenter.ts +++ b/src/main/desktop/browser/YoBrowserPresenter.ts @@ -132,8 +132,10 @@ export class YoBrowserPresenter implements IYoBrowserPresenter { hostWindowId?: number, activitySource?: YoBrowserActivitySource, agentRunId?: string, - beforeDispatch?: () => void + beforeDispatch?: () => void, + signal?: AbortSignal ): Promise { + signal?.throwIfAborted() const normalizedSessionId = sessionId.trim() if (!normalizedSessionId) { throw new Error('sessionId is required') @@ -201,17 +203,24 @@ export class YoBrowserPresenter implements IYoBrowserPresenter { normalizedSessionId, { kind: 'navigation', action: 'navigate' }, (startActivity) => - state.page.navigateUntilDomReady(url, timeoutMs ?? 30000, beforeDispatch, () => { - projectDispatch() - startActivity() - }) + state.page.navigateUntilDomReady( + url, + timeoutMs ?? 30000, + beforeDispatch, + () => { + projectDispatch() + startActivity() + }, + signal + ) ) } else { await state.page.navigateUntilDomReady( url, timeoutMs ?? 30000, beforeDispatch, - projectDispatch + projectDispatch, + signal ) } @@ -532,7 +541,8 @@ export class YoBrowserPresenter implements IYoBrowserPresenter { params?: Record, activitySource?: YoBrowserActivitySource, agentRunId?: string, - beforeDispatch?: () => void + beforeDispatch?: () => void, + signal?: AbortSignal ): Promise { const state = this.sessionBrowsers.get(sessionId) if (!state) { @@ -557,14 +567,20 @@ export class YoBrowserPresenter implements IYoBrowserPresenter { } if (activitySource === 'agent' && descriptor) { return await this.runAgentActivity(sessionId, descriptor, (startActivity) => - state.page.sendCdpCommand(method, params, beforeDispatch, () => { - projectDispatch() - startActivity() - }) + state.page.sendCdpCommand( + method, + params, + beforeDispatch, + () => { + projectDispatch() + startActivity() + }, + signal + ) ) } - return await state.page.sendCdpCommand(method, params, beforeDispatch, projectDispatch) + return await state.page.sendCdpCommand(method, params, beforeDispatch, projectDispatch, signal) } async startDownload(url: string, savePath?: string): Promise { diff --git a/src/main/desktop/browser/YoBrowserToolHandler.ts b/src/main/desktop/browser/YoBrowserToolHandler.ts index 1a814b409..2d4ff0248 100644 --- a/src/main/desktop/browser/YoBrowserToolHandler.ts +++ b/src/main/desktop/browser/YoBrowserToolHandler.ts @@ -1,4 +1,5 @@ import logger from '@shared/logger' +import { awaitWithAbort } from '@/lib/awaitWithAbort' import { getYoBrowserToolDefinitions } from '@/tool/browser/definitions' import type { YoBrowserPresenter } from './YoBrowserPresenter' import { BrowserPageStatus, type YoBrowserStatus } from '@shared/types/browser' @@ -24,9 +25,11 @@ export class YoBrowserToolHandler { args: Record, conversationId?: string, runId?: string, - beforeInvoke?: (normalizedArguments: Record) => void + beforeInvoke?: (normalizedArguments: Record) => void, + signal?: AbortSignal ): Promise { try { + signal?.throwIfAborted() const sessionId = conversationId?.trim() if (!sessionId) { throw new Error('conversationId is required for YoBrowser tools') @@ -34,7 +37,9 @@ export class YoBrowserToolHandler { switch (toolName) { case 'get_browser_status': - return JSON.stringify(await this.presenter.getBrowserStatus(sessionId)) + return JSON.stringify( + await awaitWithAbort(this.presenter.getBrowserStatus(sessionId), signal) + ) case 'load_url': { const url = typeof args.url === 'string' ? args.url : '' if (!url) { @@ -42,7 +47,7 @@ export class YoBrowserToolHandler { } const beforeDispatch = beforeInvoke ? () => beforeInvoke({ url }) : undefined return JSON.stringify( - runId || beforeDispatch + runId || beforeDispatch || signal ? await this.presenter.loadUrl( sessionId, url, @@ -50,7 +55,8 @@ export class YoBrowserToolHandler { undefined, 'agent', runId, - beforeDispatch + beforeDispatch, + signal ) : await this.presenter.loadUrl(sessionId, url, undefined, undefined, 'agent') ) @@ -61,7 +67,7 @@ export class YoBrowserToolHandler { throw new Error('CDP method is required') } - const status = await this.presenter.getBrowserStatus(sessionId) + const status = await awaitWithAbort(this.presenter.getBrowserStatus(sessionId), signal) const page = status.page if (!status.initialized || !page || page.status === BrowserPageStatus.Closed) { throw await this.createUnavailableError(sessionId, method, status) @@ -71,14 +77,15 @@ export class YoBrowserToolHandler { const params = this.normalizeCdpParams(args.params) const beforeDispatch = beforeInvoke ? () => beforeInvoke({ method, params }) : undefined const response = - runId || beforeDispatch + runId || beforeDispatch || signal ? await this.presenter.sendCdpCommand( sessionId, method, params, 'agent', runId, - beforeDispatch + beforeDispatch, + signal ) : await this.presenter.sendCdpCommand(sessionId, method, params, 'agent') return JSON.stringify(response ?? {}) diff --git a/src/main/tool/agentTools/agentToolManager.ts b/src/main/tool/agentTools/agentToolManager.ts index cadb78145..ddd4590d7 100644 --- a/src/main/tool/agentTools/agentToolManager.ts +++ b/src/main/tool/agentTools/agentToolManager.ts @@ -1005,7 +1005,8 @@ export class AgentToolManager { args, conversationId, options?.runId, - this.createAgentDispatchCommit(toolName, 'yobrowser', args, options) + this.createAgentDispatchCommit(toolName, 'yobrowser', args, options), + options?.signal ) return { content: response diff --git a/src/main/tool/codeMode/runCodeRuntimeManager.ts b/src/main/tool/codeMode/runCodeRuntimeManager.ts index 60cb10a29..ecc527fd8 100644 --- a/src/main/tool/codeMode/runCodeRuntimeManager.ts +++ b/src/main/tool/codeMode/runCodeRuntimeManager.ts @@ -7,6 +7,7 @@ import type { MCPToolDefinition, MCPToolResponse, ToolDispatchCommitInput } from import type { ToolCallOptions } from '@shared/types/tool' import { RUN_CODE_MAX_NESTED_CALLS, + RUN_CODE_DEFAULT_TIMEOUT_MS, RUN_CODE_PROTOCOL_VERSION, RUN_CODE_SOURCE_MAX_BYTES, type RunCodeFrontend, @@ -16,6 +17,7 @@ import { } from '@shared/codeModeProtocol' import { buildCanonicalToolCatalog } from '@/agent/deepchat/runtime/toolSurface' import { MAX_EXECUTION_JOURNAL_NESTED_CHILDREN } from '@/tape/domain/executionJournal' +import { awaitWithAbort } from '@/lib/awaitWithAbort' import { normalizeCodexToolName } from './toolModeTools' type CodeModeUtilityProcess = Pick & { @@ -49,6 +51,7 @@ export interface RunCodeExecutionInput { toolCallId: string frontend: RunCodeFrontend source: string + timeoutMs?: number yieldTimeMs?: number maxOutputTokens?: number executionCatalog: readonly MCPToolDefinition[] @@ -120,6 +123,7 @@ type ActiveCell = { waiter: DeferredResult | null yieldTimeMs: number maxOutputTokens: number + timeoutMs: number yieldTimer: NodeJS.Timeout | null pausedAtYield: boolean pendingYieldOutput: unknown[] | null @@ -147,7 +151,6 @@ type ActiveCell = { const READY_TIMEOUT_MS = 5_000 const HEARTBEAT_TIMEOUT_MS = 3_500 -const CELL_EXECUTION_DEADLINE_MS = 5 * 60_000 const YIELD_LEASE_MS = 60_000 const MAX_JOURNALED_NESTED_CALLS = Math.min( RUN_CODE_MAX_NESTED_CALLS, @@ -316,7 +319,13 @@ export class RunCodeRuntimeManager { } input.options.commitDispatch?.(input.outerDispatch) - const host = await this.spawnReadyHost() + const host = await this.spawnReadyHost(input.options.signal) + if (input.options.signal?.aborted) { + try { + host.kill() + } catch {} + input.options.signal.throwIfAborted() + } const cell = this.createCell({ id: cellId, @@ -329,6 +338,7 @@ export class RunCodeRuntimeManager { capabilityHash: canonicalCatalog.fullCatalogHash, yieldTimeMs: input.yieldTimeMs ?? 10_000, maxOutputTokens: input.maxOutputTokens ?? 10_000, + timeoutMs: input.timeoutMs ?? RUN_CODE_DEFAULT_TIMEOUT_MS, options: input.options, committedDispatchToolCallId: input.toolCallId }) @@ -461,6 +471,7 @@ export class RunCodeRuntimeManager { bindings: ActiveCell['bindings'] yieldTimeMs: number maxOutputTokens: number + timeoutMs: number options: ToolCallOptions capabilityHash: string committedDispatchToolCallId: string | null @@ -516,7 +527,7 @@ export class RunCodeRuntimeManager { cell.rssTimer = setInterval(() => this.inspectRss(cell), 1_000) cell.executionDeadlineTimer = setTimeout( () => this.failAndCleanup(cell, new Error('Code cell execution deadline exceeded.')), - CELL_EXECUTION_DEADLINE_MS + cell.timeoutMs ) cell.executionDeadlineTimer.unref?.() @@ -759,6 +770,7 @@ export class RunCodeRuntimeManager { } private handleCellMessage(cell: ActiveCell, rawMessage: unknown): void { + if (cell.terminal || cell.state === 'stopping') return const message = unwrapMessage(rawMessage) if (!isHostMessage(message)) return if (message.type === 'READY') return @@ -1059,10 +1071,18 @@ export class RunCodeRuntimeManager { }) } - private async spawnReadyHost(): Promise { - const host = this.managerOptions.spawnHost - ? await this.managerOptions.spawnHost() - : await this.spawnDefaultHost() + private async spawnReadyHost(signal?: AbortSignal): Promise { + signal?.throwIfAborted() + const spawning = this.managerOptions.spawnHost + ? this.managerOptions.spawnHost() + : this.spawnDefaultHost() + let host: CodeModeUtilityProcess + try { + host = await awaitWithAbort(spawning, signal) + } catch (error) { + void spawning.then((abandonedHost) => abandonedHost.kill()).catch(() => undefined) + throw error + } return await new Promise((resolve, reject) => { let settled = false const settle = (callback: () => void) => { @@ -1072,6 +1092,7 @@ export class RunCodeRuntimeManager { host.off('message', onMessage) host.off('exit', onExit) host.off('error', onError) + signal?.removeEventListener('abort', onAbort) callback() } const onMessage = (rawMessage: unknown) => { @@ -1090,6 +1111,13 @@ export class RunCodeRuntimeManager { settle(() => reject(new Error(`Code mode utility failed before ready: ${type} at ${location}`)) ) + const onAbort = () => + settle(() => { + try { + host.kill() + } catch {} + reject(signal?.reason ?? new DOMException('Aborted', 'AbortError')) + }) const timeout = setTimeout(() => { settle(() => { try { @@ -1101,6 +1129,8 @@ export class RunCodeRuntimeManager { host.on('message', onMessage) host.on('exit', onExit) host.on('error', onError) + signal?.addEventListener('abort', onAbort, { once: true }) + if (signal?.aborted) onAbort() }) } diff --git a/src/main/tool/codeMode/toolModeTools.ts b/src/main/tool/codeMode/toolModeTools.ts index f6c316a6e..3c1bab06e 100644 --- a/src/main/tool/codeMode/toolModeTools.ts +++ b/src/main/tool/codeMode/toolModeTools.ts @@ -4,7 +4,11 @@ import { formatExecCommandDescription, type ResolvedCommandShell } from '@shared/commandShell' -import { CODE_MODE_TOOL_SERVER_NAME } from '@shared/codeModeProtocol' +import { + CODE_MODE_TOOL_SERVER_NAME, + RUN_CODE_DEFAULT_TIMEOUT_MS, + RUN_CODE_MAX_TIMEOUT_MS +} from '@shared/codeModeProtocol' import { LIVE_DELEGATION_AGENT_TOOL_NAME } from '@shared/agentTools' import { UPDATE_PLAN_TOOL_NAME } from '@shared/types/agent-plan' import { QUESTION_TOOL_NAME } from '../agentTools/questionTool' @@ -73,6 +77,9 @@ const RUN_CODE_DESCRIPTION = const RUN_CODE_DESCRIPTION_PARAM_DESCRIPTION = 'Clear, concise description of what this program does in active voice, 5-10 words (shown in the UI). Examples: "Count TODO markers across packages"; "Read failing test and its fixture"; "Rename config key in every cordis.yml".' +const CODE_MODE_TIMEOUT_DESCRIPTION = + "`timeout_ms` is an optional maximum wall-clock execution time for this code cell, including awaited subtools, in milliseconds. Defaults to 300000 ms (5 minutes). Omit it for routine work; provide a larger value only when you expect the current operation to take longer than 5 minutes. It does not extend a subtool's own timeout. User cancellation still stops the cell regardless of this value." + const CODE_MODE_EXEC_DESCRIPTION = `Run JavaScript code to orchestrate/compose tool calls - Evaluates the provided JavaScript code in a fresh V8 isolate as an async module. - Invoke SDK-declared subtools only inside this program on the global \`tools\` object through \`await tools.name(args)\`. Subtool names are exposed as normalized JavaScript identifiers, for example \`await tools.mcp__ologs__get_profile(...)\`. @@ -83,6 +90,7 @@ const CODE_MODE_EXEC_DESCRIPTION = `Run JavaScript code to orchestrate/compose t - You may optionally start the tool input with a first-line pragma like \`// @exec: {"yield_time_ms": 10000, "max_output_tokens": 1000}\`. - \`yield_time_ms\` asks \`exec\` to yield early if the script is still running. Defaults to 10000 ms. - \`max_output_tokens\` sets the token budget for direct \`exec\` results. Defaults to 10000 tokens. +- ${CODE_MODE_TIMEOUT_DESCRIPTION} Set it in the first-line pragma when needed. Yielding or calling \`wait\` does not reset or extend the cell's execution timeout. - When the JS code is fully evaluated, the isolate's lifetime ends and unawaited promises are silently discarded. - Global helpers: @@ -200,7 +208,7 @@ export function createRunCodeToolDefinition(): MCPToolDefinition { providerPresentation: { type: 'function' }, function: { name: RUN_CODE_TOOL_NAME, - description: RUN_CODE_DESCRIPTION, + description: `${RUN_CODE_DESCRIPTION} ${CODE_MODE_TIMEOUT_DESCRIPTION}`, parameters: { type: 'object', properties: { @@ -211,6 +219,13 @@ export function createRunCodeToolDefinition(): MCPToolDefinition { description: { type: 'string', description: RUN_CODE_DESCRIPTION_PARAM_DESCRIPTION + }, + timeout_ms: { + type: 'integer', + minimum: 1, + maximum: RUN_CODE_MAX_TIMEOUT_MS, + default: RUN_CODE_DEFAULT_TIMEOUT_MS, + description: CODE_MODE_TIMEOUT_DESCRIPTION } }, required: ['code', 'description'] @@ -500,6 +515,8 @@ ${toolBoundary} The top-level \`exec\` starts a Code Mode cell; \`tools.exec\` inside that cell runs the selected Shell command subtool. +${CODE_MODE_TIMEOUT_DESCRIPTION} Set it in the first-line \`// @exec: {...}\` pragma when needed. Yielding or calling \`wait\` does not reset or extend the cell's execution timeout. + ${progressPrompt} \`\`\`ts @@ -517,7 +534,11 @@ declare function yield_control(): Promise ${toolBoundary} -\`run_code\` takes two required arguments: \`code\` — the body of an async TypeScript function (erasable syntax only — no \`enum\` or namespaces; type annotations are advisory, the code runs type-stripped) — and \`description\`, a short summary of what the program does. Inside the program: +\`run_code\` takes two required arguments: \`code\` — the body of an async TypeScript function (erasable syntax only — no \`enum\` or namespaces; type annotations are advisory, the code runs type-stripped) — and \`description\`, a short summary of what the program does. + +${CODE_MODE_TIMEOUT_DESCRIPTION} + +Inside the program: - Call subtools as \`await tools.name(args)\` — quoted access for exotic names: \`tools["my-tool"](args)\`. Every call resolves to the subtool's typed canonical JSON value. Subtool arguments must be lossless JSON. - A FAILED subtool call rejects with \`ToolCallError\`, whose \`toolName\` identifies the failed subtool and whose \`message\` is human-readable — \`try/catch\` it to handle and continue. diff --git a/src/main/tool/index.ts b/src/main/tool/index.ts index 325c85044..0484a4066 100644 --- a/src/main/tool/index.ts +++ b/src/main/tool/index.ts @@ -140,7 +140,11 @@ import { isCodexToolFrontend, normalizeCodexToolName } from './codeMode/toolModeTools' -import { CODE_MODE_TOOL_SERVER_NAME } from '@shared/codeModeProtocol' +import { + CODE_MODE_TOOL_SERVER_NAME, + RUN_CODE_DEFAULT_TIMEOUT_MS, + RUN_CODE_MAX_TIMEOUT_MS +} from '@shared/codeModeProtocol' type McpToolPort = Pick< McpServicePort, @@ -233,6 +237,7 @@ type InternalToolCallOptions = ToolCallOptions & { type CodexExecInput = { source: string + timeoutMs: number yieldTimeMs: number maxOutputTokens: number } @@ -1269,11 +1274,12 @@ export class ToolService implements ToolServicePort { const waitInput = isWait ? this.requireCodeModeWaitInput(request.function.arguments) : null const execInput = isExec ? this.requireCodexExecInput(request.function.arguments) : null + const codeInput = isRunCode ? this.requireRunCodeInput(request.function.arguments) : execInput const outerDispatch: ToolDispatchCommitInput = { toolName, toolSource: 'agent', normalizedArguments: execInput - ? { source: execInput.source } + ? { source: execInput.source, timeout_ms: execInput.timeoutMs } : this.parseAgentToolArguments(request.function.arguments, toolName), target: { serverName: CODE_MODE_TOOL_SERVER_NAME, originalName: toolName } } @@ -1293,7 +1299,8 @@ export class ToolService implements ToolServicePort { runId: options?.runId, toolCallId: request.id, frontend: context.frontend, - source: execInput?.source ?? this.requireRunCodeSource(request.function.arguments), + source: codeInput!.source, + timeoutMs: codeInput!.timeoutMs, ...(execInput ? { yieldTimeMs: execInput.yieldTimeMs, @@ -1321,13 +1328,22 @@ export class ToolService implements ToolServicePort { } } - private requireRunCodeSource(argumentsText: string): string { + private requireRunCodeInput(argumentsText: string): { source: string; timeoutMs: number } { const args = this.parseAgentToolArguments(argumentsText) const code = typeof args.code === 'string' ? args.code : '' const description = typeof args.description === 'string' ? args.description.trim() : '' if (!code) throw new Error('run_code requires a non-empty code string.') if (!description) throw new Error('run_code requires a non-empty description string.') - return code + return { + source: code, + timeoutMs: this.readBoundedCodeModeNumber( + args.timeout_ms, + 'timeout_ms', + RUN_CODE_DEFAULT_TIMEOUT_MS, + 1, + RUN_CODE_MAX_TIMEOUT_MS + ) + } } private requireCodexExecInput(argumentsText: string): CodexExecInput { @@ -1336,6 +1352,7 @@ export class ToolService implements ToolServicePort { if (!argumentsText.trim()) throw new Error('exec requires non-empty JavaScript source.') return { source: argumentsText, + timeoutMs: RUN_CODE_DEFAULT_TIMEOUT_MS, yieldTimeMs: CODE_MODE_DEFAULT_YIELD_TIME_MS, maxOutputTokens: CODE_MODE_DEFAULT_MAX_OUTPUT_TOKENS } @@ -1354,6 +1371,13 @@ export class ToolService implements ToolServicePort { if (!source.trim()) throw new Error('exec requires non-empty JavaScript source.') return { source, + timeoutMs: this.readBoundedCodeModeNumber( + pragma.timeout_ms, + 'timeout_ms', + RUN_CODE_DEFAULT_TIMEOUT_MS, + 1, + RUN_CODE_MAX_TIMEOUT_MS + ), yieldTimeMs: this.readBoundedCodeModeNumber( pragma.yield_time_ms, 'yield_time_ms', diff --git a/src/main/tool/runtimePorts.ts b/src/main/tool/runtimePorts.ts index 38bb333ce..d08994e06 100644 --- a/src/main/tool/runtimePorts.ts +++ b/src/main/tool/runtimePorts.ts @@ -215,7 +215,8 @@ export interface AgentBrowserToolPort { args: Record, conversationId?: string, runId?: string, - beforeInvoke?: (normalizedArguments: Record) => void + beforeInvoke?: (normalizedArguments: Record) => void, + signal?: AbortSignal ): Promise } diff --git a/src/shared/codeModeProtocol.ts b/src/shared/codeModeProtocol.ts index f5287bf44..21187e902 100644 --- a/src/shared/codeModeProtocol.ts +++ b/src/shared/codeModeProtocol.ts @@ -4,6 +4,8 @@ export const RUN_CODE_SOURCE_MAX_BYTES = 256 * 1024 export const RUN_CODE_OUTPUT_MAX_BYTES = 1024 * 1024 export const RUN_CODE_MAX_NESTED_CALLS = 128 export const RUN_CODE_MAX_NESTED_CONCURRENCY = 8 +export const RUN_CODE_DEFAULT_TIMEOUT_MS = 5 * 60_000 +export const RUN_CODE_MAX_TIMEOUT_MS = 2_147_483_647 export type RunCodeFrontend = 'codex' | 'function' diff --git a/test/main/agent/deepchat/loop/contextCoordinator.test.ts b/test/main/agent/deepchat/loop/contextCoordinator.test.ts index 8694aceea..e8651eb3b 100644 --- a/test/main/agent/deepchat/loop/contextCoordinator.test.ts +++ b/test/main/agent/deepchat/loop/contextCoordinator.test.ts @@ -3431,6 +3431,68 @@ describe('DeepChatContextCoordinator', () => { }) }) + it.each(['output', 'rejection'])( + 'cancels a silent provider without awaiting cleanup or accepting late %s', + async (lateResult) => { + vi.useFakeTimers() + const fixture = createAttemptInput() + const usageEvent: LLMCoreStreamEvent = { + type: 'usage', + usage: { prompt_tokens: 8, completion_tokens: 1, total_tokens: 9 } + } + let release!: () => void + const pending = new Promise((resolve) => { + release = resolve + }) + const waiting = vi.fn() + const cleanup = vi.fn() + fixture.input.provider.stream = async function* () { + try { + yield usageEvent + waiting() + await pending + if (lateResult === 'rejection') throw new Error('Late provider failure') + yield { type: 'tool_call_start', tool_call_id: 'late-call', tool_call_name: 'exec' } + } finally { + cleanup() + } + } + const projected: LLMCoreStreamEvent[] = [] + const execution = (async () => { + for await (const event of new DeepChatContextCoordinator().streamProviderAttempts( + fixture.input + )) { + projected.push(event) + } + })() + const cancellation = expect(execution).rejects.toMatchObject({ name: 'AbortError' }) + + await vi.advanceTimersByTimeAsync(15 * 60_000) + expect(waiting).toHaveBeenCalledOnce() + expect(fixture.outcomes).toEqual([]) + fixture.run.abortController.abort(new DOMException('Stopped', 'AbortError')) + await cancellation + + expect(projected).toEqual([usageEvent]) + expect(fixture.outcomes).toEqual([ + expectedAttemptOutcome({ + status: 'aborted', + stopReason: null, + failureClassification: 'aborted', + retryDecision: 'not_retryable', + usage: { inputTokens: 8, outputTokens: 1, totalTokens: 9 } + }) + ]) + expect(cleanup).not.toHaveBeenCalled() + + release() + await vi.advanceTimersByTimeAsync(0) + expect(cleanup).toHaveBeenCalledOnce() + expect(projected).toEqual([usageEvent]) + expect(fixture.outcomes).toHaveLength(1) + } + ) + it('settles an attempt as aborted when a provider ignores cancellation and stops cleanly', async () => { const fixture = createAttemptInput({ providerEvents: [[{ type: 'stop', stop_reason: 'complete' }]] diff --git a/test/main/desktop/browser/BrowserTab.test.ts b/test/main/desktop/browser/BrowserTab.test.ts index 6130ffc98..4460c85a6 100644 --- a/test/main/desktop/browser/BrowserTab.test.ts +++ b/test/main/desktop/browser/BrowserTab.test.ts @@ -174,6 +174,125 @@ describe('BrowserTab', () => { webContents.finishLoad() }) + it('cancels readiness waits without leaking listeners or dispatching after dom-ready', async () => { + const { tab, webContents, cdpManager } = createTab() + await makePageInteractive(tab, webContents) + webContents.emitStartNavigation('https://example.com/next') + const initialListeners = webContents.listenerCount('dom-ready') + const controller = new AbortController() + const beforeDispatch = vi.fn() + const command = tab.sendCdpCommand( + 'Runtime.evaluate', + undefined, + beforeDispatch, + undefined, + controller.signal + ) + + controller.abort() + await expect(command).rejects.toMatchObject({ name: 'AbortError' }) + expect(webContents.listenerCount('dom-ready')).toBe(initialListeners) + await vi.advanceTimersByTimeAsync(2000) + webContents.emitDomReady() + await vi.advanceTimersByTimeAsync(0) + + expect(beforeDispatch).not.toHaveBeenCalled() + expect(cdpManager.createSession).not.toHaveBeenCalled() + expect(cdpManager.evaluateScript).not.toHaveBeenCalled() + expect(webContents.debugger.sendCommand).not.toHaveBeenCalled() + }) + + it('does not dispatch a command after canceled CDP session preparation finishes', async () => { + const { tab, webContents, cdpManager } = createTab() + const controller = new AbortController() + let finishSession!: () => void + cdpManager.createSession.mockImplementationOnce( + () => + new Promise((resolve) => { + finishSession = resolve + }) + ) + const command = tab.sendCdpCommand( + 'Page.reload', + undefined, + undefined, + undefined, + controller.signal + ) + + controller.abort() + await expect(command).rejects.toMatchObject({ name: 'AbortError' }) + finishSession() + await vi.advanceTimersByTimeAsync(0) + + expect(webContents.debugger.sendCommand).not.toHaveBeenCalled() + }) + + it('cancels a silent CDP command and discards its late navigation response', async () => { + const { tab, webContents } = createTab() + await makePageInteractive(tab, webContents) + const controller = new AbortController() + let finishCommand!: (result: { loaderId: string }) => void + webContents.debugger.sendCommand.mockImplementationOnce( + () => + new Promise<{ loaderId: string }>((resolve) => { + finishCommand = resolve + }) + ) + const command = tab.sendCdpCommand( + 'Page.navigate', + { url: 'https://example.com/next' }, + undefined, + undefined, + controller.signal + ) + const failed = vi.fn() + void command.catch(failed) + await vi.advanceTimersByTimeAsync(10 * 60_000) + expect(webContents.debugger.sendCommand).toHaveBeenCalledOnce() + expect(failed).not.toHaveBeenCalled() + + controller.abort() + await expect(command).rejects.toMatchObject({ name: 'AbortError' }) + finishCommand({ loaderId: 'late-loader' }) + await vi.advanceTimersByTimeAsync(0) + + expect(tab.url).toBe('https://example.com') + expect(tab.status).toBe(BrowserPageStatus.Ready) + }) + + it('cancels a pending readiness probe without dispatching the requested command', async () => { + const { tab, webContents, cdpManager } = createTab() + await makePageInteractive(tab, webContents) + webContents.emitStartNavigation('https://example.com/next') + const controller = new AbortController() + let finishProbe!: (result: any) => void + cdpManager.evaluateScript.mockImplementationOnce( + () => + new Promise((resolve) => { + finishProbe = resolve + }) + ) + const command = tab.sendCdpCommand( + 'Runtime.evaluate', + undefined, + undefined, + undefined, + controller.signal + ) + await vi.advanceTimersByTimeAsync(2000) + expect(cdpManager.evaluateScript).toHaveBeenCalledOnce() + + controller.abort() + await expect(command).rejects.toMatchObject({ name: 'AbortError' }) + finishProbe({ readyState: 'complete', hasBody: true, href: 'https://example.com/late' }) + await vi.advanceTimersByTimeAsync(0) + + expect(webContents.debugger.sendCommand).not.toHaveBeenCalled() + expect(tab.url).toBe('https://example.com/next') + expect(tab.status).toBe(BrowserPageStatus.Loading) + }) + it('still fails immediately when the page is not loading', async () => { const { tab, webContents } = createTab() diff --git a/test/main/desktop/browser/YoBrowserPresenter.test.ts b/test/main/desktop/browser/YoBrowserPresenter.test.ts index b39a4f491..e4e27dc87 100644 --- a/test/main/desktop/browser/YoBrowserPresenter.test.ts +++ b/test/main/desktop/browser/YoBrowserPresenter.test.ts @@ -1059,6 +1059,68 @@ describe('YoBrowserPresenter', () => { ) }) + it.each(['load_url', 'cdp_send'])( + 'stops activity for a canceled %s and ignores late completion', + async (toolName) => { + const { presenter, windows, getSessionWebContents } = await setupPresenter() + windows.set(1, new MockBrowserWindow(1)) + const initialLoad = presenter.loadUrl('session-a', 'https://example.com') + const webContents = getSessionWebContents('session-a')! + webContents.emitDomReady() + await initialLoad + webContents.finishLoad() + sendToAllWindowsMock.mockClear() + + let finishCommand!: () => void + if (toolName === 'cdp_send') { + webContents.debugger.sendCommand.mockImplementation(async (method: string) => { + if (method === 'Runtime.evaluate') { + await new Promise((resolve) => { + finishCommand = resolve + }) + } + return {} + }) + } + const controller = new AbortController() + const command = presenter.toolHandler.callTool( + toolName, + toolName === 'load_url' + ? { url: 'https://example.com/next' } + : { + method: 'Runtime.evaluate', + params: { + expression: 'new Promise(() => { document.querySelector("button").click() })', + awaitPromise: true + } + }, + 'session-a', + 'run-a', + undefined, + controller.signal + ) + const cancellation = expect(command).rejects.toMatchObject({ name: 'AbortError' }) + const phases = () => + sendToAllWindowsMock.mock.calls + .map(([, envelope]) => envelope) + .filter((envelope: any) => envelope.name === 'browser.activity.changed') + .map((envelope: any) => envelope.payload.phase) + await vi.advanceTimersByTimeAsync(0) + expect(phases()).toEqual(['started']) + + controller.abort() + await cancellation + await vi.advanceTimersByTimeAsync(0) + expect(phases()).toEqual(['started', 'failed']) + + if (toolName === 'cdp_send') finishCommand() + else webContents.finishLoad() + await vi.advanceTimersByTimeAsync(0) + expect(phases()).toEqual(['started', 'failed']) + await presenter.shutdown() + } + ) + it('maps agent CDP mouse and screenshot commands to overlay activity', async () => { const { presenter, windows, getSessionWebContents } = await setupPresenter() windows.set(1, new MockBrowserWindow(1)) diff --git a/test/main/desktop/browser/YoBrowserToolHandler.test.ts b/test/main/desktop/browser/YoBrowserToolHandler.test.ts index 583345bee..5fbdf110b 100644 --- a/test/main/desktop/browser/YoBrowserToolHandler.test.ts +++ b/test/main/desktop/browser/YoBrowserToolHandler.test.ts @@ -91,7 +91,8 @@ describe('YoBrowserToolHandler', () => { undefined, 'agent', 'run-a', - expect.any(Function) + expect.any(Function), + undefined ) }) @@ -174,10 +175,39 @@ describe('YoBrowserToolHandler', () => { { type: 'mousePressed', x: 24, y: 48 }, 'agent', 'run-a', - expect.any(Function) + expect.any(Function), + undefined ) }) + it('cancels browser status resolution without dispatching a delayed CDP command', async () => { + const presenter = createPresenter() + const controller = new AbortController() + let finishStatus!: (status: typeof readyStatus) => void + presenter.getBrowserStatus.mockImplementationOnce( + () => + new Promise((resolve) => { + finishStatus = resolve + }) + ) + const handler = new YoBrowserToolHandler(presenter) + const command = handler.callTool( + 'cdp_send', + { method: 'Runtime.evaluate' }, + 'session-a', + 'run-a', + undefined, + controller.signal + ) + + controller.abort() + await expect(command).rejects.toMatchObject({ name: 'AbortError' }) + finishStatus(readyStatus) + await new Promise((resolve) => setTimeout(resolve, 0)) + + expect(presenter.sendCdpCommand).not.toHaveBeenCalled() + }) + it('rejects old tool names as unknown tools', async () => { const handler = new YoBrowserToolHandler(createPresenter()) diff --git a/test/main/tool/agentTools/agentToolManagerYoBrowser.test.ts b/test/main/tool/agentTools/agentToolManagerYoBrowser.test.ts index 8255d2ccc..3bd372c8c 100644 --- a/test/main/tool/agentTools/agentToolManagerYoBrowser.test.ts +++ b/test/main/tool/agentTools/agentToolManagerYoBrowser.test.ts @@ -65,6 +65,24 @@ describe('AgentToolManager YoBrowser routing', () => { }) }) + it('cancels a pending browser tool through the Run signal', async () => { + const controller = new AbortController() + yoBrowserCallTool.mockImplementation( + (_name, _args, _sessionId, _runId, _beforeInvoke, signal?: AbortSignal) => + new Promise((_resolve, reject) => { + signal?.addEventListener('abort', () => reject(signal.reason), { once: true }) + }) + ) + const command = manager.callTool('cdp_send', { method: 'Runtime.evaluate' }, 'session-a', { + signal: controller.signal + }) + await vi.waitFor(() => expect(yoBrowserCallTool).toHaveBeenCalledOnce()) + + controller.abort() + + await expect(command).rejects.toMatchObject({ name: 'AbortError' }) + }) + it('returns recoverable YoBrowser CDP failures as errored structured tool results', async () => { const browserStatus = { initialized: false, diff --git a/test/main/tool/codeMode/runCodeRuntimeManager.test.ts b/test/main/tool/codeMode/runCodeRuntimeManager.test.ts index 91eb12aa2..0cdd2ee89 100644 --- a/test/main/tool/codeMode/runCodeRuntimeManager.test.ts +++ b/test/main/tool/codeMode/runCodeRuntimeManager.test.ts @@ -172,7 +172,10 @@ describe('RunCodeRuntimeManager', () => { await manager.shutdown() }) - it('expires unresolved async execution despite healthy heartbeats', async () => { + it.each([ + { timeoutMs: undefined, deadlineMs: 300000 }, + { timeoutMs: 600000, deadlineMs: 600000 } + ])('expires at $deadlineMs ms despite healthy heartbeats', async ({ timeoutMs, deadlineMs }) => { vi.useFakeTimers() try { const host = new FakeUtilityProcess(false) @@ -183,6 +186,7 @@ describe('RunCodeRuntimeManager', () => { toolCallId: 'call-1', frontend: 'function', source: 'await new Promise(() => {})', + timeoutMs, executionCatalog: [tool('exec')], outerDispatch: outerDispatch(), options: {} @@ -204,7 +208,9 @@ describe('RunCodeRuntimeManager', () => { now: Date.now() }) }, 1_000) - await vi.advanceTimersByTimeAsync(5 * 60_000 + 1) + await vi.advanceTimersByTimeAsync(deadlineMs - 1) + expect(host.messages.some((message) => message.type === 'STOP')).toBe(false) + await vi.advanceTimersByTimeAsync(1) clearInterval(heartbeat) await executionFailure @@ -221,6 +227,91 @@ describe('RunCodeRuntimeManager', () => { } }) + it.each(['spawn', 'ready', 'ready race'])( + 'cancels utility startup during %s and reclaims the process without running code', + async (stage) => { + const host = new FakeUtilityProcess(false) + const controller = new AbortController() + let resolveHost!: (host: FakeUtilityProcess) => void + const spawning = new Promise((resolve) => { + resolveHost = resolve + }) + const manager = new RunCodeRuntimeManager({ + spawnHost: () => spawning, + executeNested: vi.fn() + }) + const execution = manager.execute({ + sessionId: 'session-1', + toolCallId: 'call-1', + frontend: 'function', + source: 'return true', + executionCatalog: [], + outerDispatch: outerDispatch(), + options: { signal: controller.signal } + }) + const cancellation = expect(execution).rejects.toMatchObject({ name: 'AbortError' }) + if (stage !== 'spawn') { + resolveHost(host) + await vi.waitFor(() => expect(host.listenerCount('message')).toBe(1)) + } + if (stage === 'ready race') host.becomeReady() + + controller.abort() + if (stage === 'spawn') resolveHost(host) + await cancellation + await vi.waitFor(() => expect(host.kill).toHaveBeenCalledOnce()) + + expect(host.messages).toEqual([]) + expect(host.listenerCount('message')).toBe(0) + expect(host.listenerCount('exit')).toBe(0) + expect(host.listenerCount('error')).toBe(0) + await manager.shutdown() + } + ) + + it('discards a late utility result after cancellation instead of changing the session store', async () => { + const host = new FakeUtilityProcess(false) + const nextHost = new FakeUtilityProcess(true) + const hosts = [host, nextHost] + const manager = new RunCodeRuntimeManager({ + spawnHost: async () => { + const current = hosts.shift()! + setTimeout(() => current.becomeReady(), 0) + return current + }, + executeNested: vi.fn() + }) + const controller = new AbortController() + const input = { + sessionId: 'session-1', + toolCallId: 'call-1', + frontend: 'function' as const, + source: 'return true', + executionCatalog: [], + outerDispatch: outerDispatch(), + options: { signal: controller.signal } + } + const execution = manager.execute(input) + const cancellation = expect(execution).rejects.toMatchObject({ name: 'AbortError' }) + await vi.waitFor(() => expect(host.messages[0]?.type).toBe('START')) + const start = host.messages[0] as Extract + + controller.abort() + host.emit('message', { + type: 'RESULT', + version: RUN_CODE_PROTOCOL_VERSION, + cellId: start.cellId, + output: [], + returnValue: 'late result', + store: { canceled: true } + }) + await cancellation + await manager.execute({ ...input, toolCallId: 'call-2', options: {} }) + + expect(nextHost.messages[0]).toMatchObject({ type: 'START', store: {} }) + await manager.shutdown() + }) + it('normalizes Codex binding names and serializes mutating tools', async () => { const host = new FakeUtilityProcess(true) const manager = createManager(host) @@ -808,51 +899,60 @@ describe('RunCodeRuntimeManager', () => { } }) - it('aborts an in-flight nested call when its session is cancelled', async () => { - const host = new FakeUtilityProcess(false) - let nestedSignal: AbortSignal | undefined - const executeNested = vi.fn( - async (input: RunCodeNestedExecutionInput) => - await new Promise((_resolve, reject) => { - nestedSignal = input.options.signal - input.options.signal?.addEventListener( - 'abort', - () => reject(input.options.signal?.reason), - { once: true } - ) + it.each(['session', 'run signal'])( + 'cancels an extended cell through %s even when a nested call ignores cancellation', + async (source) => { + const host = new FakeUtilityProcess(false) + let nestedSignal: AbortSignal | undefined + const controller = new AbortController() + let finishNested!: () => void + const executeNested = vi.fn(async (input: RunCodeNestedExecutionInput) => { + nestedSignal = input.options.signal + await new Promise((resolve) => { + finishNested = resolve }) - ) - const manager = createManager(host, executeNested) - const execution = manager.execute({ - sessionId: 'session-1', - toolCallId: 'call-1', - frontend: 'function', - source: 'return await tools.exec({ command: "pwd" })', - executionCatalog: [tool('exec')], - outerDispatch: outerDispatch(), - options: {} - }) - await vi.waitFor(() => - expect(host.messages.some((message) => message.type === 'START')).toBe(true) - ) - const start = host.messages.find( - (message): message is Extract => - message.type === 'START' - )! - host.emit('message', { - type: 'NESTED_CALL', - version: RUN_CODE_PROTOCOL_VERSION, - cellId: start.cellId, - callId: 'nested-1', - bindingId: start.bindings[0].id, - arguments: { command: 'pwd' } - }) - await vi.waitFor(() => expect(executeNested).toHaveBeenCalledOnce()) + return { content: 'late result', rawData: { content: 'late result' } } + }) + const manager = createManager(host, executeNested) + const execution = manager.execute({ + sessionId: 'session-1', + toolCallId: 'call-1', + frontend: 'function', + source: 'return await tools.exec({ command: "pwd" })', + timeoutMs: 600000, + executionCatalog: [tool('exec')], + outerDispatch: outerDispatch(), + options: { signal: controller.signal } + }) + const cancellation = expect(execution).rejects.toThrow( + source === 'session' ? 'test cancellation' : 'Aborted' + ) + await vi.waitFor(() => + expect(host.messages.some((message) => message.type === 'START')).toBe(true) + ) + const start = host.messages.find( + (message): message is Extract => + message.type === 'START' + )! + host.emit('message', { + type: 'NESTED_CALL', + version: RUN_CODE_PROTOCOL_VERSION, + cellId: start.cellId, + callId: 'nested-1', + bindingId: start.bindings[0].id, + arguments: { command: 'pwd' } + }) + await vi.waitFor(() => expect(executeNested).toHaveBeenCalledOnce()) - manager.cancelSession('session-1', 'test cancellation') + if (source === 'session') manager.cancelSession('session-1', 'test cancellation') + else controller.abort() - await expect(execution).rejects.toThrow('test cancellation') - expect(nestedSignal?.aborted).toBe(true) - await manager.shutdown() - }) + await cancellation + expect(nestedSignal?.aborted).toBe(true) + finishNested() + await new Promise((resolve) => setTimeout(resolve, 0)) + expect(host.messages.some((message) => message.type === 'NESTED_RESULT')).toBe(false) + await manager.shutdown() + } + ) }) diff --git a/test/main/tool/codeMode/toolModeTools.test.ts b/test/main/tool/codeMode/toolModeTools.test.ts index e0caf5647..ed99f332f 100644 --- a/test/main/tool/codeMode/toolModeTools.test.ts +++ b/test/main/tool/codeMode/toolModeTools.test.ts @@ -117,7 +117,8 @@ describe('Tool Mode provider contracts', () => { required: ['code', 'description'], properties: { code: { type: 'string' }, - description: { type: 'string' } + description: { type: 'string' }, + timeout_ms: { type: 'integer', minimum: 1, maximum: 2147483647, default: 300000 } } }) expect(runCode.function.description).toContain('`run_code` is the only code entrypoint') @@ -135,6 +136,24 @@ describe('Tool Mode provider contracts', () => { expect(sdk).toContain('Independent read-only calls MAY overlap under `Promise.all`') }) + it('explains timeout omission, extension, and cancellation in both frontends', () => { + const runCode = createRunCodeToolDefinition() + const prompts = [ + runCode.function.description, + runCode.function.parameters.properties.timeout_ms.description, + createCodexCodeModeToolDefinitions([])[0].function.description, + renderCodeModeSdk('function', []), + renderCodeModeSdk('codex', []) + ] + + for (const prompt of prompts) { + expect(prompt).toContain('300000 ms (5 minutes)') + expect(prompt).toContain('Omit it for routine work') + expect(prompt).toContain('longer than 5 minutes') + expect(prompt).toContain('User cancellation still stops the cell') + } + }) + it('omits direct Loop tools from generated SDK declarations', () => { const question = { ...nestedTool, diff --git a/test/main/tool/toolService.test.ts b/test/main/tool/toolService.test.ts index f13be439a..c49c8e812 100644 --- a/test/main/tool/toolService.test.ts +++ b/test/main/tool/toolService.test.ts @@ -1,6 +1,8 @@ import { describe, expect, it, vi } from 'vitest' import { TOOL_EXECUTION, type MCPToolDefinition } from '@shared/types/mcp' import { ToolService } from '@/tool' +import { RunCodeRuntimeManager } from '@/tool/codeMode/runCodeRuntimeManager' +import { POSIX_COMMAND_SHELL } from '../../helpers/commandShell' import { createToolCatalogPort } from '@/agent/deepchat/runtime/toolAdapters' import { AgentToolManager, @@ -1256,6 +1258,91 @@ describe('ToolService', () => { active.batch.discard() }) + describe.each([ + { providerId: 'deepseek', toolName: 'run_code' }, + { providerId: 'openai-codex', toolName: 'exec' } + ])('$toolName execution timeout', ({ providerId, toolName }) => { + const createService = () => { + const service = new ToolService({ + skillSettings: { isEnabled: () => false } as any, + mcpService: { getAllToolDefinitions: vi.fn().mockResolvedValue([]) } as any, + agentSettings: { resolveDeepChatAgentConfig: vi.fn(async () => ({})) } as any, + providerSettings: { getModelConfig: vi.fn() } as any, + settings: { get: vi.fn() }, + commandPermissionHandler: new CommandPermissionService(), + agentTools: buildAgentToolRuntimeMock() + }) + service.configureToolMode({ + conversationId: 'session-1', + mode: 'code', + providerId, + commandShell: POSIX_COMMAND_SHELL, + executionCatalog: [] + }) + return service + } + const request = (timeoutMs: unknown) => ({ + id: 'code-call', + type: 'function' as const, + function: { + name: toolName, + arguments: + toolName === 'run_code' + ? JSON.stringify({ + code: 'return true', + description: 'Return true', + timeout_ms: timeoutMs + }) + : timeoutMs === undefined + ? 'return true' + : `// @exec: ${JSON.stringify({ timeout_ms: timeoutMs })}\nreturn true` + }, + conversationId: 'session-1' + }) + + it.each([ + { timeoutMs: undefined, expected: 300000 }, + { timeoutMs: 900000, expected: 900000 }, + { timeoutMs: 1, expected: 1 }, + { timeoutMs: 2147483647, expected: 2147483647 } + ])('uses $expected ms when timeout_ms is $timeoutMs', async ({ timeoutMs, expected }) => { + const execute = vi.spyOn(RunCodeRuntimeManager.prototype, 'execute').mockResolvedValue({ + content: 'completed' + }) + try { + await createService().callTool(request(timeoutMs), { permissionMode: 'full_access' }) + + expect(execute).toHaveBeenCalledOnce() + expect(execute).toHaveBeenCalledWith( + expect.objectContaining({ source: 'return true', timeoutMs: expected }) + ) + } finally { + execute.mockRestore() + } + }) + + it.each([0, -1, 1.5, '900000', null, 2147483648])( + 'rejects invalid timeout_ms %j before starting code', + async (timeoutMs) => { + const execute = vi.spyOn(RunCodeRuntimeManager.prototype, 'execute') + const commitDispatch = vi.fn() + try { + await expect( + createService().callTool(request(timeoutMs), { + permissionMode: 'full_access', + commitDispatch + }) + ).rejects.toThrow('timeout_ms must be') + + expect(execute).not.toHaveBeenCalled() + expect(commitDispatch).not.toHaveBeenCalled() + } finally { + execute.mockRestore() + } + } + ) + }) + it('projects one execution catalog into Agent, Code, and Minimal modes', async () => { const toolService = new ToolService({ skillSettings: { isEnabled: () => false } as any,