fix(host): cancel retired metric capacity lookups (round 7)

2026-08-15 21:04:50 +00:00 · 2026-07-28 17:03:07 +08:00
parent 032cd2f72d
commit 344667cbf0
4 changed files with 287 additions and 29 deletions
--- 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()
    }
--- 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<AgentLlmTarget, 'provider' | 'model'>
@@ -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<Session, UsageState>()
  private readonly capacities = new WeakMap<Agent, CapacityState>()
+  private readonly pendingCapacities = new Set<CapacityState>()
  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
--- 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<LlmResolvedModelInfo>[] = []
+  readonly pending: {
+    result: PromiseWithResolvers<LlmResolvedModelInfo>
+    signal: AbortSignal | undefined
+  }[] = []

  constructor() {
    super('Deferred', [
@@ -73,16 +76,20 @@ class DeferredCatalogAdapter extends CatalogAdapter {
    ])
  }

-  override resolveModel(_provider: string, _model: string): Promise<LlmResolvedModelInfo> {
+  override resolveModel(
+    _provider: string,
+    _model: string,
+    signal?: AbortSignal,
+  ): Promise<LlmResolvedModelInfo> {
    const result = Promise.withResolvers<LlmResolvedModelInfo>()
-    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<undefined>()
+    const releaseCancellationWait = Promise.withResolvers<undefined>()
+    const abortObserved = Promise.withResolvers<undefined>()
+    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<void>((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<ReturnType<typeof createApiProxy>>()
+    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')
    },
--- 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<string, (contextWindow: number) => void>()
+    const resolutions = new Map<string, {
+      signal: AbortSignal | undefined
+      resolve(contextWindow: number): void
+    }>()
    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)
  })