From fdb65be4aa2d60c668594cda5ba6efd53ab065fd Mon Sep 17 00:00:00 2001 From: "haozhe.yang" Date: Wed, 8 Jul 2026 17:01:59 +0800 Subject: [PATCH] feat(ws-v1): add per-agent event subscription filter - protocol: add optional agent_filter to client_hello and subscribe - kap-server: carry per-subscription agent allowlists through the broadcaster and connection, narrowing live fan-out and replay to selected agents while keeping a single global sequence and bypassing the filter for global events - agent-core-v2: degrade MiniDbQueryStore to a no-op read model when the query-store lock is held by another process instead of crashing the host --- .../backends/minidb/miniDbQueryStore.ts | 29 +++++- .../test/persistence/miniDbQueryStore.test.ts | 23 +++++ .../ws/v1/sessionEventBroadcaster.ts | 95 +++++++++++++++---- .../src/transport/ws/v1/wsConnectionV1.ts | 57 ++++++++--- .../test/sessionEventBroadcaster.test.ts | 84 ++++++++++++++++ packages/kap-server/test/wsV1Resync.test.ts | 42 ++++++++ .../protocol/src/__tests__/ws-control.test.ts | 52 ++++++++++ packages/protocol/src/ws-control.ts | 13 +++ 8 files changed, 362 insertions(+), 33 deletions(-) diff --git a/packages/agent-core-v2/src/persistence/backends/minidb/miniDbQueryStore.ts b/packages/agent-core-v2/src/persistence/backends/minidb/miniDbQueryStore.ts index 4cabec2ec..a83f9eb1d 100644 --- a/packages/agent-core-v2/src/persistence/backends/minidb/miniDbQueryStore.ts +++ b/packages/agent-core-v2/src/persistence/backends/minidb/miniDbQueryStore.ts @@ -61,7 +61,7 @@ export class MiniDbQueryStore extends Disposable implements IQueryStore { declare readonly _serviceBrand: undefined; private readonly dir: string; - private dbPromise: Promise | undefined; + private dbPromise: Promise | undefined; private readonly ensuredIndexes = new Set(); constructor( @@ -75,7 +75,7 @@ export class MiniDbQueryStore extends Disposable implements IQueryStore { })); } - private openDb(): Promise { + private openDb(): Promise { if (this.dbPromise !== undefined) return this.dbPromise; this.dbPromise = MiniDb.openOrRebuild( { @@ -92,18 +92,33 @@ export class MiniDbQueryStore extends Disposable implements IQueryStore { }); }, }, - ); + ).catch((err) => { + // The query store is a rebuildable derived read model; authoritative data + // lives in the append-log / atomic-document stores. If it cannot be opened + // — typically because another kimi process holds the single-writer lock on + // `/query-store` — degrade to a no-op store instead of crashing + // the host. Consumers then fall back to their non-read-model paths. The + // degraded state is memoized for the process lifetime so the warning is + // emitted once and we do not hammer a contended lock. + this.log.warn('minidb query-store unavailable; disabling read model', { + dir: this.dir, + error: String(err), + }); + return undefined; + }); return this.dbPromise; } async put(collection: string, key: string, value: T): Promise { const db = await this.openDb(); + if (db === undefined) return; await db.set(physicalKey(collection, key), value); } async batch(ops: readonly WriteOp[]): Promise { if (ops.length === 0) return; const db = await this.openDb(); + if (db === undefined) return; await db.batch( ops.map((op) => op.kind === 'put' @@ -115,11 +130,13 @@ export class MiniDbQueryStore extends Disposable implements IQueryStore { async delete(collection: string, key: string): Promise { const db = await this.openDb(); + if (db === undefined) return; await db.del(physicalKey(collection, key)); } async get(collection: string, key: string): Promise { const db = await this.openDb(); + if (db === undefined) return undefined; return db.get(physicalKey(collection, key)) as T | undefined; } @@ -131,6 +148,7 @@ export class MiniDbQueryStore extends Disposable implements IQueryStore { const guard = `${collection}:${def.kind}:${def.name}`; if (this.ensuredIndexes.has(guard)) return; const db = await this.openDb(); + if (db === undefined) return; const name = indexName(collection, def.name); if (def.kind === 'value') { if (!db.listIndexes().some((i) => i.name === name)) { @@ -163,7 +181,7 @@ export class MiniDbQueryStore extends Disposable implements IQueryStore { async close(): Promise { if (this.dbPromise === undefined) return; const db = await this.dbPromise; - await db.close(); + await db?.close(); } } @@ -175,7 +193,7 @@ class MiniDbQuery implements IQuery { private skip = 0; constructor( - private readonly openDb: () => Promise, + private readonly openDb: () => Promise, private readonly collection: string, ) {} @@ -202,6 +220,7 @@ class MiniDbQuery implements IQuery { async execute(): Promise> { const db = await this.openDb(); + if (db === undefined) return { items: [] }; const prefix = `${this.collection}${SEP}`; const q: QueryOptions = { key: { prefix } }; if (Object.keys(this.filter).length > 0) q.filter = this.filter as Record; diff --git a/packages/agent-core-v2/test/persistence/miniDbQueryStore.test.ts b/packages/agent-core-v2/test/persistence/miniDbQueryStore.test.ts index 2990bd54c..b767a18db 100644 --- a/packages/agent-core-v2/test/persistence/miniDbQueryStore.test.ts +++ b/packages/agent-core-v2/test/persistence/miniDbQueryStore.test.ts @@ -9,6 +9,7 @@ import { LifecycleScope, _clearScopedRegistryForTests, registerScopedService } f import { createScopedTestHost, stubPair } from '#/_base/di/test'; import { ILogService } from '#/_base/log/log'; import { IBootstrapService } from '#/app/bootstrap/bootstrap'; +import { MiniDb } from '@moonshot-ai/minidb'; import { MiniDbQueryStore } from '#/persistence/backends/minidb/miniDbQueryStore'; import { IQueryStore } from '#/persistence/interface/queryStore'; import { stubBootstrap } from '../bootstrap/stubs'; @@ -131,6 +132,28 @@ describe('MiniDbQueryStore', () => { expect(await store.getCheckpoint('wire:abc')).toEqual({ seq: 42 }); }); + it('degrades gracefully when the database lock is held by another process', async () => { + // Simulate another kimi process holding the single-writer lock on the + // shared query-store directory. + const storeDir = join(homeDir, 'cache', 'query-store'); + const lockHolder = await MiniDb.open({ dir: storeDir, valueCodec: 'json' }); + try { + const store = build(); + // Writes must not throw: the read model is a rebuildable cache, so lock + // contention degrades to a no-op rather than crashing the host. + await expect(store.put(COLLECTION, 'a', { id: 'a' })).resolves.toBeUndefined(); + await expect(store.batch([{ kind: 'put', collection: COLLECTION, key: 'b', value: { id: 'b' } }])).resolves.toBeUndefined(); + await expect(store.ensureIndex(COLLECTION, { kind: 'value', name: 'byId', field: 'id' })).resolves.toBeUndefined(); + // Reads fall back to empty so consumers (e.g. session index) take their + // legacy path. + expect(await store.get(COLLECTION, 'a')).toBeUndefined(); + expect(await store.getCheckpoint('wire:abc')).toBeUndefined(); + expect(await store.query(COLLECTION).execute()).toEqual({ items: [] }); + } finally { + await lockHolder.close(); + } + }); + it('rebuilds a clean store after on-disk corruption', async () => { const first = build(); await first.put(COLLECTION, 'a', { v: 1 }); diff --git a/packages/kap-server/src/transport/ws/v1/sessionEventBroadcaster.ts b/packages/kap-server/src/transport/ws/v1/sessionEventBroadcaster.ts index b9499acd2..a2e74dbcf 100644 --- a/packages/kap-server/src/transport/ws/v1/sessionEventBroadcaster.ts +++ b/packages/kap-server/src/transport/ws/v1/sessionEventBroadcaster.ts @@ -92,14 +92,22 @@ export interface BroadcastTarget { send(envelope: EventEnvelope): void; } +/** + * Per-subscription agent allowlist for fine-grained v1 event delivery. + * `undefined` (or omitted) means "receive every agent" — the legacy + * session-grained behavior. A `ReadonlySet` restricts delivery to the listed + * agent ids; global events ({@link isGlobalEvent}) bypass the filter entirely. + */ +export type AgentFilter = ReadonlySet | undefined; + interface SessionState { readonly sessionId: string; readonly journal: SessionEventJournal; readonly tracker: InFlightTurnTracker; /** Recent durable envelopes for in-memory replay. */ readonly tail: Array<{ seq: number; envelope: EventEnvelope }>; - /** Connections subscribed to this session. */ - readonly targets: Set; + /** Connections subscribed to this session, each with its optional agent allowlist. */ + readonly targets: Map; /** Per-session dispatch queue — serializes stamp / journal / fan-out. */ queue: Promise; /** agentId → sink subscription. */ @@ -132,10 +140,14 @@ export class SessionEventBroadcaster { } /** Subscribe a connection to a session's stream (activates the session). */ - async subscribe(sessionId: string, target: BroadcastTarget): Promise { + async subscribe( + sessionId: string, + target: BroadcastTarget, + filter?: AgentFilter, + ): Promise { const state = await this.ensureState(sessionId); if (state === undefined) return false; - state.targets.add(target); + state.targets.set(target, filter); return true; } @@ -143,7 +155,11 @@ export class SessionEventBroadcaster { this.sessions.get(sessionId)?.targets.delete(target); } - async getBufferedSince(sessionId: string, cursor: SessionCursor): Promise { + async getBufferedSince( + sessionId: string, + cursor: SessionCursor, + filter?: AgentFilter, + ): Promise { const state = await this.ensureState(sessionId); if (state === undefined) { return { events: [], resyncRequired: 'session_recreated', currentSeq: 0, epoch: '' }; @@ -168,14 +184,24 @@ export class SessionEventBroadcaster { return { events: [], resyncRequired: 'buffer_overflow', currentSeq, epoch }; } + // Filter is a view crop over the session's single durable sequence: the + // watermark and overflow checks above stay global, only the returned + // envelopes are narrowed to the subscriber's agent allowlist. + const applyFilter = ( + entries: Array<{ seq: number; envelope: EventEnvelope }>, + ): Array<{ seq: number; envelope: EventEnvelope }> => + filter === undefined + ? entries + : entries.filter(({ envelope }) => matchesAgentFilter(envelope, filter)); + // Serve from the memory tail when it fully covers the gap; else the journal. const tailStart = tail[0]?.seq; if (tailStart !== undefined && tailStart <= cursor.seq + 1) { - const events = tail.filter((e) => e.seq > cursor.seq); + const events = applyFilter(tail.filter((e) => e.seq > cursor.seq)); return { events, resyncRequired: false, currentSeq, epoch }; } const fromDisk = await journal.readSince(cursor.seq, this.maxBufferSize); - return { events: fromDisk, resyncRequired: false, currentSeq, epoch }; + return { events: applyFilter(fromDisk), resyncRequired: false, currentSeq, epoch }; } async getCursor(sessionId: string): Promise<{ seq: number; epoch: string }> { @@ -252,7 +278,7 @@ export class SessionEventBroadcaster { journal, tracker: new InFlightTurnTracker(), tail: [], - targets: new Set(), + targets: new Map(), queue: Promise.resolve(), agentDisposables: new Map(), lifecycleDisposables: [], @@ -277,7 +303,7 @@ export class SessionEventBroadcaster { journal, tracker: new InFlightTurnTracker(), tail: [], - targets: new Set(), + targets: new Map(), queue: Promise.resolve(), agentDisposables: new Map(), lifecycleDisposables: [], @@ -526,12 +552,24 @@ export class SessionEventBroadcaster { while (tail.length > this.maxBufferSize) tail.shift(); } - const fanOut = isGlobalEvent(event.type) ? this.allTargets() : targets; - for (const target of fanOut) { - try { - target.send(envelope); - } catch { - // best-effort fan-out; a broken target is dropped, not fatal + if (isGlobalEvent(event.type)) { + // Global events (session/workspace/config/model-catalog) are not agent + // events — fan out to every subscriber regardless of any agent filter. + for (const target of this.allTargets()) { + try { + target.send(envelope); + } catch { + // best-effort fan-out; a broken target is dropped, not fatal + } + } + } else { + for (const [target, filter] of targets) { + if (!matchesAgentFilter(envelope, filter)) continue; + try { + target.send(envelope); + } catch { + // best-effort fan-out; a broken target is dropped, not fatal + } } } } @@ -554,7 +592,7 @@ export class SessionEventBroadcaster { private *allTargets(): Iterable { for (const state of this.sessions.values()) { - for (const target of state.targets) yield target; + for (const target of state.targets.keys()) yield target; } } } @@ -594,6 +632,31 @@ function isGlobalEvent(type: string): boolean { ); } +/** + * Per-subscription agent allowlist check — shared by live fan-out and replay. + * Returns `true` when the envelope should be delivered to a subscriber carrying + * `filter`: + * - `filter === undefined` → receive every agent (legacy session-grained + * behavior); + * - global events (session/workspace/config/model-catalog) are not agent + * events and always pass; + * - events without a string `agentId` (should not happen on the v1 wire, + * where the broadcaster stamps every event) pass defensively rather than + * being dropped; + * - otherwise the envelope's `payload.agentId` must be in the allowlist. + */ +function matchesAgentFilter(envelope: EventEnvelope, filter: AgentFilter): boolean { + if (filter === undefined) return true; + if (isGlobalEvent(envelope.type)) return true; + const payload = envelope.payload; + const agentId = + typeof payload === 'object' && payload !== null + ? (payload as { agentId?: unknown }).agentId + : undefined; + if (typeof agentId !== 'string') return true; + return filter.has(agentId); +} + // --------------------------------------------------------------------------- // Interaction → v1 protocol event synthesis. Event names and payload shapes // mirror v1's question/approval services diff --git a/packages/kap-server/src/transport/ws/v1/wsConnectionV1.ts b/packages/kap-server/src/transport/ws/v1/wsConnectionV1.ts index 11ca0fe23..ba4b5ebd2 100644 --- a/packages/kap-server/src/transport/ws/v1/wsConnectionV1.ts +++ b/packages/kap-server/src/transport/ws/v1/wsConnectionV1.ts @@ -29,6 +29,7 @@ import { buildServerHello, } from './protocol'; import { + type AgentFilter, type BroadcastTarget, type ResyncReason, type SessionEventBroadcaster, @@ -98,8 +99,8 @@ export class WsConnectionV1 implements BroadcastTarget { private closed = false; private gotClientHello = false; - /** Session ids this connection is currently subscribed to. */ - readonly subscriptions = new Set(); + /** Per-session subscription state, with each session's optional agent allowlist. */ + readonly subscriptions = new Map(); private pingTimer?: ReturnType; private pongTimer?: ReturnType; @@ -149,7 +150,7 @@ export class WsConnectionV1 implements BroadcastTarget { } get subscriptionSessionIds(): readonly string[] { - return Array.from(this.subscriptions).sort(); + return Array.from(this.subscriptions.keys()).sort(); } /** BroadcastTarget — forward a sequenced envelope to the socket. */ @@ -194,13 +195,21 @@ export class WsConnectionV1 implements BroadcastTarget { const payload = frame.payload ?? {}; const subscriptions = asStringArray(payload['subscriptions']); const cursors = payload['cursors'] as Record | undefined; + const agentFilter = parseAgentFilter(payload['agent_filter']); const accepted: string[] = []; const resyncRequired: string[] = []; const serverCursors: Record = {}; for (const sid of subscriptions) { - await this.attachSession(sid, cursors?.[sid], accepted, resyncRequired, serverCursors); + await this.attachSession( + sid, + cursors?.[sid], + agentFilter?.[sid], + accepted, + resyncRequired, + serverCursors, + ); } this.sendFrame( @@ -216,6 +225,7 @@ export class WsConnectionV1 implements BroadcastTarget { const payload = frame.payload ?? {}; const sessionIds = asStringArray(payload['session_ids']); const cursors = payload['cursors'] as Record | undefined; + const agentFilter = parseAgentFilter(payload['agent_filter']); const accepted: string[] = []; const notFound: string[] = []; @@ -223,16 +233,17 @@ export class WsConnectionV1 implements BroadcastTarget { const serverCursors: Record = {}; for (const sid of sessionIds) { - const ok = await this.broadcaster.subscribe(sid, this); + const filter = agentFilter?.[sid]; + const ok = await this.broadcaster.subscribe(sid, this, filter); if (!ok) { notFound.push(sid); continue; } - this.subscriptions.add(sid); + this.subscriptions.set(sid, filter); accepted.push(sid); const cursor = cursors?.[sid]; if (cursor !== undefined) { - await this.replay(sid, cursor, resyncRequired, serverCursors); + await this.replay(sid, cursor, filter, resyncRequired, serverCursors); } else { const cur = await this.broadcaster.getCursor(sid); serverCursors[sid] = cur; @@ -268,19 +279,20 @@ export class WsConnectionV1 implements BroadcastTarget { private async attachSession( sid: string, cursor: SessionCursor | undefined, + filter: AgentFilter | undefined, accepted: string[], resyncRequired: string[], serverCursors: Record, ): Promise { - const ok = await this.broadcaster.subscribe(sid, this); + const ok = await this.broadcaster.subscribe(sid, this, filter); if (!ok) { resyncRequired.push(sid); return; } - this.subscriptions.add(sid); + this.subscriptions.set(sid, filter); accepted.push(sid); if (cursor !== undefined) { - await this.replay(sid, cursor, resyncRequired, serverCursors); + await this.replay(sid, cursor, filter, resyncRequired, serverCursors); } else { const cur = await this.broadcaster.getCursor(sid); serverCursors[sid] = cur; @@ -290,10 +302,11 @@ export class WsConnectionV1 implements BroadcastTarget { private async replay( sid: string, cursor: SessionCursor, + filter: AgentFilter | undefined, resyncRequired: string[], serverCursors: Record, ): Promise { - const result = await this.broadcaster.getBufferedSince(sid, cursor); + const result = await this.broadcaster.getBufferedSince(sid, cursor, filter); if (result.resyncRequired !== false) { this.sendFrame( buildResyncRequired(sid, result.resyncRequired as ResyncReason, result.currentSeq, result.epoch), @@ -447,7 +460,7 @@ export class WsConnectionV1 implements BroadcastTarget { if (this.flushTimer !== undefined) clearTimeout(this.flushTimer); if (this.backpressureRetryTimer !== undefined) clearTimeout(this.backpressureRetryTimer); this.outbound = []; - for (const sid of this.subscriptions) this.broadcaster.unsubscribe(sid, this); + for (const sid of this.subscriptions.keys()) this.broadcaster.unsubscribe(sid, this); // registry removal is handled by registerWsV1 on the socket 'close' event. } } @@ -457,6 +470,26 @@ function asStringArray(value: unknown): string[] { return value.filter((v): v is string => typeof v === 'string'); } +/** + * Parse the wire `agent_filter` payload (`Record`) into + * a per-session allowlist map. Sessions missing from the returned map — or the + * whole field absent — fall back to "every agent" (`undefined`), the legacy + * session-grained behavior. Malformed entries (non-object, empty arrays, + * non-string ids) are dropped per-session rather than failing the whole + * handshake, so a bad entry cannot widen another session's filter. + */ +function parseAgentFilter(value: unknown): Record | undefined { + if (typeof value !== 'object' || value === null || Array.isArray(value)) return undefined; + const out: Record = {}; + for (const [sid, ids] of Object.entries(value)) { + if (!Array.isArray(ids)) continue; + const set = new Set(ids.filter((v): v is string => typeof v === 'string')); + if (set.size === 0) continue; + out[sid] = set; + } + return out; +} + function rawDataToString(data: RawData): string { if (typeof data === 'string') return data; if (Buffer.isBuffer(data)) return data.toString('utf8'); diff --git a/packages/kap-server/test/sessionEventBroadcaster.test.ts b/packages/kap-server/test/sessionEventBroadcaster.test.ts index c83ef7b75..7e88b9727 100644 --- a/packages/kap-server/test/sessionEventBroadcaster.test.ts +++ b/packages/kap-server/test/sessionEventBroadcaster.test.ts @@ -580,4 +580,88 @@ describe('SessionEventBroadcaster', () => { await bc.getCursor('s1'); expect(envelopes.map((e) => e.type)).toEqual(['event.question.answered']); }); + + // ------------------------------------------------------------------------- + // Per-agent subscription filter + // ------------------------------------------------------------------------- + + it('delivers only the allowlisted agent events on live fan-out', async () => { + const lc = new FakeLifecycle(); + const main = lc.addAgent('main'); + const sub = lc.addAgent('agent-0'); + sessions.set('s1', lc); + + const { target, envelopes } = collectingTarget(); + await bc.subscribe('s1', target, new Set(['main'])); + + main.bus.emit(agentEvent('turn.ended', { turnId: 1 })); + sub.bus.emit(agentEvent('turn.ended', { turnId: 1 })); + await bc.getCursor('s1'); + + expect(envelopes).toHaveLength(1); + expect((envelopes[0]!.payload as { agentId: string }).agentId).toBe('main'); + }); + + it('delivers every agent event when no filter is set', async () => { + const lc = new FakeLifecycle(); + const main = lc.addAgent('main'); + const sub = lc.addAgent('agent-0'); + sessions.set('s1', lc); + + const { target, envelopes } = collectingTarget(); + await bc.subscribe('s1', target); // no filter — legacy behavior + + main.bus.emit(agentEvent('turn.ended', { turnId: 1 })); + sub.bus.emit(agentEvent('turn.ended', { turnId: 1 })); + await bc.getCursor('s1'); + + const agentIds = envelopes.map((e) => (e.payload as { agentId: string }).agentId); + expect(agentIds).toEqual(['main', 'agent-0']); + }); + + it('bypasses the agent filter for global events', async () => { + const lc = new FakeLifecycle(); + lc.addAgent('main'); + sessions.set('s1', lc); + + const { target, envelopes } = collectingTarget(); + // Filter does not include 'main', yet global events must still be delivered. + await bc.subscribe('s1', target, new Set(['agent-0'])); + + eventBus.emit({ + type: 'session.meta.updated', + payload: { + agentId: 'main', + sessionId: 's1', + title: '测试', + patch: { title: '测试' }, + }, + }); + + await vi.waitFor(() => expect(envelopes).toHaveLength(1)); + expect(envelopes[0]!.type).toBe('session.meta.updated'); + }); + + it('replays only the allowlisted agent events while keeping the global sequence', async () => { + const lc = new FakeLifecycle(); + const main = lc.addAgent('main'); + const sub = lc.addAgent('agent-0'); + sessions.set('s1', lc); + + // Activate the session and journal a mixed sequence before replaying. + const warm = collectingTarget(); + await bc.subscribe('s1', warm.target); + main.bus.emit(agentEvent('turn.ended', { turnId: 1 })); // seq 1 + sub.bus.emit(agentEvent('turn.ended', { turnId: 1 })); // seq 2 + main.bus.emit(agentEvent('turn.ended', { turnId: 2 })); // seq 3 + await bc.getCursor('s1'); + + const result = await bc.getBufferedSince('s1', { seq: 0 }, new Set(['main'])); + expect(result.resyncRequired).toBe(false); + // Filtered view keeps the session's original seq offsets (with gaps). + expect(result.events.map((e) => e.seq)).toEqual([1, 3]); + expect( + result.events.every((e) => (e.envelope.payload as { agentId: string }).agentId === 'main'), + ).toBe(true); + }); }); diff --git a/packages/kap-server/test/wsV1Resync.test.ts b/packages/kap-server/test/wsV1Resync.test.ts index f49ada8de..d1a26097a 100644 --- a/packages/kap-server/test/wsV1Resync.test.ts +++ b/packages/kap-server/test/wsV1Resync.test.ts @@ -237,4 +237,46 @@ describe('server-v2 /api/v1/ws resync', () => { c.ws.close(); await c.closed; }); + + it('delivers only the allowlisted agent events via agent_filter', async () => { + const sid = await createSession(); + await ensureMainAgent(sid); + + // Add a second agent to the same session so we can distinguish sources. + const session = server!.core.accessor.get(ISessionLifecycleService).get(sid); + expect(session).toBeDefined(); + const agents = session!.accessor.get(IAgentLifecycleService); + const sub = await agents.create({ agentId: 'agent-0' }); + + const c = await openConn(wsUrl, server!.authTokenService.getToken()); + await c.next((f) => f.type === 'server_hello'); + c.send({ + type: 'client_hello', + id: 'h1', + payload: withToken({ + client_id: 'cli', + subscriptions: [sid], + agent_filter: { [sid]: ['main'] }, + }), + }); + await c.next((f) => f.type === 'ack' && f.id === 'h1'); + + // Emit one durable event per agent — only `main` is allowlisted. + agents + .getHandle('main')! + .accessor.get(IEventBus) + .publish({ type: 'turn.ended', turnId: 1 } as unknown as DomainEvent); + sub.accessor + .get(IEventBus) + .publish({ type: 'turn.ended', turnId: 2 } as unknown as DomainEvent); + + const ev = await c.next((f) => f.type === 'turn.ended'); + expect(ev.payload).toMatchObject({ agentId: 'main' }); + + // The agent-0 event is filtered out — no second turn.ended arrives. + await expect(c.next((f) => f.type === 'turn.ended', 300)).rejects.toThrow(); + + c.ws.close(); + await c.closed; + }); }); diff --git a/packages/protocol/src/__tests__/ws-control.test.ts b/packages/protocol/src/__tests__/ws-control.test.ts index 942752251..c201a4be3 100644 --- a/packages/protocol/src/__tests__/ws-control.test.ts +++ b/packages/protocol/src/__tests__/ws-control.test.ts @@ -226,6 +226,22 @@ describe('ws-control — §3.2 client_hello', () => { expect(result.success).toBe(true); }); + it('client_hello accepts an agent_filter map', () => { + const result = clientHelloMessageSchema.safeParse({ + type: 'client_hello', + id: 'c1', + payload: { + client_id: 'web_abc', + subscriptions: ['sess_1', 'sess_2'], + agent_filter: { + sess_1: ['main'], + sess_2: ['main', 'agent-0'], + }, + }, + }); + expect(result.success).toBe(true); + }); + it('client_hello rejects the v1 bare-seq cursor map', () => { const result = clientHelloMessageSchema.safeParse({ type: 'client_hello', @@ -264,6 +280,42 @@ describe('ws-control — §3.3 subscribe / unsubscribe', () => { expect(result.success).toBe(true); }); + it('subscribe accepts an agent_filter map', () => { + const result = subscribeMessageSchema.safeParse({ + type: 'subscribe', + id: 'c2', + payload: { + session_ids: ['sess_1', 'sess_2'], + agent_filter: { + sess_1: ['main'], + sess_2: ['main', 'agent-0'], + }, + }, + }); + expect(result.success).toBe(true); + }); + + it('subscribe accepts a missing agent_filter (legacy session-grained behavior)', () => { + const result = subscribeMessageSchema.safeParse({ + type: 'subscribe', + id: 'c2', + payload: { session_ids: ['sess_1'] }, + }); + expect(result.success).toBe(true); + }); + + it('subscribe rejects an empty agent_filter allowlist', () => { + const result = subscribeMessageSchema.safeParse({ + type: 'subscribe', + id: 'c2', + payload: { + session_ids: ['sess_1'], + agent_filter: { sess_1: [] }, + }, + }); + expect(result.success).toBe(false); + }); + it('subscribe rejects missing session_ids', () => { const result = subscribeMessageSchema.safeParse({ type: 'subscribe', diff --git a/packages/protocol/src/ws-control.ts b/packages/protocol/src/ws-control.ts index 8e98e5266..2098276c1 100644 --- a/packages/protocol/src/ws-control.ts +++ b/packages/protocol/src/ws-control.ts @@ -90,10 +90,22 @@ export const serverHelloMessageSchema = z.object({ export type ServerHelloMessage = z.infer; +/** + * Per-session agent allowlist for fine-grained v1 event subscriptions. Keys are + * session ids, values are the non-empty set of agent ids the client wants to + * receive events for within that session. Sessions absent from the map (or the + * whole field omitted) fall back to receiving every agent — the legacy + * session-grained behavior. + */ +export const agentFilterSchema = z.record(z.string(), z.array(z.string()).min(1)); + +export type AgentFilter = z.infer; + export const clientHelloPayloadSchema = z.object({ client_id: z.string(), subscriptions: z.array(z.string()), cursors: cursorsBySessionSchema.optional(), + agent_filter: agentFilterSchema.optional(), }); export const clientHelloMessageSchema = z.object({ @@ -124,6 +136,7 @@ export const subscribePayloadSchema = z.object({ session_ids: z.array(z.string()), cursors: cursorsBySessionSchema.optional(), watch_fs: z.record(z.string(), watchFsConfigSchema).optional(), + agent_filter: agentFilterSchema.optional(), }); export const subscribeMessageSchema = z.object({