diff --git a/packages/agent-core-v2/scripts/check-import-boundaries.mjs b/packages/agent-core-v2/scripts/check-import-boundaries.mjs index ab21965950b..3f75d4125fe 100644 --- a/packages/agent-core-v2/scripts/check-import-boundaries.mjs +++ b/packages/agent-core-v2/scripts/check-import-boundaries.mjs @@ -31,6 +31,7 @@ const HUMAN_VOCABULARY = new Set([ 'interaction/machine', 'interaction/facade', 'utils/watch', + 'xstate2', ]); const V2_ONLY_FIRST_SEGMENTS = new Set([ diff --git a/packages/agent-core-v2/src/agent/actorService/agentActorService.ts b/packages/agent-core-v2/src/agent/actorService/agentActorService.ts index 45570f2fbef..93fe6b01d2b 100644 --- a/packages/agent-core-v2/src/agent/actorService/agentActorService.ts +++ b/packages/agent-core-v2/src/agent/actorService/agentActorService.ts @@ -1,4 +1,4 @@ -import { createActor, type ActorLogic, type AnyActorRef, type Snapshot } from 'xstate'; +import { createActor, type ActorLogic, type AnyActorRef, type Snapshot } from '#human/xstate2'; import { BugIndicatingError } from '#/_base/errors/errors'; import { onUnexpectedError } from '#/_base/errors/unexpectedError'; diff --git a/packages/agent-core-v2/src/human/test/usage/machine.test.ts b/packages/agent-core-v2/src/human/test/usage/machine.test.ts index f4da64a9a03..598d2f3197a 100644 --- a/packages/agent-core-v2/src/human/test/usage/machine.test.ts +++ b/packages/agent-core-v2/src/human/test/usage/machine.test.ts @@ -15,6 +15,10 @@ import type { UsageEmitted } from '#/usage/machine'; import { createUsagePlugin } from '#/usage/plugin'; import type { UsageRecord } from '#/usage/usage'; import { createTimingPlugin } from '#/timing/plugin'; +import { + xstateInspectionCollector, + type XstateInspectionEnvelope, +} from '#/xstateInspection'; const model: LlmModel = { provider: 'test', model: 'test-model', capability: UNKNOWN_CAPABILITY }; @@ -30,32 +34,31 @@ function record( return { usage: usage(inputOther, output), model: extra?.model, turnId: extra?.turnId, at: 0 }; } -describe('usage machine', () => { - it('accumulates total, byModel and byTurn across usage.record events', () => { - const actor = createActor(createUsageMachine()); - actor.start(); - - actor.send({ type: 'usage.record', record: record(10, 2, { model, turnId: 1 }) }); - actor.send({ type: 'usage.record', record: record(5, 3, { model, turnId: 2 }) }); - actor.send({ type: 'usage.record', record: record(100, 0) }); - - const { records, summary } = actor.getSnapshot().context; - expect(records).toHaveLength(3); - expect(summary.total).toEqual({ - inputOther: 115, - output: 5, - inputCacheRead: 0, - inputCacheCreation: 0, - }); - expect(summary.byModel).toEqual({ - 'test-model': { inputOther: 15, output: 5, inputCacheRead: 0, inputCacheCreation: 0 }, - }); - expect(summary.byTurn).toEqual({ - 1: { inputOther: 10, output: 2, inputCacheRead: 0, inputCacheCreation: 0 }, - 2: { inputOther: 5, output: 3, inputCacheRead: 0, inputCacheCreation: 0 }, - }); +describe('xstate inspection collector', () => { + it('publishes JSON-safe scalar envelopes with no machine context', () => { + const envelopes: XstateInspectionEnvelope[] = []; + const unsubscribe = xstateInspectionCollector.subscribe((envelope) => envelopes.push(envelope)); + try { + const actor = createActor(createUsageMachine()); + actor.start(); + actor.send({ type: 'usage.record', record: record(10, 2, { model, turnId: 1 }) }); + } finally { + unsubscribe(); + } + const delivered = envelopes.filter((envelope) => envelope.eventType === 'usage.record'); + expect(delivered.length).toBeGreaterThan(0); + for (const envelope of delivered) { + expect(typeof envelope.actorSessionId).toBe('string'); + expect(typeof envelope.timestamp).toBe('number'); + } + expect(delivered.find((envelope) => envelope.type === '@xstate.microstep')?.stateValue).toBeDefined(); + const serialized = JSON.stringify(envelopes); + expect(serialized).not.toContain('inputOther'); + expect(JSON.parse(serialized)).toEqual(envelopes); }); +}); +describe('usage machine', () => { it('groups byModel by baseUrl + model, ignoring provider', () => { const actor = createActor(createUsageMachine()); actor.start(); diff --git a/packages/agent-core-v2/src/human/xstate2.ts b/packages/agent-core-v2/src/human/xstate2.ts index 6a163310aac..e9f709a0150 100644 --- a/packages/agent-core-v2/src/human/xstate2.ts +++ b/packages/agent-core-v2/src/human/xstate2.ts @@ -1,6 +1,8 @@ import { createActor as createXStateActor } from 'xstate'; import type { Actor, ActorOptions, AnyActorLogic, InspectionEvent } from 'xstate'; +import { xstateInspectionCollector } from '#/xstateInspection'; + export * from 'xstate'; function reportUnhandled(event: InspectionEvent): void { @@ -24,6 +26,7 @@ function createActorWithInspect( ...options, inspect: (event) => { reportUnhandled(event); + xstateInspectionCollector.publish(event); if (typeof inspect === 'function') { inspect(event); } else { diff --git a/packages/agent-core-v2/src/human/xstateInspection.ts b/packages/agent-core-v2/src/human/xstateInspection.ts new file mode 100644 index 00000000000..12671e09e4d --- /dev/null +++ b/packages/agent-core-v2/src/human/xstateInspection.ts @@ -0,0 +1,66 @@ +import type { InspectionEvent } from 'xstate'; + +export type XstateInspectionEventType = InspectionEvent['type']; + +export interface XstateInspectionEnvelope { + readonly type: XstateInspectionEventType; + readonly timestamp: number; + readonly actorSessionId: string; + readonly actorId?: string; + readonly logicId?: string; + readonly eventType?: string; + readonly stateValue?: unknown; +} + +export type XstateInspectionListener = (envelope: XstateInspectionEnvelope) => void; + +export interface XstateInspectionCollector { + subscribe(listener: XstateInspectionListener): () => void; + publish(event: InspectionEvent): void; +} + +function scalar(value: unknown): string | undefined { + return typeof value === 'string' ? value : undefined; +} + +function toEnvelope(event: InspectionEvent, now: () => number): XstateInspectionEnvelope { + const actorRef = event.actorRef as { id?: unknown; logic?: unknown }; + const logic = actorRef.logic as { id?: unknown } | undefined; + const snapshot = 'snapshot' in event ? (event.snapshot as { value?: unknown }) : undefined; + return { + type: event.type, + timestamp: now(), + actorSessionId: event.actorRef.sessionId, + actorId: scalar(actorRef.id), + logicId: scalar(logic?.id), + eventType: + 'event' in event + ? event.event.type + : event.type === '@xstate.action' + ? event.action.type + : undefined, + stateValue: snapshot?.value, + }; +} + +export function createXstateInspectionCollector(input?: { + now?: () => number; +}): XstateInspectionCollector { + const now = input?.now ?? Date.now; + const listeners = new Set(); + return { + subscribe(listener) { + listeners.add(listener); + return () => { + listeners.delete(listener); + }; + }, + publish(event) { + if (listeners.size === 0) return; + const envelope = toEnvelope(event, now); + for (const listener of listeners) listener(envelope); + }, + }; +} + +export const xstateInspectionCollector = createXstateInspectionCollector(); diff --git a/packages/kap-server/src/start.ts b/packages/kap-server/src/start.ts index d6b2892552c..e717ee0c827 100644 --- a/packages/kap-server/src/start.ts +++ b/packages/kap-server/src/start.ts @@ -60,6 +60,7 @@ import { extractWsBearerToken } from './transport/ws/bearerProtocol'; import { SessionEventBroadcaster } from './transport/ws/v1/sessionEventBroadcaster'; import type { ConfigWarningItem } from './transport/ws/v1/events'; import { registerWsV1, WS_PATH as WS_PATH_V1 } from './transport/ws/v1/registerWsV1'; +import { registerWsDebug, WS_DEBUG_PATH } from './transport/ws/debug/registerWsDebug'; import { getServerVersion } from './version'; import { classify } from './security/bindClassify'; import { @@ -303,6 +304,9 @@ export async function startServer(opts: ServerStartOptions): Promise => { + if (wssDebug !== undefined) { + for (const client of wssDebug.clients) client.terminate(); + } configChangedPublisher.close(); await remoteControlManager.close(); await app.close(); @@ -473,6 +477,7 @@ export async function startServer(opts: ServerStartOptions): Promise => { const url = req.url ?? ''; const isV1 = url === WS_PATH_V1 || url.startsWith(`${WS_PATH_V1}?`); - if (!isV1) { + const isDebug = url === WS_DEBUG_PATH || url.startsWith(`${WS_DEBUG_PATH}?`); + const wss = isV1 ? wssV1 : isDebug ? wssDebug : undefined; + if (wss === undefined) { socket.destroy(); return; } @@ -543,7 +550,7 @@ export async function startServer(opts: ServerStartOptions): Promise wssV1.emit('connection', ws, req)); + wss.handleUpgrade(req, socket, head, (ws) => wss.emit('connection', ws, req)); }; app.server.on('upgrade', (req, socket, head) => { void handleUpgrade(req, socket, head).catch((error: unknown) => @@ -554,6 +561,7 @@ export async function startServer(opts: ServerStartOptions): Promise { connectionRegistry.closeAll('server shutting down'); wssV1.close(); + wssDebug?.close(); await broadcaster.close(); }); diff --git a/packages/kap-server/src/transport/ws/debug/registerWsDebug.ts b/packages/kap-server/src/transport/ws/debug/registerWsDebug.ts new file mode 100644 index 00000000000..eb1258620ef --- /dev/null +++ b/packages/kap-server/src/transport/ws/debug/registerWsDebug.ts @@ -0,0 +1,39 @@ +import type { XstateInspectionCollector } from '@moonshot-ai/agent-core-v2/human/xstateInspection'; +import { WebSocketServer } from 'ws'; + +import { selectWsBearerProtocol } from '../bearerProtocol'; +import { WsConnectionDebug } from './wsConnectionDebug'; + +export const WS_DEBUG_PATH = '/api/v1/debug/ws'; + +export interface RegisterWsDebugOptions { + readonly collector?: XstateInspectionCollector; + readonly heartbeatIntervalMs?: number; + readonly flushIntervalMs?: number; + readonly highWaterMarkBytes?: number; +} + +export function registerWsDebug(opts: RegisterWsDebugOptions = {}): WebSocketServer { + const wss = new WebSocketServer({ noServer: true, handleProtocols: selectWsBearerProtocol }); + const connections = new Set(); + + wss.on('connection', (socket) => { + const conn = new WsConnectionDebug({ + socket, + collector: opts.collector, + heartbeatIntervalMs: opts.heartbeatIntervalMs, + flushIntervalMs: opts.flushIntervalMs, + highWaterMarkBytes: opts.highWaterMarkBytes, + }); + connections.add(conn); + socket.on('close', () => { + connections.delete(conn); + }); + }); + + wss.on('close', () => { + for (const conn of connections) conn.close(); + }); + + return wss; +} diff --git a/packages/kap-server/src/transport/ws/debug/wsConnectionDebug.ts b/packages/kap-server/src/transport/ws/debug/wsConnectionDebug.ts new file mode 100644 index 00000000000..77f6e68248c --- /dev/null +++ b/packages/kap-server/src/transport/ws/debug/wsConnectionDebug.ts @@ -0,0 +1,114 @@ +import { + xstateInspectionCollector, + type XstateInspectionCollector, + type XstateInspectionEnvelope, +} from '@moonshot-ai/agent-core-v2/human/xstateInspection'; +import type { WebSocket } from 'ws'; + +const DEFAULT_HEARTBEAT_INTERVAL_MS = 10_000; +const HEARTBEAT_MISS_LIMIT = 2; +const DEFAULT_FLUSH_INTERVAL_MS = 16; +const DEFAULT_HIGH_WATER_MARK_BYTES = 1 << 20; + +export interface WsConnectionDebugOptions { + readonly socket: WebSocket; + readonly collector?: XstateInspectionCollector; + readonly heartbeatIntervalMs?: number; + readonly flushIntervalMs?: number; + readonly highWaterMarkBytes?: number; +} + +export class WsConnectionDebug { + private readonly socket: WebSocket; + private readonly heartbeatIntervalMs: number; + private readonly flushIntervalMs: number; + private readonly highWaterMarkBytes: number; + private readonly unsubscribe: () => void; + + private closed = false; + private outbound: XstateInspectionEnvelope[] = []; + private flushTimer?: ReturnType; + private heartbeatTimer?: ReturnType; + private lastPongAt = Date.now(); + + constructor(opts: WsConnectionDebugOptions) { + this.socket = opts.socket; + this.heartbeatIntervalMs = opts.heartbeatIntervalMs ?? DEFAULT_HEARTBEAT_INTERVAL_MS; + this.flushIntervalMs = opts.flushIntervalMs ?? DEFAULT_FLUSH_INTERVAL_MS; + this.highWaterMarkBytes = opts.highWaterMarkBytes ?? DEFAULT_HIGH_WATER_MARK_BYTES; + + this.socket.on('close', () => this.onClose()); + this.socket.on('error', () => this.onClose()); + this.socket.on('pong', () => { + this.lastPongAt = Date.now(); + }); + + const collector = opts.collector ?? xstateInspectionCollector; + this.unsubscribe = collector.subscribe((envelope) => this.onEnvelope(envelope)); + + this.heartbeatTimer = setInterval(() => this.onHeartbeat(), this.heartbeatIntervalMs); + this.heartbeatTimer.unref?.(); + } + + private onEnvelope(envelope: XstateInspectionEnvelope): void { + if (this.closed) return; + if (this.socket.bufferedAmount > this.highWaterMarkBytes) return; + this.outbound.push(envelope); + this.scheduleFlush(); + } + + private scheduleFlush(): void { + if (this.flushTimer !== undefined) return; + this.flushTimer = setTimeout(() => { + this.flushTimer = undefined; + this.flush(); + }, this.flushIntervalMs); + this.flushTimer.unref?.(); + } + + private flush(): void { + if (this.outbound.length === 0) return; + if (this.closed || this.socket.readyState !== this.socket.OPEN) { + this.outbound = []; + return; + } + const envelopes = this.outbound; + this.outbound = []; + for (const envelope of envelopes) { + if (this.closed || this.socket.readyState !== this.socket.OPEN) return; + try { + this.socket.send(JSON.stringify(envelope)); + } catch { + } + } + } + + private onHeartbeat(): void { + if (Date.now() - this.lastPongAt >= this.heartbeatIntervalMs * HEARTBEAT_MISS_LIMIT) { + this.close(); + return; + } + try { + this.socket.ping(); + } catch { + } + } + + close(): void { + if (this.closed) return; + try { + this.socket.close(1000); + } catch { + } + this.onClose(); + } + + private onClose(): void { + if (this.closed) return; + this.closed = true; + if (this.flushTimer !== undefined) clearTimeout(this.flushTimer); + if (this.heartbeatTimer !== undefined) clearInterval(this.heartbeatTimer); + this.outbound = []; + this.unsubscribe(); + } +} diff --git a/packages/kap-server/test/wsUpgradeAuth.test.ts b/packages/kap-server/test/wsUpgradeAuth.test.ts index 2ca397e2f4a..3cd4c78c2a4 100644 --- a/packages/kap-server/test/wsUpgradeAuth.test.ts +++ b/packages/kap-server/test/wsUpgradeAuth.test.ts @@ -1,6 +1,16 @@ +import { mkdtemp, rm } from 'node:fs/promises'; +import { tmpdir } from 'node:os'; +import { join } from 'node:path'; + +import { IModelCatalog } from '@moonshot-ai/agent-core-v2'; +import { createActor, setup } from '@moonshot-ai/agent-core-v2/human/xstate2'; import { afterEach, describe, expect, it } from 'vitest'; import { WebSocket, type RawData } from 'ws'; +import { startServer } from '../src/start'; +import { fakeModelCatalog } from './helpers/fakeModelCatalog'; +import { fixedTokenAuth } from './helpers/fixedAuth'; +import { TEST_HOST_IDENTITY } from './helpers/hostIdentity'; import { sharedServer } from './helpers/sharedServer'; function rawToString(data: RawData): string { @@ -85,14 +95,6 @@ describe('WS upgrade auth', () => { expect(firstFrame).toMatchObject({ type: firstType }); }); - it('accepts a valid Authorization bearer header', async () => { - const { ws, firstFrame } = await openConn(url(), { - headers: { Authorization: `Bearer ${token()}` }, - }); - sockets.push(ws); - expect(firstFrame).toMatchObject({ type: firstType }); - }); - it('rejects a wrong bearer token', async () => { await expectRejected(url(), { protocols: ['kimi-code.bearer.wrong'] }); }); @@ -102,8 +104,63 @@ describe('WS upgrade auth', () => { }); }); + describe('/api/v1/debug/ws', () => { + it('streams xstate inspection envelopes to an authorized client', async () => { + const home = await mkdtemp(join(tmpdir(), 'kimi-kap-debug-ws-')); + const server = await startServer({ + hostIdentity: TEST_HOST_IDENTITY, + host: '127.0.0.1', + port: 0, + homeDir: home, + logLevel: 'silent', + debugEndpoints: true, + authTokenService: fixedTokenAuth(), + seeds: [[IModelCatalog, fakeModelCatalog()]], + }); + const ws = new WebSocket(`ws://127.0.0.1:${server.port}/api/v1/debug/ws`, { + headers: { Authorization: 'Bearer test-token' }, + }); + sockets.push(ws); + try { + const envelope = await new Promise>((resolve, reject) => { + const timer = setTimeout( + () => reject(new Error('no inspection envelope within timeout')), + 5000, + ); + ws.on('message', (data) => { + const frame = JSON.parse(rawToString(data)) as Record; + if (frame['eventType'] === 'debug.probe') { + clearTimeout(timer); + resolve(frame); + } + }); + ws.on('error', reject); + ws.once('open', () => { + const machine = setup({}).createMachine({ + id: 'debugWsProbe', + initial: 'idle', + states: { idle: { on: { 'debug.probe': 'done' } }, done: {} }, + }); + const actor = createActor(machine); + actor.start(); + actor.send({ type: 'debug.probe' }); + }); + }); + expect(envelope['type']).toBe('@xstate.event'); + expect(envelope['logicId']).toBe('debugWsProbe'); + expect(typeof envelope['actorSessionId']).toBe('string'); + expect(typeof envelope['timestamp']).toBe('number'); + } finally { + await server.close(); + await rm(home, { recursive: true, force: true }); + } + }); + }); + it('rejects upgrades to a non-WS path', async () => { const badUrl = `${v1Url().replace('/api/v1/ws', '/api/v1/other')}`; await expectRejected(badUrl, { protocols: [`kimi-code.bearer.${token()}`] }); + const debugUrl = `${v1Url().replace('/api/v1/ws', '/api/v1/debug/ws')}`; + await expectRejected(debugUrl, { protocols: [`kimi-code.bearer.${token()}`] }); }); });