Parcourir la source

fix(web): reset live metrics between connections

Hypatia May il y a 2 mois
Parent
commit
ef12895cf1

+ 7 - 2
packages/client/connection/src/client/connection.ts

@@ -46,6 +46,8 @@ export interface ConnectionSinks {
   onHostEnvelope?: (envelope: RpcRequest<HostFrame>) => void
   /** After each connection generation is established (both streams open + describe succeeded), first connect included. */
   onConnected?: () => void
+  /** After every failed generation closes and before retry starts. Not emitted when the controller is stopped. */
+  onDisconnected?: () => void
   /** Coarse state transitions (deduplicated: fires only on change). The initial pre-connect
    *  span reports nothing — the UI treats "no state yet" as connecting, not as an outage. */
   onStateChange?: (state: ConnectionState) => void
@@ -120,8 +122,8 @@ export class ConnectionController {
           if (gen === this.generation && !ac.signal.aborted) ac.abort()
           resolve()
         }
-        void this.pumpStream(this.api.events.mux({}, ac.signal, muxOpened), this.sinks.onMuxEnvelope, settle)
-        void this.pumpStream(this.api.events.host({}, ac.signal, hostOpened), this.sinks.onHostEnvelope, settle)
+        void this.pumpStream(this.api.events.mux({}, ac.signal, muxOpened), this.sinks.onMuxEnvelope, ac.signal, settle)
+        void this.pumpStream(this.api.events.host({}, ac.signal, hostOpened), this.sinks.onHostEnvelope, ac.signal, settle)
       })
 
       try {
@@ -147,6 +149,7 @@ export class ConnectionController {
 
       await failed
       if (!this.isRunning()) return
+      this.callSink(this.sinks.onDisconnected)
       this.emitState('reconnecting')
       this.attempt += 1
       console.warn(`[web-runtime] connection lost, retry #${this.attempt}`)
@@ -165,10 +168,12 @@ export class ConnectionController {
   private async pumpStream<F extends { type: string }>(
     stream: AsyncIterable<RpcRequest<F>>,
     sink: ((envelope: RpcRequest<F>) => void) | undefined,
+    signal: AbortSignal,
     onEnd: () => void,
   ): Promise<void> {
     try {
       for await (const envelope of stream) {
+        if (signal.aborted) break
         if (envelope.payload.type === 'stream/error') break
         if (sink !== undefined) this.callSink(() => { sink(envelope) })
       }

+ 50 - 2
packages/client/connection/tests/connection.spec.ts

@@ -7,7 +7,7 @@
  */
 
 import { describe, expect, it, vi } from 'vitest'
-import type { SessionId } from '../src/client/api.ts'
+import type { IApiClient, SessionId } from '../src/client/api.ts'
 import type { ConnectionState } from '../src/client/connection.ts'
 import { ConnectionController } from '../src/client/connection.ts'
 import { FakeApiClient, deferred, ok } from './fake-api.ts'
@@ -104,6 +104,51 @@ describe('connection lifecycle', () => {
     }
   })
 
+  it('drops a sibling stream frame buffered behind a generation failure', async () => {
+    const api = new FakeApiClient()
+    const lateMux = deferred<undefined>()
+    const originalEvents = api.events
+    Object.defineProperty(api, 'events', {
+      value: {
+        host: (...args: Parameters<IApiClient['events']['host']>) => originalEvents.host(...args),
+        mux: (_payload: unknown, _signal: AbortSignal, onOpen?: () => void) => (async function* () {
+          onOpen?.()
+          await lateMux.promise
+          yield { rpcId: 'late-mux' as never, payload: subscribedFrame(2) }
+        })(),
+      } satisfies IApiClient['events'],
+    })
+    const muxSeen: number[] = []
+    let connected = 0
+    let disconnected = 0
+    const warnSpy = vi.spyOn(console, 'warn').mockImplementation(() => undefined)
+    const controller = new ConnectionController(api, {
+      onMuxEnvelope: (envelope) => {
+        if (envelope.payload.type === 'session/subscribed') muxSeen.push(envelope.payload.lastSeq)
+      },
+      onConnected: () => { connected++ },
+      onDisconnected: () => {
+        disconnected++
+        lateMux.resolve(undefined)
+        controller.stop()
+      },
+    }, FAST)
+    controller.start()
+    try {
+      await vi.waitFor(() => { expect(connected).toBe(1) })
+      api.pushHost({
+        type: 'stream/error',
+        error: { code: 'internal', message: 'host stream failed', details: {} },
+      })
+      await vi.waitFor(() => { expect(disconnected).toBe(1) })
+      await new Promise(resolve => setTimeout(resolve, 0))
+      expect(muxSeen).toEqual([])
+    } finally {
+      controller.stop()
+      warnSpy.mockRestore()
+    }
+  })
+
   it('isolates sink exceptions from the pump', async () => {
     const api = new FakeApiClient()
     const seen: string[] = []
@@ -181,7 +226,7 @@ describe('connection lifecycle', () => {
     }
   })
 
-  it('deduplicates consecutive reconnecting emissions across two straight failures', async () => {
+  it('reports every failed generation while deduplicating consecutive reconnecting state', async () => {
     const api = new FakeApiClient()
     const gate = deferred<Awaited<ReturnType<FakeApiClient['onDescribe']>>>()
     let describeCalls = 0
@@ -190,10 +235,12 @@ describe('connection lifecycle', () => {
       return describeCalls <= 2 ? Promise.reject(new Error('down')) : gate.promise
     }
     const states: ConnectionState[] = []
+    let disconnected = 0
     let connected = 0
     const warnSpy = vi.spyOn(console, 'warn').mockImplementation(() => undefined)
     const controller = new ConnectionController(api, {
       onConnected: () => { connected++ },
+      onDisconnected: () => { disconnected++ },
       onStateChange: state => states.push(state),
     }, FAST)
     controller.start()
@@ -201,6 +248,7 @@ describe('connection lifecycle', () => {
       await vi.waitFor(() => { expect(describeCalls).toBe(3) })
       gate.resolve(ok({ version: '0', cwd: '/f', attachedSessions: 0 }))
       await vi.waitFor(() => { expect(connected).toBe(1) })
+      expect(disconnected).toBe(2)
       expect(states).toEqual(['reconnecting', 'connected']) // two failures, one reconnecting emission
     } finally {
       controller.stop()

+ 1 - 0
packages/client/runtime/src/client/index.ts

@@ -143,6 +143,7 @@ export function apply(ctx: Context): void {
       workspaces.handleConnected()
       ctx.emit('connection/reset')
     },
+    onDisconnected: () => { sessions.handleReconnecting() },
   })
   ctx.effect(() => () => { loop.stop() }, 'runtime: connection stream loop')
 }

+ 6 - 0
packages/client/runtime/src/client/sessions/manager.ts

@@ -430,6 +430,12 @@ export class SessionManager {
     for (const session of this.sessions.values()) void session.resync()
   }
 
+  /** Before a replacement stream generation, discard values the Host does not replay. */
+  handleReconnecting(): void {
+    this.modelRequestContextWindows.clear()
+    for (const session of this.sessions.values()) session.handleReconnecting()
+  }
+
   private buildListSnapshot(): SessionListSnapshot {
     const merged: TitledSessionSummary[] = this.summaries.map((summary) => {
       // List rows read the generic 'title' projection key (host-computed unit

+ 5 - 0
packages/client/runtime/src/client/sessions/service.ts

@@ -393,6 +393,11 @@ export class SessionsService {
     this.manager.handleConnected()
   }
 
+  /** Clear connection-local Session state before the next stream generation starts. */
+  handleReconnecting(): void {
+    this.manager.handleReconnecting()
+  }
+
   /**
    * Create a session on the host. Resolution guarantee: by the time the
    * promise resolves, the created session is in the list store and

+ 8 - 0
packages/client/runtime/src/client/sessions/session.ts

@@ -481,6 +481,14 @@ export class Session implements ObservableSnapshot<ConversationSnapshot> {
     this.notifier.markDirty()
   }
 
+  /** Connection-loss boundary: clear values that are not replayed before the next stream starts. */
+  handleReconnecting(): void {
+    if (this.metrics === null && this.contextWindow === undefined) return
+    this.metrics = null
+    this.contextWindow = undefined
+    this.notifier.markDirty()
+  }
+
   /** host/session-removed relay: flag the resident snapshot and clear connection-local capacity. */
   handleRemoved(): void {
     const changed = !this.removed || this.contextWindow !== undefined

+ 47 - 0
packages/client/runtime/tests/client-apply.spec.ts

@@ -102,6 +102,53 @@ describe('runtime client apply', () => {
     expect(bench.api.callsOf('session.create')).toHaveLength(1)
   })
 
+  it('clears connection-local Session state on disconnect but not connected', async () => {
+    const bench = await mount()
+    const sessions = bench.ctx.get('sessions') as SessionsService
+    bench.sinks?.onHostEnvelope?.({
+      rpcId: 'session' as never,
+      payload: { type: 'host/session-added', blank: true, sessionId: 's-state' } as never,
+    })
+    await Promise.resolve()
+    const session = sessions.binding('s-state' as never)?.session
+    if (session === undefined) throw new Error('session binding missing')
+    const currentMetrics = {
+      projectionRevision: 4,
+      logRevision: 10,
+      uncachedInputTokens: 10,
+      outputTokens: 4,
+      cacheReadTokens: 90,
+      cacheWriteTokens: 3,
+      contextTokens: 35,
+    }
+    bench.sinks?.onMuxEnvelope?.({
+      rpcId: 'metrics' as never,
+      payload: { type: 'session/metrics', sessionId: 's-state', metrics: currentMetrics } as never,
+    })
+    bench.sinks?.onMuxEnvelope?.({
+      rpcId: 'capacity' as never,
+      payload: {
+        type: 'session/model-request',
+        sessionId: 's-state',
+        turn: 1,
+        step: 1,
+        provider: 'test',
+        model: 'alpha',
+        contextWindow: 128_000,
+      } as never,
+    })
+
+    bench.sinks?.onConnected?.()
+    expect(session.getSnapshot()).toMatchObject({
+      metrics: currentMetrics,
+      modelRequestContextWindow: 128_000,
+    })
+
+    bench.sinks?.onDisconnected?.()
+    expect(session.getSnapshot().metrics).toBeNull()
+    expect(session.getSnapshot().modelRequestContextWindow).toBeUndefined()
+  })
+
   it('stops the stream loop when the plugin fiber unloads', async () => {
     const bench = await mount()
     const fiber = [...bench.ctx.registry.values()].find(f => f.name?.includes('client'))

+ 54 - 1
packages/client/runtime/tests/manager.spec.ts

@@ -4,7 +4,7 @@
  */
 
 import { describe, expect, it, vi } from 'vitest'
-import type { SessionId } from '@deepseek-ai/dsh-client-connection/client'
+import type { SessionId, SessionMetrics } from '@deepseek-ai/dsh-client-connection/client'
 import { SessionManager } from '../src/client/sessions/manager.ts'
 import { FakeApiClient, deferred, err, ok } from './fake-api.ts'
 import { entries, plainTurn } from './event-script.ts'
@@ -162,6 +162,59 @@ describe('instances', () => {
     expect(manager.get(S2).getSnapshot().modelRequestContextWindow).toBeUndefined()
   })
 
+  it('clears resident metrics and capacity plus lazy capacity before reconnect', () => {
+    const api = new FakeApiClient()
+    const manager = new SessionManager(api)
+    const session = manager.get(S1)
+    const currentMetrics: SessionMetrics = {
+      projectionRevision: 4,
+      logRevision: 10,
+      uncachedInputTokens: 10,
+      outputTokens: 4,
+      cacheReadTokens: 90,
+      cacheWriteTokens: 3,
+      contextTokens: 35,
+    }
+    manager.handleMuxEnvelope({
+      rpcId: 'metrics' as never,
+      payload: { type: 'session/metrics', sessionId: S1, metrics: currentMetrics },
+    })
+    manager.handleMuxEnvelope({
+      rpcId: 'resident-capacity' as never,
+      payload: {
+        type: 'session/model-request',
+        sessionId: S1,
+        turn: 1,
+        step: 1,
+        provider: 'test',
+        model: 'resident',
+        contextWindow: 128_000,
+      },
+    })
+    manager.handleMuxEnvelope({
+      rpcId: 'lazy-capacity' as never,
+      payload: {
+        type: 'session/model-request',
+        sessionId: S2,
+        turn: 1,
+        step: 1,
+        provider: 'test',
+        model: 'lazy',
+        contextWindow: 256_000,
+      },
+    })
+    expect(session.getSnapshot()).toMatchObject({
+      metrics: currentMetrics,
+      modelRequestContextWindow: 128_000,
+    })
+
+    manager.handleReconnecting()
+
+    expect(session.getSnapshot().metrics).toBeNull()
+    expect(session.getSnapshot().modelRequestContextWindow).toBeUndefined()
+    expect(manager.get(S2).getSnapshot().modelRequestContextWindow).toBeUndefined()
+  })
+
   it('caps the pending buffer at 32 keeping the newest, and drops it on session-removed', () => {
     const api = new FakeApiClient()
     const manager = new SessionManager(api)