From 513ad412608ed5efdd991575d76d2dce00120fee Mon Sep 17 00:00:00 2001 From: "haozhe.yang" Date: Sat, 11 Jul 2026 21:49:41 +0800 Subject: [PATCH] feat(klient): add WebSocket transport for calls and event streams - add WsSocket: persistent /api/v2/ws transport with hello handshake, heartbeat answers, per-call timeouts, and auto-reconnect that re-subscribes active listens; bearer token rides the kimi-code.bearer. subprotocol for browser compatibility - add WsKlient / WsChannel exposing core/session/agent scopes and listen(event, handler) over the shared socket - add Klient#ws() lazy singleton with WebSocketImpl injection - bind global fetch in HttpChannel to avoid "Illegal invocation" in browsers --- packages/klient/README.md | 27 ++- packages/klient/src/client.ts | 26 ++ packages/klient/src/httpChannel.ts | 4 +- packages/klient/src/index.ts | 17 ++ packages/klient/src/wsChannel.ts | 43 ++++ packages/klient/src/wsKlient.ts | 125 ++++++++++ packages/klient/src/wsSocket.ts | 370 +++++++++++++++++++++++++++++ packages/klient/test/ws.test.ts | 289 ++++++++++++++++++++++ 8 files changed, 898 insertions(+), 3 deletions(-) create mode 100644 packages/klient/src/wsChannel.ts create mode 100644 packages/klient/src/wsKlient.ts create mode 100644 packages/klient/src/wsSocket.ts create mode 100644 packages/klient/test/ws.test.ts diff --git a/packages/klient/README.md b/packages/klient/README.md index 8334172da..08907d88b 100644 --- a/packages/klient/README.md +++ b/packages/klient/README.md @@ -27,5 +27,28 @@ const page = await index.list({ workspaceId: 'w1' }); ``` Service interfaces and tokens are imported directly from `agent-core-v2` leaf -subpaths; the channel and proxy live in this package. WebSocket events -(`listen`) are out of scope for this scaffold — HTTP only carries `call`. +subpaths; the channel and proxy live in this package. + +## WebSocket transport (calls + events) + +`Klient#ws()` returns a lazily-created `WsKlient` over the persistent +`/api/v2/ws` socket: the same scope entries and typed proxies (one socket +multiplexes every `call`), plus `listen(event, handler)` on each scope for the +server's event streams — core `events`, session `interactions` / +`interactions:resolved`, agent `events`: + +```ts +const ws = client.ws(); +const sub = ws.session('s1').agent('main').listen('events', (event) => { + console.log('agent event', event); +}); +const pending = await ws.session('s1').service(ISessionApprovalService).listPending(); +sub.dispose(); +ws.close(); +``` + +The socket answers heartbeats, applies per-call timeouts, and reconnects +automatically after an unexpected close (active `listen`s are re-subscribed; +in-flight calls reject). The bearer token rides the +`kimi-code.bearer.` subprotocol, so the transport works unchanged in +browsers. diff --git a/packages/klient/src/client.ts b/packages/klient/src/client.ts index b566bcc23..be1a925a4 100644 --- a/packages/klient/src/client.ts +++ b/packages/klient/src/client.ts @@ -19,6 +19,8 @@ import type { ServiceIdentifier } from '@moonshot-ai/agent-core-v2/_base/di/inst import type { IChannel } from './channel.js'; import { HttpChannel, type HttpChannelOptions } from './httpChannel.js'; import { makeProxy } from './proxy.js'; +import { WsKlient } from './wsKlient.js'; +import type { WsLikeCtor } from './wsSocket.js'; export interface KlientOptions { /** Base URL of the server, e.g. `http://127.0.0.1:58627`. */ @@ -27,17 +29,22 @@ export interface KlientOptions { readonly token?: string; /** `fetch` implementation; defaults to the global `fetch`. */ readonly fetch?: typeof fetch; + /** WebSocket implementation for `ws()`; defaults to the global `WebSocket`. */ + readonly WebSocketImpl?: WsLikeCtor; } export class Klient { private readonly url: string; private readonly token?: string; private readonly fetchImpl?: typeof fetch; + private readonly wsImpl?: WsLikeCtor; + private wsKlient?: WsKlient; constructor(opts: KlientOptions) { this.url = opts.url.replace(/\/$/, ''); this.token = opts.token; this.fetchImpl = opts.fetch; + this.wsImpl = opts.WebSocketImpl; } private channelOptions(baseUrl: string): HttpChannelOptions { @@ -55,6 +62,25 @@ export class Klient { session(sessionId: string): SessionClient { return new SessionClient(this.url, this.token, this.fetchImpl, sessionId); } + + /** + * WebSocket counterpart of this client — same scopes and typed proxies over + * the persistent `/api/v2/ws` socket, plus event `listen`s. Lazily created + * on first call so one `Klient` holds at most one live socket; close it with + * `client.ws().close()`. After a close, the next `ws()` call lazily creates a + * fresh `WsKlient` (so React StrictMode's mount → unmount → mount cycle, + * whose cleanup closes the socket, recovers on the second mount). + */ + ws(): WsKlient { + if (this.wsKlient === undefined || this.wsKlient.state === 'closed') { + this.wsKlient = new WsKlient({ + url: this.url, + token: this.token, + WebSocketImpl: this.wsImpl, + }); + } + return this.wsKlient; + } } export class SessionClient { diff --git a/packages/klient/src/httpChannel.ts b/packages/klient/src/httpChannel.ts index 024c5d3d7..b7d1561e8 100644 --- a/packages/klient/src/httpChannel.ts +++ b/packages/klient/src/httpChannel.ts @@ -36,7 +36,9 @@ export class HttpChannel implements IChannel { constructor(opts: HttpChannelOptions) { this.baseUrl = opts.baseUrl.replace(/\/$/, ''); this.token = opts.token; - this.fetchImpl = opts.fetch ?? fetch; + // Bind the global fetch: browsers throw "Illegal invocation" when the + // native function is invoked with a non-Window receiver. + this.fetchImpl = opts.fetch ?? fetch.bind(globalThis); } async call(command: string, arg?: unknown): Promise { diff --git a/packages/klient/src/index.ts b/packages/klient/src/index.ts index c6d3884ab..267e84d81 100644 --- a/packages/klient/src/index.ts +++ b/packages/klient/src/index.ts @@ -16,3 +16,20 @@ export { type KlientOptions, } from './client.js'; export { SessionIndexClient } from './services/sessionIndex.js'; +export { + WsSocket, + type WsLike, + type WsLikeCtor, + type WsScopeIds, + type WsScopeKind, + type WsSocketOptions, + type WsSocketState, + type WsSubscription, +} from './wsSocket.js'; +export { WsChannel, type WsChannelOptions } from './wsChannel.js'; +export { + WsAgentClient, + WsKlient, + WsSessionClient, + type WsKlientOptions, +} from './wsKlient.js'; diff --git a/packages/klient/src/wsChannel.ts b/packages/klient/src/wsChannel.ts new file mode 100644 index 000000000..ab86910fb --- /dev/null +++ b/packages/klient/src/wsChannel.ts @@ -0,0 +1,43 @@ +/** + * `WsChannel` — an `IChannel` bound to one Service that forwards `call`s over + * the shared `/api/v2/ws` socket instead of HTTP. Same VS Code shape as + * `HttpChannel` (the URL equivalent is the `{scope, service, ids}` triple the + * socket puts on each frame), so the same `makeProxy` turns it into a typed + * Service client. `listen` here takes a handler and returns a subscription + * that survives reconnects until disposed. + */ + +import type { IChannel } from './channel.js'; +import type { WsScopeIds, WsScopeKind, WsSocket, WsSubscription } from './wsSocket.js'; + +export interface WsChannelOptions { + readonly socket: WsSocket; + readonly scope: WsScopeKind; + /** Service channel name (the decorator id, `String(id)`). */ + readonly service: string; + readonly sessionId?: string; + readonly agentId?: string; +} + +export class WsChannel implements IChannel { + private readonly socket: WsSocket; + private readonly scope: WsScopeKind; + private readonly service: string; + private readonly ids: WsScopeIds; + + constructor(opts: WsChannelOptions) { + this.socket = opts.socket; + this.scope = opts.scope; + this.service = opts.service; + this.ids = { sessionId: opts.sessionId, agentId: opts.agentId }; + } + + call(command: string, arg?: unknown): Promise { + return this.socket.call(this.scope, this.service, command, arg, this.ids); + } + + /** Subscribe to an event stream in this channel's scope; dispose to unlisten. */ + listen(event: string, handler: (data: unknown) => void): WsSubscription { + return this.socket.listen(this.scope, event, this.ids, handler); + } +} diff --git a/packages/klient/src/wsKlient.ts b/packages/klient/src/wsKlient.ts new file mode 100644 index 000000000..c42a60a0d --- /dev/null +++ b/packages/klient/src/wsKlient.ts @@ -0,0 +1,125 @@ +/** + * `WsKlient` — the `/api/v2` scope-entry client over the WebSocket transport. + * + * Mirrors `Klient`'s three-level scope entry (`core` / `session` / `agent`), + * but every Service call rides the shared `WsSocket`, and each scope level + * also exposes `listen(event, handler)` for the server's event streams + * (`core` → `events`; `session` → `interactions` / `interactions:resolved`; + * `agent` → `events`): + * + * const ws = new WsKlient({ url: 'http://127.0.0.1:58627', token }); + * await ws.core(ISessionIndex).list({}); + * const sub = ws.session('s1').agent('main').listen('events', (e) => ...); + * sub.dispose(); ws.close(); + * + * Prefer `Klient#ws()` over constructing this directly so HTTP and WS share + * one configured endpoint. + */ + +import type { ServiceIdentifier } from '@moonshot-ai/agent-core-v2/_base/di/instantiation'; + +import { makeProxy } from './proxy.js'; +import { WsChannel } from './wsChannel.js'; +import { + WsSocket, + type WsSocketOptions, + type WsSocketState, + type WsSubscription, +} from './wsSocket.js'; + +export type WsKlientOptions = WsSocketOptions; + +export class WsKlient { + private readonly socket: WsSocket; + + constructor(opts: WsKlientOptions) { + this.socket = new WsSocket(opts); + } + + /** Core-scoped Service over WS, e.g. `ws.core(ISessionIndex)`. */ + core(id: ServiceIdentifier): T { + return makeProxy(new WsChannel({ socket: this.socket, scope: 'core', service: String(id) })); + } + + /** Session scope entry point. */ + session(sessionId: string): WsSessionClient { + return new WsSessionClient(this.socket, sessionId); + } + + /** Subscribe to a core-scoped event stream (e.g. `events`). */ + listen(event: string, handler: (data: unknown) => void): WsSubscription { + return this.socket.listen('core', event, {}, handler); + } + + get state(): WsSocketState { + return this.socket.currentState; + } + + onDidChangeState(listener: (state: WsSocketState) => void): WsSubscription { + return this.socket.onDidChangeState(listener); + } + + close(): void { + this.socket.close(); + } +} + +export class WsSessionClient { + constructor( + private readonly socket: WsSocket, + private readonly sessionId: string, + ) {} + + /** Session-scoped Service over WS, e.g. `.service(ISessionMetadata)`. */ + service(id: ServiceIdentifier): T { + return makeProxy( + new WsChannel({ + socket: this.socket, + scope: 'session', + service: String(id), + sessionId: this.sessionId, + }), + ); + } + + /** Subscribe to a session-scoped event stream (e.g. `interactions`). */ + listen(event: string, handler: (data: unknown) => void): WsSubscription { + return this.socket.listen('session', event, { sessionId: this.sessionId }, handler); + } + + /** Agent scope entry point. */ + agent(agentId: string): WsAgentClient { + return new WsAgentClient(this.socket, this.sessionId, agentId); + } +} + +export class WsAgentClient { + constructor( + private readonly socket: WsSocket, + private readonly sessionId: string, + private readonly agentId: string, + ) {} + + /** Agent-scoped Service over WS, e.g. `.service(IAgentProfileService)`. */ + service(id: ServiceIdentifier): T { + return makeProxy( + new WsChannel({ + socket: this.socket, + scope: 'agent', + service: String(id), + sessionId: this.sessionId, + agentId: this.agentId, + }), + ); + } + + /** Subscribe to an agent-scoped event stream (e.g. `events`). */ + listen(event: string, handler: (data: unknown) => void): WsSubscription { + return this.socket.listen( + 'agent', + event, + { sessionId: this.sessionId, agentId: this.agentId }, + handler, + ); + } +} diff --git a/packages/klient/src/wsSocket.ts b/packages/klient/src/wsSocket.ts new file mode 100644 index 000000000..9a7b03ea1 --- /dev/null +++ b/packages/klient/src/wsSocket.ts @@ -0,0 +1,370 @@ +/** + * `/api/v2/ws` socket — the persistent WebSocket transport behind `WsKlient`. + * + * Speaks the kap-server v2 JSON protocol: one socket multiplexes RPC `call`s + * and event `listen`s, correlated by client-chosen ids. Adds the client-side + * safety features a long-lived devtool connection needs: `hello` handshake, + * `ping`→`pong` heartbeat answers, per-call timeouts, and opt-out automatic + * reconnect (active `listen`s are re-subscribed after a reconnect; in-flight + * calls reject on close — the server cannot resume them). + * + * The bearer token is presented at the upgrade through the + * `kimi-code.bearer.` subprotocol (the only credential channel a browser + * WebSocket has) and again in the `hello` frame for the present-only handshake + * check. Works against the DOM WebSocket (browsers, Node ≥ 21); any compatible + * implementation can be injected for tests. + */ + +import { RPCError } from './errors.js'; + +/** Wire scope kinds, mirroring kap-server's `ScopeKind`. */ +export type WsScopeKind = 'core' | 'session' | 'agent'; + +/** Scope coordinates carried on `call` / `listen` frames. */ +export interface WsScopeIds { + readonly sessionId?: string; + readonly agentId?: string; +} + +export type WsSocketState = 'connecting' | 'open' | 'closed'; + +export interface WsSubscription { + dispose(): void; +} + +/** Minimal DOM-compatible WebSocket surface this module codes against. */ +export interface WsLike { + readonly readyState: number; + send(data: string): void; + close(code?: number, reason?: string): void; + addEventListener(type: 'open' | 'message' | 'close' | 'error', listener: (event: never) => void): void; +} + +export interface WsLikeCtor { + new (url: string, protocols?: string | string[]): WsLike; + readonly OPEN: number; +} + +export interface WsSocketOptions { + /** Server base URL (`http(s)://host:port`) or a full `ws(s)://…/api/v2/ws` URL. */ + readonly url: string; + /** Optional bearer token. */ + readonly token?: string; + /** WebSocket implementation; defaults to the global `WebSocket`. */ + readonly WebSocketImpl?: WsLikeCtor; + /** Reconnect after an unexpected close. Default `true`. */ + readonly autoReconnect?: boolean; + /** Base delay (ms) for the reconnect backoff. Default `500`. */ + readonly reconnectDelayMs?: number; + /** Per-call deadline (ms). Default `30000`. */ + readonly callTimeoutMs?: number; +} + +interface PendingCall { + readonly resolve: (data: unknown) => void; + readonly reject: (err: Error) => void; + readonly timer: ReturnType | undefined; +} + +interface ActiveListen { + readonly scope: WsScopeKind; + readonly event: string; + readonly ids: WsScopeIds; + readonly handler: (data: unknown) => void; +} + +interface ServerFrame { + readonly type: string; + readonly id?: string; + readonly data?: unknown; + readonly code?: number; + readonly msg?: string; +} + +const WS_BEARER_PROTOCOL_PREFIX = 'kimi-code.bearer.'; +const DEFAULT_CALL_TIMEOUT_MS = 30_000; + +export class WsSocket { + private readonly wsUrl: string; + private readonly token?: string; + private readonly WsCtor: WsLikeCtor; + private readonly autoReconnect: boolean; + private readonly reconnectDelayMs: number; + private readonly callTimeoutMs: number; + + private ws: WsLike | undefined; + private state: WsSocketState = 'connecting'; + private manualClose = false; + private reconnectAttempt = 0; + private reconnectTimer: ReturnType | undefined; + private readyWaiters: { resolve: () => void; reject: (err: Error) => void }[] = []; + + private readonly pending = new Map(); + private readonly listens = new Map(); + private readonly stateListeners = new Set<(state: WsSocketState) => void>(); + + private seq = 0; + private readonly idPrefix = `k${Date.now().toString(36)}`; + + constructor(opts: WsSocketOptions) { + this.wsUrl = toWsUrl(opts.url); + this.token = opts.token; + const ctor = opts.WebSocketImpl ?? (globalThis.WebSocket as unknown as WsLikeCtor | undefined); + if (ctor === undefined) { + throw new Error('no WebSocket implementation available; pass WebSocketImpl'); + } + this.WsCtor = ctor; + this.autoReconnect = opts.autoReconnect ?? true; + this.reconnectDelayMs = opts.reconnectDelayMs ?? 500; + this.callTimeoutMs = opts.callTimeoutMs ?? DEFAULT_CALL_TIMEOUT_MS; + this.connect(); + } + + get currentState(): WsSocketState { + return this.state; + } + + onDidChangeState(listener: (state: WsSocketState) => void): WsSubscription { + this.stateListeners.add(listener); + return { dispose: () => this.stateListeners.delete(listener) }; + } + + /** RPC call over the socket; rejects on `error` frame, timeout, or close. */ + async call( + scope: WsScopeKind, + service: string, + method: string, + arg?: unknown, + ids?: WsScopeIds, + ): Promise { + await this.whenReady(); + // The socket may have dropped between `whenReady` resolving and this + // continuation running; never register a call we cannot send. + if (this.state !== 'open') { + throw new Error('ws closed'); + } + const id = this.nextId(); + const promise = new Promise((resolve, reject) => { + const timer = + this.callTimeoutMs > 0 + ? setTimeout(() => { + this.pending.delete(id); + reject(new RPCError(50001, `call timed out after ${this.callTimeoutMs}ms`)); + }, this.callTimeoutMs) + : undefined; + this.pending.set(id, { + resolve: resolve as (data: unknown) => void, + reject, + timer, + }); + }); + this.send({ type: 'call', id, scope, service, method, arg, ...ids }); + return promise; + } + + /** + * Subscribe to a scope event stream. The subscription survives reconnects + * (re-sent after each reconnect) until `dispose()`d. + */ + listen( + scope: WsScopeKind, + event: string, + ids: WsScopeIds, + handler: (data: unknown) => void, + ): WsSubscription { + const id = this.nextId(); + this.listens.set(id, { scope, event, ids, handler }); + if (this.state === 'open') { + this.send({ type: 'listen', id, scope, event, ...ids }); + } + return { + dispose: () => { + if (!this.listens.delete(id)) return; + if (this.state === 'open') { + this.send({ type: 'unlisten', id }); + } + }, + }; + } + + /** Tear the socket down permanently; rejects in-flight calls. */ + close(): void { + this.manualClose = true; + if (this.reconnectTimer !== undefined) { + clearTimeout(this.reconnectTimer); + this.reconnectTimer = undefined; + } + this.setState('closed'); + this.ws?.close(); + this.ws = undefined; + this.failAll(new Error('ws closed')); + this.rejectReadyWaiters(new Error('ws closed')); + } + + // ------------------------------------------------------------------------- + // Internals + // ------------------------------------------------------------------------- + + private nextId(): string { + this.seq += 1; + return `${this.idPrefix}_${this.seq}`; + } + + private connect(): void { + this.setState('connecting'); + const protocols = + this.token !== undefined && this.token.length > 0 + ? [`${WS_BEARER_PROTOCOL_PREFIX}${this.token}`] + : undefined; + let ws: WsLike; + try { + ws = new this.WsCtor(this.wsUrl, protocols); + } catch (error) { + this.scheduleReconnect(error); + return; + } + this.ws = ws; + ws.addEventListener('open', () => this.onOpen()); + ws.addEventListener('message', (event: { data: unknown }) => this.onMessage(event.data)); + ws.addEventListener('close', () => this.onClose()); + ws.addEventListener('error', () => { + // The 'close' event always follows 'error'; reconnect logic lives there. + }); + } + + private onOpen(): void { + this.reconnectAttempt = 0; + this.setState('open'); + this.send({ type: 'hello', token: this.token }); + for (const [id, sub] of this.listens) { + this.send({ type: 'listen', id, scope: sub.scope, event: sub.event, ...sub.ids }); + } + const waiters = this.readyWaiters; + this.readyWaiters = []; + for (const w of waiters) w.resolve(); + } + + private onMessage(raw: unknown): void { + let frame: ServerFrame; + try { + frame = JSON.parse(typeof raw === 'string' ? raw : String(raw)) as ServerFrame; + } catch { + return; + } + switch (frame.type) { + case 'ready': + case 'server_hello': + return; + case 'ping': + this.send({ type: 'pong' }); + return; + case 'result': { + const p = this.take(frame.id); + p?.resolve(frame.data); + return; + } + case 'error': { + const p = this.take(frame.id); + p?.reject(new RPCError(frame.code ?? 50001, frame.msg ?? 'error')); + return; + } + case 'event': { + const sub = this.listens.get(frame.id ?? ''); + sub?.handler(frame.data); + return; + } + } + } + + private onClose(): void { + this.ws = undefined; + this.failAll(new Error('ws closed')); + if (this.manualClose || !this.autoReconnect) { + this.setState('closed'); + this.rejectReadyWaiters(new Error('ws closed')); + return; + } + // Transient drop: queued calls keep waiting for the reconnect. + this.scheduleReconnect(undefined); + } + + private scheduleReconnect(_cause: unknown): void { + if (this.manualClose || !this.autoReconnect) { + this.setState('closed'); + return; + } + this.reconnectAttempt += 1; + const delay = Math.min(this.reconnectDelayMs * 2 ** (this.reconnectAttempt - 1), 10_000); + this.setState('connecting'); + this.reconnectTimer = setTimeout(() => { + this.reconnectTimer = undefined; + this.connect(); + }, delay); + this.reconnectTimer.unref?.(); + } + + private whenReady(): Promise { + if (this.state === 'open') return Promise.resolve(); + if (this.state === 'closed' && this.manualClose) { + return Promise.reject(new Error('ws closed')); + } + return new Promise((resolve, reject) => { + this.readyWaiters.push({ resolve, reject }); + }); + } + + private rejectReadyWaiters(err: Error): void { + const waiters = this.readyWaiters; + this.readyWaiters = []; + for (const w of waiters) w.reject(err); + } + + private take(id: string | undefined): PendingCall | undefined { + const p = this.pending.get(id ?? ''); + if (p !== undefined) { + this.pending.delete(id ?? ''); + if (p.timer !== undefined) clearTimeout(p.timer); + } + return p; + } + + private failAll(err: Error): void { + for (const p of this.pending.values()) { + if (p.timer !== undefined) clearTimeout(p.timer); + p.reject(err); + } + this.pending.clear(); + } + + private send(frame: Record): void { + const ws = this.ws; + if (ws === undefined || ws.readyState !== this.WsCtor.OPEN) return; + try { + ws.send(JSON.stringify(frame)); + } catch { + // best-effort; the close handler handles teardown + } + } + + private setState(next: WsSocketState): void { + if (this.state === next) return; + this.state = next; + for (const listener of this.stateListeners) listener(next); + } +} + +/** Derive the `/api/v2/ws` WebSocket URL from a server base URL (or pass a full ws URL through). */ +function toWsUrl(base: string): string { + const url = new URL(base); + if (url.protocol === 'http:') url.protocol = 'ws:'; + else if (url.protocol === 'https:') url.protocol = 'wss:'; + if (url.protocol !== 'ws:' && url.protocol !== 'wss:') { + throw new Error(`unsupported URL scheme for WS transport: ${base}`); + } + if (!url.pathname.endsWith('/api/v2/ws')) { + url.pathname = `${url.pathname.replace(/\/$/, '')}/api/v2/ws`; + } + url.search = ''; + url.hash = ''; + return url.toString(); +} diff --git a/packages/klient/test/ws.test.ts b/packages/klient/test/ws.test.ts new file mode 100644 index 000000000..9d70a7f11 --- /dev/null +++ b/packages/klient/test/ws.test.ts @@ -0,0 +1,289 @@ +import { describe, expect, it } from 'vitest'; + +import { ISessionIndex } from '@moonshot-ai/agent-core-v2/app/sessionIndex/sessionIndex'; +import { ISessionMetadata } from '@moonshot-ai/agent-core-v2/session/sessionMetadata/sessionMetadata'; + +import { Klient } from '../src/client.js'; +import { WsKlient } from '../src/wsKlient.js'; +import type { WsLike, WsLikeCtor, WsSocketState } from '../src/wsSocket.js'; + +const tick = (ms = 0): Promise => new Promise((resolve) => setTimeout(resolve, ms)); + +type Listener = (event: never) => void; + +/** + * In-memory emulation of the kap-server `/api/v2/ws` endpoint: answers `call` + * with an echo `result`, `boom` with an `error`, pushes `event`s to active + * `listen`s, and can drop the socket to exercise reconnect. + */ +class FakeServer { + readonly frames: Record[] = []; + readonly listens = new Set(); + pongs = 0; + helloCount = 0; + lastUrl = ''; + lastProtocols: string[] | undefined; + private socket: FakeClientSocket | undefined; + + attach(socket: FakeClientSocket): void { + this.socket = socket; + queueMicrotask(() => { + socket.readyState = FakeClientSocket.OPEN; + socket.fire('open'); + this.send({ type: 'ready', heartbeatMs: 30_000 }); + }); + } + + receive(raw: string): void { + const frame = JSON.parse(raw) as Record; + this.frames.push(frame); + switch (frame['type']) { + case 'hello': + this.helloCount += 1; + return; + case 'call': + if (frame['method'] === 'boom') { + this.send({ type: 'error', id: frame['id'], code: 40001, msg: 'boom' }); + } else { + this.send({ + type: 'result', + id: frame['id'], + data: { + scope: frame['scope'], + service: frame['service'], + method: frame['method'], + arg: frame['arg'] ?? null, + sessionId: frame['sessionId'] ?? null, + agentId: frame['agentId'] ?? null, + }, + }); + } + return; + case 'listen': + this.listens.add(frame['id'] as string); + return; + case 'unlisten': + this.listens.delete(frame['id'] as string); + return; + case 'pong': + this.pongs += 1; + return; + } + } + + pushEvent(id: string, data: unknown): void { + this.send({ type: 'event', id, data }); + } + + ping(): void { + this.send({ type: 'ping' }); + } + + drop(): void { + this.socket?.dropFromServer(); + } + + private send(frame: Record): void { + this.socket?.deliver(frame); + } +} + +class FakeClientSocket implements WsLike { + static readonly OPEN = 1; + readyState = 0; + private readonly handlers = new Map>(); + + constructor( + private readonly server: FakeServer, + url: string, + protocols?: string | string[], + ) { + server.lastUrl = url; + server.lastProtocols = Array.isArray(protocols) ? protocols : protocols ? [protocols] : undefined; + server.attach(this); + } + + addEventListener(type: string, listener: Listener): void { + const set = this.handlers.get(type) ?? new Set(); + set.add(listener); + this.handlers.set(type, set); + } + + send(data: string): void { + this.server.receive(data); + } + + close(): void { + this.readyState = 3; + this.fire('close'); + } + + fire(type: string): void { + for (const handler of this.handlers.get(type) ?? []) handler(undefined as never); + } + + deliver(frame: Record): void { + queueMicrotask(() => { + for (const handler of this.handlers.get('message') ?? []) { + handler({ data: JSON.stringify(frame) } as never); + } + }); + } + + dropFromServer(): void { + this.readyState = 3; + this.fire('close'); + } +} + +function fakeCtor(server: FakeServer): WsLikeCtor { + class BoundFakeSocket extends FakeClientSocket { + constructor(url: string, protocols?: string | string[]) { + super(server, url, protocols); + } + } + return BoundFakeSocket as unknown as WsLikeCtor; +} + +async function openKlient(server: FakeServer, opts: { token?: string } = {}): Promise { + const ws = new WsKlient({ + url: 'http://127.0.0.1:58627', + token: opts.token, + WebSocketImpl: fakeCtor(server), + reconnectDelayMs: 10, + }); + await tick(5); + return ws; +} + +describe('WsKlient', () => { + it('routes calls by scope / service / method with scope ids', async () => { + const server = new FakeServer(); + const ws = await openKlient(server); + + const core = await ws.core(ISessionIndex).list({ workspaceId: 'w1' }); + const session = await ws.session('s1').service(ISessionMetadata).read(); + const agent = await ws.session('s1').agent('a1').service(ISessionMetadata).read(); + + expect(core).toMatchObject({ scope: 'core', service: 'sessionIndex', method: 'list' }); + expect(session).toMatchObject({ scope: 'session', service: 'sessionMetadata', sessionId: 's1' }); + expect(agent).toMatchObject({ + scope: 'agent', + service: 'sessionMetadata', + sessionId: 's1', + agentId: 'a1', + }); + ws.close(); + }); + + it('rejects calls with RPCError on error frames', async () => { + const server = new FakeServer(); + const ws = await openKlient(server); + const meta = ws.session('s1').service(ISessionMetadata) as unknown as { + boom(): Promise; + }; + await expect(meta.boom()).rejects.toMatchObject({ + name: 'RPCError', + code: 40001, + }); + ws.close(); + }); + + it('delivers events to listen handlers and sends unlisten on dispose', async () => { + const server = new FakeServer(); + const ws = await openKlient(server); + const seen: unknown[] = []; + const sub = ws.session('s1').listen('interactions', (data) => seen.push(data)); + await tick(5); + const listenId = [...server.listens][0]!; + server.pushEvent(listenId, [{ id: 'a1' }]); + await tick(5); + expect(seen).toEqual([[{ id: 'a1' }]]); + + sub.dispose(); + expect(server.listens.size).toBe(0); + ws.close(); + }); + + it('answers heartbeat pings with pong', async () => { + const server = new FakeServer(); + const ws = await openKlient(server); + server.ping(); + await tick(5); + expect(server.pongs).toBe(1); + ws.close(); + }); + + it('reconnects after a drop: calls reject, listens re-subscribe, state is observable', async () => { + const server = new FakeServer(); + const ws = await openKlient(server); + const states: WsSocketState[] = []; + ws.onDidChangeState((s) => states.push(s)); + const seen: unknown[] = []; + ws.session('s1').agent('a1').listen('events', (data) => seen.push(data)); + await tick(5); + + const inFlight = ws.core(ISessionIndex).countActive('w1'); + server.drop(); + await expect(inFlight).rejects.toThrow('ws closed'); + expect(ws.state).toBe('connecting'); + + await tick(50); // backoff 10ms → reconnect + expect(ws.state).toBe('open'); + expect(server.helloCount).toBe(2); + expect(server.listens.size).toBe(1); + + const listenId = [...server.listens][0]!; + server.pushEvent(listenId, { type: 'turn.started' }); + await tick(5); + expect(seen).toEqual([{ type: 'turn.started' }]); + + const data = await ws.core(ISessionIndex).countActive('w1'); + expect(data).toMatchObject({ method: 'countActive' }); + expect(states).toContain('connecting'); + ws.close(); + expect(ws.state).toBe('closed'); + }); + + it('Klient.ws() is a lazy singleton deriving the ws URL and bearer subprotocol', async () => { + const server = new FakeServer(); + const client = new Klient({ + url: 'http://127.0.0.1:58627', + token: 'tok', + WebSocketImpl: fakeCtor(server), + }); + const ws = client.ws(); + expect(client.ws()).toBe(ws); + await tick(5); + expect(server.lastUrl).toBe('ws://127.0.0.1:58627/api/v2/ws'); + expect(server.lastProtocols).toEqual(['kimi-code.bearer.tok']); + const data = await ws.core(ISessionIndex).list({}); + expect(data).toMatchObject({ service: 'sessionIndex' }); + ws.close(); + }); + + it('rejects calls made after close', async () => { + const server = new FakeServer(); + const ws = await openKlient(server); + ws.close(); + await expect(ws.core(ISessionIndex).list({})).rejects.toThrow('ws closed'); + }); + + it('Klient.ws() recreates a fresh socket after the previous one was closed', async () => { + const server = new FakeServer(); + const client = new Klient({ + url: 'http://127.0.0.1:58627', + WebSocketImpl: fakeCtor(server), + }); + const first = client.ws(); + await tick(5); + first.close(); + const second = client.ws(); + expect(second).not.toBe(first); + await tick(5); + expect(second.state).toBe('open'); + const data = await second.core(ISessionIndex).list({}); + expect(data).toMatchObject({ service: 'sessionIndex' }); + second.close(); + }); +});