diff --git a/packages/host/apiproxy/src/api-proxy.ts b/packages/host/apiproxy/src/api-proxy.ts index 149143b900..46ed485d45 100644 --- a/packages/host/apiproxy/src/api-proxy.ts +++ b/packages/host/apiproxy/src/api-proxy.ts @@ -488,6 +488,10 @@ export function createApiProxy(ctx: Context, defaults: ApiProxyDefaults): ApiPro ctx.on('session/disposed', (session: Session) => { pendingMetricSessions.delete(session) }), ctx.on('internal/status', (fiber) => { if (metricsDisposed) return + if (fiber.state === FiberState.UNLOADING) { + metricsProjector.invalidateCapacities() + return + } if (fiber.state !== FiberState.ACTIVE && fiber.state !== FiberState.FAILED && fiber.state !== FiberState.DISPOSED) return @@ -499,6 +503,7 @@ export function createApiProxy(ctx: Context, defaults: ApiProxyDefaults): ApiPro ] return () => { metricsDisposed = true + metricsProjector.dispose() pendingMetricSessions.clear() for (const dispose of disposers) dispose() } diff --git a/packages/host/apiproxy/src/session-metrics.ts b/packages/host/apiproxy/src/session-metrics.ts index 1ad1f488af..1eebc7b7d8 100644 --- a/packages/host/apiproxy/src/session-metrics.ts +++ b/packages/host/apiproxy/src/session-metrics.ts @@ -26,6 +26,7 @@ interface CapacityState { epoch: number status: 'pending' | 'ready' | 'retryable' contextWindow?: number + controller?: AbortController } type CapacityTarget = Pick @@ -35,7 +36,7 @@ interface TokenMeterLike { } interface LlmLike { - resolveModelInfo(provider: string, model: string): Promise<{ + resolveModelInfo(provider: string, model: string, signal?: AbortSignal): Promise<{ context?: { contextWindow: number } }> } @@ -89,7 +90,9 @@ function routeKeyFor(target: CapacityTarget | undefined): string | undefined { export class SessionMetricsProjector { private readonly usage = new WeakMap() private readonly capacities = new WeakMap() + private readonly pendingCapacities = new Set() private capacityEpoch = 0 + private disposed = false /** * @param ctx - Host context providing optional token-meter and LLM services. @@ -105,6 +108,13 @@ export class SessionMetricsProjector { /** Retire adapter-owned metadata and fence every resolution already in flight. */ invalidateCapacities(): void { this.capacityEpoch++ + for (const pending of this.pendingCapacities) this.abortCapacityResolution(pending) + } + + /** Permanently retire capacity projection and cancel every adapter-owned lookup. */ + dispose(): void { + this.disposed = true + this.invalidateCapacities() } /** @@ -163,6 +173,7 @@ export class SessionMetricsProjector { } private capacityFor(agent: Agent): number | undefined { + if (this.disposed) return undefined const target = this.targetFor(agent) const routeKey = routeKeyFor(target) let state = this.capacities.get(agent) @@ -170,6 +181,7 @@ export class SessionMetricsProjector { || state.routeKey !== routeKey || state.epoch !== this.capacityEpoch || state.status === 'retryable') { + this.abortCapacityResolution(state) state = { routeKey, generation: (state?.generation ?? 0) + 1, @@ -192,21 +204,42 @@ export class SessionMetricsProjector { pending.status = 'retryable' return } + const controller = new AbortController() + pending.controller = controller + this.pendingCapacities.add(pending) void Promise.resolve() - .then(() => llm.resolveModelInfo(target.provider, target.model)) + .then(() => { + controller.signal.throwIfAborted() + return llm.resolveModelInfo(target.provider, target.model, controller.signal) + }) .then( (resolved) => { + this.finishCapacityResolution(pending, controller) if (this.capacityResolutionIsStale(agent, pending)) return pending.status = 'ready' if (resolved.context !== undefined) pending.contextWindow = resolved.context.contextWindow this.onCapacityResolved(agent) }, () => { + this.finishCapacityResolution(pending, controller) if (!this.capacityResolutionIsStale(agent, pending)) pending.status = 'retryable' }, ) } + private abortCapacityResolution(pending: CapacityState | undefined): void { + if (pending === undefined || pending.controller === undefined) return + const controller = pending.controller + delete pending.controller + this.pendingCapacities.delete(pending) + controller.abort() + } + + private finishCapacityResolution(pending: CapacityState, controller: AbortController): void { + this.pendingCapacities.delete(pending) + if (pending.controller === controller) delete pending.controller + } + private capacityResolutionIsStale(agent: Agent, pending: CapacityState): boolean { if (pending.epoch !== this.capacityEpoch) return true if (this.capacities.get(agent)?.generation !== pending.generation) return true diff --git a/packages/host/apiproxy/tests/api-proxy-models.spec.ts b/packages/host/apiproxy/tests/api-proxy-models.spec.ts index 758669f5db..7b2595e528 100644 --- a/packages/host/apiproxy/tests/api-proxy-models.spec.ts +++ b/packages/host/apiproxy/tests/api-proxy-models.spec.ts @@ -65,7 +65,10 @@ class CatalogAdapter extends LlmAdapter { } class DeferredCatalogAdapter extends CatalogAdapter { - readonly pending: PromiseWithResolvers[] = [] + readonly pending: { + result: PromiseWithResolvers + signal: AbortSignal | undefined + }[] = [] constructor() { super('Deferred', [ @@ -73,16 +76,20 @@ class DeferredCatalogAdapter extends CatalogAdapter { ]) } - override resolveModel(_provider: string, _model: string): Promise { + override resolveModel( + _provider: string, + _model: string, + signal?: AbortSignal, + ): Promise { const result = Promise.withResolvers() - this.pending.push(result) + this.pending.push({ result, signal }) return result.promise } resolve(index: number, contextWindow: number): void { const pending = this.pending[index] if (pending === undefined) throw new Error(`no pending resolution at index ${String(index)}`) - pending.resolve({ + pending.result.resolve({ provider: 'deferred', id: 'lifecycle-model', name: 'Lifecycle model', @@ -563,6 +570,140 @@ describe('Web session model selection', () => { await ctx.fiber.dispose() }) + it('aborts pending capacity during adapter UNLOADING and refreshes after settlement', async () => { + const ctx = await hostContext() + const retiredAdapter = new DeferredCatalogAdapter() + const releaseUnload = Promise.withResolvers() + const releaseCancellationWait = Promise.withResolvers() + const abortObserved = Promise.withResolvers() + const retiredFiber = await ctx.plugin(Object.assign((inner: Context) => { + inner.llm.registerAdapter(['deferred'], retiredAdapter) + inner.effect( + () => () => releaseUnload.promise, + 'test: hold adapter unload', + ) + inner.effect(() => () => { + const signal = retiredAdapter.pending[0]?.signal + if (signal === undefined) throw new Error('pending capacity signal missing') + const cancellation = new Promise((resolve) => { + const finish = () => { + abortObserved.resolve(undefined) + resolve() + } + if (signal.aborted) finish() + else signal.addEventListener('abort', finish, { once: true }) + }) + return Promise.race([cancellation, releaseCancellationWait.promise]) + }, 'test: await capacity cancellation') + }, { inject: ['llm'] })) + const lifecycle = attachLifecycleSession(ctx, SessionId('capacity-adapter-unloading')) + const detachAgent = attachLifecycleAgent(ctx, lifecycle.session) + const api = createApiProxy(ctx, { + provider: 'deepseek', + model: 'deepseek-chat', + cwd: '/tmp', + workspaceRoot: '/tmp', + }) + const controller = new AbortController() + const iterator = api.events.mux(request({}), controller.signal)[Symbol.asyncIterator]() + const resolveModelInfo = vi.spyOn(ctx.llm, 'resolveModelInfo') + + expect((await nextMetrics(iterator)).contextWindow).toBeUndefined() + await vi.waitFor(() => { expect(retiredAdapter.pending).toHaveLength(1) }) + expect(resolveModelInfo).toHaveBeenCalledOnce() + const listSessions = vi.spyOn(ctx.sessions, 'list') + listSessions.mockClear() + const pendingFrame = iterator.next() + const disposing = retiredFiber.dispose() + try { + await vi.waitFor(() => { + expect(retiredAdapter.pending[0]?.signal?.aborted).toBe(true) + }) + await abortObserved.promise + expect(retiredFiber.state).toBe(FiberState.UNLOADING) + retiredAdapter.resolve(0, 64_000) + await settleCapacityCompletion() + expect(resolveModelInfo).toHaveBeenCalledOnce() + expect(listSessions).not.toHaveBeenCalled() + const outcome = await Promise.race([ + pendingFrame.then(() => 'frame' as const), + new Promise<'idle'>((resolve) => { setImmediate(() => { resolve('idle') }) }), + ]) + expect(outcome).toBe('idle') + } finally { + releaseCancellationWait.resolve(undefined) + releaseUnload.resolve(undefined) + await disposing + } + + expect(retiredFiber.state).toBe(FiberState.DISPOSED) + const settledFrame = await pendingFrame + if (settledFrame.done || settledFrame.value.payload.type !== 'session/metrics') { + throw new Error('expected settled metrics refresh') + } + expect(settledFrame.value.payload.metrics.contextWindow).toBeUndefined() + await settleCapacityCompletion() + expect(resolveModelInfo).toHaveBeenCalledTimes(2) + + const replacementAdapter = new DeferredCatalogAdapter() + const replacementFiber = await installDeferredAdapter(ctx, replacementAdapter) + expect((await nextMetrics(iterator)).contextWindow).toBeUndefined() + await vi.waitFor(() => { expect(replacementAdapter.pending).toHaveLength(1) }) + expect(resolveModelInfo).toHaveBeenCalledTimes(3) + replacementAdapter.resolve(0, 128_000) + await settleCapacityCompletion() + expect((await nextMetrics(iterator)).contextWindow).toBe(128_000) + + controller.abort() + await iterator.return?.() + detachAgent() + lifecycle.detach() + await replacementFiber.dispose() + await ctx.fiber.dispose() + }) + + it('aborts pending capacity when the API proxy fiber is disposed', async () => { + const ctx = await hostContext() + const deferred = new DeferredCatalogAdapter() + ctx.llm.registerAdapter(['deferred'], deferred) + const lifecycle = attachLifecycleSession(ctx, SessionId('capacity-api-proxy-teardown')) + const detachAgent = attachLifecycleAgent(ctx, lifecycle.session) + const proxy = Promise.withResolvers>() + const proxyFiber = await ctx.plugin(Object.assign((inner: Context) => { + proxy.resolve(createApiProxy(inner, { + provider: 'deepseek', + model: 'deepseek-chat', + cwd: '/tmp', + workspaceRoot: '/tmp', + })) + }, { inject: ['agents', 'sessions', 'userInteraction'] })) + const api = await proxy.promise + const controller = new AbortController() + const iterator = api.events.mux(request({}), controller.signal)[Symbol.asyncIterator]() + + expect((await nextMetrics(iterator)).contextWindow).toBeUndefined() + await vi.waitFor(() => { expect(deferred.pending).toHaveLength(1) }) + expect(deferred.pending[0]?.signal?.aborted).toBe(false) + + await proxyFiber.dispose() + expect(deferred.pending[0]?.signal?.aborted).toBe(true) + deferred.resolve(0, 64_000) + await settleCapacityCompletion() + const pendingFrame = iterator.next() + const outcome = await Promise.race([ + pendingFrame.then(() => 'frame' as const), + new Promise<'idle'>((resolve) => { setImmediate(() => { resolve('idle') }) }), + ]) + expect(outcome).toBe('idle') + + controller.abort() + await expect(pendingFrame).resolves.toMatchObject({ done: true }) + await iterator.return?.() + detachAgent() + lifecycle.detach() + await ctx.fiber.dispose() + }) + it('does not read the sessions service after its disposal status', async () => { let sessionsFiber: Fiber | undefined const ctx = await hostContext((fiber) => { sessionsFiber = fiber }) @@ -720,7 +861,7 @@ describe('Web session model selection', () => { detachAgent() lifecycle.detach() await ctx.fiber.dispose() - expect(agentReads).toBe(serviceName === 'sessions' ? 1 : 0) + expect(agentReads).toBe(0) expect(sessionReads).toBe(0) expect(outcome).toBe('idle') }, diff --git a/packages/host/apiproxy/tests/session-metrics.spec.ts b/packages/host/apiproxy/tests/session-metrics.spec.ts index 3b3653aad8..deb5f740fe 100644 --- a/packages/host/apiproxy/tests/session-metrics.spec.ts +++ b/packages/host/apiproxy/tests/session-metrics.spec.ts @@ -159,12 +159,20 @@ describe('SessionMetricsProjector', () => { it('publishes only the selected route capacity when asynchronous resolutions race', async () => { const ctx = new Context() - const resolutions = new Map void>() + const resolutions = new Map() ctx.provide('tokenMeter', { measure: () => ({ totalTokens: 35_000 }) }) ctx.provide('llm', { - resolveModelInfo(_provider: string, model: string) { + resolveModelInfo(_provider: string, model: string, signal?: AbortSignal) { return new Promise<{ context: { contextWindow: number } }>((resolve) => { - resolutions.set(model, (contextWindow) => { resolve({ context: { contextWindow } }) }) + resolutions.set(model, { + signal, + resolve(contextWindow) { + resolve({ context: { contextWindow } }) + }, + }) }) }, }) @@ -176,14 +184,17 @@ describe('SessionMetricsProjector', () => { expect(projector.snapshot(session, attached).contextWindow).toBeUndefined() await vi.waitFor(() => { expect(resolutions.has('alpha')).toBe(true) }) + expect(resolutions.get('alpha')?.signal?.aborted).toBe(false) current = { provider: 'test', model: 'beta' } expect(projector.snapshot(session, attached).contextWindow).toBeUndefined() + expect(resolutions.get('alpha')?.signal?.aborted).toBe(true) await vi.waitFor(() => { expect(resolutions.has('beta')).toBe(true) }) + expect(resolutions.get('beta')?.signal?.aborted).toBe(false) - resolutions.get('alpha')?.(64_000) + resolutions.get('alpha')?.resolve(64_000) await Promise.resolve() expect(resolved).not.toHaveBeenCalled() - resolutions.get('beta')?.(128_000) + resolutions.get('beta')?.resolve(128_000) await vi.waitFor(() => { expect(resolved).toHaveBeenCalledOnce() }) expect(projector.snapshot(session, attached)).toMatchObject({ contextTokens: 35_000, @@ -225,17 +236,74 @@ describe('SessionMetricsProjector', () => { expect(projector.snapshot(session, attached).contextWindow).toBe(128_000) }) - it('invalidates same-route capacity and fences the prior epoch in flight', async () => { + it('aborts every active capacity on invalidation and resolves fresh generations', async () => { const ctx = new Context() - const attempts: PromiseWithResolvers<{ context: { contextWindow: number } }>[] = [] + const attempts: { + result: PromiseWithResolvers<{ context: { contextWindow: number } }> + signal: AbortSignal | undefined + }[] = [] ctx.provide('llm', { - resolveModelInfo() { - const attempt = Promise.withResolvers<{ context: { contextWindow: number } }>() - attempts.push(attempt) - return attempt.promise + resolveModelInfo(_provider: string, _model: string, signal?: AbortSignal) { + const result = Promise.withResolvers<{ context: { contextWindow: number } }>() + attempts.push({ result, signal }) + return result.promise }, }) - const session = new Session(SessionId('capacity-invalidation')) + const firstSession = new Session(SessionId('capacity-invalidation-first')) + const secondSession = new Session(SessionId('capacity-invalidation-second')) + const firstAgent = agent(firstSession) + const secondAgent = agent(secondSession) + const resolved = vi.fn() + const projector = new SessionMetricsProjector( + ctx, + () => ({ provider: 'test', model: 'alpha' }), + resolved, + ) + + expect(projector.snapshot(firstSession, firstAgent).contextWindow).toBeUndefined() + expect(projector.snapshot(secondSession, secondAgent).contextWindow).toBeUndefined() + await vi.waitFor(() => { expect(attempts).toHaveLength(2) }) + expect(attempts.map(attempt => attempt.signal?.aborted)).toEqual([false, false]) + + projector.invalidateCapacities() + expect(attempts.map(attempt => attempt.signal?.aborted)).toEqual([true, true]) + attempts[0]?.result.resolve({ context: { contextWindow: 32_000 } }) + attempts[1]?.result.resolve({ context: { contextWindow: 64_000 } }) + await settleAsyncWork() + expect(resolved).not.toHaveBeenCalled() + + expect(projector.snapshot(firstSession, firstAgent).contextWindow).toBeUndefined() + expect(projector.snapshot(secondSession, secondAgent).contextWindow).toBeUndefined() + await vi.waitFor(() => { expect(attempts).toHaveLength(4) }) + expect(attempts.slice(2).map(attempt => attempt.signal?.aborted)).toEqual([false, false]) + attempts[2]?.result.resolve({ context: { contextWindow: 128_000 } }) + attempts[3]?.result.resolve({ context: { contextWindow: 256_000 } }) + await vi.waitFor(() => { expect(resolved).toHaveBeenCalledTimes(2) }) + expect(projector.snapshot(firstSession, firstAgent).contextWindow).toBe(128_000) + expect(projector.snapshot(secondSession, secondAgent).contextWindow).toBe(256_000) + + projector.dispose() + expect(projector.snapshot(firstSession, firstAgent).contextWindow).toBeUndefined() + expect(attempts).toHaveLength(4) + }) + + it('skips adapter work invalidated before its deferred invocation', async () => { + const ctx = new Context() + const attempts: { + result: PromiseWithResolvers<{ context: { contextWindow: number } }> + signal: AbortSignal | undefined + }[] = [] + const resolveModelInfo = vi.fn(( + _provider: string, + _model: string, + signal?: AbortSignal, + ) => { + const result = Promise.withResolvers<{ context: { contextWindow: number } }>() + attempts.push({ result, signal }) + return result.promise + }) + ctx.provide('llm', { resolveModelInfo }) + const session = new Session(SessionId('capacity-pre-invocation-invalidation')) const attached = agent(session) const resolved = vi.fn() const projector = new SessionMetricsProjector( @@ -245,26 +313,34 @@ describe('SessionMetricsProjector', () => { ) expect(projector.snapshot(session, attached).contextWindow).toBeUndefined() - await vi.waitFor(() => { expect(attempts).toHaveLength(1) }) projector.invalidateCapacities() - attempts[0]?.resolve({ context: { contextWindow: 64_000 } }) await settleAsyncWork() + expect(resolveModelInfo).not.toHaveBeenCalled() expect(resolved).not.toHaveBeenCalled() expect(projector.snapshot(session, attached).contextWindow).toBeUndefined() - await vi.waitFor(() => { expect(attempts).toHaveLength(2) }) - attempts[1]?.resolve({ context: { contextWindow: 128_000 } }) + await vi.waitFor(() => { expect(attempts).toHaveLength(1) }) + expect(attempts[0]?.signal?.aborted).toBe(false) + attempts[0]?.result.resolve({ context: { contextWindow: 128_000 } }) await vi.waitFor(() => { expect(resolved).toHaveBeenCalledOnce() }) expect(projector.snapshot(session, attached).contextWindow).toBe(128_000) }) it('starts a fresh capacity generation when an unavailable route returns', async () => { const ctx = new Context() - const resolutions: ((contextWindow: number) => void)[] = [] + const resolutions: { + signal: AbortSignal | undefined + resolve(contextWindow: number): void + }[] = [] ctx.provide('llm', { - resolveModelInfo() { + resolveModelInfo(_provider: string, _model: string, signal?: AbortSignal) { return new Promise<{ context: { contextWindow: number } }>((resolve) => { - resolutions.push((contextWindow) => { resolve({ context: { contextWindow } }) }) + resolutions.push({ + signal, + resolve(contextWindow) { + resolve({ context: { contextWindow } }) + }, + }) }) }, }) @@ -278,13 +354,16 @@ describe('SessionMetricsProjector', () => { expect(projector.snapshot(session, attached).contextWindow).toBeUndefined() await vi.waitFor(() => { expect(resolutions).toHaveLength(1) }) current = undefined - resolutions[0]?.(64_000) - await vi.waitFor(() => { expect(targetFor).toHaveBeenCalledTimes(2) }) + expect(projector.snapshot(session, attached).contextWindow).toBeUndefined() + expect(resolutions[0]?.signal?.aborted).toBe(true) + resolutions[0]?.resolve(64_000) + await settleAsyncWork() expect(resolved).not.toHaveBeenCalled() current = { provider: 'test', model: 'alpha' } expect(projector.snapshot(session, attached).contextWindow).toBeUndefined() await vi.waitFor(() => { expect(resolutions).toHaveLength(2) }) - resolutions[1]?.(128_000) + expect(resolutions[1]?.signal?.aborted).toBe(false) + resolutions[1]?.resolve(128_000) await vi.waitFor(() => { expect(resolved).toHaveBeenCalledOnce() }) expect(projector.snapshot(session, attached).contextWindow).toBe(128_000) })