mirror of
https://github.com/MoonshotAI/kimi-code.git
synced 2026-07-10 01:39:25 +00:00
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
This commit is contained in:
parent
5eef0a6736
commit
fdb65be4aa
8 changed files with 362 additions and 33 deletions
|
|
@ -61,7 +61,7 @@ export class MiniDbQueryStore extends Disposable implements IQueryStore {
|
|||
declare readonly _serviceBrand: undefined;
|
||||
|
||||
private readonly dir: string;
|
||||
private dbPromise: Promise<MiniDb> | undefined;
|
||||
private dbPromise: Promise<MiniDb | undefined> | undefined;
|
||||
private readonly ensuredIndexes = new Set<string>();
|
||||
|
||||
constructor(
|
||||
|
|
@ -75,7 +75,7 @@ export class MiniDbQueryStore extends Disposable implements IQueryStore {
|
|||
}));
|
||||
}
|
||||
|
||||
private openDb(): Promise<MiniDb> {
|
||||
private openDb(): Promise<MiniDb | undefined> {
|
||||
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
|
||||
// `<cacheDir>/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<T>(collection: string, key: string, value: T): Promise<void> {
|
||||
const db = await this.openDb();
|
||||
if (db === undefined) return;
|
||||
await db.set(physicalKey(collection, key), value);
|
||||
}
|
||||
|
||||
async batch(ops: readonly WriteOp[]): Promise<void> {
|
||||
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<void> {
|
||||
const db = await this.openDb();
|
||||
if (db === undefined) return;
|
||||
await db.del(physicalKey(collection, key));
|
||||
}
|
||||
|
||||
async get<T>(collection: string, key: string): Promise<T | undefined> {
|
||||
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<void> {
|
||||
if (this.dbPromise === undefined) return;
|
||||
const db = await this.dbPromise;
|
||||
await db.close();
|
||||
await db?.close();
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -175,7 +193,7 @@ class MiniDbQuery<T> implements IQuery<T> {
|
|||
private skip = 0;
|
||||
|
||||
constructor(
|
||||
private readonly openDb: () => Promise<MiniDb>,
|
||||
private readonly openDb: () => Promise<MiniDb | undefined>,
|
||||
private readonly collection: string,
|
||||
) {}
|
||||
|
||||
|
|
@ -202,6 +220,7 @@ class MiniDbQuery<T> implements IQuery<T> {
|
|||
|
||||
async execute(): Promise<Page<T>> {
|
||||
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<string, unknown>;
|
||||
|
|
|
|||
|
|
@ -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 });
|
||||
|
|
|
|||
|
|
@ -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<string> | 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<BroadcastTarget>;
|
||||
/** Connections subscribed to this session, each with its optional agent allowlist. */
|
||||
readonly targets: Map<BroadcastTarget, AgentFilter>;
|
||||
/** Per-session dispatch queue — serializes stamp / journal / fan-out. */
|
||||
queue: Promise<void>;
|
||||
/** 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<boolean> {
|
||||
async subscribe(
|
||||
sessionId: string,
|
||||
target: BroadcastTarget,
|
||||
filter?: AgentFilter,
|
||||
): Promise<boolean> {
|
||||
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<BufferedSinceResult> {
|
||||
async getBufferedSince(
|
||||
sessionId: string,
|
||||
cursor: SessionCursor,
|
||||
filter?: AgentFilter,
|
||||
): Promise<BufferedSinceResult> {
|
||||
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<BroadcastTarget> {
|
||||
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
|
||||
|
|
|
|||
|
|
@ -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<string>();
|
||||
/** Per-session subscription state, with each session's optional agent allowlist. */
|
||||
readonly subscriptions = new Map<string, AgentFilter>();
|
||||
|
||||
private pingTimer?: ReturnType<typeof setInterval>;
|
||||
private pongTimer?: ReturnType<typeof setTimeout>;
|
||||
|
|
@ -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<string, SessionCursor> | undefined;
|
||||
const agentFilter = parseAgentFilter(payload['agent_filter']);
|
||||
|
||||
const accepted: string[] = [];
|
||||
const resyncRequired: string[] = [];
|
||||
const serverCursors: Record<string, { seq: number; epoch?: string }> = {};
|
||||
|
||||
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<string, SessionCursor> | 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<string, { seq: number; epoch?: string }> = {};
|
||||
|
||||
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<string, { seq: number; epoch?: string }>,
|
||||
): Promise<void> {
|
||||
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<string, { seq: number; epoch?: string }>,
|
||||
): Promise<void> {
|
||||
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<session_id, agent_id[]>`) 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<string, AgentFilter> | undefined {
|
||||
if (typeof value !== 'object' || value === null || Array.isArray(value)) return undefined;
|
||||
const out: Record<string, AgentFilter> = {};
|
||||
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');
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -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',
|
||||
|
|
|
|||
|
|
@ -90,10 +90,22 @@ export const serverHelloMessageSchema = z.object({
|
|||
|
||||
export type ServerHelloMessage = z.infer<typeof serverHelloMessageSchema>;
|
||||
|
||||
/**
|
||||
* 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<typeof agentFilterSchema>;
|
||||
|
||||
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({
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue