From 2b9ab7401525c346c516cfb6428d0122023b6a85 Mon Sep 17 00:00:00 2001 From: qer Date: Wed, 22 Jul 2026 18:29:17 +0800 Subject: [PATCH] feat(kap-server): register global WS targets --- .changeset/ws-global-target.md | 5 + .../ws/v1/sessionEventBroadcaster.ts | 29 ++- .../src/transport/ws/v1/wsConnectionV1.ts | 14 ++ .../test/sessionEventBroadcaster.test.ts | 45 +++++ .../kap-server/test/wsConnectionV1.test.ts | 184 ++++++++++++++++++ 5 files changed, 276 insertions(+), 1 deletion(-) create mode 100644 .changeset/ws-global-target.md diff --git a/.changeset/ws-global-target.md b/.changeset/ws-global-target.md new file mode 100644 index 0000000000..503b4061db --- /dev/null +++ b/.changeset/ws-global-target.md @@ -0,0 +1,5 @@ +--- +"@moonshot-ai/kimi-code": patch +--- + +Allow WebSocket connections to receive global session events without per-session subscriptions. diff --git a/packages/kap-server/src/transport/ws/v1/sessionEventBroadcaster.ts b/packages/kap-server/src/transport/ws/v1/sessionEventBroadcaster.ts index 200f7311d4..ce3a7152aa 100644 --- a/packages/kap-server/src/transport/ws/v1/sessionEventBroadcaster.ts +++ b/packages/kap-server/src/transport/ws/v1/sessionEventBroadcaster.ts @@ -30,6 +30,10 @@ * A session is activated (journaling starts) on first `subscribe` / * `getSnapshotState` / `getCursor` and stays active for the process lifetime so * the journal is continuous from first activation onward. + * + * Global events (`isGlobalEvent`) fan out to the union of every session's + * subscribers and the connection-level global targets registered after + * `client_hello`. Per-session frames still reach subscribers only. */ import type { @@ -205,6 +209,8 @@ export class SessionEventBroadcaster { * per-chunk `AABBCC` stream while every seq and offset still looks valid. */ private readonly pendingStates = new Map>(); + /** Connections that completed `client_hello`, independent of session subscriptions. */ + private readonly globalTargets = new Set(); private readonly maxBufferSize: number; private readonly coreEventSubscription: IDisposable; private closed = false; @@ -229,6 +235,16 @@ export class SessionEventBroadcaster { .subscribe((event) => this.onCoreEvent(event)); } + /** Register a connection for global fan-out until it closes. */ + registerGlobalTarget(target: BroadcastTarget): void { + if (this.closed) return; + this.globalTargets.add(target); + } + + unregisterGlobalTarget(target: BroadcastTarget): void { + this.globalTargets.delete(target); + } + /** * Subscribe a connection to a session's stream (activates the session). * @@ -618,6 +634,7 @@ export class SessionEventBroadcaster { if (this.closed) return; this.closed = true; this.coreEventSubscription.dispose(); + this.globalTargets.clear(); for (const [sessionId, state] of this.sessions) { await disposeSessionState(state); // Transcript bindings die with the session stream (its store @@ -1165,9 +1182,19 @@ export class SessionEventBroadcaster { }; } + /** Yield global targets and session subscribers without duplicates. */ private *allTargets(): Iterable { + const seen = new Set(); + for (const target of this.globalTargets) { + seen.add(target); + yield target; + } for (const state of this.sessions.values()) { - for (const target of state.targets.keys()) yield target; + for (const target of state.targets.keys()) { + if (seen.has(target)) continue; + seen.add(target); + yield target; + } } } } diff --git a/packages/kap-server/src/transport/ws/v1/wsConnectionV1.ts b/packages/kap-server/src/transport/ws/v1/wsConnectionV1.ts index 1536873993..fdac222f00 100644 --- a/packages/kap-server/src/transport/ws/v1/wsConnectionV1.ts +++ b/packages/kap-server/src/transport/ws/v1/wsConnectionV1.ts @@ -9,6 +9,9 @@ * `{seq, epoch}` cursor, or sends `resync_required` when the gap cannot be * served incrementally. * + * A successful `client_hello` also registers the connection for global event + * fan-out, independent of per-session subscriptions, until the socket closes. + * * The server never initiates a disconnect: unlike v1's `WsConnection` * (`packages/server/src/ws/connection.ts`) there is no ping/pong heartbeat — * a connection stays open until the client closes it or the process shuts @@ -191,6 +194,7 @@ export class WsConnectionV1 implements BroadcastTarget { private async onClientHello(frame: InboundFrame): Promise { if (!(await this.authorize(frame))) return; + if (this.closed) return; this.gotClientHello = true; const payload = frame.payload ?? {}; @@ -215,6 +219,11 @@ export class WsConnectionV1 implements BroadcastTarget { ); } + // Register after requested cursor replays so live global events cannot + // arrive before older durable backlog for the same session. + if (this.closed) return; + this.broadcaster.registerGlobalTarget(this); + this.sendFrame( buildAck(frame.id ?? '', 0, 'success', { accepted_subscriptions: accepted, @@ -329,6 +338,10 @@ export class WsConnectionV1 implements BroadcastTarget { const ok = await this.broadcaster.subscribe(sid, this, filter, transcriptGrades, { deferTranscriptReset: cursor !== undefined, }); + if (this.closed) { + if (ok) this.broadcaster.unsubscribe(sid, this); + return; + } if (!ok) { resyncRequired.push(sid); return; @@ -478,6 +491,7 @@ export class WsConnectionV1 implements BroadcastTarget { if (this.flushTimer !== undefined) clearTimeout(this.flushTimer); if (this.backpressureRetryTimer !== undefined) clearTimeout(this.backpressureRetryTimer); this.outbound = []; + this.broadcaster.unregisterGlobalTarget(this); for (const sid of this.subscriptions.keys()) this.broadcaster.unsubscribe(sid, this); this.fsWatchBridge?.detachConnection(this); // registry removal is handled by registerWsV1 on the socket 'close' event. diff --git a/packages/kap-server/test/sessionEventBroadcaster.test.ts b/packages/kap-server/test/sessionEventBroadcaster.test.ts index 54720a8279..92023f5918 100644 --- a/packages/kap-server/test/sessionEventBroadcaster.test.ts +++ b/packages/kap-server/test/sessionEventBroadcaster.test.ts @@ -1154,6 +1154,51 @@ describe('SessionEventBroadcaster', () => { expect(s1View.envelopes.some((e) => e.type === 'turn.started')).toBe(false); }); + it('fans global events out to hello-only targets without leaking session frames', async () => { + const lc = new FakeLifecycle(); + const main = lc.addAgent('main'); + sessions.set('s1', lc); + + // A regular subscriber activates the session producer. The global target + // models a connection that completed client_hello but did not subscribe. + const subscribed = collectingTarget(); + const global = collectingTarget(); + await bc.subscribe('s1', subscribed.target); + bc.registerGlobalTarget(global.target); + + main.bus.emit(agentEvent('turn.started', { turnId: 1 })); + await bc.getCursor('s1'); + + expect(global.envelopes.map((e) => e.type)).toEqual(['event.session.work_changed']); + expect(global.envelopes[0]).toMatchObject({ + session_id: 's1', + payload: { busy: true }, + }); + expect(global.envelopes.some((e) => e.type === 'turn.started')).toBe(false); + + bc.unregisterGlobalTarget(global.target); + main.bus.emit(agentEvent('turn.ended', { turnId: 1, reason: 'completed' })); + await bc.getCursor('s1'); + expect(global.envelopes).toHaveLength(1); + }); + + it('sends a global event once to a target subscribed to multiple sessions', async () => { + const lc = new FakeLifecycle(); + const main = lc.addAgent('main'); + sessions.set('s1', lc); + sessions.set('s2', new FakeLifecycle()); + + const view = collectingTarget(); + await bc.subscribe('s1', view.target); + await bc.subscribe('s2', view.target); + + main.bus.emit(agentEvent('turn.started', { turnId: 1 })); + await bc.getCursor('s1'); + + expect(view.envelopes.filter((e) => e.type === 'event.session.work_changed')).toHaveLength(1); + expect(view.envelopes.filter((e) => e.type === 'turn.started')).toHaveLength(1); + }); + it('does not re-announce interactions already pending at activation, but still broadcasts their resolution', async () => { const lc = new FakeLifecycle(); lc.addAgent('main'); diff --git a/packages/kap-server/test/wsConnectionV1.test.ts b/packages/kap-server/test/wsConnectionV1.test.ts index 3c91224b8f..8a95c9e66a 100644 --- a/packages/kap-server/test/wsConnectionV1.test.ts +++ b/packages/kap-server/test/wsConnectionV1.test.ts @@ -62,6 +62,8 @@ function makeBroadcaster(): SessionEventBroadcaster { return { subscribe: async () => true, unsubscribe: () => {}, + registerGlobalTarget: () => {}, + unregisterGlobalTarget: () => {}, getCursor: async () => ({ seq: 0, epoch: '' }), getBufferedSince: async () => ({ events: [], @@ -226,6 +228,8 @@ describe('WsConnectionV1 transcript subscriptions', () => { return true; }, unsubscribe: () => {}, + registerGlobalTarget: () => {}, + unregisterGlobalTarget: () => {}, getCursor: async () => ({ seq: 0, epoch: '' }), getBufferedSince: async () => ({ events: [], @@ -315,6 +319,8 @@ describe('WsConnectionV1 transcript subscriptions', () => { target.send({ type: 'transcript.reset', seq: 10, session_id: 's1', payload: {} }); }, unsubscribe: () => {}, + registerGlobalTarget: () => {}, + unregisterGlobalTarget: () => {}, getCursor: async () => ({ seq: 10, epoch: 'e1' }), getBufferedSince: async () => ({ events: backlog.map((envelope) => ({ seq: envelope.seq, envelope })), @@ -372,6 +378,184 @@ describe('WsConnectionV1 transcript subscriptions', () => { }); }); +// --------------------------------------------------------------------------- +// WsConnectionV1 — global target registration +// --------------------------------------------------------------------------- + +describe('WsConnectionV1 global target registration', () => { + function makeRegisteringBroadcaster(): { + broadcaster: SessionEventBroadcaster; + registered: unknown[]; + unregistered: unknown[]; + } { + const registered: unknown[] = []; + const unregistered: unknown[] = []; + const broadcaster = { + subscribe: async () => true, + unsubscribe: () => {}, + registerGlobalTarget: (target: unknown) => registered.push(target), + unregisterGlobalTarget: (target: unknown) => unregistered.push(target), + getCursor: async () => ({ seq: 0, epoch: '' }), + getBufferedSince: async () => ({ + events: [], + resyncRequired: false, + currentSeq: 0, + epoch: '', + }), + } as unknown as SessionEventBroadcaster; + return { broadcaster, registered, unregistered }; + } + + it('registers after client_hello and unregisters on close', async () => { + const socket = new FakeSocket(); + const { broadcaster, registered, unregistered } = makeRegisteringBroadcaster(); + const conn = makeConn(socket, { broadcaster }); + + socket.emit( + 'message', + JSON.stringify({ type: 'client_hello', id: 'h1', payload: { client_id: 'c1' } }), + ); + await vi.waitFor(() => expect(registered).toEqual([conn])); + + socket.emit('close'); + expect(unregistered).toEqual([conn]); + }); + + it('registers only after requested cursor replay completes', async () => { + const socket = new FakeSocket(); + const order: string[] = []; + const broadcaster = { + subscribe: async () => { + order.push('subscribe'); + return true; + }, + unsubscribe: () => {}, + registerGlobalTarget: () => order.push('register-global'), + unregisterGlobalTarget: () => {}, + getCursor: async () => ({ seq: 1, epoch: 'e1' }), + getBufferedSince: async () => { + order.push('replay'); + return { + events: [], + resyncRequired: false, + currentSeq: 1, + epoch: 'e1', + }; + }, + flushTranscriptSeed: async () => order.push('flush-transcript'), + } as unknown as SessionEventBroadcaster; + const conn = makeConn(socket, { broadcaster }); + + socket.emit( + 'message', + JSON.stringify({ + type: 'client_hello', + id: 'h1', + payload: { + client_id: 'c1', + subscriptions: ['s1'], + cursors: { s1: { seq: 0, epoch: 'e1' } }, + }, + }), + ); + await vi.waitFor(() => expect(order).toContain('register-global')); + + expect(order).toEqual(['subscribe', 'replay', 'flush-transcript', 'register-global']); + conn.close(); + }); + + it('stops client_hello when the socket closes during authorization', async () => { + const socket = new FakeSocket(); + let releaseAuthorization!: () => void; + const authorizationGate = new Promise((resolve) => { + releaseAuthorization = resolve; + }); + const authorizationFinished = vi.fn(); + const subscribe = vi.fn(async () => true); + const registered: unknown[] = []; + const broadcaster = { + subscribe, + unsubscribe: () => {}, + registerGlobalTarget: (target: unknown) => registered.push(target), + unregisterGlobalTarget: () => {}, + getCursor: async () => ({ seq: 0, epoch: '' }), + getBufferedSince: async () => ({ + events: [], + resyncRequired: false, + currentSeq: 0, + epoch: '', + }), + } as unknown as SessionEventBroadcaster; + makeConn(socket, { + broadcaster, + validateCredential: async () => { + await authorizationGate; + authorizationFinished(); + return true; + }, + }); + + socket.emit( + 'message', + JSON.stringify({ + type: 'client_hello', + id: 'h1', + payload: { client_id: 'c1', token: 'token', subscriptions: ['s1'] }, + }), + ); + socket.emit('close'); + releaseAuthorization(); + await vi.waitFor(() => expect(authorizationFinished).toHaveBeenCalledOnce()); + + expect(subscribe).not.toHaveBeenCalled(); + expect(registered).toHaveLength(0); + }); + + it('undoes a subscription that resolves after the socket closes', async () => { + const socket = new FakeSocket(); + let releaseSubscription!: () => void; + const subscriptionGate = new Promise((resolve) => { + releaseSubscription = resolve; + }); + const subscriptionStarted = vi.fn(); + const unsubscribed: string[] = []; + const registered: unknown[] = []; + const broadcaster = { + subscribe: async () => { + subscriptionStarted(); + await subscriptionGate; + return true; + }, + unsubscribe: (sessionId: string) => unsubscribed.push(sessionId), + registerGlobalTarget: (target: unknown) => registered.push(target), + unregisterGlobalTarget: () => {}, + getCursor: async () => ({ seq: 0, epoch: '' }), + getBufferedSince: async () => ({ + events: [], + resyncRequired: false, + currentSeq: 0, + epoch: '', + }), + } as unknown as SessionEventBroadcaster; + makeConn(socket, { broadcaster }); + + socket.emit( + 'message', + JSON.stringify({ + type: 'client_hello', + id: 'h1', + payload: { client_id: 'c1', subscriptions: ['s1'] }, + }), + ); + await vi.waitFor(() => expect(subscriptionStarted).toHaveBeenCalledOnce()); + socket.emit('close'); + releaseSubscription(); + await vi.waitFor(() => expect(unsubscribed).toEqual(['s1'])); + + expect(registered).toHaveLength(0); + }); +}); + // --------------------------------------------------------------------------- // WsConnectionV1 — flush / backpressure / close // ---------------------------------------------------------------------------