From 10922fc70fc9b69bb7b5806a1e545854c50ac72a Mon Sep 17 00:00:00 2001 From: 7Sageer <7sageer@djwcb.cn> Date: Mon, 6 Jul 2026 21:54:12 +0800 Subject: [PATCH] feat(telemetry): add system metrics collection (#1435) * feat(telemetry): add system metrics collection Add periodic CPU and memory telemetry sampling with warmup capture, lifecycle cleanup, and tests. * fix(telemetry): attach prompt session to system metrics --- apps/kimi-code/src/cli/run-prompt.ts | 1 + apps/kimi-code/src/cli/telemetry.ts | 2 + apps/kimi-code/test/cli/export.test.ts | 1 + apps/kimi-code/test/cli/run-prompt.test.ts | 3 + apps/kimi-code/test/cli/run-shell.test.ts | 1 + packages/telemetry/src/bootstrap.ts | 6 + packages/telemetry/src/client.ts | 22 ++++ packages/telemetry/src/systemMetrics.ts | 107 ++++++++++++++++ packages/telemetry/test/telemetry.test.ts | 135 +++++++++++++++++++++ 9 files changed, 278 insertions(+) create mode 100644 packages/telemetry/src/systemMetrics.ts diff --git a/apps/kimi-code/src/cli/run-prompt.ts b/apps/kimi-code/src/cli/run-prompt.ts index 66315ff3e..a9221d215 100644 --- a/apps/kimi-code/src/cli/run-prompt.ts +++ b/apps/kimi-code/src/cli/run-prompt.ts @@ -179,6 +179,7 @@ export async function runPrompt( version, uiMode: PROMPT_UI_MODE, model: telemetryModel, + sessionId: session.id, }); setCrashPhase('runtime'); diff --git a/apps/kimi-code/src/cli/telemetry.ts b/apps/kimi-code/src/cli/telemetry.ts index 2e7f49b4b..b228c913e 100644 --- a/apps/kimi-code/src/cli/telemetry.ts +++ b/apps/kimi-code/src/cli/telemetry.ts @@ -32,6 +32,7 @@ export interface InitializeCliTelemetryOptions { readonly version: string; readonly uiMode: string; readonly model?: string; + readonly sessionId?: string; } export function createCliTelemetryBootstrap(): CliTelemetryBootstrap { @@ -54,6 +55,7 @@ export function initializeCliTelemetry(options: InitializeCliTelemetryOptions): version: options.version, uiMode: options.uiMode, model: options.model ?? options.config.defaultModel, + sessionId: options.sessionId, getAccessToken: async () => (await options.harness.auth.getCachedAccessToken(KIMI_CODE_PROVIDER_NAME)) ?? null, }); diff --git a/apps/kimi-code/test/cli/export.test.ts b/apps/kimi-code/test/cli/export.test.ts index 4518cf3e7..14fdc0190 100644 --- a/apps/kimi-code/test/cli/export.test.ts +++ b/apps/kimi-code/test/cli/export.test.ts @@ -412,6 +412,7 @@ describe('kimi export', () => { version: expect.any(String), uiMode: 'shell', model: 'k2', + sessionId: undefined, getAccessToken: expect.any(Function), }); expect(mocks.initializeTelemetry.mock.invocationCallOrder[0]).toBeLessThan( diff --git a/apps/kimi-code/test/cli/run-prompt.test.ts b/apps/kimi-code/test/cli/run-prompt.test.ts index d75a6b3a5..aafb99af9 100644 --- a/apps/kimi-code/test/cli/run-prompt.test.ts +++ b/apps/kimi-code/test/cli/run-prompt.test.ts @@ -217,6 +217,9 @@ describe('runPrompt', () => { expect(mocks.session.prompt).toHaveBeenCalledWith('say hello'); expect(stdout.text()).toBe('• hello world\n\n'); expect(stderr.text()).toBe('To resume this session: kimi -r ses_prompt\n'); + expect(mocks.initializeTelemetry).toHaveBeenCalledWith( + expect.objectContaining({ sessionId: 'ses_prompt' }), + ); expect(mocks.shutdownTelemetry).toHaveBeenCalled(); expect(mocks.harnessClose).toHaveBeenCalled(); }); diff --git a/apps/kimi-code/test/cli/run-shell.test.ts b/apps/kimi-code/test/cli/run-shell.test.ts index b582c8543..3fa673ab9 100644 --- a/apps/kimi-code/test/cli/run-shell.test.ts +++ b/apps/kimi-code/test/cli/run-shell.test.ts @@ -213,6 +213,7 @@ describe('runShell', () => { version: '1.2.3-test', uiMode: 'shell', model: 'k2', + sessionId: undefined, getAccessToken: expect.any(Function), }); expect(mocks.setCrashPhase).toHaveBeenCalledWith('runtime'); diff --git a/packages/telemetry/src/bootstrap.ts b/packages/telemetry/src/bootstrap.ts index 041e650d4..f5d79f063 100644 --- a/packages/telemetry/src/bootstrap.ts +++ b/packages/telemetry/src/bootstrap.ts @@ -1,5 +1,6 @@ import { getDefaultTelemetryClient } from './client'; import { EventSink } from './sink'; +import { SystemMetricsCollector } from './systemMetrics'; import { AsyncTransport } from './transport'; export const TELEMETRY_DISABLE_ENV = 'KIMI_DISABLE_TELEMETRY'; @@ -65,5 +66,10 @@ export function initializeTelemetry(options: TelemetryBootstrapOptions): void { client.attachSink(sink); sink.startPeriodicFlush(); + + const systemMetricsCollector = new SystemMetricsCollector({ client }); + client.setSystemMetricsCollector(systemMetricsCollector); + systemMetricsCollector.start(); + void sink.retryDiskEvents().catch(() => {}); } diff --git a/packages/telemetry/src/client.ts b/packages/telemetry/src/client.ts index d4904e8f2..3710278a1 100644 --- a/packages/telemetry/src/client.ts +++ b/packages/telemetry/src/client.ts @@ -13,6 +13,10 @@ export interface TelemetryShutdownOptions { readonly timeoutMs?: number; } +export interface SystemMetricsCollectorHandle { + stop(): void; +} + const MAX_QUEUE_SIZE = 1000; interface PendingTelemetryEvent extends TelemetryEvent { @@ -25,6 +29,7 @@ interface PendingTelemetryEvent extends TelemetryEvent { export class TelemetryClient { private queue: PendingTelemetryEvent[] = []; private sink: EventSink | null = null; + private systemMetricsCollector: SystemMetricsCollectorHandle | null = null; private deviceId: string | null = null; private sessionId: string | null = null; private disabled = false; @@ -38,6 +43,13 @@ export class TelemetryClient { return new ScopedTelemetryClient(this, input); } + setSystemMetricsCollector(collector: SystemMetricsCollectorHandle): void { + if (this.systemMetricsCollector !== null && this.systemMetricsCollector !== collector) { + this.systemMetricsCollector.stop(); + } + this.systemMetricsCollector = collector; + } + attachSink(sink: EventSink): void { if (this.sink !== null && this.sink !== sink) { this.sink.stopPeriodicFlush(); @@ -60,6 +72,8 @@ export class TelemetryClient { disable(): void { this.disabled = true; this.queue = []; + this.systemMetricsCollector?.stop(); + this.systemMetricsCollector = null; if (this.sink !== null) { this.sink.stopPeriodicFlush(); this.sink.clearBuffer(); @@ -116,6 +130,8 @@ export class TelemetryClient { } async shutdown(options: TelemetryShutdownOptions = {}): Promise { + this.systemMetricsCollector?.stop(); + this.systemMetricsCollector = null; const sink = this.sink; if (sink === null) return; sink.stopPeriodicFlush(); @@ -139,6 +155,8 @@ export class TelemetryClient { resetForTests(): void { this.sink?.stopPeriodicFlush(); + this.systemMetricsCollector?.stop(); + this.systemMetricsCollector = null; this.queue = []; this.sink = null; this.deviceId = null; @@ -163,6 +181,10 @@ class ScopedTelemetryClient extends TelemetryClient { return new ScopedTelemetryClient(this.parent, mergeContext(this.context, input)); } + override setSystemMetricsCollector(collector: SystemMetricsCollectorHandle): void { + this.parent.setSystemMetricsCollector(collector); + } + override attachSink(sink: EventSink): void { this.parent.attachSink(sink); } diff --git a/packages/telemetry/src/systemMetrics.ts b/packages/telemetry/src/systemMetrics.ts new file mode 100644 index 000000000..a00fe03d4 --- /dev/null +++ b/packages/telemetry/src/systemMetrics.ts @@ -0,0 +1,107 @@ +import { cpus, freemem, loadavg, totalmem } from 'node:os'; + +import type { TelemetryProperties } from './types'; + +const DEFAULT_INTERVAL_MS = 30_000; +const DEFAULT_WARMUP_SAMPLE_MS = 1_500; +const SYSTEM_METRICS_EVENT = 'system_metrics'; + +export interface SystemMetricsTrackClient { + track(event: string, properties?: TelemetryProperties): void; +} + +export interface SystemMetricsCollectorOptions { + readonly client: SystemMetricsTrackClient; + readonly intervalMs?: number; + readonly warmupSampleMs?: number | null; +} + +export class SystemMetricsCollector { + private readonly client: SystemMetricsTrackClient; + private readonly intervalMs: number; + private readonly warmupSampleMs: number | null; + private intervalTimer: ReturnType | null = null; + private warmupTimer: ReturnType | null = null; + private previousCpuUsage = process.cpuUsage(); + private previousHrtime = process.hrtime.bigint(); + private readonly processStartedAtSeconds = Math.floor(Date.now() / 1000 - process.uptime()); + + constructor(options: SystemMetricsCollectorOptions) { + this.client = options.client; + this.intervalMs = options.intervalMs ?? DEFAULT_INTERVAL_MS; + this.warmupSampleMs = + options.warmupSampleMs === undefined ? DEFAULT_WARMUP_SAMPLE_MS : options.warmupSampleMs; + } + + start(): void { + if (this.intervalTimer !== null) return; + + if (this.warmupSampleMs !== null && this.warmupSampleMs > 0) { + this.warmupTimer = setTimeout(() => { + this.warmupTimer = null; + this.sampleSafely(); + }, this.warmupSampleMs); + this.warmupTimer.unref?.(); + } + + this.intervalTimer = setInterval(() => { + this.sampleSafely(); + }, this.intervalMs); + this.intervalTimer.unref?.(); + } + + stop(): void { + if (this.warmupTimer !== null) { + clearTimeout(this.warmupTimer); + this.warmupTimer = null; + } + if (this.intervalTimer !== null) { + clearInterval(this.intervalTimer); + this.intervalTimer = null; + } + } + + private sampleSafely(): void { + try { + this.sample(); + } catch { + this.stop(); + } + } + + private sample(): void { + const now = process.hrtime.bigint(); + const elapsedUs = Number(now - this.previousHrtime) / 1_000; + + const cpu = process.cpuUsage(this.previousCpuUsage); + this.previousCpuUsage = process.cpuUsage(); + this.previousHrtime = now; + + const mem = process.memoryUsage(); + const constrainedMemory = getConstrainedMemoryBytes(); + + this.client.track(SYSTEM_METRICS_EVENT, { + process_started_at: this.processStartedAtSeconds, + process_uptime_ms: Math.round(process.uptime() * 1000), + rss_bytes: mem.rss, + heap_used_bytes: mem.heapUsed, + heap_total_bytes: mem.heapTotal, + external_bytes: mem.external, + array_buffers_bytes: mem.arrayBuffers, + constrained_memory_bytes: constrainedMemory, + cpu_user_us: cpu.user, + cpu_system_us: cpu.system, + cpu_elapsed_us: Math.round(elapsedUs), + load_avg_1m: loadavg()[0], + free_mem_bytes: freemem(), + total_mem_bytes: totalmem(), + cpu_count: cpus().length, + }); + } +} + +function getConstrainedMemoryBytes(): number | undefined { + if (typeof process.constrainedMemory !== 'function') return undefined; + const value = process.constrainedMemory(); + return Number.isFinite(value) ? value : undefined; +} diff --git a/packages/telemetry/test/telemetry.test.ts b/packages/telemetry/test/telemetry.test.ts index a19082b7c..8e048f635 100644 --- a/packages/telemetry/test/telemetry.test.ts +++ b/packages/telemetry/test/telemetry.test.ts @@ -12,6 +12,7 @@ import { isTelemetryDisabledByEnv } from '../src/bootstrap'; import { TelemetryClient, resetDefaultTelemetryClientForTests } from '../src/client'; import { installCrashHandlersForClient, setCrashPhase, uninstallCrashHandlers } from '../src/crash'; import { EventSink } from '../src/sink'; +import { SystemMetricsCollector } from '../src/systemMetrics'; import { AsyncTransport, DISK_EVENT_MAX_AGE_MS, @@ -30,6 +31,7 @@ afterEach(() => { uninstallCrashHandlers(); setCrashPhase('startup'); resetDefaultTelemetryClientForTests(); + vi.useRealTimers(); for (const dir of tempDirs.splice(0)) { rmSync(dir, { recursive: true, force: true }); } @@ -158,6 +160,19 @@ describe('TelemetryClient', () => { expect(transport.sent).toHaveLength(0); }); + it('stops the previous system metrics collector when replacing it', () => { + const client = new TelemetryClient(); + const first = { stop: vi.fn() }; + const second = { stop: vi.fn() }; + + client.setSystemMetricsCollector(first); + client.setSystemMetricsCollector(second); + client.disable(); + + expect(first.stop).toHaveBeenCalledTimes(1); + expect(second.stop).toHaveBeenCalledTimes(1); + }); + it('flushes the previous sink synchronously when replacing sinks', () => { const client = new TelemetryClient(); const first = new RecordingTransport(); @@ -204,6 +219,77 @@ describe('TelemetryClient', () => { }); }); +describe('SystemMetricsCollector', () => { + it('emits a numeric system_metrics sample after the warmup delay', () => { + vi.useFakeTimers(); + const tracked: Array<{ + event: string; + properties: Record; + }> = []; + const client = { + track( + event: string, + properties: Record = {}, + ): void { + tracked.push({ event, properties }); + }, + }; + const collector = new SystemMetricsCollector({ + client, + intervalMs: 30_000, + warmupSampleMs: 1_500, + }); + + collector.start(); + vi.advanceTimersByTime(1_499); + expect(tracked).toHaveLength(0); + + vi.advanceTimersByTime(1); + collector.stop(); + + expect(tracked).toHaveLength(1); + const event = tracked[0]; + if (event === undefined) throw new Error('Expected a system_metrics event'); + expect(event.event).toBe('system_metrics'); + expect(numberProperty(event.properties, 'process_started_at')).toBeGreaterThan(0); + expect(numberProperty(event.properties, 'process_uptime_ms')).toBeGreaterThanOrEqual(0); + expect(numberProperty(event.properties, 'rss_bytes')).toBeGreaterThan(0); + expect(numberProperty(event.properties, 'heap_used_bytes')).toBeGreaterThan(0); + expect(numberProperty(event.properties, 'heap_total_bytes')).toBeGreaterThan(0); + expect(numberProperty(event.properties, 'external_bytes')).toBeGreaterThanOrEqual(0); + expect(numberProperty(event.properties, 'array_buffers_bytes')).toBeGreaterThanOrEqual(0); + expect(numberProperty(event.properties, 'cpu_user_us')).toBeGreaterThanOrEqual(0); + expect(numberProperty(event.properties, 'cpu_system_us')).toBeGreaterThanOrEqual(0); + expect(numberProperty(event.properties, 'cpu_elapsed_us')).toBeGreaterThan(0); + expect(numberProperty(event.properties, 'load_avg_1m')).toBeGreaterThanOrEqual(0); + expect(numberProperty(event.properties, 'free_mem_bytes')).toBeGreaterThanOrEqual(0); + expect(numberProperty(event.properties, 'total_mem_bytes')).toBeGreaterThan(0); + expect(numberProperty(event.properties, 'cpu_count')).toBeGreaterThanOrEqual(1); + }); + + it('does not duplicate interval sampling when started twice', () => { + vi.useFakeTimers(); + const tracked: string[] = []; + const client = { + track(event: string): void { + tracked.push(event); + }, + }; + const collector = new SystemMetricsCollector({ + client, + intervalMs: 30_000, + warmupSampleMs: null, + }); + + collector.start(); + collector.start(); + vi.advanceTimersByTime(30_000); + collector.stop(); + + expect(tracked).toEqual(['system_metrics']); + }); +}); + describe('EventSink', () => { it('enriches context without mutating the original event', () => { const transport = new RecordingTransport(); @@ -692,6 +778,44 @@ describe('telemetry bootstrap', () => { const file = readFileSync(join(telemetryDir, readdirOne(telemetryDir)), 'utf-8'); expect(file).toContain('"event":"sync_flush"'); }); + + it('writes system metrics with the singleton session context', async () => { + vi.useFakeTimers(); + const homeDir = await tempHome(); + initializeTelemetry({ + homeDir, + deviceId: 'dev', + sessionId: 'ses', + appName: 'kimi-code-cli', + version: '1.2.3', + }); + + vi.advanceTimersByTime(1_500); + flushTelemetrySync(); + + const telemetryDir = join(homeDir, 'telemetry'); + const file = readFileSync(join(telemetryDir, readdirOne(telemetryDir)), 'utf-8'); + const events = file + .trim() + .split('\n') + .map( + (line) => + JSON.parse(line) as { + event: string; + session_id: string | null; + properties: Record; + }, + ); + const metrics = events.find((event) => event.event === 'system_metrics'); + if (metrics === undefined) throw new Error('Expected a system_metrics event'); + + expect(metrics.session_id).toBe('ses'); + expect(Number.isFinite(metrics.properties['process_started_at'])).toBe(true); + expect(metrics.properties['process_started_at']).toBeGreaterThan(0); + expect(Number.isFinite(metrics.properties['process_uptime_ms'])).toBe(true); + expect(metrics.properties['process_uptime_ms']).toBeGreaterThanOrEqual(0); + expect(metrics.properties['rss_bytes']).toBeGreaterThan(0); + }); }); describe('crash handler', () => { @@ -826,6 +950,17 @@ function readdirOne(dir: string): string { return entry; } +function numberProperty( + properties: Record, + key: string, +): number { + const value = properties[key]; + if (typeof value !== 'number' || !Number.isFinite(value)) { + throw new Error(`Expected property ${key} to be a finite number, got ${String(value)}`); + } + return value; +} + function requestInitFrom( fetchImpl: { readonly mock: { readonly calls: readonly unknown[][] } }, index = 0,