fix(web): reset live metrics between connections

2026-07-28 20:28:01 +08:00
parent 47205cbb82
commit ef12895cf1
8 changed files with 178 additions and 5 deletions
@@ -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) })
      }
@@ -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()
@@ -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')
 }
@@ -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
@@ -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
@@ -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
@@ -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'))
@@ -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)