feat(tools): require cancellation signal on every invocation

This commit is contained in:
Tianyi Cui
2026-07-19 23:38:54 +08:00
parent a99750f341
commit e8b95c8754
77 changed files with 1129 additions and 446 deletions

View File

@@ -54,8 +54,7 @@ export function apply(ctx: Context): void {
using d = deadline(exec.signal, timeoutMs, TOOL_TIMEOUT)
// Swap the derived deadline onto exec for dispatch, then restore the
// caller's own signal so post-execute listeners never see this plugin's
// (possibly already-aborted) timeout signal. `undefined` is not assignable to
// the optional `signal` under exactOptionalPropertyTypes, so branch on it.
// (possibly already-aborted) timeout signal.
const upstream = exec.signal
exec.signal = d.signal
try {
@@ -69,8 +68,7 @@ export function apply(ctx: Context): void {
}
return result
} finally {
if (upstream === undefined) delete exec.signal
else exec.signal = upstream
exec.signal = upstream
}
})
}

View File

@@ -11,10 +11,12 @@ import { Context } from 'cordis'
import Loader from '@cordisjs/plugin-loader'
import { CallId, HarnessError } from '@deepseek-ai/dsh-llm'
import SystemPrompt from '@deepseek-ai/dsh-system-prompt'
import ToolRegistry, { defineTool, type ToolExecutionInput, type PostToolDecision } from '@deepseek-ai/dsh-tools'
import ToolRegistry, { defineTool, TOOL_ABORTED, type ToolExecutionInput, type PostToolDecision } from '@deepseek-ai/dsh-tools'
import * as timeoutPolicy from '@deepseek-ai/dsh-timeout-policy'
import { TOOL_TIMEOUT } from '@deepseek-ai/dsh-timeout-policy'
const testToolSignal = new AbortController().signal
/** Mount the registry + the zero-config timeout-policy enforcer. */
async function setup() {
const ctx = new Context()
@@ -29,8 +31,8 @@ const cooperativeTool = defineTool({
name: 'slow', description: 'stops when aborted', parameters: {}, timeoutMs: 100,
execute(_args, exec): Promise<{ type: 'text'; text: string }[]> {
const done = [{ type: 'text' as const, text: 'stopped cooperatively' }]
if (exec.signal?.aborted) return Promise.resolve(done)
return new Promise((resolve) => { exec.signal?.addEventListener('abort', () => { resolve(done) }) })
if (exec.signal.aborted) return Promise.resolve(done)
return new Promise((resolve) => { exec.signal.addEventListener('abort', () => { resolve(done) }) })
},
})
@@ -38,8 +40,8 @@ const cooperativeTool = defineTool({
const abortThrowingTool = defineTool({
name: 'aborter', description: 'throws WEB_ABORTED when aborted', parameters: {}, timeoutMs: 100,
execute(_args, exec): Promise<never> {
if (exec.signal?.aborted) return Promise.reject(new HarnessError('web fetch aborted', 'WEB_ABORTED'))
return new Promise((_resolve, reject) => { exec.signal?.addEventListener('abort', () => { reject(new HarnessError('web fetch aborted', 'WEB_ABORTED')) }) })
if (exec.signal.aborted) return Promise.reject(new HarnessError('web fetch aborted', 'WEB_ABORTED'))
return new Promise((_resolve, reject) => { exec.signal.addEventListener('abort', () => { reject(new HarnessError('web fetch aborted', 'WEB_ABORTED')) }) })
},
})
@@ -59,7 +61,7 @@ describe('timeout-policy delegation (unconfigured / fast)', () => {
const ctx = await setup()
ctx.tools.register(defineTool({ name: 'fast', description: 'd', parameters: {}, timeoutMs: 10_000,
async execute() { return [{ type: 'text' as const, text: 'ok' }] } }))
const result = await ctx.tools.execute({ callId: CallId('c1'), name: 'fast', arguments: {} })
const result = await ctx.tools.execute({ signal: testToolSignal, callId: CallId('c1'), name: 'fast', arguments: {} })
expect(result).toEqual({ content: [{ type: 'text', text: 'ok' }], isError: false })
})
@@ -86,16 +88,6 @@ describe('timeout-policy signal restoration', () => {
await ctx.tools.execute({ callId: CallId('c1'), name: 'fast', arguments: {}, signal: upstream })
expect(postSignal).toBe(upstream)
})
it('deletes exec.signal again when the caller passed none', async () => {
const ctx = await setup()
ctx.tools.register(defineTool({ name: 'fast', description: 'd', parameters: {}, timeoutMs: 10_000,
async execute() { return [{ type: 'text' as const, text: 'ok' }] } }))
let hadSignal: boolean | undefined
ctx.on('tools/post-execute', async (exec, _result, next): Promise<PostToolDecision> => { hadSignal = 'signal' in exec && exec.signal !== undefined; return next() })
await ctx.tools.execute({ callId: CallId('c1'), name: 'fast', arguments: {} })
expect(hadSignal).toBe(false)
})
})
describe('timeout-policy TOOL_TIMEOUT replacement (deadline wins)', () => {
@@ -105,7 +97,7 @@ describe('timeout-policy TOOL_TIMEOUT replacement (deadline wins)', () => {
it('replaces a cooperative tool result with TOOL_TIMEOUT when its own deadline fires', async () => {
const ctx = await setup()
ctx.tools.register(cooperativeTool)
const pending = ctx.tools.execute({ callId: CallId('c1'), name: 'slow', arguments: {} })
const pending = ctx.tools.execute({ signal: testToolSignal, callId: CallId('c1'), name: 'slow', arguments: {} })
await vi.advanceTimersByTimeAsync(150)
const result = await pending
expect(result).toEqual({
@@ -118,7 +110,7 @@ describe('timeout-policy TOOL_TIMEOUT replacement (deadline wins)', () => {
it('replaces a provider-owned abort ERROR result with TOOL_TIMEOUT when the signal was ours', async () => {
const ctx = await setup()
ctx.tools.register(abortThrowingTool)
const pending = ctx.tools.execute({ callId: CallId('c1'), name: 'aborter', arguments: {} })
const pending = ctx.tools.execute({ signal: testToolSignal, callId: CallId('c1'), name: 'aborter', arguments: {} })
await vi.advanceTimersByTimeAsync(150)
const result = await pending
expect(result.isError).toBe(true)
@@ -128,14 +120,26 @@ describe('timeout-policy TOOL_TIMEOUT replacement (deadline wins)', () => {
it('preserves registry ABORTED when the caller aborts first (upstream cancel, not our timeout)', async () => {
const ctx = await setup()
ctx.tools.register(cooperativeTool)
const entered = Promise.withResolvers<undefined>()
ctx.tools.register(defineTool({
name: 'slow', description: 'stops when aborted', parameters: {}, timeoutMs: 100,
execute(_args, exec) {
entered.resolve(undefined)
const done = [{ type: 'text' as const, text: 'stopped cooperatively' }]
if (exec.signal.aborted) return Promise.resolve(done)
return new Promise((resolve) => {
exec.signal.addEventListener('abort', () => { resolve(done) }, { once: true })
})
},
}))
const upstream = new AbortController()
const pending = ctx.tools.execute({ callId: CallId('c1'), name: 'slow', arguments: {}, signal: upstream.signal })
await entered.promise
upstream.abort('user cancelled')
await vi.advanceTimersByTimeAsync(0)
const result = await pending
expect(result.isError).toBe(true)
expect(result.error).toEqual({ name: 'AbortError', code: 'ABORTED' })
expect(result.error).toEqual({ name: 'AbortError', code: TOOL_ABORTED })
expect(result.content[0]).toMatchObject({ text: 'Error: tool call aborted' })
})
@@ -146,9 +150,9 @@ describe('timeout-policy TOOL_TIMEOUT replacement (deadline wins)', () => {
ctx.tools.register(defineTool({
name: 'slow-cleanup', description: 'settles after abort cleanup', parameters: {}, timeoutMs: 100,
async execute(_args, exec) {
if (!exec.signal?.aborted) {
if (!exec.signal.aborted) {
await new Promise<undefined>((resolve) => {
exec.signal?.addEventListener('abort', () => { resolve(undefined) }, { once: true })
exec.signal.addEventListener('abort', () => { resolve(undefined) }, { once: true })
})
}
sawAbort.resolve(undefined)
@@ -217,7 +221,7 @@ describe('dsh-timeout-policy real-load-path guard', () => {
const loader = Object.create(Loader.prototype) as Loader
const unwrapped = loader.unwrapExports(timeoutPolicy) as Parameters<Context['plugin']>[0]
const fiber = await ctx.plugin(unwrapped)
const result = await ctx.tools.execute({ callId: CallId('c1'), name: 'fast', arguments: {} } satisfies ToolExecutionInput)
const result = await ctx.tools.execute({ signal: testToolSignal, callId: CallId('c1'), name: 'fast', arguments: {} } satisfies ToolExecutionInput)
expect(result.isError).toBe(false)
await fiber.dispose()
})