From ef12895cf1c82cf2e6d233924bfc9df03020df57 Mon Sep 17 00:00:00 2001 From: Hypatia May Date: Tue, 28 Jul 2026 20:28:01 +0800 Subject: [PATCH] fix(web): reset live metrics between connections --- .../connection/src/client/connection.ts | 9 ++- .../connection/tests/connection.spec.ts | 52 +++++++++++++++++- packages/client/runtime/src/client/index.ts | 1 + .../runtime/src/client/sessions/manager.ts | 6 ++ .../runtime/src/client/sessions/service.ts | 5 ++ .../runtime/src/client/sessions/session.ts | 8 +++ .../client/runtime/tests/client-apply.spec.ts | 47 ++++++++++++++++ packages/client/runtime/tests/manager.spec.ts | 55 ++++++++++++++++++- 8 files changed, 178 insertions(+), 5 deletions(-) diff --git a/packages/client/connection/src/client/connection.ts b/packages/client/connection/src/client/connection.ts index 6eb6491e2f..a15312c9f9 100644 --- a/packages/client/connection/src/client/connection.ts +++ b/packages/client/connection/src/client/connection.ts @@ -46,6 +46,8 @@ export interface ConnectionSinks { onHostEnvelope?: (envelope: RpcRequest) => 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( stream: AsyncIterable>, sink: ((envelope: RpcRequest) => void) | undefined, + signal: AbortSignal, onEnd: () => void, ): Promise { 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) }) } diff --git a/packages/client/connection/tests/connection.spec.ts b/packages/client/connection/tests/connection.spec.ts index 4de4a31f25..9c54fcbcff 100644 --- a/packages/client/connection/tests/connection.spec.ts +++ b/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() + const originalEvents = api.events + Object.defineProperty(api, 'events', { + value: { + host: (...args: Parameters) => 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>>() 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() diff --git a/packages/client/runtime/src/client/index.ts b/packages/client/runtime/src/client/index.ts index 48514a9df4..f971331ba8 100644 --- a/packages/client/runtime/src/client/index.ts +++ b/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') } diff --git a/packages/client/runtime/src/client/sessions/manager.ts b/packages/client/runtime/src/client/sessions/manager.ts index b6daecb97c..3443f60dd2 100644 --- a/packages/client/runtime/src/client/sessions/manager.ts +++ b/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 diff --git a/packages/client/runtime/src/client/sessions/service.ts b/packages/client/runtime/src/client/sessions/service.ts index 19e022d450..e342450b94 100644 --- a/packages/client/runtime/src/client/sessions/service.ts +++ b/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 diff --git a/packages/client/runtime/src/client/sessions/session.ts b/packages/client/runtime/src/client/sessions/session.ts index 2eec667e46..49035736dc 100644 --- a/packages/client/runtime/src/client/sessions/session.ts +++ b/packages/client/runtime/src/client/sessions/session.ts @@ -481,6 +481,14 @@ export class Session implements ObservableSnapshot { 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 diff --git a/packages/client/runtime/tests/client-apply.spec.ts b/packages/client/runtime/tests/client-apply.spec.ts index d5b29f10a9..7399c46f6b 100644 --- a/packages/client/runtime/tests/client-apply.spec.ts +++ b/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')) diff --git a/packages/client/runtime/tests/manager.spec.ts b/packages/client/runtime/tests/manager.spec.ts index 410058228d..874e0be477 100644 --- a/packages/client/runtime/tests/manager.spec.ts +++ b/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)