fix(tools): single ordered driver lane for the sub-dispatch scheduler; validate the cap
Responding to ds-review-bot round 2 on #658 (three critical findings, one warning — all rooted in the pump/commit split racing ordered stages): - ONE driver lane now owns every ordered stage: the start append, prepare (pre-execute/guards), and the head-of-line commit (post-execute, context deferral, settle append). start() is awaited before the next entry can start, so concurrent submissions can no longer run pre-execute pipelines concurrently; only the around-dispatch/body stage overlaps, matching the native loop's fillPool sequencing. - An exclusive call's barrier now holds through its COMMIT: later starts wait for the exclusive pipeline (post-execute included) to finish, the native exclusive-group semantics. - drainDispatches() awaits the driver run itself, so a commit already mid-flight when the program returns is drained before run_code closes the turn — the settle event and deferred contexts land inside it. - maxParallelSubCalls is resolved and validated at construction (positive integer), so direct construction can no longer wedge the pool with 0. New tests: overlapping-submission ordered-prepare, barrier-through-commit, drain-mid-commit, cap rejection. 96 keyless snapshots replay unchanged; Agent Note updated (both languages).
This commit is contained in:
Tianyi Cui
@@ -249,100 +249,114 @@ export function createRunCodeTool(registry: ToolRegistry, requireRuntime: () =>
|
||||
|
||||
let dispatches = 0
|
||||
// The per-run scheduler, reusing the NATIVE concurrency contract through
|
||||
// the registry's staged view (the loop scheduler's own seam): submitted
|
||||
// calls START strictly in submission order; only the around-dispatch/body
|
||||
// stage overlaps — ordered pre-execute runs at start time and ordered
|
||||
// post-execute/context commitment runs in submission order through the
|
||||
// commit cursor below, so stateful policy listeners observe submission
|
||||
// order exactly as they do under the native loop. Consecutive
|
||||
// parallel-classified calls overlap up to maxParallel; an exclusive call
|
||||
// waits for the pool to drain, runs alone, and bars later calls.
|
||||
// Classification is re-read via executionMode() immediately before each
|
||||
// start (a registry mutation while queued can flip a call exclusive),
|
||||
// matching the native scheduler's lazy reclassification.
|
||||
// the registry's staged view (the loop scheduler's own seam) — and the
|
||||
// native loop's SEQUENCING: every ordered stage (the dispatch-start
|
||||
// append, prepare = pre-execute/guards, finalize/finish = post-execute,
|
||||
// context deferral, the settle append) runs inside ONE driver lane, so
|
||||
// ordered policy stages never overlap each other and only the
|
||||
// around-dispatch/body stage runs concurrently. Starts are strictly
|
||||
// submission-ordered; results commit in submission order through the
|
||||
// head-of-line cursor. Consecutive parallel-classified calls overlap up
|
||||
// to maxParallel; an exclusive call waits for the pool to drain, runs
|
||||
// alone, and holds its barrier until its COMMIT (post-execute included)
|
||||
// completes, exactly like a native exclusive group. Classification is
|
||||
// re-read via executionMode() immediately before each start (a registry
|
||||
// mutation while queued can flip a call exclusive), matching the native
|
||||
// scheduler's lazy reclassification.
|
||||
interface PendingDispatch {
|
||||
/** Ordered stage: append the start event, prepare, dispatch (body overlaps), park for commit. */
|
||||
/** Ordered stage: append the start event, await prepare (pre-execute/guards), launch the body into `flight`. */
|
||||
start(): Promise<void>
|
||||
classify(): 'parallel' | 'exclusive'
|
||||
abandon(): void
|
||||
/** Ordered stage: post-execute + context deferral + settle event, in submission order. */
|
||||
commit(): Promise<void>
|
||||
/** Set once the dispatch stage settles; commit() runs after this resolves. */
|
||||
dispatched?: Promise<void>
|
||||
/** The launched around-dispatch/body stage; resolved until start() replaces it. */
|
||||
flight: Promise<void>
|
||||
/** True once the dispatch stage parked its outcome; the commit cursor waits on it. */
|
||||
settled: boolean
|
||||
/** The classification this entry started under; an exclusive holds its barrier through commit(). */
|
||||
mode?: 'parallel' | 'exclusive'
|
||||
}
|
||||
const pendingQueue: PendingDispatch[] = []
|
||||
const inFlight = new Set<Promise<void>>()
|
||||
const commitQueue: PendingDispatch[] = []
|
||||
let committing = false
|
||||
let exclusiveActive = false
|
||||
let pumping = false
|
||||
/** Ordered commit cursor: drain the head-of-line settled dispatches one at a time. */
|
||||
const commitReady = async (): Promise<void> => {
|
||||
if (committing) return
|
||||
committing = true
|
||||
try {
|
||||
while (commitQueue.length > 0) {
|
||||
const head = commitQueue[0]
|
||||
/* v8 ignore next -- the loop condition bounds the index. */
|
||||
if (head === undefined) break
|
||||
/* v8 ignore next -- entries join commitQueue only after start() set dispatched (see pump). */
|
||||
if (head.dispatched === undefined) break
|
||||
await head.dispatched
|
||||
commitQueue.shift()
|
||||
await head.commit()
|
||||
}
|
||||
} finally {
|
||||
committing = false
|
||||
}
|
||||
let driving = false
|
||||
let driverRun: Promise<void> = Promise.resolve()
|
||||
let wake: (() => void) | undefined
|
||||
const wakeup = (): void => {
|
||||
const release = wake
|
||||
wake = undefined
|
||||
release?.()
|
||||
}
|
||||
const pump = (): void => {
|
||||
// Defensive re-entry guard: today every caller (binding submission,
|
||||
// flight.finally, drain) runs off promise callbacks, never while pump
|
||||
// is on the stack, so this cannot fire — kept against a future
|
||||
// synchronous caller.
|
||||
/* v8 ignore next -- see the re-entry note above. */
|
||||
if (pumping) return
|
||||
pumping = true
|
||||
try {
|
||||
for (;;) {
|
||||
const head = pendingQueue[0]
|
||||
if (head === undefined) return
|
||||
if (runController.signal.aborted) {
|
||||
pendingQueue.shift()
|
||||
head.abandon()
|
||||
continue
|
||||
/**
|
||||
* The single ordered lane. Each pass commits the head-of-line settled
|
||||
* dispatch (ordered post-execute), then starts the next queued entry if
|
||||
* its slot is free (ordered pre-execute), and otherwise sleeps until a
|
||||
* body settles or a new submission arrives. One run reaching the
|
||||
* empty-queues/empty-pool state is quiescence.
|
||||
*/
|
||||
const drive = (): Promise<void> => {
|
||||
if (driving) return driverRun
|
||||
driving = true
|
||||
driverRun = (async () => {
|
||||
try {
|
||||
for (;;) {
|
||||
// Arm before inspecting state so a settle or submission landing
|
||||
// between the checks and the await below cannot be lost.
|
||||
const signal = new Promise<void>((resolve) => { wake = resolve })
|
||||
const commitHead = commitQueue[0]
|
||||
if (commitHead !== undefined && commitHead.settled) {
|
||||
commitQueue.shift()
|
||||
await commitHead.commit()
|
||||
// The barrier covers post-execute: later starts wait for the
|
||||
// exclusive call's full pipeline, as under the native loop.
|
||||
if (commitHead.mode === 'exclusive') exclusiveActive = false
|
||||
continue
|
||||
}
|
||||
const head = pendingQueue[0]
|
||||
if (head !== undefined) {
|
||||
if (runController.signal.aborted) {
|
||||
pendingQueue.shift()
|
||||
head.abandon()
|
||||
continue
|
||||
}
|
||||
// Reclassify at start time (fail-closed on registry changes).
|
||||
const mode = head.classify()
|
||||
const capacity = !exclusiveActive
|
||||
&& (mode === 'exclusive' ? inFlight.size === 0 : inFlight.size < maxParallel)
|
||||
if (capacity) {
|
||||
if (mode === 'exclusive') exclusiveActive = true
|
||||
head.mode = mode
|
||||
pendingQueue.shift()
|
||||
// Joined before start() so the commit cursor sees submission
|
||||
// order; nothing commits it until `settled` flips.
|
||||
commitQueue.push(head)
|
||||
await head.start()
|
||||
const flight: Promise<void> = head.flight.finally(() => {
|
||||
inFlight.delete(flight)
|
||||
wakeup()
|
||||
})
|
||||
inFlight.add(flight)
|
||||
continue
|
||||
}
|
||||
}
|
||||
if (pendingQueue.length === 0 && commitQueue.length === 0 && inFlight.size === 0) return
|
||||
await signal
|
||||
}
|
||||
// Reclassify at start time (fail-closed on registry changes).
|
||||
const mode = head.classify()
|
||||
if (exclusiveActive || inFlight.size >= (mode === 'exclusive' ? 1 : maxParallel)) return
|
||||
// The guard above already returned for an exclusive head with any
|
||||
// in-flight sibling, so claiming the barrier here is race-free.
|
||||
if (mode === 'exclusive') exclusiveActive = true
|
||||
pendingQueue.shift()
|
||||
const flight = head.start().finally(() => {
|
||||
inFlight.delete(flight)
|
||||
if (mode === 'exclusive') exclusiveActive = false
|
||||
// Commit ordering and slot refill are independent: the cursor
|
||||
// may wait head-of-line on an earlier dispatch while later
|
||||
// slots keep starting.
|
||||
void commitReady()
|
||||
pump()
|
||||
})
|
||||
// Joined AFTER start() ran synchronously, so every commitQueue
|
||||
// entry already carries its `dispatched` promise.
|
||||
commitQueue.push(head)
|
||||
inFlight.add(flight)
|
||||
} finally {
|
||||
driving = false
|
||||
wake = undefined
|
||||
}
|
||||
} finally {
|
||||
pumping = false
|
||||
}
|
||||
})()
|
||||
return driverRun
|
||||
}
|
||||
/** Every in-flight dispatch settled AND committed; nothing can start (the run is aborted at call time). */
|
||||
/** Every dispatch settled AND committed; nothing can start (the run is aborted at call time). */
|
||||
const drainDispatches = async (): Promise<void> => {
|
||||
// Abandon queued-unstarted tasks first, then await the live set until quiescent.
|
||||
pump()
|
||||
while (inFlight.size > 0) await Promise.allSettled([...inFlight])
|
||||
await commitReady()
|
||||
// The abort already fired: the driver abandons queued-unstarted
|
||||
// entries, awaits the live pool, and drains the ordered commit lane —
|
||||
// including a commit already in progress when the program returned.
|
||||
await drive()
|
||||
}
|
||||
|
||||
// Read through a call, not a bare property: the abort state genuinely
|
||||
@@ -368,7 +382,7 @@ export function createRunCodeTool(registry: ToolRegistry, requireRuntime: () =>
|
||||
type DispatchOutcome = { isError: true; message: string } | { isError: false; value: JsonValue }
|
||||
const scheduler = registry[TOOL_REGISTRY_SCHEDULER]
|
||||
const outcome = await new Promise<DispatchOutcome>((resolve, reject) => {
|
||||
// Set by start(): what commit() finalizes in submission order.
|
||||
// Set by the dispatch stage (or start() for a pre-settled result): what commit() finalizes in submission order.
|
||||
let parked:
|
||||
| { kind: 'post-result' | 'final-result'; exec: ToolRunContext; result: ToolExecutionResult }
|
||||
| undefined
|
||||
@@ -391,34 +405,37 @@ export function createRunCodeTool(registry: ToolRegistry, requireRuntime: () =>
|
||||
: { isError: false, value: result.value })
|
||||
}
|
||||
pendingQueue.push({
|
||||
// Re-read per pump pass against the same agent view the SDK
|
||||
flight: Promise.resolve(),
|
||||
settled: false,
|
||||
// Re-read per driver pass against the same agent view the SDK
|
||||
// declared; fail-closed exclusive when undeclared/invalid.
|
||||
classify: () => registry.executionMode(input).kind,
|
||||
abandon: () => {
|
||||
reject(new Error(`run_code run is over (${String(runController.signal.reason)}); ${name} tool call abandoned`))
|
||||
},
|
||||
start(): Promise<void> {
|
||||
async start(): Promise<void> {
|
||||
exec.agent?.session.append('tool/code-dispatch-start', {
|
||||
parentCallId: exec.callId,
|
||||
subCallId,
|
||||
name,
|
||||
arguments: normalized.logged,
|
||||
})
|
||||
// Ordered prepare (pre-execute/guards) runs here — starts are
|
||||
// strictly submission-ordered; only dispatch overlaps.
|
||||
this.dispatched = (async () => {
|
||||
const prepared = await scheduler.prepare(input)
|
||||
if (prepared.kind === 'dispatch') {
|
||||
const dispatchOutcome = await scheduler.dispatch(prepared.exec)
|
||||
// Ordered prepare runs INSIDE the driver lane: the next entry's
|
||||
// pre-execute waits for this resolution, as under the native
|
||||
// scheduler. Only the launched body below overlaps.
|
||||
const prepared = await scheduler.prepare(input)
|
||||
if (prepared.kind === 'dispatch') {
|
||||
this.flight = scheduler.dispatch(prepared.exec).then((dispatchOutcome) => {
|
||||
parked = { kind: dispatchOutcome.kind, exec: prepared.exec, result: dispatchOutcome.result }
|
||||
return
|
||||
}
|
||||
parked = { kind: prepared.kind, exec: prepared.exec, result: prepared.result }
|
||||
})()
|
||||
return this.dispatched
|
||||
this.settled = true
|
||||
})
|
||||
return
|
||||
}
|
||||
parked = { kind: prepared.kind, exec: prepared.exec, result: prepared.result }
|
||||
this.settled = true
|
||||
},
|
||||
async commit(): Promise<void> {
|
||||
/* v8 ignore next -- commit() runs only after this.dispatched resolved, which set parked. */
|
||||
/* v8 ignore next -- commit() runs only after `settled` flipped, which set parked. */
|
||||
if (parked === undefined) return
|
||||
const result = parked.kind === 'post-result'
|
||||
? await scheduler.finalize(parked.exec, parked.result)
|
||||
@@ -429,7 +446,8 @@ export function createRunCodeTool(registry: ToolRegistry, requireRuntime: () =>
|
||||
settle(result)
|
||||
},
|
||||
})
|
||||
pump()
|
||||
wakeup()
|
||||
void drive()
|
||||
})
|
||||
// A budget expiry or outer cancel that lands while this call was in
|
||||
// flight already aborted the dispatch; stop the program now rather
|
||||
|
||||
@@ -635,6 +635,15 @@ interface FusedToolSignal {
|
||||
dispose(): void
|
||||
}
|
||||
|
||||
/** Resolve the run_code overlap cap at the owning config boundary (direct construction bypasses the Loader schema). */
|
||||
function resolveMaxParallelSubCalls(value: number | undefined): number {
|
||||
const maxParallelSubCalls = value ?? 10
|
||||
if (!Number.isInteger(maxParallelSubCalls) || maxParallelSubCalls < 1) {
|
||||
throw new Error('maxParallelSubCalls must be a positive integer')
|
||||
}
|
||||
return maxParallelSubCalls
|
||||
}
|
||||
|
||||
/**
|
||||
* Tool registry and execution pipeline. Scoped registrations shadow globals;
|
||||
* one visibility resolver feeds presentation, lookup, and dispatch.
|
||||
@@ -681,7 +690,7 @@ export class ToolRegistry extends Service {
|
||||
// the filterable global/scoped capability layers.
|
||||
this.codeTransport = this.mode === 'native'
|
||||
? undefined
|
||||
: createRunCodeTool(this, () => this.requireCodeRuntime(), config.maxParallelSubCalls ?? 10)
|
||||
: createRunCodeTool(this, () => this.requireCodeRuntime(), resolveMaxParallelSubCalls(config.maxParallelSubCalls))
|
||||
ctx.systemPrompt.tools(context => this.wireSchemas(context.scope))
|
||||
if (this.mode !== 'native') {
|
||||
ctx.systemPrompt.section({
|
||||
|
||||
@@ -510,6 +510,113 @@ describe('the sub-dispatch scheduler (native concurrency contract)', () => {
|
||||
expect(calls).toEqual([])
|
||||
})
|
||||
|
||||
it('ordered pre-execute never overlaps: a slow policy on one call delays the next start', async () => {
|
||||
const { ctx, runtime } = await setup({ mode: 'code' })
|
||||
const gated = registerGated(ctx, 'safe_read', true)
|
||||
const stages: string[] = []
|
||||
let releaseGate: (() => void) | undefined
|
||||
ctx.on('tools/pre-execute', async (preExec, next) => {
|
||||
if (preExec.name !== 'safe_read') return next()
|
||||
stages.push(`pre-enter:${String(preExec.callId)}`)
|
||||
if (releaseGate === undefined) {
|
||||
// The FIRST call's policy awaits an asynchronous decision.
|
||||
await new Promise<void>((resolve) => { releaseGate = resolve })
|
||||
}
|
||||
stages.push(`pre-exit:${String(preExec.callId)}`)
|
||||
return next()
|
||||
})
|
||||
runtime.behavior = async (request) => {
|
||||
const tools = request.bindings[0]!.functions
|
||||
const all = Promise.all([tools.safe_read!({ id: 'a' }), tools.safe_read!({ id: 'b' })])
|
||||
// Both submissions are in; the second pre-execute must NOT have entered
|
||||
// while the first is still awaiting its policy decision.
|
||||
await expect.poll(() => stages.length).toBeGreaterThanOrEqual(1)
|
||||
expect(stages).toEqual(['pre-enter:call-1:code:1'])
|
||||
releaseGate!()
|
||||
await expect.poll(() => gated.pending()).toBe(2)
|
||||
gated.releaseAll()
|
||||
await all
|
||||
return { logs: [], value: 'ordered-prepare' }
|
||||
}
|
||||
const result = await runCode(ctx, 'program')
|
||||
expect(result.isError).toBe(false)
|
||||
expect(stages).toEqual([
|
||||
'pre-enter:call-1:code:1', 'pre-exit:call-1:code:1',
|
||||
'pre-enter:call-1:code:2', 'pre-exit:call-1:code:2',
|
||||
])
|
||||
})
|
||||
|
||||
it('an exclusive call holds its barrier through post-execute: the next start waits for the commit', async () => {
|
||||
const { ctx, runtime } = await setup({ mode: 'code' })
|
||||
const writer = registerGated(ctx, 'writer', false)
|
||||
const reader = registerGated(ctx, 'safe_read', true)
|
||||
const stages: string[] = []
|
||||
let releasePost: (() => void) | undefined
|
||||
ctx.on('tools/post-execute', async (postExec, _result, next): Promise<PostToolDecision> => {
|
||||
if (postExec.name === 'writer') {
|
||||
stages.push('post-enter:writer')
|
||||
await new Promise<void>((resolve) => { releasePost = resolve })
|
||||
stages.push('post-exit:writer')
|
||||
}
|
||||
return next()
|
||||
})
|
||||
runtime.behavior = async (request) => {
|
||||
const tools = request.bindings[0]!.functions
|
||||
const w = tools.writer!({ id: 'w' })
|
||||
const r = tools.safe_read!({ id: 'r' })
|
||||
await expect.poll(() => writer.pending()).toBe(1)
|
||||
writer.release()
|
||||
// The writer's body is done and its async post-execute is running; the
|
||||
// parallel read must not have STARTED (no pre/body) while the exclusive
|
||||
// call's pipeline is still open.
|
||||
await expect.poll(() => stages).toContain('post-enter:writer')
|
||||
expect(reader.pending()).toBe(0)
|
||||
releasePost!()
|
||||
await w
|
||||
await expect.poll(() => reader.pending()).toBe(1)
|
||||
reader.releaseAll()
|
||||
await r
|
||||
return { logs: [], value: 'barrier-through-commit' }
|
||||
}
|
||||
const result = await runCode(ctx, 'program')
|
||||
expect(result.isError).toBe(false)
|
||||
expect(stages).toEqual(['post-enter:writer', 'post-exit:writer'])
|
||||
})
|
||||
|
||||
it('run settlement drains a commit already in progress: the settle event lands inside the turn', async () => {
|
||||
const { ctx, runtime } = await setup({ mode: 'code' })
|
||||
const gated = registerGated(ctx, 'safe_read', true)
|
||||
const { agent, events } = fakeAgent()
|
||||
let releasePost: (() => void) | undefined
|
||||
ctx.on('tools/post-execute', async (postExec, _result, next): Promise<PostToolDecision> => {
|
||||
if (postExec.name === 'safe_read') {
|
||||
await new Promise<void>((resolve) => { releasePost = resolve })
|
||||
}
|
||||
return next()
|
||||
})
|
||||
runtime.behavior = async (request) => {
|
||||
// Fire-and-forget: the program returns while the sub-call's async
|
||||
// post-execute commit is mid-flight.
|
||||
request.bindings[0]!.functions.safe_read!({ id: 'a' }).catch(() => 'run-over')
|
||||
await expect.poll(() => gated.pending()).toBe(1)
|
||||
gated.release()
|
||||
await expect.poll(() => releasePost !== undefined).toBe(true)
|
||||
queueMicrotask(() => { releasePost!() })
|
||||
return { logs: [], value: 'returned-early' }
|
||||
}
|
||||
const result = await runCode(ctx, 'program', { agent })
|
||||
expect(result.isError).toBe(false)
|
||||
// The drain awaited the in-progress commit: the settle event exists and
|
||||
// preceded the run_code turn closing (all appends happen inside
|
||||
// execute()). The run's settlement aborted the sub-call's signal while
|
||||
// its post-execute was mid-flight, so the native cancellation contract
|
||||
// replaces the successful outcome with the aborted result — the event is
|
||||
// still durable and in-turn, which is the invariant under test.
|
||||
const settles = events.filter(event => event.type === 'tool/code-dispatch')
|
||||
expect(settles).toHaveLength(1)
|
||||
expect(settles[0]?.data).toMatchObject({ name: 'safe_read', isError: true })
|
||||
})
|
||||
|
||||
it('post-execute and context commitment stay in submission order under out-of-order completion', async () => {
|
||||
const { ctx, runtime } = await setup({ mode: 'code' })
|
||||
const gated = registerGated(ctx, 'safe_read', true)
|
||||
@@ -1293,6 +1400,13 @@ describe('the run_code dispatch bridge', () => {
|
||||
expect(derived[0]?.role).toBe('user')
|
||||
})
|
||||
|
||||
it('direct construction rejects a non-positive parallel sub-call cap at load', async () => {
|
||||
const ctx = new Context()
|
||||
await ctx.plugin(SystemPrompt, {})
|
||||
expect(() => new ToolRegistry(ctx, { mode: 'code', maxParallelSubCalls: 0 }))
|
||||
.toThrow('maxParallelSubCalls must be a positive integer')
|
||||
})
|
||||
|
||||
it('direct construction in code mode defaults the parallel sub-call cap', async () => {
|
||||
const ctx = new Context()
|
||||
await ctx.plugin(SystemPrompt, {})
|
||||
|
||||
Reference in New Issue
Block a user