diff --git a/packages/daemon/src/server/__tests__/request-execution-locks.test.ts b/packages/daemon/src/server/__tests__/request-execution-locks.test.ts index 65200672cb..7cb2d5c9cd 100644 --- a/packages/daemon/src/server/__tests__/request-execution-locks.test.ts +++ b/packages/daemon/src/server/__tests__/request-execution-locks.test.ts @@ -279,3 +279,22 @@ test('a nested request retaining a lock its enclosing request holds re-enters it }), ).resolves.toBe('reentered'); }); + +test('a nested request whose initial key its parent retained re-enters it instead of waiting on its parent', async () => { + const locks = new Map>(); + const parent = createRequestExecutionLocks({ locks, initialKeys: ['session:replay'] }); + + await expect( + parent.run(async () => { + // A fresh replay retains its device once its first open selects it. + await parent.retainDevice('SIM-001'); + // A later step runs as a full request whose plan names the same device. + const step = createRequestExecutionLocks({ + locks, + initialKeys: ['session:replay', 'device:SIM-001'], + }); + return await step.run(async () => 'reentered'); + }), + ).resolves.toBe('reentered'); + expect(locks.size).toBe(0); +}); diff --git a/packages/daemon/src/server/request-execution-locks.ts b/packages/daemon/src/server/request-execution-locks.ts index 65a0d061fb..a6c025b774 100644 --- a/packages/daemon/src/server/request-execution-locks.ts +++ b/packages/daemon/src/server/request-execution-locks.ts @@ -1,29 +1,18 @@ -import { AsyncLocalStorage } from 'node:async_hooks'; -import { withKeyedLock } from '@agent-device/kernel/keyed-lock'; +import { + heldKeyedLockKeys, + withKeyedLock, + withKeyedLockHolder, + type KeyedLockHolder, +} from '@agent-device/kernel/keyed-lock'; import { hostDeviceIdOf } from '@agent-device/kernel/device'; import { AppError } from '@agent-device/kernel/errors'; import type { RequestExecutionLockKey } from './request-binding.ts'; -type RetainedExecutionLock = { - acquired: Promise; - completion: Promise; - release(): void; -}; - export type RequestExecutionLocks = { run(task: () => Promise): Promise; retainDevice(deviceId: string): Promise; }; -/** - * The locks each enclosing request holds, innermost last. A request nested in another (a batch - * step) takes its own locks while the enclosing request still holds its own, so the same-host rule - * below applies to the locks it takes up front as much as to the ones it retains later. - */ -const enclosingRequestLocks = new AsyncLocalStorage< - readonly ReadonlySet[] ->(); - /** * Owns the stable execution locks for one request and any device lock discovered * after admission. Dynamic device retention is used by a fresh replay whose @@ -35,72 +24,50 @@ export function createRequestExecutionLocks(params: { scope?: Readonly<{ signal: AbortSignal }>; }): RequestExecutionLocks { const { locks, initialKeys, scope } = params; - const initialKeySet = new Set(initialKeys); - const retainedLocks = new Map(); - // Every lock this request holds, for the requests nested in it to see. - const heldLocks = new Set(); - const enclosing = (): readonly ReadonlySet[] => - enclosingRequestLocks.getStore() ?? []; + let holder: KeyedLockHolder | undefined; let running = false; - const releaseRetainedLocks = async (): Promise => { - const held = [...retainedLocks.values()]; - for (const lock of held) lock.release(); - await Promise.all(held.map(async (lock) => await lock.completion)); - retainedLocks.clear(); - }; - return { run: async (task) => { if (running) { throw new Error('Request execution locks cannot run the same scope concurrently.'); } running = true; - const outer = enclosing(); - const heldOutside = outer.flatMap((held) => [...held]); + // What the requests this one is nested in hold (a batch step runs inside its batch). Its own + // initial keys are taken in sorted order, waiting for each in turn. + const heldOutside = heldRequestLockKeys(locks); try { return await withRequestExecutionLockKeys( locks, initialKeys, - async () => await enclosingRequestLocks.run([...outer, heldLocks], task), + async () => + await withKeyedLockHolder(async (taken) => { + holder = taken; + return await task(); + }), scope?.signal, - (key) => { - // An enclosing request already holding this lock re-enters it without waiting. - if (!heldOutside.includes(key)) assertSameHostLockFree(locks, key, heldOutside); - heldLocks.add(key); - }, + (key) => assertSameHostLockFree(locks, key, heldOutside), ); } finally { - heldLocks.clear(); - try { - await releaseRetainedLocks(); - } finally { - running = false; - } + holder = undefined; + running = false; } }, retainDevice: async (deviceId) => { - if (!running) { + if (!holder) { throw new Error('A device execution lock can only be retained by a running request.'); } const key = `device:${deviceId}` as const; - if (initialKeySet.has(key)) return; - // An enclosing request already holds it, and this request runs inside that hold. - if (enclosing().some((held) => held.has(key))) return; - - let retained = retainedLocks.get(key); - if (!retained) { - const heldKeys = [...heldLocks, ...enclosing().flatMap((held) => [...held])]; - assertSameHostLockFree(locks, key, heldKeys); - retained = retainExecutionLock(locks, key, scope?.signal); - retainedLocks.set(key, retained); - heldLocks.add(key); - } - await retained.acquired; + assertSameHostLockFree(locks, key, heldRequestLockKeys(locks)); + await holder.hold(locks, key, scope?.signal); }, }; } +function heldRequestLockKeys(locks: Map>): RequestExecutionLockKey[] { + return heldKeyedLockKeys(locks) as RequestExecutionLockKey[]; +} + /** * Whether the request already holds a lock on `deviceId`'s host or one of its app devices * (ADR 0035). Waiting for a second lock on the same host while holding one could cross another @@ -113,7 +80,8 @@ function assertSameHostLockFree( key: RequestExecutionLockKey, heldKeys: readonly RequestExecutionLockKey[], ): void { - if (!key.startsWith('device:')) return; + // A held key re-enters, and only device keys belong to a host. + if (heldKeys.includes(key) || !key.startsWith('device:')) return; const deviceId = key.slice('device:'.length); if (holdsDeviceOfSameHost(deviceId, heldKeys) && locks.has(key)) { throw deviceExecutionLockBusy(deviceId); @@ -164,31 +132,3 @@ export async function withRequestExecutionLockKeys( signal, ); } - -function retainExecutionLock( - locks: Map>, - key: RequestExecutionLockKey, - signal?: AbortSignal, -): RetainedExecutionLock { - let markAcquired: () => void = () => {}; - let rejectAcquired: (error: unknown) => void = () => {}; - let release: () => void = () => {}; - const acquired = new Promise((resolve, reject) => { - markAcquired = resolve; - rejectAcquired = reject; - }); - const released = new Promise((resolve) => { - release = resolve; - }); - const completion = withKeyedLock( - locks, - key, - async () => { - markAcquired(); - await released; - }, - signal, - ); - void completion.catch((error) => rejectAcquired(error)); - return { acquired, completion, release }; -} diff --git a/packages/daemon/src/sessions/handlers/__tests__/session-device-claims.test.ts b/packages/daemon/src/sessions/handlers/__tests__/session-device-claims.test.ts index eb02617791..61f215d6ab 100644 --- a/packages/daemon/src/sessions/handlers/__tests__/session-device-claims.test.ts +++ b/packages/daemon/src/sessions/handlers/__tests__/session-device-claims.test.ts @@ -64,6 +64,11 @@ import { handleSessionOpenCommands as handleProductionOpenCommand } from '../../ import { hostAppDevice, type DeviceInfo } from '@agent-device/kernel/device'; import { makeAuthoringSession } from '../../../__tests__/support/session-factories.ts'; import { AppError } from '@agent-device/kernel/errors'; +import { resolveMacOsAppDeviceOpenTarget } from '@agent-device/platform-apple/app-resolution'; +import { + createLocalAppleToolProvider, + withAppleToolProvider, +} from '@agent-device/platform-apple/tool-provider'; import { bindProviderLifecycleRuntime, bindLifecycleRuntime, @@ -671,11 +676,24 @@ test('closing another app from an app device session is refused before teardown, }); mockDispatch.mockResolvedValue(undefined); // The platform resolves the target as open does and refuses another app of the Mac. - resolveOpenTargetFixture.mockImplementationOnce(async () => { - throw new AppError('INVALID_ARGS', 'This session drives com.apple.TextEdit.', { - reason: 'app-device-scope', - }); + // The production platform refusal, over a stubbed list of installed apps. + const provider = createLocalAppleToolProvider({ + runCommand: async () => ({ exitCode: 0, stdout: '', stderr: '' }), }); + const macosHost = { + ...provider.macosHost!, + listApps: async () => [ + { bundleId: 'com.apple.TextEdit', name: 'TextEdit' }, + { bundleId: 'com.apple.calculator', name: 'Calculator' }, + ], + }; + resolveOpenTargetFixture.mockImplementationOnce( + async (device, _directResolve, input) => + await withAppleToolProvider( + { ...provider, macosHost }, + async () => await resolveMacOsAppDeviceOpenTarget(device, input.target), + ), + ); const outcome = await handleCloseCommand({ host: createTestHostEnvironment(), diff --git a/packages/kernel/src/keyed-lock.test.ts b/packages/kernel/src/keyed-lock.test.ts index 302be03472..9b3a5cad29 100644 --- a/packages/kernel/src/keyed-lock.test.ts +++ b/packages/kernel/src/keyed-lock.test.ts @@ -1,6 +1,11 @@ import { test, expect, vi } from 'vitest'; import assert from 'node:assert/strict'; -import { withKeyedLock } from './keyed-lock.ts'; +import { + heldKeyedLockKeys, + withKeyedLock, + withKeyedLockHolder, + type KeyedLockHolder, +} from './keyed-lock.ts'; test('canceled queued work returns before its owner and never runs or admits a competing task early', async () => { const locks = new Map>(); @@ -214,3 +219,162 @@ test.each([false, true])( assert.deepEqual(order, ['descendant', 'finished', 'nested', 'competing']); }, ); + +test('a lock a holder takes partway through is held by later sibling scopes until the holder ends', async () => { + const locks = new Map>(); + const order: string[] = []; + let markHeld: () => void = () => {}; + const held = new Promise((resolve) => { + markHeld = resolve; + }); + let finishHolder: () => void = () => {}; + const holderMayFinish = new Promise((resolve) => { + finishHolder = resolve; + }); + + const holderRun = withKeyedLockHolder(async (holder) => { + // A first nested scope takes the hold, as a replay action retains its device. + await withKeyedLock(locks, 'session', async () => { + await holder.hold(locks, 'device:app'); + }); + markHeld(); + // A later sibling scope re-enters the hold instead of queueing behind it. + await withKeyedLock(locks, 'device:app', async () => { + order.push('sibling re-entered'); + expect(heldKeyedLockKeys(locks)).toEqual(['device:app']); + }); + await holderMayFinish; + order.push('holder done'); + }); + await held; + // Another context queues behind the live hold. + let waiterEntered = false; + const waiter = withKeyedLock(locks, 'device:app', async () => { + waiterEntered = true; + order.push('waiter entered'); + }); + await new Promise((resolve) => setTimeout(resolve, 10)); + expect(waiterEntered).toBe(false); + finishHolder(); + + await Promise.all([holderRun, waiter]); + expect(order).toEqual(['sibling re-entered', 'holder done', 'waiter entered']); + expect(locks.size).toBe(0); +}); + +test('a hold cancelled while queued rejects with the signal reason and leaves nothing held', async () => { + const locks = new Map>(); + let releaseOwner: () => void = () => {}; + const owner = withKeyedLock( + locks, + 'device:app', + async () => + await new Promise((resolve) => { + releaseOwner = resolve; + }), + ); + const controller = new AbortController(); + const reason = new Error('request canceled'); + + await withKeyedLockHolder(async (holder) => { + const pending = holder.hold(locks, 'device:app', controller.signal); + // A key the holder is still taking already reads as held, so no second same-host wait starts. + expect(heldKeyedLockKeys(locks)).toEqual(['device:app']); + controller.abort(reason); + await expect(pending).rejects.toBe(reason); + expect(heldKeyedLockKeys(locks)).toEqual([]); + }); + releaseOwner(); + await owner; + expect(locks.size).toBe(0); +}); + +test('a holder refuses a lock once it has ended', async () => { + const locks = new Map>(); + let late: KeyedLockHolder | undefined; + await withKeyedLockHolder(async (holder) => { + late = holder; + }); + await expect(late!.hold(locks, 'device:app')).rejects.toThrow(/after it ends/); + expect(locks.size).toBe(0); +}); + +test('heldKeyedLockKeys lists only the locks of the map held by the current context', async () => { + const locks = new Map>(); + const other = new Map>(); + await withKeyedLock(other, 'other', async () => { + await withKeyedLock(locks, 'host', async () => { + await withKeyedLock(locks, 'host:app:one', async () => { + expect(heldKeyedLockKeys(locks).sort()).toEqual(['host', 'host:app:one']); + }); + expect(heldKeyedLockKeys(locks)).toEqual(['host']); + }); + }); +}); + +test('a scope wanting a key its enclosing holder is still taking re-enters it once taken', async () => { + const locks = new Map>(); + let releaseOwner: () => void = () => {}; + const owner = withKeyedLock( + locks, + 'device:app', + async () => + await new Promise((resolve) => { + releaseOwner = resolve; + }), + ); + + const outcome = withKeyedLockHolder(async (holder) => { + const taking = holder.hold(locks, 'device:app'); + // A nested step asks for the same key while the hold still waits for the owner. + const step = withKeyedLock(locks, 'device:app', async () => 'step re-entered'); + await new Promise((resolve) => setTimeout(resolve, 10)); + releaseOwner(); + await taking; + return await step; + }); + + await expect(outcome).resolves.toBe('step re-entered'); + await owner; + expect(locks.size).toBe(0); +}); + +test('taking or holding a free key queues it synchronously, so a check just before stays true', async () => { + const locks = new Map>(); + const taken = withKeyedLock(locks, 'device:a', async () => {}); + expect(locks.has('device:a')).toBe(true); + await taken; + + await withKeyedLockHolder(async (holder) => { + const hold = holder.hold(locks, 'device:b'); + expect(locks.has('device:b')).toBe(true); + await hold; + }); + expect(locks.size).toBe(0); +}); + +test('a scope waiting on its holder for a key fails with the hold when the hold is cancelled', async () => { + const locks = new Map>(); + let releaseOwner: () => void = () => {}; + const owner = withKeyedLock( + locks, + 'device:app', + async () => + await new Promise((resolve) => { + releaseOwner = resolve; + }), + ); + const controller = new AbortController(); + const reason = new Error('request canceled'); + + await withKeyedLockHolder(async (holder) => { + const taking = holder.hold(locks, 'device:app', controller.signal); + const step = withKeyedLock(locks, 'device:app', async () => 'ran'); + controller.abort(reason); + await expect(taking).rejects.toBe(reason); + await expect(step).rejects.toBe(reason); + }); + releaseOwner(); + await owner; + expect(locks.size).toBe(0); +}); diff --git a/packages/kernel/src/keyed-lock.ts b/packages/kernel/src/keyed-lock.ts index b7ffcb0e55..14f16138dc 100644 --- a/packages/kernel/src/keyed-lock.ts +++ b/packages/kernel/src/keyed-lock.ts @@ -7,32 +7,122 @@ type LockFrame = { descendants: Set>; }; -const keyedLockStorage = new AsyncLocalStorage(); +/** + * The locks an async context holds: its own frames and, through `parent`, every enclosing scope's. + * Parents are live, so a lock a holder scope retains later is visible to scopes nested in it. + */ +type LockScope = { + parent: LockScope | undefined; + frames: LockFrame[]; + /** Set on a holder scope while it can still take locks, including the keys it is still taking. */ + holder?: { open: boolean; pending: Set }; +}; -export async function withKeyedLock( +type PendingHold = { + locks: Map>; + key: string; + acquired: Promise; +}; + +function findPendingHold( locks: Map>, key: string, - task: () => Promise, - signal?: AbortSignal, -): Promise { - signal?.throwIfAborted(); - const activeLocks = keyedLockStorage.getStore() ?? []; - const held = activeLocks.find( - (entry) => entry.active && entry.locks === locks && entry.key === key, +): PendingHold | undefined { + for (let current = keyedLockStorage.getStore(); current; current = current.parent) { + for (const pending of current.holder?.pending ?? []) { + if (pending.locks === locks && pending.key === key) return pending; + } + } + return undefined; +} + +/** + * The frame an enclosing holder is still taking for `key`, once it lands. Queueing behind that hold + * instead would wait on the holder's own request. Only this rare path awaits: with no pending hold, + * callers reach `locks` synchronously, so a check made just before cannot go stale. A hold that + * fails (it was cancelled) fails the caller too, rather than letting it queue unchecked. + */ +async function pendingFrameOnceTaken(pending: PendingHold): Promise { + await pending.acquired; + return findHeldFrame(pending.locks, pending.key); +} + +const keyedLockStorage = new AsyncLocalStorage(); + +/** + * The nearest scope that still matters to a new nested scope: one holding an active lock, or an open + * holder scope that may take one. Skipping the rest keeps chains short in contexts that outlive their + * locks, such as a callback scheduled from inside a lock task. + */ +function liveScope(scope: LockScope | undefined): LockScope | undefined { + let current = scope; + while (current && !current.holder?.open && !current.frames.some((frame) => frame.active)) { + current = current.parent; + } + return current; +} + +function activeFrames(scope: LockScope | undefined): LockFrame[] { + const frames: LockFrame[] = []; + for (let current = scope; current; current = current.parent) { + for (const frame of current.frames) if (frame.active) frames.push(frame); + } + return frames; +} + +function findHeldFrame(locks: Map>, key: string): LockFrame | undefined { + return activeFrames(keyedLockStorage.getStore()).find( + (frame) => frame.locks === locks && frame.key === key, ); - if (held) { - let markSettled!: () => void; - const settled = new Promise((resolve) => { - markSettled = resolve; - }); - held.descendants.add(settled); - try { - return await task(); - } finally { - held.descendants.delete(settled); - markSettled(); +} + +/** + * The keys of `locks` the current async context holds through any enclosing scope, together with + * the keys an enclosing holder scope is still taking, which it will hold. + */ +export function heldKeyedLockKeys(locks: Map>): string[] { + const store = keyedLockStorage.getStore(); + const keys = activeFrames(store) + .filter((frame) => frame.locks === locks) + .map((frame) => frame.key); + for (let current = store; current; current = current.parent) { + for (const pending of current.holder?.pending ?? []) { + if (pending.locks === locks && !keys.includes(pending.key)) keys.push(pending.key); } } + return keys; +} + +async function runInHeldFrame(frame: LockFrame, task: () => Promise): Promise { + let markSettled!: () => void; + const settled = new Promise((resolve) => { + markSettled = resolve; + }); + frame.descendants.add(settled); + try { + return await task(); + } finally { + frame.descendants.delete(settled); + markSettled(); + } +} + +async function settleFrame(frame: LockFrame): Promise { + while (frame.descendants.size > 0) await Promise.all(frame.descendants); + frame.active = false; +} + +/** + * Queues `enter` behind the current holder of `key` and records the queue position in `locks` + * synchronously, before this returns, so a caller that checked `locks.has(key)` just before cannot + * see a stale answer. A request that gave up while queued rejects with the signal's reason. + */ +async function acquireKeyedLock( + locks: Map>, + key: string, + enter: () => Promise, + signal: AbortSignal | undefined, +): Promise { const previous = locks.get(key) ?? Promise.resolve(); let entered = false; const current = previous @@ -40,16 +130,7 @@ export async function withKeyedLock( .then(async () => { entered = true; signal?.throwIfAborted(); - const frame: LockFrame = { locks, key, active: true, descendants: new Set() }; - try { - return await keyedLockStorage.run( - [...activeLocks.filter((entry) => entry.active), frame], - task, - ); - } finally { - while (frame.descendants.size > 0) await Promise.all(frame.descendants); - frame.active = false; - } + return await enter(); }); locks.set(key, current); const completion = current.finally(() => { @@ -67,3 +148,128 @@ export async function withKeyedLock( if (signal.aborted) onAbort(); }); } + +export async function withKeyedLock( + locks: Map>, + key: string, + task: () => Promise, + signal?: AbortSignal, +): Promise { + signal?.throwIfAborted(); + const held = findHeldFrame(locks, key); + if (held) return await runInHeldFrame(held, task); + const pending = findPendingHold(locks, key); + const landed = pending ? await pendingFrameOnceTaken(pending) : undefined; + if (landed) return await runInHeldFrame(landed, task); + const parent = liveScope(keyedLockStorage.getStore()); + return await acquireKeyedLock( + locks, + key, + async () => { + const frame: LockFrame = { locks, key, active: true, descendants: new Set() }; + try { + return await keyedLockStorage.run({ parent, frames: [frame] }, task); + } finally { + await settleFrame(frame); + } + }, + signal, + ); +} + +/** Takes locks for the rest of a {@link withKeyedLockHolder} scope. */ +export type KeyedLockHolder = Readonly<{ + /** + * Takes `key` until the holder scope ends. Once taken, every scope nested in the holder scope sees + * it as held and re-enters it, including scopes that started before it was taken. Holding a key the + * context already holds, or one this holder is already taking, waits for that same hold. + */ + hold(locks: Map>, key: string, signal?: AbortSignal): Promise; +}>; + +type Hold = { + locks: Map>; + key: string; + acquired: Promise; + completion: Promise; + release(): void; +}; + +/** + * Runs `task` in a scope that can take locks partway through and keep them until `task` ends, when + * they are released after every scope that re-entered them settles. + */ +export async function withKeyedLockHolder( + task: (holder: KeyedLockHolder) => Promise, +): Promise { + const state = { open: true, pending: new Set() }; + const scope: LockScope = { + parent: liveScope(keyedLockStorage.getStore()), + frames: [], + holder: state, + }; + const holds: Hold[] = []; + const holder: KeyedLockHolder = Object.freeze({ + hold: async (locks, key, signal) => { + signal?.throwIfAborted(); + if (!state.open) throw new Error('A keyed lock holder cannot take a lock after it ends.'); + const existing = holds.find((hold) => hold.locks === locks && hold.key === key); + if (existing) return await existing.acquired; + if (findHeldFrame(locks, key)) return; + const enclosing = findPendingHold(locks, key); + if (enclosing && (await pendingFrameOnceTaken(enclosing))) return; + const hold = takeHold(scope, locks, key, signal); + const pending: PendingHold = { locks, key, acquired: hold.acquired }; + state.pending.add(pending); + holds.push(hold); + try { + await hold.acquired; + } finally { + state.pending.delete(pending); + } + }, + }); + try { + return await keyedLockStorage.run(scope, async () => await task(holder)); + } finally { + state.open = false; + for (const hold of holds) hold.release(); + await Promise.all(holds.map(async (hold) => await hold.completion)); + } +} + +function takeHold( + scope: LockScope, + locks: Map>, + key: string, + signal: AbortSignal | undefined, +): Hold { + let markAcquired: () => void = () => {}; + let rejectAcquired: (error: unknown) => void = () => {}; + let release: () => void = () => {}; + const acquired = new Promise((resolve, reject) => { + markAcquired = resolve; + rejectAcquired = reject; + }); + const released = new Promise((resolve) => { + release = resolve; + }); + const completion = acquireKeyedLock( + locks, + key, + async () => { + const frame: LockFrame = { locks, key, active: true, descendants: new Set() }; + scope.frames.push(frame); + markAcquired(); + try { + await released; + } finally { + await settleFrame(frame); + scope.frames.splice(scope.frames.indexOf(frame), 1); + } + }, + signal, + ); + completion.catch(rejectAcquired); + return { locks, key, acquired, completion: completion.catch(() => {}), release }; +}