diff --git a/apps/api/src/common/interfaces/storage-port.interface.ts b/apps/api/src/common/interfaces/storage-port.interface.ts index 6319f4c8c..bf0648d38 100644 --- a/apps/api/src/common/interfaces/storage-port.interface.ts +++ b/apps/api/src/common/interfaces/storage-port.interface.ts @@ -438,6 +438,12 @@ export interface KvCacheFootprintQuery extends KvCacheRangeQuery { export interface KvCacheSampleQuery extends KvCacheRangeQuery { engineId?: string; modelName?: string; + bucketMs?: number; +} + +export interface KvCacheSampleTotals { + requestedTokens: number; + hitTokens: number; } export interface MemorySnapshotQueryOptions { @@ -709,7 +715,9 @@ export interface StoragePort { pruneOldKvCacheFootprintSnapshots(cutoffTimestamp: number, connectionId?: string): Promise; saveKvCacheEngineSamples(samples: KvCacheEngineSample[]): Promise; getKvCacheEngineSamples(query: KvCacheSampleQuery): Promise; + getKvCacheSampleTotals(query: Omit): Promise; pruneOldKvCacheEngineSamples(cutoffTimestamp: number, connectionId?: string): Promise; + deleteKvCacheEngineSamples(engineId: string): Promise; saveKvCacheEngine(engine: StoredKvCacheEngine): Promise; getKvCacheEngines(connectionId?: string): Promise; getKvCacheEngine(id: string): Promise; diff --git a/apps/api/src/common/utils/__tests__/outbound-url-guard.spec.ts b/apps/api/src/common/utils/__tests__/outbound-url-guard.spec.ts new file mode 100644 index 000000000..b91ce01ee --- /dev/null +++ b/apps/api/src/common/utils/__tests__/outbound-url-guard.spec.ts @@ -0,0 +1,294 @@ +import { BadRequestException, Logger } from '@nestjs/common'; +import { promises as dns } from 'dns'; +import { assertSafeOutboundUrl, BlockedAddressError, guardedLookup, isBlockedIp } from '../outbound-url-guard'; + +describe('outbound url guard', () => { + const env = process.env.NODE_ENV; + afterEach(() => { + process.env.NODE_ENV = env; + jest.restoreAllMocks(); + }); + + it('classifies addresses', () => { + expect(isBlockedIp('10.0.0.1')).toBe(true); + expect(isBlockedIp('10.0.0.1', true)).toBe(false); + expect(isBlockedIp('169.254.169.254', true)).toBe(true); + expect(isBlockedIp('fe80::1', true)).toBe(true); + expect(isBlockedIp('8.8.8.8')).toBe(false); + expect(isBlockedIp('[::1]', true)).toBe(false); + expect(isBlockedIp('127.0.0.1', true)).toBe(false); + expect(isBlockedIp('[::1]')).toBe(true); + expect(isBlockedIp('127.0.0.1')).toBe(true); + expect(isBlockedIp('0.0.0.0')).toBe(true); + expect(isBlockedIp('100.100.100.200')).toBe(true); + expect(isBlockedIp('100.63.255.255')).toBe(false); + expect(isBlockedIp('100.128.0.1')).toBe(false); + }); + + it('rejects non-http protocols with the existing message', async () => { + await expect(assertSafeOutboundUrl('ftp://x', { label: 'metrics URL' })).rejects.toThrow('Only HTTP and HTTPS protocols are allowed'); + }); + + it('allows localhost outside production and labels the production error', async () => { + process.env.NODE_ENV = 'development'; + await expect(assertSafeOutboundUrl('http://localhost:9400/metrics', { label: 'metrics URL' })).resolves.toBeInstanceOf(URL); + process.env.NODE_ENV = 'production'; + await expect(assertSafeOutboundUrl('http://localhost:9400/metrics', { label: 'metrics URL' })).rejects.toThrow( + 'Cannot use localhost or 0.0.0.0 as metrics URL in production', + ); + }); + + it('lets the scraper reach private networks but never link-local', async () => { + process.env.NODE_ENV = 'production'; + jest.spyOn(dns, 'lookup').mockResolvedValue([{ address: '10.1.2.3', family: 4 }] as any); + await expect(assertSafeOutboundUrl('http://10.1.2.3:8000/metrics', { label: 'metrics URL', allowPrivateNetworks: true })).resolves.toBeInstanceOf(URL); + await expect(assertSafeOutboundUrl('http://10.1.2.3:8000/metrics', { label: 'metrics URL' })).rejects.toBeInstanceOf(BadRequestException); + await expect(assertSafeOutboundUrl('http://169.254.169.254/latest', { label: 'metrics URL', allowPrivateNetworks: true })).rejects.toBeInstanceOf(BadRequestException); + }); + + it('accepts private addresses that merely contain 127 when private networks are allowed', async () => { + process.env.NODE_ENV = 'production'; + jest.spyOn(dns, 'lookup').mockResolvedValue([{ address: '10.127.0.5', family: 4 }] as any); + await expect(assertSafeOutboundUrl('http://10.127.0.5:8000/metrics', { label: 'metrics URL', allowPrivateNetworks: true })).resolves.toBeInstanceOf(URL); + await expect(assertSafeOutboundUrl('http://192.168.127.10/metrics', { label: 'metrics URL', allowPrivateNetworks: true })).resolves.toBeInstanceOf(URL); + }); + + it('lets loopback follow the private network setting in production', async () => { + process.env.NODE_ENV = 'production'; + const resolve = jest.spyOn(dns, 'lookup').mockResolvedValue([{ address: '127.0.0.1', family: 4 }] as any); + await expect(assertSafeOutboundUrl('http://localhost:8000/metrics', { label: 'metrics URL', allowPrivateNetworks: true })).resolves.toBeInstanceOf(URL); + await expect(assertSafeOutboundUrl('http://127.0.0.2/metrics', { label: 'metrics URL', allowPrivateNetworks: true })).resolves.toBeInstanceOf(URL); + await expect(assertSafeOutboundUrl('http://[::1]:9400/metrics', { label: 'metrics URL', allowPrivateNetworks: true })).resolves.toBeInstanceOf(URL); + await expect(assertSafeOutboundUrl('http://localhost:8000/metrics', { label: 'metrics URL' })).rejects.toBeInstanceOf(BadRequestException); + await expect(assertSafeOutboundUrl('http://127.0.0.2/metrics', { label: 'metrics URL' })).rejects.toBeInstanceOf(BadRequestException); + await expect(assertSafeOutboundUrl('http://[fe80::1]/metrics', { label: 'metrics URL', allowPrivateNetworks: true })).rejects.toBeInstanceOf(BadRequestException); + expect(resolve).not.toHaveBeenCalled(); + }); + + it('lets IPv6 unique local addresses follow the private network setting', () => { + for (const ip of ['fc00::1', 'fdff:ffff::1', 'fd12:3456::8']) { + expect(isBlockedIp(ip, true)).toBe(false); + expect(isBlockedIp(ip)).toBe(true); + } + }); + + it('treats the whole 0.0.0.0/8 range like loopback', async () => { + for (const ip of ['0.0.0.0', '0.1.2.3', '0.255.255.255', '::ffff:0.1.2.3']) { + expect(isBlockedIp(ip)).toBe(true); + expect(isBlockedIp(ip, true)).toBe(false); + } + process.env.NODE_ENV = 'production'; + await expect(assertSafeOutboundUrl('http://0.1.2.3:8000/metrics', { label: 'metrics URL' })).rejects.toBeInstanceOf(BadRequestException); + }); + + it('always blocks the IPv6 metadata address', async () => { + for (const ip of ['fd00:ec2::254', '[fd00:ec2::254]', 'FD00:EC2::254', 'fd00:0ec2:0000:0000:0000:0000:0000:0254']) { + expect(isBlockedIp(ip, true)).toBe(true); + } + expect(isBlockedIp('fe00::1', true)).toBe(false); + expect(isBlockedIp('fc0::1', true)).toBe(false); + expect(isBlockedIp('2001:db8::1', true)).toBe(false); + process.env.NODE_ENV = 'production'; + await expect(assertSafeOutboundUrl('http://[fd00:ec2::254]/latest/meta-data', { label: 'metrics URL', allowPrivateNetworks: true })).rejects.toThrow('Cannot use private IP addresses as metrics URL'); + await expect(assertSafeOutboundUrl('http://[fd00:ec2:0:0:0:0:0:254]/latest', { label: 'metrics URL', allowPrivateNetworks: true })).rejects.toBeInstanceOf(BadRequestException); + jest.spyOn(dns, 'lookup').mockResolvedValue([{ address: 'fd00:ec2::254', family: 6 }] as any); + await expect(assertSafeOutboundUrl('http://metrics.example.com/x', { label: 'metrics URL', allowPrivateNetworks: true })).rejects.toThrow('Metrics URL resolves to a blocked address'); + }); + + it('always blocks other cloud metadata endpoints', () => { + for (const ip of ['100.100.100.200', '192.0.0.192', '::ffff:100.100.100.200']) { + expect(isBlockedIp(ip, true)).toBe(true); + } + expect(isBlockedIp('100.100.100.201', true)).toBe(false); + expect(isBlockedIp('192.0.0.193', true)).toBe(false); + }); + + it('unwraps IPv4-mapped IPv6 addresses before classifying them', async () => { + expect(isBlockedIp('::ffff:169.254.169.254', true)).toBe(true); + expect(isBlockedIp('[::ffff:a9fe:a9fe]', true)).toBe(true); + expect(isBlockedIp('::ffff:7f00:1')).toBe(true); + expect(isBlockedIp('::ffff:0a00:0001')).toBe(true); + expect(isBlockedIp('::ffff:0a00:0001', true)).toBe(false); + expect(isBlockedIp('::ffff:808:808')).toBe(false); + expect(isBlockedIp('fd12::1')).toBe(true); + expect(isBlockedIp('febf::1', true)).toBe(true); + process.env.NODE_ENV = 'production'; + await expect(assertSafeOutboundUrl('http://[::ffff:169.254.169.254]/latest', { label: 'metrics URL', allowPrivateNetworks: true })).rejects.toBeInstanceOf(BadRequestException); + await expect(assertSafeOutboundUrl('http://[::ffff:7f00:1]/metrics', { label: 'metrics URL' })).rejects.toBeInstanceOf(BadRequestException); + }); + + it('treats IPv6 loopback as localhost outside production', async () => { + process.env.NODE_ENV = 'development'; + await expect(assertSafeOutboundUrl('http://[::1]:9400/metrics', { label: 'metrics URL' })).resolves.toBeInstanceOf(URL); + }); + + it('does not resolve DNS for IP literals', async () => { + process.env.NODE_ENV = 'production'; + const resolve = jest.spyOn(dns, 'lookup'); + await expect(assertSafeOutboundUrl('http://10.1.2.3:8000/metrics', { label: 'metrics URL', allowPrivateNetworks: true })).resolves.toBeInstanceOf(URL); + await expect(assertSafeOutboundUrl('http://8.8.8.8/x', { label: 'webhook URL' })).resolves.toBeInstanceOf(URL); + expect(resolve).not.toHaveBeenCalled(); + }); + + it('treats a DNS answer in loopback space like a private address', async () => { + process.env.NODE_ENV = 'production'; + jest.spyOn(dns, 'lookup').mockResolvedValue([{ address: '127.0.0.1', family: 4 }] as any); + await expect(assertSafeOutboundUrl('http://metrics.example.com/x', { label: 'metrics URL', allowPrivateNetworks: true })).resolves.toBeInstanceOf(URL); + await expect(assertSafeOutboundUrl('http://metrics.example.com/x', { label: 'metrics URL' })).rejects.toThrow('Metrics URL resolves to a blocked address'); + }); + + it('resolves webhook hostnames only in production', async () => { + process.env.NODE_ENV = 'development'; + const resolve = jest.spyOn(dns, 'lookup').mockResolvedValue([{ address: '172.18.0.5', family: 4 }] as any); + await expect(assertSafeOutboundUrl('http://receiver:8080/hook', { label: 'webhook URL', resolveDns: 'production' })).resolves.toBeInstanceOf(URL); + expect(resolve).not.toHaveBeenCalled(); + await expect(assertSafeOutboundUrl('http://10.0.0.1/hook', { label: 'webhook URL', resolveDns: 'production' })).rejects.toThrow('Cannot use private IP addresses as webhook URL'); + process.env.NODE_ENV = 'production'; + await expect(assertSafeOutboundUrl('http://receiver:8080/hook', { label: 'webhook URL', resolveDns: 'production' })).rejects.toThrow('Webhook URL resolves to a blocked address'); + }); + + it('resolves DNS by default in every environment', async () => { + for (const nodeEnv of ['development', 'production']) { + process.env.NODE_ENV = nodeEnv; + const resolve = jest.spyOn(dns, 'lookup').mockResolvedValue([{ address: '10.9.9.9', family: 4 }] as any); + await expect(assertSafeOutboundUrl('http://metrics.example.com/x', { label: 'metrics URL' })).rejects.toThrow('Metrics URL resolves to a blocked address'); + expect(resolve).toHaveBeenCalled(); + resolve.mockRestore(); + } + }); + + it('only skips the DNS pre-check for an explicit never', async () => { + process.env.NODE_ENV = 'development'; + const resolve = jest.spyOn(dns, 'lookup').mockResolvedValue([{ address: '10.9.9.9', family: 4 }] as any); + for (const resolveDns of [undefined, 'always', 'sometimes']) { + await expect(assertSafeOutboundUrl('http://metrics.example.com/x', { label: 'metrics URL', resolveDns: resolveDns as any })).rejects.toBeInstanceOf(BadRequestException); + } + expect(resolve).toHaveBeenCalledTimes(3); + }); + + it('skips the DNS pre-check when resolveDns is never', async () => { + process.env.NODE_ENV = 'production'; + const resolve = jest.spyOn(dns, 'lookup').mockResolvedValue([{ address: '169.254.169.254', family: 4 }] as any); + await expect(assertSafeOutboundUrl('http://metrics.example.com/x', { label: 'metrics URL', resolveDns: 'never' })).resolves.toBeInstanceOf(URL); + expect(resolve).not.toHaveBeenCalled(); + await expect(assertSafeOutboundUrl('http://10.0.0.1/x', { label: 'metrics URL', resolveDns: 'never' })).rejects.toBeInstanceOf(BadRequestException); + }); + + it('applies the private IP patterns only to IP literals, not hostnames', async () => { + process.env.NODE_ENV = 'production'; + for (const url of ['http://10.metrics.example.com/x', 'http://192.168.example.com/x', 'http://fc00-metrics.example.com/x']) { + await expect(assertSafeOutboundUrl(url, { label: 'metrics URL', resolveDns: 'never' })).resolves.toBeInstanceOf(URL); + } + }); + + it('skips the suspicious hostname check when private networks are allowed', async () => { + process.env.NODE_ENV = 'production'; + await expect(assertSafeOutboundUrl('http://vllm-localhost-proxy:8000/x', { label: 'metrics URL', allowPrivateNetworks: true, resolveDns: 'never' })).resolves.toBeInstanceOf(URL); + await expect(assertSafeOutboundUrl('http://app.localhost:8000/x', { label: 'metrics URL', resolveDns: 'never' })).rejects.toThrow('Suspicious hostname detected'); + }); + + it('only treats localhost and its subdomains as suspicious hostnames', async () => { + process.env.NODE_ENV = 'production'; + for (const url of ['http://notlocalhost.io/x', 'http://vllm-localhost-proxy:8000/x', 'http://localhost.example.com/x']) { + await expect(assertSafeOutboundUrl(url, { label: 'metrics URL', resolveDns: 'never' })).resolves.toBeInstanceOf(URL); + } + for (const url of ['http://app.localhost/x', 'http://LOCALHOST./x', 'http://a.b.localhost./x']) { + await expect(assertSafeOutboundUrl(url, { label: 'metrics URL', resolveDns: 'never' })).rejects.toThrow(BadRequestException); + } + }); + + it('keeps the blocked address out of the error and logs it instead', async () => { + process.env.NODE_ENV = 'production'; + const warn = jest.spyOn(Logger.prototype, 'warn').mockImplementation(() => undefined); + jest.spyOn(dns, 'lookup').mockResolvedValue([{ address: '10.9.9.9', family: 4 }] as any); + const error = await assertSafeOutboundUrl('http://db.internal/x', { label: 'metrics URL' }).catch((e: Error) => e); + expect((error as Error).message).toBe('Metrics URL resolves to a blocked address'); + expect(warn.mock.calls.flat().join(' ')).toContain('10.9.9.9'); + const { error: lookupError } = await new Promise<{ error: Error | null }>((resolve) => + guardedLookup({ label: 'metrics URL' })('db.internal', { all: false } as any, (e: NodeJS.ErrnoException | null) => resolve({ error: e })), + ); + expect(lookupError?.message).toBe('Metrics URL resolves to a blocked address'); + }); + + it('rejects a DNS answer in link-local space even when private networks are allowed', async () => { + process.env.NODE_ENV = 'production'; + jest.spyOn(dns, 'lookup').mockResolvedValue([{ address: '169.254.169.254', family: 4 }] as any); + await expect(assertSafeOutboundUrl('http://metrics.example.com/x', { label: 'metrics URL', allowPrivateNetworks: true })).rejects.toThrow('Metrics URL resolves to a blocked address'); + }); + + it('uses the label in the remaining messages', async () => { + process.env.NODE_ENV = 'production'; + await expect(assertSafeOutboundUrl('http://10.0.0.1/x', { label: 'webhook URL' })).rejects.toThrow('Cannot use private IP addresses as webhook URL'); + await expect(assertSafeOutboundUrl('not a url', { label: 'webhook URL' })).rejects.toThrow('Invalid webhook URL'); + jest.spyOn(dns, 'lookup').mockResolvedValue([{ address: '10.9.9.9', family: 4 }] as any); + await expect(assertSafeOutboundUrl('http://example.com/x', { label: 'webhook URL' })).rejects.toThrow('Webhook URL resolves to a blocked address'); + jest.spyOn(dns, 'lookup').mockRejectedValue(new Error('nx')); + await expect(assertSafeOutboundUrl('http://example.com/x', { label: 'webhook URL' })).rejects.toThrow('Failed to resolve webhook URL hostname'); + }); + + it('checks DNS answers outside production too', async () => { + process.env.NODE_ENV = 'development'; + jest.spyOn(dns, 'lookup').mockResolvedValue([{ address: '169.254.169.254', family: 4 }] as any); + await expect(assertSafeOutboundUrl('http://metadata.example.com/x', { label: 'metrics URL', allowPrivateNetworks: true })).rejects.toThrow('Metrics URL resolves to a blocked address'); + }); + + it('rejects a hostname whose AAAA answer is blocked', async () => { + process.env.NODE_ENV = 'production'; + jest.spyOn(dns, 'lookup').mockResolvedValue([ + { address: '8.8.8.8', family: 4 }, + { address: 'fe80::1', family: 6 }, + ] as any); + await expect(assertSafeOutboundUrl('http://metrics.example.com/x', { label: 'metrics URL', allowPrivateNetworks: true })).rejects.toThrow('Metrics URL resolves to a blocked address'); + }); + + describe('guardedLookup', () => { + const run = (lookup: ReturnType, host: string, all: boolean) => + new Promise<{ error: NodeJS.ErrnoException | null; address: unknown; family?: number }>((resolve) => + lookup(host, { all } as any, (error: NodeJS.ErrnoException | null, address: unknown, family?: number) => resolve({ error, address, family })), + ); + + it('returns the validated address in both callback shapes', async () => { + jest.spyOn(dns, 'lookup').mockResolvedValue([{ address: '8.8.8.8', family: 4 }] as any); + const lookup = guardedLookup({ label: 'metrics URL' }); + await expect(run(lookup, 'metrics.example.com', false)).resolves.toEqual({ error: null, address: '8.8.8.8', family: 4 }); + await expect(run(lookup, 'metrics.example.com', true)).resolves.toEqual({ error: null, address: [{ address: '8.8.8.8', family: 4 }], family: undefined }); + }); + + it('fails the connection when any answer is blocked', async () => { + jest.spyOn(dns, 'lookup').mockResolvedValue([ + { address: '8.8.8.8', family: 4 }, + { address: '169.254.169.254', family: 4 }, + ] as any); + const { error } = await run(guardedLookup({ label: 'metrics URL', allowPrivateNetworks: true }), 'metrics.example.com', true); + expect(error).toBeInstanceOf(BlockedAddressError); + expect(error?.message).toBe('Metrics URL resolves to a blocked address'); + }); + + it('applies the private network setting', async () => { + jest.spyOn(dns, 'lookup').mockResolvedValue([{ address: '10.1.2.3', family: 4 }] as any); + expect((await run(guardedLookup({ label: 'metrics URL' }), 'h', false)).error).toBeInstanceOf(BlockedAddressError); + expect((await run(guardedLookup({ label: 'metrics URL', allowPrivateNetworks: true }), 'h', false)).address).toBe('10.1.2.3'); + }); + + it('agrees with the pre-check on loopback outside production', async () => { + jest.spyOn(dns, 'lookup').mockResolvedValue([{ address: '127.0.0.1', family: 4 }] as any); + process.env.NODE_ENV = 'development'; + await expect(assertSafeOutboundUrl('http://localhost:8000/metrics', { label: 'metrics URL' })).resolves.toBeInstanceOf(URL); + expect((await run(guardedLookup({ label: 'metrics URL' }), 'localhost', false)).address).toBe('127.0.0.1'); + process.env.NODE_ENV = 'production'; + expect((await run(guardedLookup({ label: 'metrics URL' }), 'localhost', false)).error).toBeInstanceOf(BlockedAddressError); + }); + + it('still blocks other private answers outside production', async () => { + process.env.NODE_ENV = 'development'; + jest.spyOn(dns, 'lookup').mockResolvedValue([{ address: '10.1.2.3', family: 4 }] as any); + expect((await run(guardedLookup({ label: 'metrics URL' }), 'h', false)).error).toBeInstanceOf(BlockedAddressError); + }); + + it('passes DNS errors through', async () => { + const failure = Object.assign(new Error('getaddrinfo ENOTFOUND h'), { code: 'ENOTFOUND' }); + jest.spyOn(dns, 'lookup').mockRejectedValue(failure); + expect((await run(guardedLookup({ label: 'metrics URL' }), 'h', false)).error).toBe(failure); + }); + }); +}); diff --git a/apps/api/src/common/utils/outbound-url-guard.ts b/apps/api/src/common/utils/outbound-url-guard.ts new file mode 100644 index 000000000..bd8d387b1 --- /dev/null +++ b/apps/api/src/common/utils/outbound-url-guard.ts @@ -0,0 +1,158 @@ +import { BadRequestException, Logger } from '@nestjs/common'; +import { promises as dns } from 'dns'; +import { isIP, type LookupFunction } from 'net'; + +export interface OutboundUrlOptions { + label: string; + allowPrivateNetworks?: boolean; + resolveDns?: 'always' | 'production' | 'never'; +} + +export class BlockedAddressError extends Error {} + +const logger = new Logger('OutboundUrlGuard'); + +const BLOCKED_IP_PATTERNS = [ + /^127\./, + /^10\./, + /^172\.(1[6-9]|2[0-9]|3[01])\./, + /^192\.168\./, + /^100\.(6[4-9]|[7-9][0-9]|1[01][0-9]|12[0-7])\./, + /^169\.254\./, + /^::1?$/, + /^0\./, + /^fe[89ab][0-9a-f]:/i, + /^f[cd][0-9a-f]{2}:/i, +]; + +const ALWAYS_BLOCKED = [/^169\.254\./, /^fe[89ab][0-9a-f]:/i, /^fd00:ec2::254$/, /^100\.100\.100\.200$/, /^192\.0\.0\.192$/]; + +const LOOPBACK = [/^127\./, /^0\./, /^::1?$/]; + +const MAPPED_DOTTED = /^::ffff:(?:0:)?(\d{1,3}(?:\.\d{1,3}){3})$/; +const MAPPED_HEX = /^::ffff:(?:0:)?([0-9a-f]{1,4}):([0-9a-f]{1,4})$/; + +const bareHost = (hostname: string) => hostname.replace(/^\[|\]$/g, '').toLowerCase(); + +function canonicalIp(ip: string): string { + const address = bareHost(ip); + if (isIP(address) !== 6) return address; + try { + return bareHost(new URL(`http://[${address}]`).hostname); + } catch { + return address; + } +} + +function unmapIpv4(address: string): string { + const dotted = MAPPED_DOTTED.exec(address); + if (dotted) return dotted[1]; + const hex = MAPPED_HEX.exec(address); + if (!hex) return address; + const high = parseInt(hex[1], 16); + const low = parseInt(hex[2], 16); + return `${high >> 8}.${high & 255}.${low >> 8}.${low & 255}`; +} + +const capitalize = (text: string) => text.charAt(0).toUpperCase() + text.slice(1); + +const isLocalhostName = (hostname: string) => /(^|\.)localhost\.?$/.test(hostname); + +function blockedAddressMessage(label: string, hostname: string, address: string): string { + logger.warn(`${capitalize(label)} ${hostname} resolves to blocked IP address ${address}`); + return `${capitalize(label)} resolves to a blocked address`; +} + +export function isBlockedIp(ip: string, allowPrivateNetworks = false): boolean { + const address = unmapIpv4(canonicalIp(ip)); + if (ALWAYS_BLOCKED.some((pattern) => pattern.test(address))) return true; + if (allowPrivateNetworks) return false; + return BLOCKED_IP_PATTERNS.some((pattern) => pattern.test(address)); +} + +function isLoopback(ip: string): boolean { + const address = unmapIpv4(canonicalIp(ip)); + return LOOPBACK.some((pattern) => pattern.test(address)); +} + +export async function assertSafeOutboundUrl(rawUrl: string, options: OutboundUrlOptions): Promise { + const { label, allowPrivateNetworks, resolveDns = 'always' } = options; + try { + const parsed = new URL(rawUrl); + + if (!['http:', 'https:'].includes(parsed.protocol)) { + throw new BadRequestException('Only HTTP and HTTPS protocols are allowed'); + } + + if (parsed.username || parsed.password) { + logger.warn(`${capitalize(label)} contains credentials, consider using custom headers instead: ${parsed.hostname}`); + } + + const isProduction = process.env.NODE_ENV === 'production'; + const host = unmapIpv4(bareHost(parsed.hostname)); + const isLocalhost = host === 'localhost' || LOOPBACK.some((pattern) => pattern.test(host)); + + if (isLocalhost && (!isProduction || allowPrivateNetworks)) { + logger.debug(`Allowing localhost ${label} in ${process.env.NODE_ENV || 'development'} mode: ${rawUrl}`); + return parsed; + } + + if (parsed.hostname === 'localhost' || parsed.hostname === '0.0.0.0') { + throw new BadRequestException(`Cannot use localhost or 0.0.0.0 as ${label} in production`); + } + + if (isIP(host) && isBlockedIp(host, allowPrivateNetworks)) { + throw new BadRequestException(`Cannot use private IP addresses as ${label}`); + } + + if (!allowPrivateNetworks && (parsed.hostname.startsWith('127.') || isLocalhostName(host))) { + throw new BadRequestException('Suspicious hostname detected'); + } + + if (!isIP(host) && resolveDns !== 'never' && (resolveDns !== 'production' || isProduction)) { + try { + const answers = await dns.lookup(host, { all: true, verbatim: true }); + for (const { address: addr } of answers) { + if (isBlockedIp(addr, allowPrivateNetworks)) { + throw new BadRequestException(blockedAddressMessage(label, host, addr)); + } + } + } catch (dnsError: any) { + if (dnsError instanceof BadRequestException) { + throw dnsError; + } + logger.warn(`Failed to resolve DNS for ${label}: ${parsed.hostname}`); + throw new BadRequestException(`Failed to resolve ${label} hostname`); + } + } + + return parsed; + } catch (error) { + if (error instanceof BadRequestException) { + throw error; + } + throw new BadRequestException(`Invalid ${label}`); + } +} + +export function guardedLookup(options: Omit): LookupFunction { + const { label, allowPrivateNetworks } = options; + return (hostname, lookupOptions, callback) => { + dns.lookup(hostname, { family: lookupOptions.family, hints: lookupOptions.hints, all: true, verbatim: true }).then( + (answers) => { + const allowLoopback = allowPrivateNetworks || process.env.NODE_ENV !== 'production'; + const blocked = answers.find(({ address }) => isBlockedIp(address, allowPrivateNetworks) && !(allowLoopback && isLoopback(address))); + if (blocked) { + callback(new BlockedAddressError(blockedAddressMessage(label, hostname, blocked.address)), '', 0); + } else if (answers.length === 0) { + callback(Object.assign(new Error(`getaddrinfo ENOTFOUND ${hostname}`), { code: 'ENOTFOUND' }), '', 0); + } else if (lookupOptions.all) { + (callback as unknown as (error: null, addresses: typeof answers) => void)(null, answers); + } else { + callback(null, answers[0].address, answers[0].family); + } + }, + (error: NodeJS.ErrnoException) => callback(error, '', 0), + ); + }; +} diff --git a/apps/api/src/common/utils/pinned-request.ts b/apps/api/src/common/utils/pinned-request.ts new file mode 100644 index 000000000..6740d86d0 --- /dev/null +++ b/apps/api/src/common/utils/pinned-request.ts @@ -0,0 +1,58 @@ +import { request as httpRequest, type IncomingMessage } from 'http'; +import { request as httpsRequest } from 'https'; +import type { LookupFunction } from 'net'; +import { Readable } from 'stream'; + +export interface PinnedRequestInit { + method?: string; + headers: Record; + body?: string; + signal: AbortSignal; + lookup?: LookupFunction; +} + +export class UnexpectedStatusError extends Error {} + +const NULL_BODY_STATUSES = new Set([204, 205, 304]); + +function toResponse(message: IncomingMessage): Response { + const headers = new Headers(); + for (const [name, value] of Object.entries(message.headers)) { + if (value === undefined) continue; + for (const item of Array.isArray(value) ? value : [value]) headers.append(name, item); + } + const status = message.statusCode ?? 502; + if (status < 200 || status > 599) throw new UnexpectedStatusError(`unexpected HTTP status ${status}`); + if (NULL_BODY_STATUSES.has(status)) { + message.resume(); + return new Response(null, { status, headers }); + } + return new Response(Readable.toWeb(message) as ReadableStream, { status, headers }); +} + +export const pinnedRequest = (url: string, init: PinnedRequestInit): Promise => + new Promise((resolve, reject) => { + const target = new URL(url); + const send = target.protocol === 'https:' ? httpsRequest : httpRequest; + const headers = init.body === undefined ? init.headers : { 'Content-Length': String(Buffer.byteLength(init.body)), ...init.headers }; + const req = send( + target, + { + method: init.method ?? 'GET', + agent: false, + headers, + signal: init.signal, + ...(init.lookup ? { lookup: init.lookup } : {}), + }, + (message) => { + try { + resolve(toResponse(message)); + } catch (error) { + message.destroy(); + reject(error); + } + }, + ); + req.on('error', reject); + req.end(init.body); + }); diff --git a/apps/api/src/external-metrics/__tests__/otel-metrics-ingest.controller.spec.ts b/apps/api/src/external-metrics/__tests__/otel-metrics-ingest.controller.spec.ts index 4e1c8ec05..ca2183aba 100644 --- a/apps/api/src/external-metrics/__tests__/otel-metrics-ingest.controller.spec.ts +++ b/apps/api/src/external-metrics/__tests__/otel-metrics-ingest.controller.spec.ts @@ -19,6 +19,7 @@ const clean: IngestResult = { unmapped_metric: 0, invalid_value: 0, cardinality_limit: 0, + unknown_engine: 0, }, }; diff --git a/apps/api/src/external-metrics/__tests__/otel-metrics-ingest.service.spec.ts b/apps/api/src/external-metrics/__tests__/otel-metrics-ingest.service.spec.ts index 1590099b4..2e4b68b56 100644 --- a/apps/api/src/external-metrics/__tests__/otel-metrics-ingest.service.spec.ts +++ b/apps/api/src/external-metrics/__tests__/otel-metrics-ingest.service.spec.ts @@ -120,6 +120,7 @@ describe('OtelMetricsIngestService', () => { unmapped_metric: 1, invalid_value: 1, cardinality_limit: 0, + unknown_engine: 0, }); expect(toPartialSuccess(result)).toEqual({ rejectedDataPoints: 8, @@ -420,4 +421,78 @@ describe('OtelMetricsIngestService discovery', () => { expect(discovered.list(NOW_MS)).toEqual([]); expect(result.dropped.unknown_instance).toBe(1); }); + + describe('KV cache routing', () => { + const lmcacheSum = (name: string): OtlpMetric => ({ + name, + sum: { dataPoints: [{ attributes: [], asInt: '5' }, { attributes: [], asInt: '6' }], aggregationTemporality: 2 }, + }); + + function buildWithSink(sinkResult = { accepted: 2, dropped: { unmapped_metric: 3 } as Record }) { + const registry = { findByHostPort: jest.fn() } as unknown as ConnectionRegistry; + const store = new ExternalMetricsStore(); + const applySpy = jest.spyOn(store, 'apply'); + const sink = { ingest: jest.fn().mockReturnValue(sinkResult) }; + const service = new OtelMetricsIngestService(registry, store, undefined, undefined, sink); + return { service, registry, applySpy, sink }; + } + + it('routes a resource with the engine attribute to the sink', () => { + const { service, registry, applySpy, sink } = buildWithSink(); + const metrics = [lmcacheSum('lmcache:num_hit_tokens_total')]; + const result = service.ingest(resource([str('betterdb.lmcache.engine', 'eng-1')], metrics), NOW_MS); + expect(sink.ingest).toHaveBeenCalledWith({ 'betterdb.lmcache.engine': 'eng-1' }, metrics, NOW_MS); + expect(registry.findByHostPort).not.toHaveBeenCalled(); + expect(applySpy).not.toHaveBeenCalled(); + expect(result.accepted).toBe(2); + expect(result.dropped.unmapped_metric).toBe(3); + }); + + it('routes a resource whose metrics are all lmcache without the attribute', () => { + const { service, sink } = buildWithSink(); + service.ingest(resource([], [lmcacheSum('lmcache:num_hit_tokens_total')]), NOW_MS); + expect(sink.ingest).toHaveBeenCalledTimes(1); + }); + + it('drops an lmcache resource as unknown_engine when no sink is present', () => { + const { service } = build(); + const result = service.ingest( + resource([str('betterdb.lmcache.engine', 'eng-1')], [lmcacheSum('lmcache:num_hit_tokens_total')]), + NOW_MS, + ); + expect(result.dropped.unknown_engine).toBe(2); + expect(result.accepted).toBe(0); + }); + + it('keeps Valkey metrics of a tagged resource on the Valkey path and routes the rest to the sink', () => { + const { service, registry, sink } = buildWithSink({ accepted: 2, dropped: {} }); + (registry.findByHostPort as jest.Mock).mockReturnValue({ id: 'ext', connectionType: 'external' }); + const kv = lmcacheSum('lmcache:num_hit_tokens_total'); + const result = service.ingest( + resource([str('betterdb.lmcache.engine', 'eng-1'), ...identity], [gauge('redis.memory.used', 1024), kv]), + NOW_MS, + ); + expect(sink.ingest).toHaveBeenCalledWith(expect.objectContaining({ 'betterdb.lmcache.engine': 'eng-1' }), [kv], NOW_MS); + expect(registry.findByHostPort).toHaveBeenCalledWith('cache.internal', 6379); + expect(result.accepted).toBe(3); + }); + + it('does not call the sink for a tagged resource that only carries Valkey metrics', () => { + const { service, registry, sink } = buildWithSink(); + (registry.findByHostPort as jest.Mock).mockReturnValue({ id: 'ext', connectionType: 'external' }); + const result = service.ingest( + resource([str('betterdb.lmcache.engine', 'eng-1'), ...identity], [gauge('valkey.memory.used', 1024)]), + NOW_MS, + ); + expect(sink.ingest).not.toHaveBeenCalled(); + expect(result.accepted).toBe(1); + }); + + it('keeps the Valkey path when a sink is present', () => { + const { service, sink } = buildWithSink(); + const result = service.ingest(resource([], [gauge('redis.memory.used', 1)]), NOW_MS); + expect(sink.ingest).not.toHaveBeenCalled(); + expect(result.dropped.unidentified).toBe(1); + }); + }); }); diff --git a/apps/api/src/external-metrics/kv-cache-otlp-sink.ts b/apps/api/src/external-metrics/kv-cache-otlp-sink.ts new file mode 100644 index 000000000..54d2a3e3b --- /dev/null +++ b/apps/api/src/external-metrics/kv-cache-otlp-sink.ts @@ -0,0 +1,18 @@ +import type { DropReason, OtlpMetric } from './otlp-metrics-types'; + +export const KV_CACHE_OTLP_SINK = Symbol('KV_CACHE_OTLP_SINK'); +export const LMCACHE_ENGINE_ATTRIBUTE = 'betterdb.lmcache.engine'; + +export interface KvCacheOtlpSinkResult { + accepted: number; + dropped: Partial>; +} + +export interface KvCacheOtlpSink { + ingest(resourceAttributes: Record, metrics: OtlpMetric[], nowMs: number): KvCacheOtlpSinkResult; +} + +export function isKvCacheResource(resourceAttributes: Record, metrics: OtlpMetric[]): boolean { + if (resourceAttributes[LMCACHE_ENGINE_ATTRIBUTE] !== undefined) return true; + return metrics.length > 0 && metrics.every((metric) => (metric.name ?? '').startsWith('lmcache:')); +} diff --git a/apps/api/src/external-metrics/otel-metrics-ingest.service.ts b/apps/api/src/external-metrics/otel-metrics-ingest.service.ts index d4f93ea70..2a47ab391 100644 --- a/apps/api/src/external-metrics/otel-metrics-ingest.service.ts +++ b/apps/api/src/external-metrics/otel-metrics-ingest.service.ts @@ -1,13 +1,17 @@ -import { Injectable, Logger, Optional } from '@nestjs/common'; +import { Inject, Injectable, Logger, Optional } from '@nestjs/common'; import { ConnectionRegistry } from '../connections/connection-registry.service'; import { PrometheusService } from '../prometheus/prometheus.service'; import { DiscoveredInstancesStore } from './discovered-instances.store'; import { ExternalMetricsStore } from './external-metrics-store'; +import { isKvCacheResource, KV_CACHE_OTLP_SINK, KvCacheOtlpSink, LMCACHE_ENGINE_ATTRIBUTE } from './kv-cache-otlp-sink'; import { attrsToRecord, + isDeltaTemporality, mapDataPoint, + metricPointCount, metricVocabulary, nanosToMs, + NO_RECORDED_VALUE, pointValue, resolveInstanceKeys, explicitInstanceKey, @@ -35,21 +39,14 @@ function emptyDropCounts(): Record { return Object.fromEntries(DROP_REASONS.map((reason) => [reason, 0])) as Record; } -const NO_RECORDED_VALUE = 1; -function isDelta(temporality: number | string | undefined): boolean { - return temporality === 1 || temporality === 'AGGREGATION_TEMPORALITY_DELTA'; -} +const isValkeyMetric = (metric: OtlpMetric) => metricVocabulary(metric.name ?? '') !== null; -function metricPointCount(metric: OtlpMetric): number { - const points = - metric.gauge?.dataPoints ?? - metric.sum?.dataPoints ?? - metric.histogram?.dataPoints ?? - metric.exponentialHistogram?.dataPoints ?? - metric.summary?.dataPoints ?? - []; - return points.length; +function keepValkeyMetrics(resourceMetrics: OtlpResourceMetrics): OtlpResourceMetrics { + return { + ...resourceMetrics, + scopeMetrics: (resourceMetrics.scopeMetrics ?? []).map((scope) => ({ ...scope, metrics: (scope.metrics ?? []).filter(isValkeyMetric) })), + }; } function resourcePointCount(resourceMetrics: OtlpResourceMetrics): number { @@ -79,6 +76,7 @@ export class OtelMetricsIngestService { private readonly store: ExternalMetricsStore, @Optional() private readonly prometheus?: PrometheusService, @Optional() private readonly discovered?: DiscoveredInstancesStore, + @Optional() @Inject(KV_CACHE_OTLP_SINK) private readonly kvCacheSink?: KvCacheOtlpSink, ) {} ingest(request: OtlpMetricsRequest, nowMs: number = Date.now()): IngestResult { @@ -90,8 +88,17 @@ export class OtelMetricsIngestService { return result; } - private ingestResource(resourceMetrics: OtlpResourceMetrics, nowMs: number, result: IngestResult): void { + private ingestResource(input: OtlpResourceMetrics, nowMs: number, result: IngestResult): void { + let resourceMetrics = input; const attrs = attrsToRecord(resourceMetrics.resource?.attributes); + const metrics = (resourceMetrics.scopeMetrics ?? []).flatMap((scope) => scope.metrics ?? []); + if (isKvCacheResource(attrs, metrics)) { + const tagged = attrs[LMCACHE_ENGINE_ATTRIBUTE] !== undefined; + const kvMetrics = tagged ? metrics.filter((metric) => !isValkeyMetric(metric)) : metrics; + if (kvMetrics.length > 0) this.routeToKvCache(attrs, kvMetrics, nowMs, result); + if (kvMetrics.length === metrics.length) return; + resourceMetrics = keepValkeyMetrics(resourceMetrics); + } const keys = resolveInstanceKeys(attrs); if (keys.length === 0) { this.drop(result, 'unidentified', resourcePointCount(resourceMetrics), 'unidentified resource', nowMs); @@ -133,6 +140,24 @@ export class OtelMetricsIngestService { this.drop(result, 'cardinality_limit', applied.rejected, instance, nowMs); } + private routeToKvCache( + attrs: Record, + metrics: OtlpMetric[], + nowMs: number, + result: IngestResult, + ): void { + if (!this.kvCacheSink) { + const points = metrics.reduce((sum, metric) => sum + metricPointCount(metric), 0); + this.drop(result, 'unknown_engine', points, 'kv cache engine', nowMs); + return; + } + const routed = this.kvCacheSink.ingest(attrs, metrics, nowMs); + result.accepted += routed.accepted; + for (const [reason, count] of Object.entries(routed.dropped)) { + this.drop(result, reason as DropReason, count ?? 0, 'kv cache engine', nowMs); + } + } + private collectMetric( metric: OtlpMetric, instance: string, @@ -144,7 +169,7 @@ export class OtelMetricsIngestService { this.drop(result, 'unsupported_type', metricPointCount(metric), instance, nowMs); return false; } - if (metric.sum && isDelta(metric.sum.aggregationTemporality)) { + if (metric.sum && isDeltaTemporality(metric.sum.aggregationTemporality)) { this.drop(result, 'unsupported_temporality', metricPointCount(metric), instance, nowMs); return false; } diff --git a/apps/api/src/external-metrics/otlp-metric-map.ts b/apps/api/src/external-metrics/otlp-metric-map.ts index fd5babaab..e22045fe4 100644 --- a/apps/api/src/external-metrics/otlp-metric-map.ts +++ b/apps/api/src/external-metrics/otlp-metric-map.ts @@ -1,4 +1,6 @@ -import type { InfoTarget, InstanceKey, OtlpKeyValue, OtlpNumberDataPoint } from './otlp-metrics-types'; +import type { InfoTarget, InstanceKey, OtlpKeyValue, OtlpMetric, OtlpNumberDataPoint } from './otlp-metrics-types'; + +export const NO_RECORDED_VALUE = 1; export interface MappedPoint { target: InfoTarget; @@ -168,6 +170,10 @@ export function resolveInstanceKeys(attrs: Record): InstanceKey[ return keys; } +export function isDeltaTemporality(temporality: number | string | undefined): boolean { + return temporality === 1 || temporality === 'AGGREGATION_TEMPORALITY_DELTA'; +} + export function pointValue(dp: OtlpNumberDataPoint): string | null { if (dp.asInt !== undefined && dp.asInt !== null) { const s = String(dp.asInt); @@ -187,3 +193,14 @@ export function nanosToMs(v: string | number | undefined): number | null { const ms = Number(BigInt(v) / 1_000_000n); return ms > 0 ? ms : null; } + +export function metricPointCount(metric: OtlpMetric): number { + const points = + metric.gauge?.dataPoints ?? + metric.sum?.dataPoints ?? + metric.histogram?.dataPoints ?? + metric.exponentialHistogram?.dataPoints ?? + metric.summary?.dataPoints ?? + []; + return points.length; +} diff --git a/apps/api/src/external-metrics/otlp-metrics-types.ts b/apps/api/src/external-metrics/otlp-metrics-types.ts index d431b02a4..f92978ccb 100644 --- a/apps/api/src/external-metrics/otlp-metrics-types.ts +++ b/apps/api/src/external-metrics/otlp-metrics-types.ts @@ -65,6 +65,7 @@ export const DROP_REASONS = [ 'unmapped_metric', 'invalid_value', 'cardinality_limit', + 'unknown_engine', ] as const; export type DropReason = (typeof DROP_REASONS)[number]; diff --git a/apps/api/src/otel-telemetry/__tests__/semconv-map.spec.ts b/apps/api/src/otel-telemetry/__tests__/semconv-map.spec.ts index c4ac7a2e8..c48ebbafe 100644 --- a/apps/api/src/otel-telemetry/__tests__/semconv-map.spec.ts +++ b/apps/api/src/otel-telemetry/__tests__/semconv-map.spec.ts @@ -120,6 +120,9 @@ describe('planSemconvInstruments', () => { ['betterdb_cve_findings', 'betterdb.cve.findings', 'gauge', '{finding}'], ['betterdb_cve_kev', 'betterdb.cve.kev', 'gauge', '{finding}'], ['betterdb_cve_dataset_stale', 'betterdb.cve.dataset_stale', 'gauge', '1'], + ['betterdb_kv_cache_hit_rate', 'betterdb.kv_cache.hit_rate', 'gauge', '1'], + ['betterdb_kv_cache_chunks', 'betterdb.kv_cache.chunks', 'gauge', '{chunk}'], + ['betterdb_kv_cache_bytes', 'betterdb.kv_cache.bytes', 'gauge', 'By'], ])('%s → %s (%s, %s)', (prom, name, kind, unit) => { expect(planSemconvInstruments([family(prom, [])])).toEqual([ { name, kind, unit, description: `${prom} help` }, diff --git a/apps/api/src/otel-telemetry/semconv-map.ts b/apps/api/src/otel-telemetry/semconv-map.ts index 42be0625d..fdc7760d4 100644 --- a/apps/api/src/otel-telemetry/semconv-map.ts +++ b/apps/api/src/otel-telemetry/semconv-map.ts @@ -279,6 +279,9 @@ export const SEMCONV_RULES: Readonly> = { betterdb_cve_findings: gauge('betterdb.cve.findings', '{finding}'), betterdb_cve_kev: gauge('betterdb.cve.kev', '{finding}'), betterdb_cve_dataset_stale: gauge('betterdb.cve.dataset_stale', '1'), + betterdb_kv_cache_hit_rate: gauge('betterdb.kv_cache.hit_rate', '1'), + betterdb_kv_cache_chunks: gauge('betterdb.kv_cache.chunks', '{chunk}'), + betterdb_kv_cache_bytes: gauge('betterdb.kv_cache.bytes', 'By'), betterdb_otlp_metric_points_accepted_total: counter('betterdb.otlp.metric_points.accepted', '{point}'), betterdb_otlp_metric_points_dropped_total: counter('betterdb.otlp.metric_points.dropped', '{point}'), betterdb_process_cpu_user_seconds_total: described( diff --git a/apps/api/src/prometheus/__tests__/prometheus.kv-cache.spec.ts b/apps/api/src/prometheus/__tests__/prometheus.kv-cache.spec.ts new file mode 100644 index 000000000..de6f208d3 --- /dev/null +++ b/apps/api/src/prometheus/__tests__/prometheus.kv-cache.spec.ts @@ -0,0 +1,148 @@ +import { ConfigService } from '@nestjs/config'; +import { PrometheusService } from '../prometheus.service'; +import { ConnectionRegistry } from '../../connections/connection-registry.service'; +import { RuntimeCapabilityTracker } from '../../connections/runtime-capability-tracker.service'; +import { SlowLogAnalyticsService } from '../../slowlog-analytics/slowlog-analytics.service'; +import { CommandLogAnalyticsService } from '../../commandlog-analytics/commandlog-analytics.service'; +import { HealthService } from '../../health/health.service'; +import type { StoragePort } from '../../common/interfaces/storage-port.interface'; + +const INSTANCES: Record = { + 'conn-1': { host: '10.0.0.1', port: 6379 }, + 'conn-2': { host: '10.0.0.2', port: 6379 }, + 'agent-1': { host: 'agent', port: 0 }, + 'agent-2': { host: 'agent', port: 0 }, +}; + +const footprint = (model: string, chunksEst: number, bytesEst: number | null) => ({ model, dtype: 'bfloat16', chunksEst, bytesEst }); + +describe('PrometheusService KV cache metrics', () => { + let service: PrometheusService; + + beforeEach(() => { + const registry = { + getConfig: jest.fn().mockImplementation((id: string) => INSTANCES[id]), + list: jest.fn().mockReturnValue([]), + getName: jest.fn().mockReturnValue('primary'), + } as unknown as ConnectionRegistry; + service = new PrometheusService( + {} as StoragePort, + registry, + { get: jest.fn().mockReturnValue(5000) } as unknown as ConfigService, + {} as RuntimeCapabilityTracker, + {} as SlowLogAnalyticsService, + {} as CommandLogAnalyticsService, + {} as HealthService, + undefined, + undefined, + undefined, + undefined, + undefined, + undefined, + ); + }); + + async function exposition(): Promise { + return (service as unknown as { registry: { metrics: () => Promise } }).registry.metrics(); + } + + it('exports chunks and bytes per model', async () => { + service.setKvCacheFootprint('conn-1', [footprint('llama', 10, 1000), footprint('qwen', 4, 400)]); + const text = await exposition(); + expect(text).toContain('betterdb_kv_cache_chunks{connection="10.0.0.1:6379",connection_id="conn-1",model="llama"} 10'); + expect(text).toContain('betterdb_kv_cache_bytes{connection="10.0.0.1:6379",connection_id="conn-1",model="qwen"} 400'); + }); + + it('sums a model served with more than one dtype', async () => { + service.setKvCacheFootprint('conn-1', [footprint('llama', 10, 1000), { ...footprint('llama', 3, 600), dtype: 'float16' }]); + const text = await exposition(); + expect(text).toContain('betterdb_kv_cache_chunks{connection="10.0.0.1:6379",connection_id="conn-1",model="llama"} 13'); + expect(text).toContain('betterdb_kv_cache_bytes{connection="10.0.0.1:6379",connection_id="conn-1",model="llama"} 1600'); + }); + + it('exports chunks but no byte series when the byte estimate is unknown', async () => { + service.setKvCacheFootprint('conn-1', [footprint('llama', 10, 1000)]); + service.setKvCacheFootprint('conn-1', [footprint('llama', 12, null)]); + const text = await exposition(); + expect(text).toContain('betterdb_kv_cache_chunks{connection="10.0.0.1:6379",connection_id="conn-1",model="llama"} 12'); + expect(text).not.toContain('betterdb_kv_cache_bytes{connection="10.0.0.1:6379",connection_id="conn-1",model="llama"}'); + }); + + it('drops the series of a model that is no longer in the footprint', async () => { + service.setKvCacheFootprint('conn-1', [footprint('llama', 10, 1000), footprint('qwen', 4, 400)]); + service.setKvCacheFootprint('conn-1', [footprint('llama', 12, 1200)]); + const text = await exposition(); + expect(text).toContain('betterdb_kv_cache_chunks{connection="10.0.0.1:6379",connection_id="conn-1",model="llama"} 12'); + expect(text).not.toContain('model="qwen"'); + }); + + it('replaces hit rate series for a connection', async () => { + service.setKvCacheHitRates('conn-1', [ + { engine: 'vllm-a', model: 'llama', hitRate: 0.5 }, + { engine: 'vllm-b', model: 'llama', hitRate: 0.1 }, + ]); + service.setKvCacheHitRates('conn-1', [{ engine: 'vllm-a', model: 'llama', hitRate: 0.75 }]); + const text = await exposition(); + expect(text).toContain('betterdb_kv_cache_hit_rate{connection="10.0.0.1:6379",connection_id="conn-1",engine="vllm-a",model="llama"} 0.75'); + expect(text).not.toContain('engine="vllm-b"'); + }); + + it('removes every hit rate series when the list is empty', async () => { + service.setKvCacheHitRates('conn-1', [{ engine: 'vllm-a', model: 'llama', hitRate: 0.5 }]); + service.setKvCacheHitRates('conn-1', []); + expect(await exposition()).not.toContain('betterdb_kv_cache_hit_rate{'); + }); + + it('clears all three families for one connection only', async () => { + service.setKvCacheFootprint('conn-1', [footprint('llama', 10, 1000)]); + service.setKvCacheHitRates('conn-1', [{ engine: 'vllm-a', model: 'llama', hitRate: 0.5 }]); + service.setKvCacheFootprint('conn-2', [footprint('qwen', 4, 400)]); + service.setKvCacheHitRates('conn-2', [{ engine: 'vllm-b', model: 'qwen', hitRate: 0.3 }]); + service.clearKvCache('conn-1'); + const text = await exposition(); + expect(text).not.toContain('connection="10.0.0.1:6379",connection_id="conn-1",model="llama"'); + expect(text).not.toContain('engine="vllm-a"'); + expect(text).toContain('betterdb_kv_cache_chunks{connection="10.0.0.2:6379",connection_id="conn-2",model="qwen"} 4'); + expect(text).toContain('betterdb_kv_cache_bytes{connection="10.0.0.2:6379",connection_id="conn-2",model="qwen"} 400'); + expect(text).toContain('betterdb_kv_cache_hit_rate{connection="10.0.0.2:6379",connection_id="conn-2",engine="vllm-b",model="qwen"} 0.3'); + }); + + it('clears KV cache series for every connection', async () => { + service.setKvCacheFootprint('conn-1', [footprint('llama', 10, 1000)]); + service.setKvCacheHitRates('conn-1', [{ engine: 'vllm-a', model: 'llama', hitRate: 0.5 }]); + service.setKvCacheFootprint('conn-2', [footprint('qwen', 4, 400)]); + service.clearAllKvCache(); + const text = await exposition(); + expect(text).not.toContain('betterdb_kv_cache_chunks{'); + expect(text).not.toContain('betterdb_kv_cache_bytes{'); + expect(text).not.toContain('betterdb_kv_cache_hit_rate{'); + }); + + it('clears KV cache series when the connection is cleaned up', async () => { + service.setKvCacheFootprint('conn-1', [footprint('llama', 10, 1000)]); + service.setKvCacheHitRates('conn-1', [{ engine: 'vllm-a', model: 'llama', hitRate: 0.5 }]); + service.cleanupConnectionMetrics('conn-1'); + const text = await exposition(); + expect(text).not.toContain('betterdb_kv_cache_chunks{'); + expect(text).not.toContain('betterdb_kv_cache_hit_rate{'); + }); + + it('keeps the series of connections that share a host:port label apart', async () => { + service.setKvCacheFootprint('agent-1', [footprint('llama', 10, 1000)]); + service.setKvCacheHitRates('agent-1', [{ engine: 'vllm', model: 'llama', hitRate: 0.5 }]); + service.setKvCacheFootprint('agent-2', [footprint('llama', 4, 400)]); + service.setKvCacheHitRates('agent-2', [{ engine: 'vllm', model: 'llama', hitRate: 0.2 }]); + service.clearKvCache('agent-2'); + const text = await exposition(); + expect(text).toContain('betterdb_kv_cache_chunks{connection="agent:0",connection_id="agent-1",model="llama"} 10'); + expect(text).toContain('betterdb_kv_cache_hit_rate{connection="agent:0",connection_id="agent-1",engine="vllm",model="llama"} 0.5'); + expect(text).not.toContain('connection_id="agent-2"'); + }); + + it('ignores unknown connections', async () => { + service.setKvCacheFootprint('missing', [footprint('llama', 10, 1000)]); + service.setKvCacheHitRates('missing', [{ engine: 'a', model: 'm', hitRate: 0.1 }]); + service.clearKvCache('missing'); + expect(await exposition()).not.toContain('betterdb_kv_cache_chunks{'); + }); +}); diff --git a/apps/api/src/prometheus/prometheus.service.ts b/apps/api/src/prometheus/prometheus.service.ts index 84d53e936..412634340 100644 --- a/apps/api/src/prometheus/prometheus.service.ts +++ b/apps/api/src/prometheus/prometheus.service.ts @@ -34,7 +34,7 @@ import { ConnectionContext, } from '../common/services/multi-connection-poller'; import { MetricForecastingService } from '../metric-forecasting/metric-forecasting.service'; -import { ALL_METRIC_KINDS } from '@betterdb/shared'; +import { ALL_METRIC_KINDS, type KvCacheModelFootprint } from '@betterdb/shared'; import { isCveEnabled } from '../cve/cve.constants'; import { OtelEventDispatcherService } from '../otel-telemetry/otel-event-dispatcher.service'; import { @@ -377,6 +377,11 @@ export class PrometheusService extends MultiConnectionPoller implements OnModule private cveKev: Gauge; private cveDatasetStale: Gauge; + private kvCacheHitRate: Gauge; + private kvCacheChunks: Gauge; + private kvCacheBytes: Gauge; + private readonly kvCacheSeries = new Map; rates: Set }>(); + // OTLP Ingest Metrics private otlpPointsAccepted: Counter; private otlpPointsDropped: Counter; @@ -1134,6 +1139,20 @@ export class PrometheusService extends MultiConnectionPoller implements OnModule 'Whether the CVE scan is partial or sources are missing: 1 stale, 0 ok', ); + this.kvCacheHitRate = this.createGauge( + 'kv_cache_hit_rate', + 'LMCache token hit rate over the last 15 minutes', + ['connection_id', 'engine', 'model'], + ); + this.kvCacheChunks = this.createGauge('kv_cache_chunks', 'Estimated LMCache chunks stored', [ + 'connection_id', + 'model', + ]); + this.kvCacheBytes = this.createGauge('kv_cache_bytes', 'Estimated bytes held by LMCache keys', [ + 'connection_id', + 'model', + ]); + // OTLP Ingest Metrics this.otlpPointsAccepted = new Counter({ name: 'betterdb_otlp_metric_points_accepted_total', @@ -1371,6 +1390,84 @@ export class PrometheusService extends MultiConnectionPoller implements OnModule } } + setKvCacheFootprint(connectionId: string, perModel: KvCacheModelFootprint[]): void { + const series = this.kvCacheSeriesFor(connectionId); + if (!series) return; + const totals = new Map(); + for (const entry of perModel) { + const total = totals.get(entry.model) ?? { chunks: 0, bytes: 0 }; + total.chunks += entry.chunksEst; + total.bytes = total.bytes === null || entry.bytesEst === null ? null : total.bytes + entry.bytesEst; + totals.set(entry.model, total); + } + for (const model of series.models) { + if (!totals.has(model)) { + this.kvCacheChunks.remove(series.label, connectionId, model); + this.kvCacheBytes.remove(series.label, connectionId, model); + } + } + for (const [model, total] of totals) { + this.kvCacheChunks.labels(series.label, connectionId, model).set(total.chunks); + if (total.bytes === null) this.kvCacheBytes.remove(series.label, connectionId, model); + else this.kvCacheBytes.labels(series.label, connectionId, model).set(total.bytes); + } + series.models = new Set(totals.keys()); + } + + setKvCacheHitRates( + connectionId: string, + rates: Array<{ engine: string; model: string; hitRate: number }>, + ): void { + const series = this.kvCacheSeriesFor(connectionId); + if (!series) return; + const next = new Set(rates.map((rate) => `${rate.engine}\u0000${rate.model}`)); + for (const key of series.rates) { + if (!next.has(key)) { + const [engine, model] = key.split('\u0000'); + this.kvCacheHitRate.remove(series.label, connectionId, engine, model); + } + } + for (const rate of rates) { + this.kvCacheHitRate.labels(series.label, connectionId, rate.engine, rate.model).set(rate.hitRate); + } + series.rates = next; + } + + clearAllKvCache(): void { + for (const connectionId of [...this.kvCacheSeries.keys()]) this.clearKvCache(connectionId); + } + + clearKvCache(connectionId: string): void { + const series = this.kvCacheSeries.get(connectionId); + if (!series) return; + this.kvCacheSeries.delete(connectionId); + for (const model of series.models) { + this.kvCacheChunks.remove(series.label, connectionId, model); + this.kvCacheBytes.remove(series.label, connectionId, model); + } + for (const key of series.rates) { + const [engine, model] = key.split('\u0000'); + this.kvCacheHitRate.remove(series.label, connectionId, engine, model); + } + } + + private kvCacheSeriesFor( + connectionId: string, + ): { label: string; models: Set; rates: Set } | null { + const label = this.getConnectionLabel(connectionId); + if (label === null) return null; + const existing = this.kvCacheSeries.get(connectionId); + if (existing && existing.label !== label) { + this.clearKvCache(connectionId); + } + let series = this.kvCacheSeries.get(connectionId); + if (!series) { + series = { label, models: new Set(), rates: new Set() }; + this.kvCacheSeries.set(connectionId, series); + } + return series; + } + private async updateMetricForecastMetrics( connectionId: string, connLabel: string, @@ -3708,6 +3805,7 @@ export class PrometheusService extends MultiConnectionPoller implements OnModule if (cveLabel !== null) { this.removeCveSeries(cveLabel); } + this.clearKvCache(connectionId); this.perConnectionState.delete(connectionId); // Drop any in-flight metric update so a reused connection ID cannot join // stale work. The promise itself still settles; its compare-and-delete diff --git a/apps/api/src/storage/adapters/__tests__/kv-cache-storage.spec.ts b/apps/api/src/storage/adapters/__tests__/kv-cache-storage.spec.ts index 62cfd20cb..227c4a559 100644 --- a/apps/api/src/storage/adapters/__tests__/kv-cache-storage.spec.ts +++ b/apps/api/src/storage/adapters/__tests__/kv-cache-storage.spec.ts @@ -225,11 +225,68 @@ describe.each<[string, AdapterFactory]>([ expect(rows).toEqual([sample()]); }); + it('sums every counter per engine, model and epoch-aligned bucket when bucketing', async () => { + await storage.saveKvCacheEngineSamples([ + sample({ timestamp: 60_000 }), + sample({ timestamp: 120_000 }), + sample({ timestamp: 180_000 }), + sample({ timestamp: 120_000, engineId: 'eng-2' }), + sample({ timestamp: 240_000, modelName: 'other' }), + sample({ timestamp: 999_999 }), + ]); + const doubled = { + requestedTokens: 2000, + hitTokens: 1400, + lookupTokens: 2000, + lookupHits: 6, + remoteReadBytes: 20, + remoteWriteBytes: 40, + remoteReadRequests: 2, + remoteWriteRequests: 4, + remotePingErrors: 0, + }; + const query = { connectionId: 'conn-a', from: 0, to: 300_000, bucketMs: 120_000 }; + expect(await storage.getKvCacheEngineSamples(query)).toEqual([ + sample({ timestamp: 0 }), + sample({ timestamp: 120_000, ...doubled }), + sample({ timestamp: 120_000, engineId: 'eng-2' }), + sample({ timestamp: 240_000, modelName: 'other' }), + ]); + expect(await storage.getKvCacheEngineSamples({ ...query, limit: 1 })).toEqual([sample({ timestamp: 240_000, modelName: 'other' })]); + expect(await storage.getKvCacheEngineSamples({ ...query, engineId: 'eng-2' })).toEqual([sample({ timestamp: 120_000, engineId: 'eng-2' })]); + }); + + it('sums requested and hit tokens over the whole filtered range', async () => { + await storage.saveKvCacheEngineSamples([ + sample({ timestamp: 1, requestedTokens: 100, hitTokens: 10 }), + sample({ timestamp: 2, requestedTokens: 200, hitTokens: 20 }), + sample({ timestamp: 2, engineId: 'eng-2', requestedTokens: 1000, hitTokens: 1000 }), + sample({ timestamp: 999_999, requestedTokens: 5000, hitTokens: 0 }), + ]); + expect(await storage.getKvCacheSampleTotals({ connectionId: 'conn-a', engineId: 'eng-1', from: 0, to: 100_000 })).toEqual({ + requestedTokens: 300, + hitTokens: 30, + }); + expect(await storage.getKvCacheSampleTotals({ connectionId: 'conn-a', from: 0, to: 100_000 })).toEqual({ requestedTokens: 1300, hitTokens: 1030 }); + expect(await storage.getKvCacheSampleTotals({ connectionId: 'conn-b' })).toEqual({ requestedTokens: 0, hitTokens: 0 }); + }); + it('prunes samples', async () => { await storage.saveKvCacheEngineSamples([sample({ timestamp: 1 }), sample({ timestamp: 10 })]); expect(await storage.pruneOldKvCacheEngineSamples(5)).toBe(1); }); + it('deletes the samples of one engine only', async () => { + await storage.saveKvCacheEngineSamples([ + sample({ timestamp: 60_000 }), + sample({ timestamp: 120_000 }), + sample({ engineId: 'eng-2', timestamp: 60_000 }), + ]); + await storage.deleteKvCacheEngineSamples('eng-1'); + const rows = await storage.getKvCacheEngineSamples({ connectionId: 'conn-a' }); + expect(rows.map((r) => r.engineId)).toEqual(['eng-2']); + }); + it('upserts, lists, reads and deletes engines', async () => { await storage.saveKvCacheEngine(engine()); await storage.saveKvCacheEngine(engine({ id: 'eng-2', source: 'otlp', scrapeUrl: null, otlpEngineId: 'lmc-1', createdAt: 2 })); diff --git a/apps/api/src/storage/adapters/kv-cache-sample-columns.ts b/apps/api/src/storage/adapters/kv-cache-sample-columns.ts new file mode 100644 index 000000000..f28dd5b93 --- /dev/null +++ b/apps/api/src/storage/adapters/kv-cache-sample-columns.ts @@ -0,0 +1,25 @@ +import type { KvCacheSampleCounters } from '@betterdb/shared'; + +export const KV_CACHE_SAMPLE_COUNTER_COLUMNS = [ + 'requested_tokens', + 'hit_tokens', + 'lookup_tokens', + 'lookup_hits', + 'remote_read_bytes', + 'remote_write_bytes', + 'remote_read_requests', + 'remote_write_requests', + 'remote_ping_errors', +] as const; + +export const KV_CACHE_SAMPLE_COUNTER_FIELDS: ReadonlyArray = [ + 'requestedTokens', + 'hitTokens', + 'lookupTokens', + 'lookupHits', + 'remoteReadBytes', + 'remoteWriteBytes', + 'remoteReadRequests', + 'remoteWriteRequests', + 'remotePingErrors', +]; diff --git a/apps/api/src/storage/adapters/memory.adapter.ts b/apps/api/src/storage/adapters/memory.adapter.ts index ad9a52529..5c1ac841d 100644 --- a/apps/api/src/storage/adapters/memory.adapter.ts +++ b/apps/api/src/storage/adapters/memory.adapter.ts @@ -1,4 +1,5 @@ import { randomUUID } from 'crypto'; +import { KV_CACHE_SAMPLE_COUNTER_FIELDS } from './kv-cache-sample-columns'; import type { RawDatabaseHandle, RawDatabaseHandleProvider } from '../raw-database-handle'; import { approxOtelSpanBytes, @@ -106,6 +107,7 @@ import type { KvCacheRangeQuery, KvCacheFootprintQuery, KvCacheSampleQuery, + KvCacheSampleTotals, } from '../../common/interfaces/storage-port.interface'; import type { ReadinessBand, @@ -1278,16 +1280,44 @@ export class MemoryAdapter implements StoragePort, RawDatabaseHandleProvider { } async getKvCacheEngineSamples(query: KvCacheSampleQuery): Promise { - return this.kvCacheSamples - .filter( - (r) => - this.inKvCacheRange(r, query) && - (query.engineId === undefined || r.engineId === query.engineId) && - (query.modelName === undefined || r.modelName === query.modelName), - ) - .sort((a, b) => a.timestamp - b.timestamp) - .slice(-(query.limit ?? 50_000)) - .map((r) => ({ ...r })); + const rows = query.bucketMs + ? this.bucketKvCacheSamples(this.matchingKvCacheSamples(query), Math.ceil(query.bucketMs)) + : this.matchingKvCacheSamples(query).sort((a, b) => a.timestamp - b.timestamp); + return rows.slice(-(query.limit ?? 50_000)).map((r) => ({ ...r })); + } + + private bucketKvCacheSamples(rows: KvCacheEngineSample[], bucketMs: number): KvCacheEngineSample[] { + const buckets = new Map(); + for (const row of rows) { + const timestamp = Math.floor(row.timestamp / bucketMs) * bucketMs; + const key = JSON.stringify([row.engineId, row.connectionId, row.modelName, timestamp]); + const bucket = buckets.get(key); + if (!bucket) { + buckets.set(key, { ...row, timestamp }); + continue; + } + for (const field of KV_CACHE_SAMPLE_COUNTER_FIELDS) bucket[field] += row[field]; + } + const compare = (a: string, b: string) => (a < b ? -1 : a > b ? 1 : 0); + return [...buckets.values()].sort( + (a, b) => a.timestamp - b.timestamp || compare(a.engineId, b.engineId) || compare(a.modelName, b.modelName), + ); + } + + async getKvCacheSampleTotals(query: Omit): Promise { + return this.matchingKvCacheSamples(query).reduce( + (totals, r) => ({ requestedTokens: totals.requestedTokens + r.requestedTokens, hitTokens: totals.hitTokens + r.hitTokens }), + { requestedTokens: 0, hitTokens: 0 }, + ); + } + + private matchingKvCacheSamples(query: Omit): KvCacheEngineSample[] { + return this.kvCacheSamples.filter( + (r) => + this.inKvCacheRange(r, query) && + (query.engineId === undefined || r.engineId === query.engineId) && + (query.modelName === undefined || r.modelName === query.modelName), + ); } async pruneOldKvCacheEngineSamples(cutoffTimestamp: number, connectionId?: string): Promise { @@ -1298,6 +1328,10 @@ export class MemoryAdapter implements StoragePort, RawDatabaseHandleProvider { return before - this.kvCacheSamples.length; } + async deleteKvCacheEngineSamples(engineId: string): Promise { + this.kvCacheSamples = this.kvCacheSamples.filter((r) => r.engineId !== engineId); + } + async saveKvCacheEngine(engine: StoredKvCacheEngine): Promise { const existing = this.kvCacheEngines.get(engine.id); const stored = existing diff --git a/apps/api/src/storage/adapters/postgres.adapter.ts b/apps/api/src/storage/adapters/postgres.adapter.ts index 3aa2fc008..bf711d416 100644 --- a/apps/api/src/storage/adapters/postgres.adapter.ts +++ b/apps/api/src/storage/adapters/postgres.adapter.ts @@ -1,4 +1,5 @@ import { Client, ClientBase, ClientConfig, Pool, PoolConfig } from 'pg'; +import { KV_CACHE_SAMPLE_COUNTER_COLUMNS } from './kv-cache-sample-columns'; import { isTrueFlag } from '../../config/env-normalize'; import { chunkedPostgresDelete } from './postgres-chunked-delete'; import { @@ -15,6 +16,7 @@ import type { KvCacheRangeQuery, KvCacheFootprintQuery, KvCacheSampleQuery, + KvCacheSampleTotals, } from '../../common/interfaces/storage-port.interface'; import type { ReadinessDimension, @@ -4333,6 +4335,49 @@ export class PostgresAdapter implements StoragePort, RawDatabaseHandleProvider { async getKvCacheEngineSamples(query: KvCacheSampleQuery): Promise { if (!this.pool) throw new Error('Database not initialized'); + if (query.bucketMs) return this.getBucketedKvCacheEngineSamples(query, Math.ceil(query.bucketMs)); + const { where, params } = this.kvCacheSampleFilter(query); + params.push(query.limit ?? 50_000); + const result = await this.pool.query( + `SELECT * FROM ( + SELECT * FROM kv_cache_engine_samples WHERE ${where.join(' AND ')} + ORDER BY timestamp DESC, id DESC LIMIT $${params.length} + ) AS recent ORDER BY timestamp ASC, id ASC`, + params, + ); + return result.rows.map((r) => this.kvCacheSampleFromRow(r)); + } + + private async getBucketedKvCacheEngineSamples(query: KvCacheSampleQuery, bucketMs: number): Promise { + const { where, params } = this.kvCacheSampleFilter(query); + params.push(bucketMs); + const bucket = `$${params.length}::bigint`; + params.push(query.limit ?? 50_000); + const sums = KV_CACHE_SAMPLE_COUNTER_COLUMNS.map((column) => `SUM(${column}) AS ${column}`).join(', '); + const result = await this.pool!.query( + `SELECT * FROM ( + SELECT engine_id, connection_id, model_name, (timestamp / ${bucket}) * ${bucket} AS bucket_ts, ${sums} + FROM kv_cache_engine_samples WHERE ${where.join(' AND ')} + GROUP BY engine_id, connection_id, model_name, bucket_ts + ORDER BY bucket_ts DESC, engine_id DESC, model_name DESC LIMIT $${params.length} + ) AS recent ORDER BY bucket_ts ASC, engine_id ASC, model_name ASC`, + params, + ); + return result.rows.map((r) => this.kvCacheSampleFromRow({ ...r, timestamp: r.bucket_ts })); + } + + async getKvCacheSampleTotals(query: Omit): Promise { + if (!this.pool) throw new Error('Database not initialized'); + const { where, params } = this.kvCacheSampleFilter(query); + const result = await this.pool.query( + `SELECT COALESCE(SUM(requested_tokens), 0) AS requested, COALESCE(SUM(hit_tokens), 0) AS hit + FROM kv_cache_engine_samples WHERE ${where.join(' AND ')}`, + params, + ); + return { requestedTokens: Number(result.rows[0].requested), hitTokens: Number(result.rows[0].hit) }; + } + + private kvCacheSampleFilter(query: Omit): { where: string[]; params: unknown[] } { const params: unknown[] = [query.connectionId, query.from ?? 0, query.to ?? Number.MAX_SAFE_INTEGER]; const where = ['connection_id = $1', 'timestamp >= $2', 'timestamp <= $3']; if (query.engineId !== undefined) { @@ -4343,15 +4388,7 @@ export class PostgresAdapter implements StoragePort, RawDatabaseHandleProvider { params.push(query.modelName); where.push(`model_name = $${params.length}`); } - params.push(query.limit ?? 50_000); - const result = await this.pool.query( - `SELECT * FROM ( - SELECT * FROM kv_cache_engine_samples WHERE ${where.join(' AND ')} - ORDER BY timestamp DESC, id DESC LIMIT $${params.length} - ) AS recent ORDER BY timestamp ASC, id ASC`, - params, - ); - return result.rows.map((r) => this.kvCacheSampleFromRow(r)); + return { where, params }; } async pruneOldKvCacheEngineSamples(cutoffTimestamp: number, connectionId?: string): Promise { @@ -4362,6 +4399,11 @@ export class PostgresAdapter implements StoragePort, RawDatabaseHandleProvider { return chunkedPostgresDelete(this.pool, 'kv_cache_engine_samples', 'timestamp < $1', [cutoffTimestamp]); } + async deleteKvCacheEngineSamples(engineId: string): Promise { + if (!this.pool) throw new Error('Database not initialized'); + await chunkedPostgresDelete(this.pool, 'kv_cache_engine_samples', 'engine_id = $1', [engineId]); + } + async saveKvCacheEngine(e: StoredKvCacheEngine): Promise { if (!this.pool) throw new Error('Database not initialized'); const result = await this.pool.query( diff --git a/apps/api/src/storage/adapters/sqlite.adapter.ts b/apps/api/src/storage/adapters/sqlite.adapter.ts index 0cd6aa0cc..772ed8a50 100644 --- a/apps/api/src/storage/adapters/sqlite.adapter.ts +++ b/apps/api/src/storage/adapters/sqlite.adapter.ts @@ -1,4 +1,5 @@ import type Database from 'better-sqlite3'; +import { KV_CACHE_SAMPLE_COUNTER_COLUMNS } from './kv-cache-sample-columns'; import { chunkedSqliteDelete } from './sqlite-chunked-delete'; import { otelSpanCap, @@ -16,6 +17,7 @@ import type { KvCacheRangeQuery, KvCacheFootprintQuery, KvCacheSampleQuery, + KvCacheSampleTotals, } from '../../common/interfaces/storage-port.interface'; import type { ReadinessBand, @@ -4264,16 +4266,8 @@ export class SqliteAdapter implements StoragePort, RawDatabaseHandleProvider { async getKvCacheEngineSamples(query: KvCacheSampleQuery): Promise { if (!this.db) throw new Error('Database not initialized'); - const where = ['connection_id = ?', 'timestamp >= ?', 'timestamp <= ?']; - const params: unknown[] = [query.connectionId, query.from ?? 0, query.to ?? Number.MAX_SAFE_INTEGER]; - if (query.engineId !== undefined) { - where.push('engine_id = ?'); - params.push(query.engineId); - } - if (query.modelName !== undefined) { - where.push('model_name = ?'); - params.push(query.modelName); - } + if (query.bucketMs) return this.getBucketedKvCacheEngineSamples(query, Math.ceil(query.bucketMs)); + const { where, params } = this.kvCacheSampleFilter(query); params.push(query.limit ?? 50_000); const rows = this.db .prepare( @@ -4286,6 +4280,48 @@ export class SqliteAdapter implements StoragePort, RawDatabaseHandleProvider { return rows.map((r) => this.kvCacheSampleFromRow(r)); } + private getBucketedKvCacheEngineSamples(query: KvCacheSampleQuery, bucketMs: number): KvCacheEngineSample[] { + const { where, params } = this.kvCacheSampleFilter(query); + const sums = KV_CACHE_SAMPLE_COUNTER_COLUMNS.map((column) => `SUM(${column}) AS ${column}`).join(', '); + const rows = this.db! + .prepare( + `SELECT * FROM ( + SELECT engine_id, connection_id, model_name, (timestamp / CAST(? AS INTEGER)) * CAST(? AS INTEGER) AS bucket_ts, ${sums} + FROM kv_cache_engine_samples WHERE ${where.join(' AND ')} + GROUP BY engine_id, connection_id, model_name, bucket_ts + ORDER BY bucket_ts DESC, engine_id DESC, model_name DESC LIMIT ? + ) ORDER BY bucket_ts ASC, engine_id ASC, model_name ASC`, + ) + .all(bucketMs, bucketMs, ...params, query.limit ?? 50_000) as Record[]; + return rows.map((r) => this.kvCacheSampleFromRow({ ...r, timestamp: r.bucket_ts })); + } + + async getKvCacheSampleTotals(query: Omit): Promise { + if (!this.db) throw new Error('Database not initialized'); + const { where, params } = this.kvCacheSampleFilter(query); + const row = this.db + .prepare( + `SELECT COALESCE(SUM(requested_tokens), 0) AS requested, COALESCE(SUM(hit_tokens), 0) AS hit + FROM kv_cache_engine_samples WHERE ${where.join(' AND ')}`, + ) + .get(...params) as { requested: number; hit: number }; + return { requestedTokens: Number(row.requested), hitTokens: Number(row.hit) }; + } + + private kvCacheSampleFilter(query: Omit): { where: string[]; params: unknown[] } { + const where = ['connection_id = ?', 'timestamp >= ?', 'timestamp <= ?']; + const params: unknown[] = [query.connectionId, query.from ?? 0, query.to ?? Number.MAX_SAFE_INTEGER]; + if (query.engineId !== undefined) { + where.push('engine_id = ?'); + params.push(query.engineId); + } + if (query.modelName !== undefined) { + where.push('model_name = ?'); + params.push(query.modelName); + } + return { where, params }; + } + async pruneOldKvCacheEngineSamples(cutoffTimestamp: number, connectionId?: string): Promise { if (!this.db) throw new Error('Database not initialized'); if (connectionId) { @@ -4294,6 +4330,11 @@ export class SqliteAdapter implements StoragePort, RawDatabaseHandleProvider { return chunkedSqliteDelete(this.db, 'kv_cache_engine_samples', 'timestamp < ?', [cutoffTimestamp]); } + async deleteKvCacheEngineSamples(engineId: string): Promise { + if (!this.db) throw new Error('Database not initialized'); + await chunkedSqliteDelete(this.db, 'kv_cache_engine_samples', 'engine_id = ?', [engineId]); + } + async saveKvCacheEngine(e: StoredKvCacheEngine): Promise { if (!this.db) throw new Error('Database not initialized'); const row = this.db diff --git a/apps/api/src/webhooks/__tests__/webhook-dispatcher.service.spec.ts b/apps/api/src/webhooks/__tests__/webhook-dispatcher.service.spec.ts index b7bfddc78..98f20a8d7 100644 --- a/apps/api/src/webhooks/__tests__/webhook-dispatcher.service.spec.ts +++ b/apps/api/src/webhooks/__tests__/webhook-dispatcher.service.spec.ts @@ -10,6 +10,9 @@ import { type WebhookPayload, } from '@betterdb/shared'; import { ConfigService } from '@nestjs/config'; +import { createServer, type Server } from 'http'; +import type { AddressInfo } from 'net'; +import { promises as dns } from 'dns'; describe('WebhookDispatcherService', () => { let service: WebhookDispatcherService; @@ -865,4 +868,102 @@ describe('WebhookDispatcherService', () => { ); }); }); + + describe('Delivery transport', () => { + let server: Server; + let port: number; + let hits: Record }[]>; + const env = process.env.NODE_ENV; + + const hook = (url: string) => ({ + id: '1', + name: 'Receiver', + url, + enabled: true, + events: [WebhookEventType.INSTANCE_DOWN], + headers: {}, + retryPolicy: { maxRetries: 3, backoffMultiplier: 2, initialDelayMs: 1000, maxDelayMs: 60000 }, + createdAt: Date.now(), + updatedAt: Date.now(), + }); + + const payload = { + id: 'p1', + event: WebhookEventType.INSTANCE_DOWN, + timestamp: Date.now(), + instance: { host: 'localhost', port: 6379 }, + data: { reason: 'test' }, + } as WebhookPayload; + + beforeAll(async () => { + server = createServer((req, res) => { + let body = ''; + req.on('data', (chunk) => (body += chunk)); + req.on('end', () => { + (hits[req.url ?? ''] ??= []).push({ body, headers: req.headers }); + if (req.url === '/redirect') { + res.writeHead(302, { Location: '/target' }).end(); + } else { + res.writeHead(200).end('ok'); + } + }); + }); + await new Promise((resolve) => server.listen(0, '127.0.0.1', resolve)); + port = (server.address() as AddressInfo).port; + }); + + afterAll(() => new Promise((resolve) => server.close(() => resolve()))); + + beforeEach(() => { + hits = {}; + storageClient.getDelivery.mockResolvedValue({ id: 'd1', attempts: 0 } as any); + storageClient.updateDelivery.mockResolvedValue(undefined as any); + webhooksService.generateSignature.mockReturnValue('test-signature'); + }); + + afterEach(() => { + process.env.NODE_ENV = env; + jest.restoreAllMocks(); + }); + + it('delivers the signed JSON body', async () => { + await expect(service.sendWebhook(hook(`http://127.0.0.1:${port}/ok`), 'd1', payload)).resolves.toBe(DeliveryStatus.SUCCESS); + expect(hits['/ok']).toHaveLength(1); + expect(JSON.parse(hits['/ok'][0].body)).toMatchObject({ id: 'p1', event: WebhookEventType.INSTANCE_DOWN }); + expect(hits['/ok'][0].headers['content-type']).toBe('application/json'); + expect(hits['/ok'][0].headers['x-webhook-signature']).toBe('test-signature'); + }); + + it('does not follow a redirect on delivery', async () => { + await expect(service.sendWebhook(hook(`http://127.0.0.1:${port}/redirect`), 'd1', payload)).resolves.toBe(DeliveryStatus.FAILED); + expect(hits['/target']).toBeUndefined(); + expect(storageClient.updateDelivery).toHaveBeenCalledWith('d1', expect.objectContaining({ status: DeliveryStatus.FAILED, statusCode: 302 })); + }); + + it('does not follow a redirect on a test delivery', async () => { + const result = await service.testWebhook(hook(`http://127.0.0.1:${port}/redirect`)); + expect(result).toMatchObject({ success: false, statusCode: 302 }); + expect(hits['/target']).toBeUndefined(); + }); + + it('refuses to connect when the hostname resolves to a blocked address in production', async () => { + process.env.NODE_ENV = 'production'; + jest.spyOn(dns, 'lookup').mockResolvedValue([{ address: '10.1.2.3', family: 4 }] as any); + await expect(service.sendWebhook(hook(`http://hooks.example.com:${port}/ok`), 'd1', payload)).resolves.toBe(DeliveryStatus.FAILED); + expect(storageClient.updateDelivery).toHaveBeenCalledWith( + 'd1', + expect.objectContaining({ status: DeliveryStatus.FAILED, responseBody: 'Webhook URL resolves to a blocked address' }), + ); + const result = await service.testWebhook(hook(`http://hooks.example.com:${port}/ok`)); + expect(result).toMatchObject({ success: false, error: 'Webhook URL resolves to a blocked address' }); + expect(hits['/ok']).toBeUndefined(); + }); + + it('connects to the pinned address in production', async () => { + process.env.NODE_ENV = 'production'; + jest.spyOn(dns, 'lookup').mockResolvedValue([{ address: '8.8.8.8', family: 4 }, { address: '10.1.2.3', family: 4 }] as any); + await expect(service.sendWebhook(hook(`http://hooks.example.com:${port}/ok`), 'd1', payload)).resolves.toBe(DeliveryStatus.FAILED); + expect(hits['/ok']).toBeUndefined(); + }); + }); }); diff --git a/apps/api/src/webhooks/webhook-dispatcher.service.ts b/apps/api/src/webhooks/webhook-dispatcher.service.ts index 70c50669c..6ad4b90d1 100644 --- a/apps/api/src/webhooks/webhook-dispatcher.service.ts +++ b/apps/api/src/webhooks/webhook-dispatcher.service.ts @@ -21,6 +21,8 @@ import { StoragePort } from '../common/interfaces/storage-port.interface'; import { WebhooksService } from './webhooks.service'; import { ConnectionRegistry } from '../connections/connection-registry.service'; import { formatWebhookBody } from './webhook-payload-formatter'; +import { BlockedAddressError, guardedLookup } from '../common/utils/outbound-url-guard'; +import { pinnedRequest } from '../common/utils/pinned-request'; interface AlertState { fired: boolean; @@ -646,6 +648,11 @@ export class WebhookDispatcherService implements OnModuleDestroy { return delivery.id; } + private post(url: string, headers: Record, body: string, signal: AbortSignal): Promise { + const lookup = process.env.NODE_ENV === 'production' ? guardedLookup({ label: 'webhook URL' }) : undefined; + return pinnedRequest(url, { method: 'POST', headers, body, signal, lookup }); + } + /** * Send webhook HTTP request * @returns The status for this attempt (SUCCESS only on 2xx). @@ -694,12 +701,7 @@ export class WebhookDispatcherService implements OnModuleDestroy { const timeoutId = setTimeout(() => controller.abort(), timeoutMs); try { - const response = await fetch(webhook.url, { - method: 'POST', - headers, - body: payloadString, - signal: controller.signal, - }); + const response = await this.post(webhook.url, headers, payloadString, controller.signal); clearTimeout(timeoutId); @@ -710,6 +712,11 @@ export class WebhookDispatcherService implements OnModuleDestroy { if (response.ok) { status = DeliveryStatus.SUCCESS; this.logger.log(`Webhook delivered successfully: ${webhook.id} -> ${webhook.url}`); + } else if (statusCode >= 300 && statusCode < 400) { + status = DeliveryStatus.FAILED; + this.logger.warn( + `Webhook delivery redirect not followed (HTTP ${statusCode}): ${webhook.id} -> ${webhook.url}`, + ); } else if (statusCode >= 400 && statusCode < 500) { // 4xx errors are client errors - don't retry status = DeliveryStatus.FAILED; @@ -727,7 +734,11 @@ export class WebhookDispatcherService implements OnModuleDestroy { clearTimeout(timeoutId); // Handle specific errors - if (fetchError instanceof Error && fetchError.name === 'AbortError') { + if (fetchError instanceof BlockedAddressError) { + status = DeliveryStatus.FAILED; + responseBody = fetchError.message; + this.logger.warn(`Webhook delivery blocked: ${webhook.id} -> ${webhook.url}`); + } else if (fetchError instanceof Error && fetchError.name === 'AbortError') { status = DeliveryStatus.RETRYING; responseBody = 'Request timeout'; this.logger.warn(`Webhook delivery timeout: ${webhook.id} -> ${webhook.url}`); @@ -884,12 +895,7 @@ export class WebhookDispatcherService implements OnModuleDestroy { const controller = new AbortController(); const timeoutId = setTimeout(() => controller.abort(), timeoutMs); - const response = await fetch(webhook.url, { - method: 'POST', - headers, - body: payloadString, - signal: controller.signal, - }); + const response = await this.post(webhook.url, headers, payloadString, controller.signal); clearTimeout(timeoutId); diff --git a/apps/api/src/webhooks/webhooks.service.ts b/apps/api/src/webhooks/webhooks.service.ts index 4c9ee8c5d..0524c38cc 100644 --- a/apps/api/src/webhooks/webhooks.service.ts +++ b/apps/api/src/webhooks/webhooks.service.ts @@ -1,9 +1,9 @@ import { Injectable, Inject, Logger, BadRequestException, NotFoundException, ForbiddenException, Optional } from '@nestjs/common'; import { randomBytes, createHmac } from 'crypto'; -import { promises as dns } from 'dns'; import type { Webhook, WebhookDelivery, WebhookEventType, DEFAULT_RETRY_POLICY } from '@betterdb/shared'; import { Tier, validateEventsForTier, getRequiredTierForEvent, getEventsForTier, getLockedEventsForTier } from '@betterdb/shared'; import { WebhookPayloadFormat } from '@betterdb/shared'; +import { assertSafeOutboundUrl } from '../common/utils/outbound-url-guard'; import { StoragePort } from '../common/interfaces/storage-port.interface'; import { CreateWebhookDto, UpdateWebhookDto } from '../common/dto/webhook.dto'; import { ConnectionRegistry } from '../connections/connection-registry.service'; @@ -13,18 +13,6 @@ import { LicenseService } from '@proprietary/licenses'; export class WebhooksService { private readonly logger = new Logger(WebhooksService.name); - // SSRF Protection: Private IP ranges to block - private readonly BLOCKED_IP_PATTERNS = [ - /^127\./, // localhost - /^10\./, // 10.0.0.0/8 - /^172\.(1[6-9]|2[0-9]|3[01])\./, // 172.16.0.0/12 - /^192\.168\./, // 192.168.0.0/16 - /^169\.254\./, // link-local - /^::1$/, // IPv6 localhost - /^fe80:/, // IPv6 link-local - /^fc00:/, // IPv6 unique local - ]; - constructor( @Inject('STORAGE_CLIENT') private readonly storageClient: StoragePort, private readonly connectionRegistry: ConnectionRegistry, @@ -39,90 +27,6 @@ export class WebhooksService { return connectionId || this.connectionRegistry.getDefaultId() || undefined; } - /** - * Check if an IP address is blocked - */ - private isBlockedIp(ip: string): boolean { - for (const pattern of this.BLOCKED_IP_PATTERNS) { - if (pattern.test(ip)) { - return true; - } - } - return false; - } - - /** - * Validate webhook URL for SSRF protection - */ - private async validateUrl(url: string): Promise { - try { - const parsed = new URL(url); - - // Only allow http and https - if (!['http:', 'https:'].includes(parsed.protocol)) { - throw new BadRequestException('Only HTTP and HTTPS protocols are allowed'); - } - - // Warn if URL contains credentials - if (parsed.username || parsed.password) { - this.logger.warn(`Webhook URL contains credentials, consider using custom headers instead: ${parsed.hostname}`); - } - - // Allow localhost in development/non-production environments - const isProduction = process.env.NODE_ENV === 'production'; - const isLocalhost = parsed.hostname === 'localhost' || - parsed.hostname === '0.0.0.0' || - parsed.hostname === '127.0.0.1' || - parsed.hostname.startsWith('127.'); - - if (isLocalhost && !isProduction) { - // Allow localhost in development - this.logger.debug(`Allowing localhost webhook URL in ${process.env.NODE_ENV || 'development'} mode: ${url}`); - return; - } - - // Block localhost in production - if (parsed.hostname === 'localhost' || parsed.hostname === '0.0.0.0') { - throw new BadRequestException('Cannot use localhost or 0.0.0.0 as webhook URL in production'); - } - - // Check if hostname is already an IP - if (this.isBlockedIp(parsed.hostname)) { - throw new BadRequestException('Cannot use private IP addresses as webhook URL'); - } - - // Additional checks for common bypass attempts - if (parsed.hostname.includes('127.') || parsed.hostname.includes('localhost')) { - throw new BadRequestException('Suspicious hostname detected'); - } - - // DNS resolution to prevent DNS rebinding attacks - if (isProduction) { - try { - const addresses = await dns.resolve(parsed.hostname); - for (const addr of addresses) { - if (this.isBlockedIp(addr)) { - throw new BadRequestException(`Webhook URL resolves to blocked IP address: ${addr}`); - } - } - } catch (dnsError: any) { - // If DNS resolution fails, it might be unreachable but not necessarily malicious - if (dnsError instanceof BadRequestException) { - throw dnsError; - } - this.logger.warn(`Failed to resolve DNS for webhook URL: ${parsed.hostname}`); - throw new BadRequestException('Failed to resolve webhook URL hostname'); - } - } - - } catch (error) { - if (error instanceof BadRequestException) { - throw error; - } - throw new BadRequestException('Invalid webhook URL'); - } - } - /** * Generate a secure webhook secret */ @@ -176,7 +80,7 @@ export class WebhooksService { */ async createWebhook(dto: CreateWebhookDto & { connectionId?: string }): Promise { // Validate URL for SSRF - await this.validateUrl(dto.url); + await assertSafeOutboundUrl(dto.url, { label: 'webhook URL', resolveDns: 'production' }); // Validate event tier access this.validateEventTiers(dto.events); @@ -261,7 +165,7 @@ export class WebhooksService { async updateWebhook(id: string, dto: UpdateWebhookDto): Promise { // Validate URL if provided if (dto.url) { - await this.validateUrl(dto.url); + await assertSafeOutboundUrl(dto.url, { label: 'webhook URL', resolveDns: 'production' }); } // Validate event tier access if events are being updated diff --git a/apps/web/src/components/webhooks/WebhookForm.tsx b/apps/web/src/components/webhooks/WebhookForm.tsx index 802924f4b..e40804afa 100644 --- a/apps/web/src/components/webhooks/WebhookForm.tsx +++ b/apps/web/src/components/webhooks/WebhookForm.tsx @@ -82,6 +82,8 @@ const EVENT_LABELS: Record = { 'config.changed': 'Config Changed', 'metric_forecast.limit': 'Metric Forecast Limit', 'scaling_readiness.low': 'Scaling Readiness Low', + 'kv_cache.hit_rate_low': 'KV cache hit rate low', + 'kv_cache.eviction_risk': 'KV cache eviction risk', 'inference.sla.breach': 'Inference SLA Breach', 'cve.critical_detected': 'CVE Critical Detected', 'cve.kev_detected': 'CVE KEV Detected', diff --git a/packages/shared/src/webhooks/types.ts b/packages/shared/src/webhooks/types.ts index 5b0ca853c..b4ba6439d 100644 --- a/packages/shared/src/webhooks/types.ts +++ b/packages/shared/src/webhooks/types.ts @@ -25,6 +25,8 @@ export enum WebhookEventType { COMPLIANCE_ALERT = 'compliance.alert', METRIC_FORECAST_LIMIT = 'metric_forecast.limit', SCALING_READINESS_LOW = 'scaling_readiness.low', + KV_CACHE_HIT_RATE_LOW = 'kv_cache.hit_rate_low', + KV_CACHE_EVICTION_RISK = 'kv_cache.eviction_risk', INFERENCE_SLA_BREACH = 'inference.sla.breach', CVE_CRITICAL_DETECTED = 'cve.critical_detected', CVE_KEV_DETECTED = 'cve.kev_detected', @@ -65,6 +67,8 @@ export const PRO_EVENTS: WebhookEventType[] = [ WebhookEventType.CONNECTION_SPIKE, WebhookEventType.METRIC_FORECAST_LIMIT, WebhookEventType.SCALING_READINESS_LOW, + WebhookEventType.KV_CACHE_HIT_RATE_LOW, + WebhookEventType.KV_CACHE_EVICTION_RISK, WebhookEventType.FAILOVER_STARTED, WebhookEventType.FAILOVER_COMPLETED, WebhookEventType.DATA_LOSS_DETECTED, @@ -91,6 +95,7 @@ export const ENTERPRISE_EVENTS: WebhookEventType[] = [ import { Tier } from '../license/types'; import type { CveWebhookFindingSummary } from '../types/cve'; import type { MetricKind } from '../types/metric-forecasting.types'; +import type { KvCacheEvictionRiskData, KvCacheHitRateLowData } from '../types/kv-cache.types'; import type { ScalingReadinessLowData } from '../types/scaling-readiness.types'; export { Tier }; @@ -122,6 +127,8 @@ export const WEBHOOK_EVENT_TIERS: Record = { [WebhookEventType.CONNECTION_SPIKE]: Tier.pro, [WebhookEventType.METRIC_FORECAST_LIMIT]: Tier.pro, [WebhookEventType.SCALING_READINESS_LOW]: Tier.pro, + [WebhookEventType.KV_CACHE_HIT_RATE_LOW]: Tier.pro, + [WebhookEventType.KV_CACHE_EVICTION_RISK]: Tier.pro, [WebhookEventType.FAILOVER_STARTED]: Tier.pro, [WebhookEventType.FAILOVER_COMPLETED]: Tier.pro, [WebhookEventType.DATA_LOSS_DETECTED]: Tier.pro, @@ -517,6 +524,10 @@ export interface IWebhookEventsProService { dispatchScalingReadinessLow(data: ScalingReadinessLowData): Promise; + dispatchKvCacheHitRateLow(data: KvCacheHitRateLowData): Promise; + + dispatchKvCacheEvictionRisk(data: KvCacheEvictionRiskData): Promise; + dispatchInferenceSlaBreach(data: { indexName: string; currentP99Us: number; diff --git a/proprietary/kv-cache/__tests__/counter-deltas.spec.ts b/proprietary/kv-cache/__tests__/counter-deltas.spec.ts new file mode 100644 index 000000000..64fa07251 --- /dev/null +++ b/proprietary/kv-cache/__tests__/counter-deltas.spec.ts @@ -0,0 +1,213 @@ +import { addCounters, BASELINE_IDLE_MS, CounterDeltaTracker, labelSeries, METRIC_FIELDS, type CounterObservation } from '../counter-deltas'; + +const obs = (o: Partial): CounterObservation => ({ + engineId: 'e1', + connectionId: 'c1', + modelName: 'm', + series: '0|worker', + metric: 'num_hit_tokens', + value: 0, + cumulative: true, + ...o, +}); + +describe('CounterDeltaTracker', () => { + it('uses the first point as a baseline and sums growth', () => { + const t = new CounterDeltaTracker(); + t.observe(obs({ value: 100 }), 1_000); + t.observe(obs({ value: 150 }), 2_000); + t.observe(obs({ value: 175 }), 3_000); + const [bucket] = t.drain(60_000); + expect(bucket).toMatchObject({ engineId: 'e1', connectionId: 'c1', modelName: 'm', timestamp: 0, hitTokens: 75, requestedTokens: 0 }); + }); + + it('treats a lower value as a reset', () => { + const t = new CounterDeltaTracker(); + t.observe(obs({ value: 100 }), 1_000); + t.observe(obs({ value: 30 }), 2_000); + expect(t.drain(60_000)[0].hitTokens).toBe(30); + }); + + it('ignores an out-of-order point from the same series start', () => { + const t = new CounterDeltaTracker(); + t.observe(obs({ value: 100, startMs: 5 }), 1_000); + t.observe(obs({ value: 150, startMs: 5 }), 2_000); + t.observe(obs({ value: 120, startMs: 5 }), 3_000); + t.observe(obs({ value: 160, startMs: 5 }), 4_000); + expect(t.drain(60_000)[0].hitTokens).toBe(60); + }); + + it('takes the first point after a start-time change as a new baseline', () => { + const t = new CounterDeltaTracker(); + t.observe(obs({ value: 100, startMs: 5 }), 1_000); + t.observe(obs({ value: 140, startMs: 9 }), 2_000); + t.observe(obs({ value: 150, startMs: 9 }), 3_000); + expect(t.drain(60_000)[0].hitTokens).toBe(10); + }); + + it('does not count lifetime totals when a collector restart moves the start time', () => { + const t = new CounterDeltaTracker(); + t.observe(obs({ value: 5_000_000, startMs: 5 }), 1_000); + t.observe(obs({ value: 5_000_100, startMs: 5 }), 2_000); + t.observe(obs({ value: 5_000_400, startMs: 30 }), 3_000); + t.observe(obs({ value: 5_000_500, startMs: 30 }), 4_000); + expect(t.drain(60_000)[0].hitTokens).toBe(200); + }); + + it('ignores a late point from an older series start', () => { + const t = new CounterDeltaTracker(); + t.observe(obs({ value: 1000, startMs: 1 }), 1_000); + t.observe(obs({ value: 50, startMs: 2 }), 2_000); + expect(t.drain(200_000, true)).toEqual([]); + t.observe(obs({ value: 1000, startMs: 1 }), 3_000); + expect(t.drain(200_000, true)).toEqual([]); + t.observe(obs({ value: 70, startMs: 2 }), 4_000); + expect(t.drain(200_000, true)[0].hitTokens).toBe(20); + }); + + it('keeps the newest start when a point arrives without one', () => { + const t = new CounterDeltaTracker(); + t.observe(obs({ value: 100, startMs: 9 }), 1_000); + t.observe(obs({ value: 120 }), 2_000); + t.observe(obs({ value: 500, startMs: 5 }), 3_000); + t.observe(obs({ value: 130, startMs: 9 }), 4_000); + expect(t.drain(200_000, true)[0].hitTokens).toBe(30); + }); + + it('does not flap between interleaved series that share a key', () => { + const t = new CounterDeltaTracker(); + for (let i = 0; i < 4; i += 1) { + t.observe(obs({ value: 1000 + i * 10, startMs: 1 }), 1_000 + i); + t.observe(obs({ value: 2000 + i * 10, startMs: 2 }), 1_000 + i); + } + expect(t.drain(200_000, true)[0].hitTokens).toBe(30); + }); + + it('sums workers into one bucket per engine and model', () => { + const t = new CounterDeltaTracker(); + for (const series of ['0|worker', '1|worker']) { + t.observe(obs({ series, value: 0 }), 1_000); + t.observe(obs({ series, value: 10 }), 2_000); + } + const buckets = t.drain(60_000); + expect(buckets).toHaveLength(1); + expect(buckets[0].hitTokens).toBe(20); + }); + + it('adds delta-temporality points directly', () => { + const t = new CounterDeltaTracker(); + t.observe(obs({ value: 40, cumulative: false, metric: 'num_requested_tokens' }), 1_000); + t.observe(obs({ value: 60, cumulative: false, metric: 'num_requested_tokens' }), 2_000); + expect(t.drain(60_000)[0].requestedTokens).toBe(100); + }); + + it('keeps the open minute until it closes', () => { + const t = new CounterDeltaTracker(); + t.observe(obs({ value: 1 }), 61_000); + t.observe(obs({ value: 2 }), 62_000); + expect(t.drain(90_000)).toEqual([]); + expect(t.drain(120_000)).toHaveLength(1); + expect(t.drain(120_000)).toEqual([]); + }); + + it('splits buckets by minute and rejects unknown metrics', () => { + const t = new CounterDeltaTracker(); + expect(t.observe(obs({ metric: 'lookup_hit_rate' }), 1)).toBe(false); + t.observe(obs({ value: 0 }), 1_000); + t.observe(obs({ value: 5 }), 59_000); + t.observe(obs({ value: 9 }), 61_000); + expect(t.drain(200_000).map((b) => [b.timestamp, b.hitTokens])).toEqual([[0, 5], [60_000, 4]]); + }); + + it('rejects inherited metric names', () => { + const t = new CounterDeltaTracker(); + expect(t.observe(obs({ metric: 'constructor', value: 1 }), 1_000)).toBe(false); + expect(t.observe(obs({ metric: 'toString', value: 2, cumulative: false }), 2_000)).toBe(false); + expect(t.drain(200_000, true)).toEqual([]); + }); + + it('rejects non-finite values without disturbing the next delta', () => { + const t = new CounterDeltaTracker(); + t.observe(obs({ value: 100 }), 1_000); + expect(t.observe(obs({ value: NaN }), 2_000)).toBe(false); + t.observe(obs({ value: 150 }), 3_000); + expect(t.drain(60_000)[0].hitTokens).toBe(50); + }); + + it('drops baselines of series that have gone idle', () => { + const t = new CounterDeltaTracker(); + t.observe(obs({ series: 'old-worker', value: 100 }), 1_000); + t.observe(obs({ series: 'live', value: 10 }), 1_000); + t.observe(obs({ series: 'live', value: 20 }), BASELINE_IDLE_MS); + t.drain(BASELINE_IDLE_MS + 2_000); + expect(t.baselineCount()).toBe(1); + t.observe(obs({ series: 'old-worker', value: 500 }), BASELINE_IDLE_MS + 3_000); + t.observe(obs({ series: 'live', value: 25 }), BASELINE_IDLE_MS + 3_000); + expect(t.drain(10 * BASELINE_IDLE_MS, true).reduce((sum, b) => sum + b.hitTokens, 0)).toBe(15); + }); + + it('forgets an engine', () => { + const t = new CounterDeltaTracker(); + t.observe(obs({ value: 0 }), 1_000); + t.observe(obs({ value: 5 }), 2_000); + t.forgetEngine('e1'); + expect(t.drain(200_000, true)).toEqual([]); + t.observe(obs({ value: 50 }), 3_000); + expect(t.drain(200_000, true)).toEqual([]); + }); + + it('forgets baselines but keeps pending buckets', () => { + const t = new CounterDeltaTracker(); + t.observe(obs({ value: 0 }), 1_000); + t.observe(obs({ value: 5 }), 2_000); + t.forgetBaselines('e1'); + t.observe(obs({ value: 500 }), 3_000); + t.observe(obs({ value: 510 }), 4_000); + expect(t.drain(200_000, true)[0].hitTokens).toBe(15); + }); + + it('folds a late observation for an already drained minute into the open minute', () => { + const t = new CounterDeltaTracker(); + t.observe(obs({ value: 100 }), 50_000); + t.observe(obs({ value: 150 }), 55_000); + expect(t.drain(61_000)).toEqual([expect.objectContaining({ timestamp: 0, hitTokens: 50 })]); + t.observe(obs({ value: 170 }), 59_000); + expect(t.drain(61_000)).toEqual([]); + expect(t.drain(125_000)).toEqual([expect.objectContaining({ timestamp: 60_000, hitTokens: 20 })]); + }); + + it('never recreates a minute that was drained with the open bucket', () => { + const t = new CounterDeltaTracker(); + t.observe(obs({ value: 100 }), 10_000); + t.observe(obs({ value: 110 }), 20_000); + expect(t.drain(30_000, true)).toEqual([expect.objectContaining({ timestamp: 0, hitTokens: 10 })]); + t.observe(obs({ value: 115 }), 40_000); + expect(t.drain(200_000, true)).toEqual([expect.objectContaining({ timestamp: 60_000, hitTokens: 5 })]); + }); +}); + +describe('labelSeries', () => { + it('keeps the worker and role key when there are no other labels', () => { + expect(labelSeries({ model_name: 'm', worker_id: '0', role: 'worker' })).toBe('0|worker'); + expect(labelSeries({})).toBe('|'); + }); + + it('separates series that differ only by other labels, in a stable order', () => { + const a = labelSeries({ model_name: 'm', worker_id: '0', served_model_name: 'a', zone: 'z' }); + const b = labelSeries({ zone: 'z', served_model_name: 'a', worker_id: '0', model_name: 'm' }); + const c = labelSeries({ model_name: 'm', worker_id: '0', served_model_name: 'b', zone: 'z' }); + expect(a).toBe(b); + expect(a).not.toBe(c); + expect(a).toBe('0||[["served_model_name","a"],["zone","z"]]'); + }); +}); + +describe('addCounters', () => { + it('adds every mapped counter field into the target', () => { + const fields = Object.values(METRIC_FIELDS); + const target = Object.fromEntries(fields.map((field) => [field, 1])) as Record; + const source = Object.fromEntries(fields.map((field) => [field, 2])) as Record; + addCounters(target as any, source as any); + for (const field of fields) expect(target[field]).toBe(3); + }); +}); diff --git a/proprietary/kv-cache/__tests__/fixtures/lmcache-metrics.txt b/proprietary/kv-cache/__tests__/fixtures/lmcache-metrics.txt new file mode 100644 index 000000000..662eddcef --- /dev/null +++ b/proprietary/kv-cache/__tests__/fixtures/lmcache-metrics.txt @@ -0,0 +1,11 @@ +# HELP lmcache:num_requested_tokens_total Total number of tokens requested +# TYPE lmcache:num_requested_tokens_total counter +lmcache:num_requested_tokens_total{model_name="Qwen/Qwen2.5-0.5B-Instruct",role="worker",served_model_name="qwen",worker_id="0"} 1000.0 +lmcache:num_requested_tokens_created{model_name="Qwen/Qwen2.5-0.5B-Instruct",role="worker",served_model_name="qwen",worker_id="0"} 1.7599e+09 +# HELP lmcache:num_hit_tokens_total Total number of hit tokens +# TYPE lmcache:num_hit_tokens_total counter +lmcache:num_hit_tokens_total{model_name="Qwen/Qwen2.5-0.5B-Instruct",role="worker",served_model_name="qwen",worker_id="0"} 700.0 +lmcache:lookup_hit_rate{model_name="Qwen/Qwen2.5-0.5B-Instruct",role="worker",served_model_name="qwen",worker_id="0"} NaN +vllm:num_requests_running{model_name="Qwen/Qwen2.5-0.5B-Instruct"} 3.0 +lmcache:remote_ping_errors_total{model_name="a\"b\\c\nd",role="worker",served_model_name="",worker_id="1"} 2 1759900000000 +lmcache:num_remote_read_bytes_total{model_name="m",role="worker",served_model_name="m",worker_id="0"} +Inf diff --git a/proprietary/kv-cache/__tests__/kv-cache-alert-rules.spec.ts b/proprietary/kv-cache/__tests__/kv-cache-alert-rules.spec.ts new file mode 100644 index 000000000..cf1336e1c --- /dev/null +++ b/proprietary/kv-cache/__tests__/kv-cache-alert-rules.spec.ts @@ -0,0 +1,137 @@ +import type { KvCacheEngineSample, KvCacheFootprintSnapshot, KvCacheSettings } from '@betterdb/shared'; +import { evictionConditions, shouldCheckHitRate, windowHitRates, type HitRateWindow } from '../kv-cache-alert-rules'; + +const settings = (over: Partial = {}): KvCacheSettings => ({ + connectionId: 'c1', + hitRateAlertEnabled: true, + hitRateThreshold: 0.2, + evictionAlertEnabled: true, + updatedAt: 0, + ...over, +}); + +const window = (over: Partial = {}): HitRateWindow => ({ + engineId: 'e1', + model: 'm', + requestedTokens: 10000, + hitTokens: 1000, + hitRate: 0.1, + ...over, +}); + +const sample = (over: Partial = {}): KvCacheEngineSample => ({ + connectionId: 'c1', + engineId: 'e1', + modelName: 'm', + timestamp: 0, + requestedTokens: 100, + hitTokens: 50, + lookupTokens: 0, + lookupHits: 0, + remoteReadBytes: 0, + remoteWriteBytes: 0, + remoteReadRequests: 0, + remoteWriteRequests: 0, + remotePingErrors: 0, + ...over, +}); + +const snapshot = (over: Partial = {}): KvCacheFootprintSnapshot => ({ + connectionId: 'c1', + timestamp: 0, + detected: true, + layout: null, + scannedKeys: 0, + matchedKeys: 0, + sampledKeys: 0, + scanComplete: true, + primariesScanned: 1, + primariesTotal: 1, + chunksEst: 0, + bytesEst: 0, + usedMemory: 50, + maxmemory: 100, + maxmemoryPolicy: 'allkeys-lru', + lmcacheMemoryShare: 0.1, + noTtlRatio: 0, + orphanRatio: null, + evictedKeysDelta: 0, + evictedKeysIntervalMs: 300_000, + otherDbs: [], + perModel: [], + ...over, +}); + +describe('shouldCheckHitRate', () => { + it('requires the token floor', () => { + expect(shouldCheckHitRate(window({ requestedTokens: 9999 }), settings())).toBe(false); + expect(shouldCheckHitRate(window({ requestedTokens: 10000 }), settings())).toBe(true); + }); + + it('is false when the alert is disabled', () => { + expect(shouldCheckHitRate(window(), settings({ hitRateAlertEnabled: false }))).toBe(false); + }); + + it('is false when the hit rate is null', () => { + expect(shouldCheckHitRate(window({ hitRate: null }), settings())).toBe(false); + }); +}); + +describe('windowHitRates', () => { + it('sums minutes per engine and model and keeps models apart', () => { + const result = windowHitRates([ + sample({ timestamp: 1, requestedTokens: 100, hitTokens: 10 }), + sample({ timestamp: 2, requestedTokens: 300, hitTokens: 30 }), + sample({ modelName: 'other', requestedTokens: 50, hitTokens: 50 }), + ]); + expect(result).toEqual([ + { engineId: 'e1', model: 'm', requestedTokens: 400, hitTokens: 40, hitRate: 0.1 }, + { engineId: 'e1', model: 'other', requestedTokens: 50, hitTokens: 50, hitRate: 1 }, + ]); + }); + + it('reports a null hit rate when nothing was requested', () => { + expect(windowHitRates([sample({ requestedTokens: 0, hitTokens: 0 })])[0].hitRate).toBeNull(); + }); +}); + +describe('evictionConditions', () => { + it('flags unevictable keys under a volatile policy', () => { + expect(evictionConditions(snapshot({ maxmemoryPolicy: 'volatile-lru', noTtlRatio: 0.95 })).unevictable).toBe(true); + }); + + it('does not flag unevictable keys without TTL data', () => { + expect(evictionConditions(snapshot({ maxmemoryPolicy: 'volatile-lru', noTtlRatio: null })).unevictable).toBe(false); + }); + + it('does not flag allkeys or noeviction policies', () => { + expect(evictionConditions(snapshot({ maxmemoryPolicy: 'allkeys-lru', noTtlRatio: 1 })).unevictable).toBe(false); + expect(evictionConditions(snapshot({ maxmemoryPolicy: 'noeviction', noTtlRatio: 1 })).unevictable).toBe(false); + }); + + const evicting = { evictedKeysDelta: 5, usedMemory: 95, maxmemory: 100, lmcacheMemoryShare: 0.6 }; + + it('flags evicting when all three conditions hold', () => { + expect(evictionConditions(snapshot(evicting)).evicting).toBe(true); + }); + + it('flags evicting when the LMCache share is unknown', () => { + expect(evictionConditions(snapshot({ ...evicting, lmcacheMemoryShare: null })).evicting).toBe(true); + }); + + it.each([ + ['no evictions', { evictedKeysDelta: 0 }], + ['null evictions', { evictedKeysDelta: null }], + ['low memory pressure', { usedMemory: 80 }], + ['no maxmemory', { maxmemory: 0 }], + ['low lmcache share', { lmcacheMemoryShare: 0.4 }], + ])('does not flag evicting with %s', (_name, override) => { + expect(evictionConditions(snapshot({ ...evicting, ...override })).evicting).toBe(false); + }); + + it('flags nothing when LMCache is not detected', () => { + expect( + evictionConditions(snapshot({ ...evicting, detected: false, maxmemoryPolicy: 'volatile-lru', noTtlRatio: 1 })), + ).toEqual({ unevictable: false, evicting: false }); + }); +}); diff --git a/proprietary/kv-cache/__tests__/kv-cache-alerts.service.spec.ts b/proprietary/kv-cache/__tests__/kv-cache-alerts.service.spec.ts new file mode 100644 index 000000000..4e1c61dbb --- /dev/null +++ b/proprietary/kv-cache/__tests__/kv-cache-alerts.service.spec.ts @@ -0,0 +1,476 @@ +import { Logger } from '@nestjs/common'; +import type { KvCacheFootprintSnapshot } from '@betterdb/shared'; +import { KvCacheAlertsService } from '../kv-cache-alerts.service'; + +const snapshot = (over: Partial = {}): KvCacheFootprintSnapshot => + ({ + connectionId: 'c1', + timestamp: 5, + detected: true, + scanComplete: true, + usedMemory: 95, + maxmemory: 100, + maxmemoryPolicy: 'volatile-lru', + lmcacheMemoryShare: 0.6, + noTtlRatio: 0.95, + evictedKeysDelta: 0, + ...over, + }) as KvCacheFootprintSnapshot; + +const engine = (over: Record = {}) => ({ id: 'e1', connectionId: 'c1', name: 'vllm', enabled: true, ...over }); + +const row = (over: Record = {}) => ({ + connectionId: 'c1', + engineId: 'e1', + modelName: 'm', + timestamp: 1, + requestedTokens: 20000, + hitTokens: 2000, + ...over, +}); + +function setup(opts: { licensed?: boolean; engines?: unknown[]; stored?: unknown; rows?: unknown } = {}) { + const storage = { + getKvCacheSettings: jest.fn().mockResolvedValue(opts.stored ?? null), + saveKvCacheSettings: jest.fn().mockImplementation(async (s) => s), + getKvCacheEngineSamples: jest.fn().mockResolvedValue(opts.rows ?? [row()]), + }; + const connectionRegistry = { getConfig: jest.fn().mockReturnValue({ host: 'h', port: 6379 }) }; + let listener: ((s: KvCacheFootprintSnapshot) => void) | undefined; + let removalListener: ((id: string) => void) | undefined; + const footprint = { + onSnapshot: jest.fn().mockImplementation((l) => (listener = l)), + onConnectionRemoval: jest.fn().mockImplementation((l) => (removalListener = l)), + }; + const prometheus = { setKvCacheFootprint: jest.fn(), setKvCacheHitRates: jest.fn(), clearKvCache: jest.fn(), clearAllKvCache: jest.fn() }; + const registry = { + list: jest.fn().mockReturnValue(opts.engines ?? [engine()]), + get: jest.fn().mockImplementation((id: string) => (opts.engines ?? [engine()]).find((e: any) => e.id === id) ?? null), + }; + const license = { hasFeature: jest.fn().mockReturnValue(opts.licensed ?? true) }; + const pro = { + dispatchKvCacheHitRateLow: jest.fn().mockResolvedValue(undefined), + dispatchKvCacheEvictionRisk: jest.fn().mockResolvedValue(undefined), + }; + const service = new KvCacheAlertsService( + storage as any, + connectionRegistry as any, + footprint as any, + registry as any, + license as any, + pro as any, + prometheus as any, + ); + return { + storage, + connectionRegistry, + footprint, + registry, + license, + pro, + prometheus, + service, + emit: (s: KvCacheFootprintSnapshot) => listener?.(s), + remove: (id: string) => removalListener?.(id), + }; +} + +describe('KvCacheAlertsService settings', () => { + it('returns defaults when nothing is stored', async () => { + const { service } = setup(); + expect(await service.getSettings('c1')).toEqual({ + connectionId: 'c1', + hitRateAlertEnabled: true, + hitRateThreshold: 0.2, + evictionAlertEnabled: true, + updatedAt: 0, + }); + }); + + it('keeps the stored value when an update field is undefined', async () => { + const { service } = setup({ + stored: { connectionId: 'c1', hitRateAlertEnabled: true, hitRateThreshold: 0.3, evictionAlertEnabled: true, updatedAt: 1 }, + }); + const saved = await service.updateSettings('c1', { hitRateThreshold: undefined }); + expect(saved.hitRateThreshold).toBe(0.3); + }); + + it('merges updates onto the current settings and stamps the time', async () => { + const { service, storage } = setup({ + stored: { connectionId: 'c1', hitRateAlertEnabled: true, hitRateThreshold: 0.2, evictionAlertEnabled: true, updatedAt: 1 }, + }); + const saved = await service.updateSettings('c1', { hitRateThreshold: 0.4 }); + expect(saved).toMatchObject({ connectionId: 'c1', hitRateAlertEnabled: true, hitRateThreshold: 0.4, evictionAlertEnabled: true }); + expect(saved.updatedAt).toBeGreaterThan(1); + expect(storage.saveKvCacheSettings).toHaveBeenCalledWith(saved); + }); +}); + +describe('KvCacheAlertsService eviction risk', () => { + const flush = () => new Promise((resolve) => setImmediate(resolve)); + + it('subscribes to footprint snapshots on init and dispatches both reasons', async () => { + const { service, pro, emit } = setup(); + service.onModuleInit(); + emit(snapshot()); + await flush(); + service.onModuleDestroy(); + expect(pro.dispatchKvCacheEvictionRisk).toHaveBeenCalledTimes(2); + expect(pro.dispatchKvCacheEvictionRisk).toHaveBeenCalledWith( + expect.objectContaining({ + connectionId: 'c1', + reason: 'unevictable', + active: true, + policy: 'volatile-lru', + instance: { host: 'h', port: 6379 }, + timestamp: 5, + }), + ); + expect(pro.dispatchKvCacheEvictionRisk).toHaveBeenCalledWith(expect.objectContaining({ reason: 'evicting', active: false })); + }); + + it('does not hold the snapshot on webhook delivery and logs a failed dispatch', async () => { + const { service, pro } = setup(); + const warn = jest.spyOn(Logger.prototype, 'warn').mockImplementation(() => undefined); + pro.dispatchKvCacheEvictionRisk.mockReturnValueOnce(new Promise(() => undefined)).mockRejectedValueOnce(new Error('receiver down')); + await expect(service.onSnapshot(snapshot())).resolves.toBeUndefined(); + expect(pro.dispatchKvCacheEvictionRisk).toHaveBeenCalledTimes(2); + await flush(); + expect(warn).toHaveBeenCalledWith(expect.stringContaining('receiver down')); + warn.mockRestore(); + }); + + it('holds the evicting alert until three clean snapshots in a row', async () => { + const { service, pro } = setup(); + const evictingStates = () => + pro.dispatchKvCacheEvictionRisk.mock.calls.filter(([e]) => e.reason === 'evicting').map(([e]) => e.active); + for (const delta of [50, 0, 30, 0, 0, 0, 0]) await service.onSnapshot(snapshot({ evictedKeysDelta: delta })); + expect(evictingStates()).toEqual([true, true, true, true, true, false, false]); + }); + + it('forgets the evicting hold when the connection is removed', async () => { + const { service, pro, remove } = setup(); + service.onModuleInit(); + await service.onSnapshot(snapshot({ evictedKeysDelta: 50 })); + remove('c1'); + await service.onSnapshot(snapshot({ evictedKeysDelta: 0 })); + expect(pro.dispatchKvCacheEvictionRisk).toHaveBeenLastCalledWith(expect.objectContaining({ reason: 'evicting', active: false })); + service.onModuleDestroy(); + }); + + it('omits the instance when the connection is unknown', async () => { + const { service, pro, connectionRegistry } = setup(); + connectionRegistry.getConfig.mockReturnValue(null); + await service.onSnapshot(snapshot()); + expect(pro.dispatchKvCacheEvictionRisk.mock.calls[0][0]).not.toHaveProperty('instance'); + }); + + it('clears both eviction alerts and the evicting hold when eviction alerts are disabled', async () => { + const { service, pro, storage } = setup(); + const settings = (evictionAlertEnabled: boolean) => ({ + connectionId: 'c1', + hitRateAlertEnabled: true, + hitRateThreshold: 0.2, + evictionAlertEnabled, + updatedAt: 1, + }); + await service.onSnapshot(snapshot({ evictedKeysDelta: 50 })); + pro.dispatchKvCacheEvictionRisk.mockClear(); + storage.getKvCacheSettings.mockResolvedValue(settings(false)); + await service.onSnapshot(snapshot({ evictedKeysDelta: 50 })); + expect(pro.dispatchKvCacheEvictionRisk.mock.calls.map(([e]) => [e.reason, e.active])).toEqual([ + ['unevictable', false], + ['evicting', false], + ]); + pro.dispatchKvCacheEvictionRisk.mockClear(); + storage.getKvCacheSettings.mockResolvedValue(settings(true)); + await service.onSnapshot(snapshot({ evictedKeysDelta: 0 })); + expect(pro.dispatchKvCacheEvictionRisk).toHaveBeenCalledWith(expect.objectContaining({ reason: 'evicting', active: false })); + }); + + it('clears the evicting alert at once when LMCache is no longer detected', async () => { + const { service, pro } = setup(); + await service.onSnapshot(snapshot({ evictedKeysDelta: 50 })); + await service.onSnapshot(snapshot({ detected: false, evictedKeysDelta: 0 })); + expect(pro.dispatchKvCacheEvictionRisk).toHaveBeenLastCalledWith(expect.objectContaining({ reason: 'evicting', active: false })); + await service.onSnapshot(snapshot({ evictedKeysDelta: 0 })); + expect(pro.dispatchKvCacheEvictionRisk).toHaveBeenLastCalledWith(expect.objectContaining({ reason: 'evicting', active: false })); + }); + + it('keeps the unevictable state when a snapshot has no TTL data', async () => { + const { service, pro } = setup(); + const unevictableStates = () => + pro.dispatchKvCacheEvictionRisk.mock.calls.filter(([e]) => e.reason === 'unevictable').map(([e]) => e.active); + for (const noTtlRatio of [null, 0.95, null, 0.1, null]) await service.onSnapshot(snapshot({ noTtlRatio })); + expect(unevictableStates()).toEqual([false, true, true, false, false]); + }); + + it('forgets the unevictable state when the connection is removed', async () => { + const { service, pro, remove } = setup(); + service.onModuleInit(); + await service.onSnapshot(snapshot({ noTtlRatio: 0.95 })); + remove('c1'); + await service.onSnapshot(snapshot({ noTtlRatio: null })); + expect(pro.dispatchKvCacheEvictionRisk).toHaveBeenCalledWith(expect.objectContaining({ reason: 'unevictable', active: false })); + service.onModuleDestroy(); + }); + + it('sends nothing when unlicensed', async () => { + const { service, pro } = setup({ licensed: false }); + await service.onSnapshot(snapshot()); + expect(pro.dispatchKvCacheEvictionRisk).not.toHaveBeenCalled(); + }); +}); + +describe('KvCacheAlertsService hit rate', () => { + it('dispatches for a window below the floor of tokens with a low rate', async () => { + const { service, pro, storage } = setup(); + await service.evaluateHitRates(1_000_000); + expect(storage.getKvCacheEngineSamples).toHaveBeenCalledWith({ connectionId: 'c1', from: 1_000_000 - 900_000, to: 1_000_000 }); + expect(pro.dispatchKvCacheHitRateLow).toHaveBeenCalledWith({ + connectionId: 'c1', + engineId: 'e1', + engineName: 'vllm', + model: 'm', + hitRate: 0.1, + threshold: 0.2, + requestedTokens: 20000, + windowMs: 900_000, + timestamp: 1_000_000, + instance: { host: 'h', port: 6379 }, + }); + }); + + it('skips windows under the token floor', async () => { + const { service, pro } = setup({ rows: [row({ requestedTokens: 5000, hitTokens: 100 })] }); + await service.evaluateHitRates(1); + expect(pro.dispatchKvCacheHitRateLow).not.toHaveBeenCalled(); + }); + + it('skips disabled engines', async () => { + const { service, pro, prometheus } = setup({ engines: [engine({ enabled: false })] }); + await service.evaluateHitRates(1); + expect(pro.dispatchKvCacheHitRateLow).not.toHaveBeenCalled(); + expect(prometheus.setKvCacheHitRates).toHaveBeenCalledWith('c1', []); + }); + + it('skips samples of engines that are no longer registered', async () => { + const { service, pro } = setup({ rows: [row({ engineId: 'gone' })] }); + await service.evaluateHitRates(1); + expect(pro.dispatchKvCacheHitRateLow).not.toHaveBeenCalled(); + }); + + it('does nothing when unlicensed', async () => { + const { service, pro, storage } = setup({ licensed: false }); + await service.evaluateHitRates(1); + expect(storage.getKvCacheEngineSamples).not.toHaveBeenCalled(); + expect(pro.dispatchKvCacheHitRateLow).not.toHaveBeenCalled(); + }); + + it('continues with the next connection after a storage error', async () => { + const engines = [engine(), engine({ id: 'e2', connectionId: 'c2', name: 'other' })]; + const { service, pro, storage } = setup({ engines }); + storage.getKvCacheEngineSamples + .mockRejectedValueOnce(new Error('boom')) + .mockResolvedValueOnce([row({ connectionId: 'c2', engineId: 'e2' })]); + await service.evaluateHitRates(1); + expect(pro.dispatchKvCacheHitRateLow).toHaveBeenCalledTimes(1); + expect(pro.dispatchKvCacheHitRateLow).toHaveBeenCalledWith(expect.objectContaining({ connectionId: 'c2', engineId: 'e2' })); + }); + + it('only dispatches for the enabled engine when a sibling is disabled', async () => { + const engines = [engine(), engine({ id: 'e2', name: 'off', enabled: false })]; + const { service, pro } = setup({ engines, rows: [row(), row({ engineId: 'e2' })] }); + await service.evaluateHitRates(1); + expect(pro.dispatchKvCacheHitRateLow).toHaveBeenCalledTimes(1); + expect(pro.dispatchKvCacheHitRateLow).toHaveBeenCalledWith(expect.objectContaining({ engineId: 'e1' })); + }); + + it('continues with the next window after a failed dispatch', async () => { + const engines = [engine(), engine({ id: 'e2', name: 'second' })]; + const { service, pro } = setup({ engines, rows: [row(), row({ engineId: 'e2' })] }); + pro.dispatchKvCacheHitRateLow.mockRejectedValueOnce(new Error('boom')); + await service.evaluateHitRates(1); + expect(pro.dispatchKvCacheHitRateLow).toHaveBeenCalledTimes(2); + }); + + it('keeps the alert state while a checked window has too little traffic', async () => { + const { service, pro, storage } = setup(); + await service.evaluateHitRates(1); + storage.getKvCacheEngineSamples.mockResolvedValue([row({ requestedTokens: 5000, hitTokens: 100 })]); + await service.evaluateHitRates(2); + await service.evaluateHitRates(3); + expect(pro.dispatchKvCacheHitRateLow).toHaveBeenCalledTimes(1); + storage.getKvCacheEngineSamples.mockResolvedValue([row({ requestedTokens: 20000, hitTokens: 19000 })]); + await service.evaluateHitRates(4); + expect(pro.dispatchKvCacheHitRateLow).toHaveBeenCalledTimes(2); + expect(pro.dispatchKvCacheHitRateLow).toHaveBeenLastCalledWith(expect.objectContaining({ hitRate: 0.95, timestamp: 4 })); + }); + + it('recovers a checked window once hit rate alerts are disabled', async () => { + const { service, pro, storage } = setup(); + await service.evaluateHitRates(1); + storage.getKvCacheSettings.mockResolvedValue({ connectionId: 'c1', hitRateAlertEnabled: false, hitRateThreshold: 0.3, evictionAlertEnabled: true, updatedAt: 1 }); + await service.evaluateHitRates(2); + expect(pro.dispatchKvCacheHitRateLow).toHaveBeenCalledTimes(2); + expect(pro.dispatchKvCacheHitRateLow).toHaveBeenLastCalledWith( + expect.objectContaining({ engineId: 'e1', model: 'm', hitRate: 1, requestedTokens: 0, timestamp: 2 }), + ); + await service.evaluateHitRates(3); + expect(pro.dispatchKvCacheHitRateLow).toHaveBeenCalledTimes(2); + }); + + it('recovers a checked window when its engine is gone and no samples remain', async () => { + const { service, pro, storage, registry } = setup(); + await service.evaluateHitRates(1); + registry.list.mockReturnValue([]); + registry.get.mockReturnValue(null); + storage.getKvCacheEngineSamples.mockResolvedValue([]); + await service.evaluateHitRates(2); + expect(storage.getKvCacheEngineSamples).toHaveBeenLastCalledWith(expect.objectContaining({ connectionId: 'c1' })); + expect(pro.dispatchKvCacheHitRateLow).toHaveBeenLastCalledWith(expect.objectContaining({ engineName: 'vllm', hitRate: 1 })); + }); + + it('skips a run while the previous one is still in flight', async () => { + const { service, pro, storage } = setup(); + let release: () => void = () => undefined; + pro.dispatchKvCacheHitRateLow.mockImplementationOnce(() => new Promise((resolve) => (release = resolve))); + const first = service.evaluateHitRates(1_000_000); + await new Promise((resolve) => setImmediate(resolve)); + await service.evaluateHitRates(1_060_000); + expect(storage.getKvCacheEngineSamples).toHaveBeenCalledTimes(1); + release(); + await first; + await service.evaluateHitRates(1_120_000); + expect(storage.getKvCacheEngineSamples).toHaveBeenCalledTimes(2); + }); + + it('never rejects when the registry throws', async () => { + const { service, registry } = setup(); + registry.list.mockImplementation(() => { + throw new Error('boom'); + }); + await expect(service.evaluateHitRates(1)).resolves.toBeUndefined(); + }); +}); + +describe('KvCacheAlertsService Prometheus export', () => { + const flush = () => new Promise((resolve) => setImmediate(resolve)); + const perModel = [{ model: 'm', dtype: 'bfloat16', chunksEst: 3, bytesEst: 300 }]; + + it('sets footprint gauges from a snapshot', async () => { + const { service, prometheus, emit } = setup(); + service.onModuleInit(); + emit(snapshot({ perModel })); + await flush(); + service.onModuleDestroy(); + expect(prometheus.setKvCacheFootprint).toHaveBeenCalledWith('c1', perModel); + }); + + it('sets footprint gauges even when eviction alerts are disabled', async () => { + const { service, prometheus } = setup({ + stored: { connectionId: 'c1', hitRateAlertEnabled: true, hitRateThreshold: 0.2, evictionAlertEnabled: false, updatedAt: 1 }, + }); + await service.onSnapshot(snapshot({ perModel })); + expect(prometheus.setKvCacheFootprint).toHaveBeenCalledWith('c1', perModel); + }); + + it('does not export footprint gauges from a partial snapshot', async () => { + const { service, prometheus, pro } = setup(); + await service.onSnapshot(snapshot({ perModel, scanComplete: false })); + expect(prometheus.setKvCacheFootprint).not.toHaveBeenCalled(); + expect(pro.dispatchKvCacheEvictionRisk).toHaveBeenCalled(); + }); + + it('does not set footprint gauges when unlicensed', async () => { + const { service, prometheus } = setup({ licensed: false }); + await service.onSnapshot(snapshot({ perModel })); + expect(prometheus.setKvCacheFootprint).not.toHaveBeenCalled(); + }); + + it('clears exported series once the license lapses', async () => { + const { service, prometheus, license } = setup(); + await service.onSnapshot(snapshot({ perModel })); + await service.evaluateHitRates(1); + expect(prometheus.clearAllKvCache).not.toHaveBeenCalled(); + license.hasFeature.mockReturnValue(false); + await service.evaluateHitRates(2); + expect(prometheus.clearAllKvCache).toHaveBeenCalledTimes(1); + await service.onSnapshot(snapshot({ perModel })); + expect(prometheus.clearAllKvCache).toHaveBeenCalledTimes(2); + }); + + it('sets hit rates with the engine name', async () => { + const { service, prometheus } = setup(); + await service.evaluateHitRates(1); + expect(prometheus.setKvCacheHitRates).toHaveBeenCalledWith('c1', [{ engine: 'vllm', model: 'm', hitRate: 0.1 }]); + }); + + it('sets hit rates when hit rate alerts are disabled', async () => { + const { service, prometheus, pro } = setup({ + stored: { connectionId: 'c1', hitRateAlertEnabled: false, hitRateThreshold: 0.2, evictionAlertEnabled: true, updatedAt: 1 }, + }); + await service.evaluateHitRates(1); + expect(pro.dispatchKvCacheHitRateLow).not.toHaveBeenCalled(); + expect(prometheus.setKvCacheHitRates).toHaveBeenCalledWith('c1', [{ engine: 'vllm', model: 'm', hitRate: 0.1 }]); + }); + + it('sets hit rates even when dispatch fails', async () => { + const { service, prometheus, pro } = setup(); + pro.dispatchKvCacheHitRateLow.mockRejectedValue(new Error('boom')); + await service.evaluateHitRates(1); + expect(prometheus.setKvCacheHitRates).toHaveBeenCalledWith('c1', [{ engine: 'vllm', model: 'm', hitRate: 0.1 }]); + }); + + it('exports an empty list when there are no windows', async () => { + const { service, prometheus } = setup({ rows: [] }); + await service.evaluateHitRates(1); + expect(prometheus.setKvCacheHitRates).toHaveBeenCalledWith('c1', []); + }); + + it('excludes windows with a null hit rate', async () => { + const { service, prometheus } = setup({ rows: [row({ requestedTokens: 0, hitTokens: 0 })] }); + await service.evaluateHitRates(1); + expect(prometheus.setKvCacheHitRates).toHaveBeenCalledWith('c1', []); + }); + + it('excludes windows of disabled engines', async () => { + const engines = [engine(), engine({ id: 'e2', name: 'off', enabled: false })]; + const { service, prometheus } = setup({ engines, rows: [row(), row({ engineId: 'e2' })] }); + await service.evaluateHitRates(1); + expect(prometheus.setKvCacheHitRates).toHaveBeenCalledWith('c1', [{ engine: 'vllm', model: 'm', hitRate: 0.1 }]); + }); + + it('excludes windows of engines that are no longer registered', async () => { + const { service, prometheus } = setup({ rows: [row({ engineId: 'gone' })] }); + await service.evaluateHitRates(1); + expect(prometheus.setKvCacheHitRates).toHaveBeenCalledWith('c1', []); + }); + + it('exports an empty list once the last engine of a connection is removed', async () => { + const { service, prometheus, registry } = setup(); + await service.evaluateHitRates(1); + registry.list.mockReturnValue([]); + registry.get.mockReturnValue(null); + prometheus.setKvCacheHitRates.mockClear(); + await service.evaluateHitRates(2); + expect(prometheus.setKvCacheHitRates).toHaveBeenCalledWith('c1', []); + prometheus.setKvCacheHitRates.mockClear(); + await service.evaluateHitRates(3); + expect(prometheus.setKvCacheHitRates).not.toHaveBeenCalled(); + }); + + it('does not export hit rates when unlicensed', async () => { + const { service, prometheus } = setup({ licensed: false }); + await service.evaluateHitRates(1); + expect(prometheus.setKvCacheHitRates).not.toHaveBeenCalled(); + }); + + it('clears the connection series on connection removal', () => { + const { service, prometheus, remove } = setup(); + service.onModuleInit(); + remove('c1'); + service.onModuleDestroy(); + expect(prometheus.clearKvCache).toHaveBeenCalledWith('c1'); + }); +}); diff --git a/proprietary/kv-cache/__tests__/kv-cache-dto.spec.ts b/proprietary/kv-cache/__tests__/kv-cache-dto.spec.ts new file mode 100644 index 000000000..607ca2e65 --- /dev/null +++ b/proprietary/kv-cache/__tests__/kv-cache-dto.spec.ts @@ -0,0 +1,57 @@ +import { plainToInstance } from 'class-transformer'; +import { validate } from 'class-validator'; +import { CreateKvCacheEngineDto, UpdateKvCacheEngineDto } from '../dto/kv-cache-engine.dto'; +import { UpdateKvCacheSettingsDto } from '../dto/kv-cache-settings.dto'; + +async function errorsFor(cls: new () => T, body: Record): Promise { + const errors = await validate(plainToInstance(cls, body)); + return errors.map((error) => error.property); +} + +describe('KV cache DTO null handling', () => { + it('rejects null for optional engine update fields', async () => { + for (const field of ['name', 'scrapeUrl', 'enabled']) { + expect(await errorsFor(UpdateKvCacheEngineDto, { [field]: null })).toEqual([field]); + } + }); + + it('still allows clearing the scrape auth header with null and omitting fields', async () => { + expect(await errorsFor(UpdateKvCacheEngineDto, { scrapeAuthHeader: null })).toEqual([]); + expect(await errorsFor(UpdateKvCacheEngineDto, {})).toEqual([]); + }); + + it('rejects null for optional engine create fields', async () => { + const base = { name: 'e', source: 'otlp' }; + for (const field of ['scrapeAuthHeader', 'otlpEngineId', 'enabled']) { + expect(await errorsFor(CreateKvCacheEngineDto, { ...base, [field]: null })).toEqual([field]); + } + expect(await errorsFor(CreateKvCacheEngineDto, base)).toEqual([]); + }); + + it('rejects null for settings fields', async () => { + for (const field of ['hitRateAlertEnabled', 'hitRateThreshold', 'evictionAlertEnabled']) { + expect(await errorsFor(UpdateKvCacheSettingsDto, { [field]: null })).toEqual([field]); + } + expect(await errorsFor(UpdateKvCacheSettingsDto, {})).toEqual([]); + }); + + it('rejects credentials in the scrape URL and points at the auth header', async () => { + const url = 'http://admin:secret@vllm:8000/metrics'; + for (const [cls, body] of [ + [CreateKvCacheEngineDto, { name: 'e', source: 'scrape', scrapeUrl: url }], + [UpdateKvCacheEngineDto, { scrapeUrl: url }], + [UpdateKvCacheEngineDto, { scrapeUrl: 'http://token@vllm/metrics' }], + ] as const) { + const errors = await validate(plainToInstance(cls as new () => object, body)); + expect(errors.map((error) => error.property)).toEqual(['scrapeUrl']); + expect(Object.values(errors[0].constraints ?? {}).join(' ')).toContain('scrapeAuthHeader'); + } + expect(await errorsFor(UpdateKvCacheEngineDto, { scrapeUrl: 'http://vllm:8000/metrics' })).toEqual([]); + }); + + it('accepts scrape hosts with underscores such as compose service names', async () => { + const scrapeUrl = 'http://vllm_engine:8000/metrics'; + expect(await errorsFor(CreateKvCacheEngineDto, { name: 'e', source: 'scrape', scrapeUrl })).toEqual([]); + expect(await errorsFor(UpdateKvCacheEngineDto, { scrapeUrl })).toEqual([]); + }); +}); diff --git a/proprietary/kv-cache/__tests__/kv-cache-engine-registry.spec.ts b/proprietary/kv-cache/__tests__/kv-cache-engine-registry.spec.ts new file mode 100644 index 000000000..8a3aa7e65 --- /dev/null +++ b/proprietary/kv-cache/__tests__/kv-cache-engine-registry.spec.ts @@ -0,0 +1,285 @@ +import { NotFoundException } from '@nestjs/common'; +import { KV_CACHE_LAST_SEEN_PERSIST_MS, KvCacheEngineRegistry, REMOVAL_TOMBSTONE_MS } from '../kv-cache-engine-registry'; + +const engine = (id: string, connectionId: string, otlpEngineId: string | null = null) => ({ + id, connectionId, name: id, source: 'otlp' as const, scrapeUrl: null, scrapeAuthHeader: null, + scrapeAuthEncrypted: false, otlpEngineId, enabled: true, createdAt: 1, lastSeenAt: null, lastError: null, +}); + +function setup(initial = [engine('a', 'c1', 'lmc-a'), engine('b', 'c2')]) { + const storage = { + getKvCacheEngines: jest.fn().mockResolvedValue(initial), + saveKvCacheEngine: jest.fn().mockImplementation(async (e) => e), + deleteKvCacheEngine: jest.fn().mockResolvedValue(true), + } as any; + let removal: (id: string) => unknown = () => undefined; + const footprint = { onConnectionRemoval: jest.fn((l) => (removal = l)) } as any; + const samples = { forgetEngine: jest.fn(), idle: jest.fn().mockResolvedValue(undefined) } as any; + const registry = new KvCacheEngineRegistry(storage, footprint, samples); + return { registry, storage, samples, removal: (id: string) => removal(id) }; +} + +describe('KvCacheEngineRegistry', () => { + it('loads engines from storage and finds them', async () => { + const { registry } = setup(); + await registry.onModuleInit(); + expect(registry.list()).toHaveLength(2); + expect(registry.list('c1').map((e) => e.id)).toEqual(['a']); + expect(registry.get('b')?.connectionId).toBe('c2'); + expect(registry.get('zzz')).toBeNull(); + expect(registry.byOtlpId('lmc-a')?.id).toBe('a'); + expect(registry.byOtlpId('nope')).toBeNull(); + }); + + it('writes storage before memory', async () => { + const { registry, storage } = setup([]); + await registry.onModuleInit(); + storage.saveKvCacheEngine.mockRejectedValueOnce(new Error('disk')); + await expect(registry.save(engine('x', 'c1'))).rejects.toThrow('disk'); + expect(registry.get('x')).toBeNull(); + await registry.save(engine('x', 'c1')); + expect(registry.get('x')).not.toBeNull(); + }); + + it('rejects saves for a connection that has been removed', async () => { + const { registry, storage, removal } = setup(); + await registry.onModuleInit(); + await removal('c1'); + await expect(registry.save(engine('late', 'c1'))).rejects.toBeInstanceOf(NotFoundException); + expect(storage.saveKvCacheEngine).not.toHaveBeenCalled(); + expect(registry.get('late')).toBeNull(); + await expect(registry.save(engine('ok', 'c2'))).resolves.toMatchObject({ id: 'ok' }); + }); + + it('drops an engine whose save was in flight when its connection was removed', async () => { + const { registry, storage, removal } = setup(); + await registry.onModuleInit(); + let finishSave: () => void = () => undefined; + storage.saveKvCacheEngine.mockImplementationOnce((e: any) => new Promise((resolve) => (finishSave = () => resolve(e)))); + const pending = registry.save(engine('racing', 'c1')); + await new Promise((resolve) => setImmediate(resolve)); + await removal('c1'); + finishSave(); + await expect(pending).rejects.toBeInstanceOf(NotFoundException); + expect(storage.deleteKvCacheEngine).toHaveBeenCalledWith('racing'); + expect(registry.get('racing')).toBeNull(); + }); + + it('expires removal tombstones once they are older than the tombstone window', async () => { + const { registry, removal } = setup(); + await registry.onModuleInit(); + const now = jest.spyOn(Date, 'now').mockReturnValue(1_000); + try { + await registry.remove('a'); + await removal('c2'); + await expect(registry.save(engine('a', 'c1'))).rejects.toBeInstanceOf(NotFoundException); + await expect(registry.save(engine('x', 'c2'))).rejects.toBeInstanceOf(NotFoundException); + now.mockReturnValue(1_000 + REMOVAL_TOMBSTONE_MS + 1); + await registry.remove('zzz'); + await removal('c9'); + await expect(registry.save(engine('a', 'c1'))).resolves.toMatchObject({ id: 'a' }); + await expect(registry.save(engine('x', 'c2'))).resolves.toMatchObject({ id: 'x' }); + } finally { + now.mockRestore(); + } + }); + + it('records results in memory and persists them', async () => { + const { registry, storage } = setup(); + await registry.onModuleInit(); + registry.recordResult('a', { lastSeenAt: 99, lastError: 'boom' }); + expect(registry.get('a')).toMatchObject({ lastSeenAt: 99, lastError: 'boom' }); + await new Promise((resolve) => setImmediate(resolve)); + expect(storage.saveKvCacheEngine).toHaveBeenCalledWith(expect.objectContaining({ id: 'a', lastError: 'boom' })); + }); + + it('persists repeated successful results only when lastSeenAt is stale or the error changes', async () => { + const { registry, storage } = setup(); + await registry.onModuleInit(); + const flush = () => new Promise((resolve) => setImmediate(resolve)); + registry.recordResult('a', { lastSeenAt: 1_000, lastError: null }); + await flush(); + expect(storage.saveKvCacheEngine).toHaveBeenCalledTimes(1); + registry.recordResult('a', { lastSeenAt: 31_000, lastError: null }); + registry.recordResult('a', { lastSeenAt: 61_000, lastError: null }); + await flush(); + expect(storage.saveKvCacheEngine).toHaveBeenCalledTimes(1); + expect(registry.get('a')?.lastSeenAt).toBe(61_000); + registry.recordResult('a', { lastSeenAt: 91_000, lastError: 'timeout' }); + await flush(); + expect(storage.saveKvCacheEngine).toHaveBeenCalledTimes(2); + registry.recordResult('a', { lastError: 'timeout' }); + await flush(); + expect(storage.saveKvCacheEngine).toHaveBeenCalledTimes(2); + registry.recordResult('a', { lastSeenAt: 91_000 + KV_CACHE_LAST_SEEN_PERSIST_MS, lastError: 'timeout' }); + await flush(); + expect(storage.saveKvCacheEngine).toHaveBeenCalledTimes(3); + }); + + it('keeps lastSeenAt when a result omits it and survives persist failures', async () => { + const { registry, storage } = setup(); + await registry.onModuleInit(); + registry.recordResult('a', { lastSeenAt: 5, lastError: null }); + storage.saveKvCacheEngine.mockRejectedValueOnce(new Error('disk')); + registry.recordResult('a', { lastError: 'x' }); + await new Promise((resolve) => setImmediate(resolve)); + expect(registry.get('a')).toMatchObject({ lastSeenAt: 5, lastError: 'x' }); + }); + + it('persists a result raised during an update with the updated fields', async () => { + const { registry, storage } = setup(); + await registry.onModuleInit(); + let finishSave: () => void = () => undefined; + storage.saveKvCacheEngine.mockImplementationOnce((e: any) => new Promise((resolve) => (finishSave = () => resolve(e)))); + const pending = registry.save({ ...engine('a', 'c1', 'lmc-a'), name: 'renamed' }); + registry.recordResult('a', { lastSeenAt: 7, lastError: null }); + await new Promise((resolve) => setImmediate(resolve)); + finishSave(); + await pending; + await new Promise((resolve) => setImmediate(resolve)); + expect(storage.saveKvCacheEngine).toHaveBeenLastCalledWith(expect.objectContaining({ id: 'a', name: 'renamed' })); + expect(registry.get('a')?.name).toBe('renamed'); + }); + + it('keeps the probe result of a URL change over an old-URL scrape that lands during the save', async () => { + const scraped = { ...engine('s', 'c1'), source: 'scrape' as const, scrapeUrl: 'http://old/metrics' }; + const { registry, storage } = setup([scraped as any]); + await registry.onModuleInit(); + let finishSave: () => void = () => undefined; + storage.saveKvCacheEngine.mockImplementationOnce((e: any) => new Promise((resolve) => (finishSave = () => resolve(e)))); + const pending = registry.save({ ...scraped, scrapeUrl: 'http://new/metrics', lastSeenAt: 9, lastError: null } as any); + await new Promise((resolve) => setImmediate(resolve)); + registry.recordResult('s', { lastError: 'HTTP 401' }); + finishSave(); + await pending; + await new Promise((resolve) => setImmediate(resolve)); + expect(registry.get('s')).toMatchObject({ scrapeUrl: 'http://new/metrics', lastSeenAt: 9, lastError: null }); + expect(storage.saveKvCacheEngine).toHaveBeenLastCalledWith(expect.objectContaining({ scrapeUrl: 'http://new/metrics', lastError: null })); + }); + + it('keeps a scrape result that lands during an update', async () => { + const { registry, storage } = setup(); + await registry.onModuleInit(); + let finishSave: () => void = () => undefined; + storage.saveKvCacheEngine.mockImplementationOnce((e: any) => new Promise((resolve) => (finishSave = () => resolve(e)))); + const pending = registry.save({ ...engine('a', 'c1', 'lmc-a'), name: 'renamed' }); + await new Promise((resolve) => setImmediate(resolve)); + registry.recordResult('a', { lastSeenAt: 7, lastError: 'boom' }); + finishSave(); + await pending; + await new Promise((resolve) => setImmediate(resolve)); + expect(registry.get('a')).toMatchObject({ name: 'renamed', lastSeenAt: 7, lastError: 'boom' }); + expect(storage.saveKvCacheEngine).toHaveBeenLastCalledWith(expect.objectContaining({ name: 'renamed', lastSeenAt: 7, lastError: 'boom' })); + }); + + it('waits for in-flight writes on connection removal and never writes the engine again', async () => { + const { registry, storage, removal } = setup(); + await registry.onModuleInit(); + let finishSave: () => void = () => undefined; + storage.saveKvCacheEngine.mockImplementationOnce((e: any) => new Promise((resolve) => (finishSave = () => resolve(e)))); + const pending = registry.save({ ...engine('a', 'c1', 'lmc-a'), name: 'renamed' }); + await new Promise((resolve) => setImmediate(resolve)); + let removed = false; + const removing = Promise.resolve(removal('c1')).then(() => (removed = true)); + await new Promise((resolve) => setImmediate(resolve)); + expect(removed).toBe(false); + finishSave(); + await expect(pending).rejects.toBeInstanceOf(NotFoundException); + await removing; + expect(storage.deleteKvCacheEngine).toHaveBeenCalledWith('a'); + expect(registry.get('a')).toBeNull(); + await expect(registry.save(engine('a', 'c1', 'lmc-a'))).rejects.toBeInstanceOf(NotFoundException); + registry.recordResult('a', { lastSeenAt: 9, lastError: null }); + await new Promise((resolve) => setImmediate(resolve)); + expect(storage.saveKvCacheEngine).toHaveBeenCalledTimes(1); + expect(registry.get('a')).toBeNull(); + }); + + it('does not resurrect an engine removed while a result was pending', async () => { + const { registry, storage } = setup(); + await registry.onModuleInit(); + registry.recordResult('a', { lastSeenAt: 7, lastError: null }); + await registry.remove('a'); + registry.recordResult('a', { lastSeenAt: 8, lastError: null }); + await new Promise((resolve) => setImmediate(resolve)); + expect(storage.saveKvCacheEngine).not.toHaveBeenCalled(); + expect(registry.get('a')).toBeNull(); + }); + + it('rejects a save whose write was in flight when the engine was deleted', async () => { + const { registry, storage } = setup(); + await registry.onModuleInit(); + let finishSave: () => void = () => undefined; + storage.saveKvCacheEngine.mockImplementationOnce((e: any) => new Promise((resolve) => (finishSave = () => resolve(e)))); + const saving = registry.save({ ...engine('a', 'c1', 'lmc-a'), name: 'renamed' }); + await new Promise((resolve) => setImmediate(resolve)); + const removing = registry.remove('a'); + finishSave(); + await expect(saving).rejects.toBeInstanceOf(NotFoundException); + await expect(removing).resolves.toBe(true); + expect(storage.deleteKvCacheEngine).toHaveBeenCalledWith('a'); + expect(registry.get('a')).toBeNull(); + }); + + it('rejects a save queued behind a pending delete and does not re-insert the engine', async () => { + const { registry, storage } = setup(); + await registry.onModuleInit(); + let finishDelete: () => void = () => undefined; + storage.deleteKvCacheEngine.mockImplementationOnce(() => new Promise((resolve) => (finishDelete = () => resolve(true)))); + const removing = registry.remove('a'); + const saving = registry.save({ ...engine('a', 'c1', 'lmc-a'), name: 'renamed' }); + await new Promise((resolve) => setImmediate(resolve)); + finishDelete(); + await removing; + await expect(saving).rejects.toBeInstanceOf(NotFoundException); + expect(storage.saveKvCacheEngine).not.toHaveBeenCalled(); + expect(registry.get('a')).toBeNull(); + }); + + it('keeps a scrape result recorded after a save was queued behind another write', async () => { + const { registry, storage } = setup(); + await registry.onModuleInit(); + let finishPersist: () => void = () => undefined; + storage.saveKvCacheEngine.mockImplementationOnce((e: any) => new Promise((resolve) => (finishPersist = () => resolve(e)))); + registry.recordResult('a', { lastSeenAt: 5, lastError: null }); + const stale = { ...registry.get('a')!, name: 'renamed' }; + const saving = registry.save(stale); + registry.recordResult('a', { lastSeenAt: 9, lastError: 'boom' }); + await new Promise((resolve) => setImmediate(resolve)); + finishPersist(); + await saving; + await new Promise((resolve) => setImmediate(resolve)); + expect(registry.get('a')).toMatchObject({ name: 'renamed', lastSeenAt: 9, lastError: 'boom' }); + }); + + it('removes from storage and memory', async () => { + const { registry, storage } = setup(); + await registry.onModuleInit(); + await expect(registry.remove('a')).resolves.toBe(true); + expect(storage.deleteKvCacheEngine).toHaveBeenCalledWith('a'); + expect(registry.get('a')).toBeNull(); + }); + + it('keeps an engine saveable and persisting results after a failed delete', async () => { + const { registry, storage } = setup(); + await registry.onModuleInit(); + storage.deleteKvCacheEngine.mockRejectedValueOnce(new Error('disk')); + await expect(registry.remove('a')).rejects.toThrow('disk'); + expect(registry.get('a')).not.toBeNull(); + await registry.save({ ...engine('a', 'c1', 'lmc-a'), name: 'renamed' }); + expect(storage.saveKvCacheEngine).toHaveBeenLastCalledWith(expect.objectContaining({ id: 'a', name: 'renamed' })); + expect(registry.get('a')?.name).toBe('renamed'); + registry.recordResult('a', { lastSeenAt: 11, lastError: null }); + await new Promise((resolve) => setImmediate(resolve)); + expect(storage.saveKvCacheEngine).toHaveBeenLastCalledWith(expect.objectContaining({ id: 'a', lastSeenAt: 11 })); + }); + + it('drops only the removed connection engines on broadcast', async () => { + const { registry, samples, removal } = setup(); + await registry.onModuleInit(); + removal('c1'); + expect(registry.list().map((e) => e.id)).toEqual(['b']); + expect(samples.forgetEngine).toHaveBeenCalledTimes(1); + expect(samples.forgetEngine).toHaveBeenCalledWith('a'); + }); +}); diff --git a/proprietary/kv-cache/__tests__/kv-cache-engines.service.spec.ts b/proprietary/kv-cache/__tests__/kv-cache-engines.service.spec.ts new file mode 100644 index 000000000..80cb8b1be --- /dev/null +++ b/proprietary/kv-cache/__tests__/kv-cache-engines.service.spec.ts @@ -0,0 +1,324 @@ +import { BadRequestException, ConflictException, NotFoundException } from '@nestjs/common'; +import { getEncryptionService } from '@app/common/utils/encryption'; +import { fetchMetricsText, ScrapeError } from '../metrics-fetch'; +import { KvCacheEngineRegistry } from '../kv-cache-engine-registry'; +import { KvCacheEnginesService } from '../kv-cache-engines.service'; + +jest.mock('../metrics-fetch', () => ({ ...jest.requireActual('../metrics-fetch'), fetchMetricsText: jest.fn() })); +jest.mock('@app/common/utils/encryption', () => ({ ...jest.requireActual('@app/common/utils/encryption'), getEncryptionService: jest.fn() })); + +const LMCACHE_BODY = 'lmcache:num_hit_tokens_total{model_name="m",worker_id="0",role="kv_both",served_model_name="m"} 5\n'; +const fetchMock = fetchMetricsText as jest.Mock; +const encryptionMock = getEncryptionService as jest.Mock; + +function setup(encrypted = true) { + const store = new Map(); + const registry = { + list: jest.fn((connectionId?: string) => [...store.values()].filter((e) => !connectionId || e.connectionId === connectionId)), + get: jest.fn((id: string) => store.get(id) ?? null), + byOtlpId: jest.fn((otlp: string) => [...store.values()].find((e) => e.otlpEngineId === otlp) ?? null), + save: jest.fn(async (e: any) => (store.set(e.id, e), e)), + remove: jest.fn(async (id: string) => store.delete(id)), + } as any; + const encryption = encrypted ? { encrypt: (s: string) => 'enc:' + s, decrypt: (s: string) => s.slice(4) } : null; + encryptionMock.mockReturnValue(encryption); + const connections = { + getConfig: jest.fn((id: string) => (id === 'gone' ? null : { id })), + } as any; + const samples = { deleteEngine: jest.fn().mockResolvedValue(undefined), resetBaselines: jest.fn() } as any; + return { service: new KvCacheEnginesService(registry, connections, samples), registry, store, connections, samples }; +} + +describe('KvCacheEnginesService', () => { + beforeEach(() => fetchMock.mockReset().mockResolvedValue(LMCACHE_BODY)); + + it('creates a scrape engine with an encrypted header', async () => { + const { service, store } = setup(); + const created = await service.create('c1', { name: 'n', source: 'scrape', scrapeUrl: 'http://x/metrics', scrapeAuthHeader: 'Bearer t' }); + expect(fetchMock).toHaveBeenCalledWith('http://x/metrics', 'Bearer t'); + const stored = store.get(created.id); + expect(stored).toMatchObject({ scrapeAuthHeader: 'enc:Bearer t', scrapeAuthEncrypted: true, lastError: null }); + expect(stored.lastSeenAt).toEqual(expect.any(Number)); + expect(created.hasScrapeAuth).toBe(true); + expect(created).not.toHaveProperty('scrapeAuthHeader'); + }); + + it('refuses to create an engine for a connection that does not exist', async () => { + const { service, registry } = setup(); + for (const dto of [{ name: 'n', source: 'scrape', scrapeUrl: 'http://x' }, { name: 'o', source: 'otlp' }] as const) { + await expect(service.create('gone', dto as any)).rejects.toBeInstanceOf(NotFoundException); + } + expect(fetchMock).not.toHaveBeenCalled(); + expect(registry.save).not.toHaveBeenCalled(); + }); + + it('stores the header in plaintext without an encryption service', async () => { + const { service, store } = setup(false); + const created = await service.create('c1', { name: 'n', source: 'scrape', scrapeUrl: 'http://x', scrapeAuthHeader: 'Bearer t' }); + expect(store.get(created.id)).toMatchObject({ scrapeAuthHeader: 'Bearer t', scrapeAuthEncrypted: false }); + }); + + it('flags a body without lmcache series', async () => { + fetchMock.mockResolvedValue('other_metric 1\n'); + const { service, store } = setup(); + const created = await service.create('c1', { name: 'n', source: 'scrape', scrapeUrl: 'http://x' }); + expect(store.get(created.id)).toMatchObject({ lastError: 'no lmcache metrics found' }); + expect(store.get(created.id).lastSeenAt).toEqual(expect.any(Number)); + }); + + it('rejects a failed test scrape and saves nothing', async () => { + fetchMock.mockRejectedValue(new ScrapeError('HTTP 401')); + const { service, registry } = setup(); + const error = await service.create('c1', { name: 'n', source: 'scrape', scrapeUrl: 'http://x' }).catch((e) => e); + expect(error).toBeInstanceOf(BadRequestException); + expect(error.message).toBe('Test scrape failed: HTTP 401'); + expect(registry.save).not.toHaveBeenCalled(); + }); + + it('requires a URL for scrape engines', async () => { + const { service } = setup(); + await expect(service.create('c1', { name: 'n', source: 'scrape' })).rejects.toBeInstanceOf(BadRequestException); + }); + + it('rejects a concurrent create with the same otlp engine id', async () => { + const { service, registry, store } = setup(); + let finish: () => void = () => undefined; + registry.save.mockImplementationOnce((e: any) => new Promise((resolve) => (finish = () => resolve((store.set(e.id, e), e))))); + const first = service.create('c1', { name: 'a', source: 'otlp', otlpEngineId: 'shared' }); + await expect(service.create('c2', { name: 'b', source: 'otlp', otlpEngineId: 'shared' })).rejects.toBeInstanceOf(ConflictException); + finish(); + await expect(first).resolves.toMatchObject({ otlpEngineId: 'shared' }); + await expect(service.create('c2', { name: 'b', source: 'otlp', otlpEngineId: 'shared' })).rejects.toBeInstanceOf(ConflictException); + }); + + it('releases a reserved otlp engine id when the save fails', async () => { + const { service, registry } = setup(); + registry.save.mockRejectedValueOnce(new Error('disk')); + await expect(service.create('c1', { name: 'a', source: 'otlp', otlpEngineId: 'shared' })).rejects.toThrow('disk'); + await expect(service.create('c1', { name: 'a', source: 'otlp', otlpEngineId: 'shared' })).resolves.toMatchObject({ otlpEngineId: 'shared' }); + }); + + it('generates and de-duplicates otlp engine ids', async () => { + const { service, store } = setup(); + const generated = await service.create('c1', { name: 'n', source: 'otlp', scrapeUrl: 'http://ignored' }); + expect(generated.otlpEngineId).toMatch(/^lmc-[0-9a-f]{12}$/); + expect(store.get(generated.id)).toMatchObject({ scrapeUrl: null, scrapeAuthHeader: null }); + expect(fetchMock).not.toHaveBeenCalled(); + await service.create('c1', { name: 'n2', source: 'otlp', otlpEngineId: 'mine' }); + await expect(service.create('c2', { name: 'n', source: 'otlp', otlpEngineId: 'mine' })).rejects.toBeInstanceOf(ConflictException); + await expect(service.create('c1', { name: 'n3', source: 'otlp', otlpEngineId: 'bad id!' })).rejects.toBeInstanceOf(BadRequestException); + }); + + it('rejects a duplicate engine name on the same connection', async () => { + const { service, registry } = setup(); + const first = await service.create('c1', { name: 'vllm', source: 'otlp' }); + await expect(service.create('c1', { name: 'vllm', source: 'scrape', scrapeUrl: 'http://x' })).rejects.toBeInstanceOf(ConflictException); + await expect(service.create('c2', { name: 'vllm', source: 'otlp' })).resolves.toMatchObject({ name: 'vllm' }); + const second = await service.create('c1', { name: 'sglang', source: 'otlp' }); + await expect(service.update('c1', second.id, { name: 'vllm' })).rejects.toBeInstanceOf(ConflictException); + await expect(service.update('c1', first.id, { name: 'vllm' })).resolves.toMatchObject({ name: 'vllm' }); + expect(registry.save).toHaveBeenCalledTimes(4); + }); + + it('rejects a concurrent create with the same name while the first is still scraping', async () => { + let finish: (body: string) => void = () => undefined; + fetchMock.mockImplementationOnce(() => new Promise((resolve) => (finish = resolve))); + const { service, store } = setup(); + const first = service.create('c1', { name: 'vllm', source: 'scrape', scrapeUrl: 'http://x' }); + await expect(service.create('c1', { name: 'vllm', source: 'otlp' })).rejects.toBeInstanceOf(ConflictException); + finish(LMCACHE_BODY); + await expect(first).resolves.toMatchObject({ name: 'vllm' }); + expect(store.size).toBe(1); + await expect(service.create('c1', { name: 'vllm', source: 'otlp' })).rejects.toBeInstanceOf(ConflictException); + }); + + it('releases a reserved name when the create fails', async () => { + fetchMock.mockRejectedValueOnce(new ScrapeError('HTTP 401')); + const { service } = setup(); + await expect(service.create('c1', { name: 'vllm', source: 'scrape', scrapeUrl: 'http://x' })).rejects.toBeInstanceOf(BadRequestException); + await expect(service.create('c1', { name: 'vllm', source: 'otlp' })).resolves.toMatchObject({ name: 'vllm' }); + }); + + it('resolves the stored auth header', () => { + const base = { id: 'e', scrapeAuthHeader: 'Bearer t', scrapeAuthEncrypted: false } as any; + expect(setup(false).service.authHeaderFor(base)).toBe('Bearer t'); + expect(setup(true).service.authHeaderFor({ ...base, scrapeAuthHeader: 'enc:Bearer t', scrapeAuthEncrypted: true })).toBe('Bearer t'); + expect(setup(true).service.authHeaderFor({ ...base, scrapeAuthHeader: null })).toBeNull(); + expect(setup(false).service.authHeaderFor({ ...base, scrapeAuthEncrypted: true })).toBeNull(); + }); + + it('returns null when decryption throws', () => { + const { service } = setup(); + encryptionMock.mockReturnValue({ decrypt: () => { throw new Error('bad'); } }); + expect(service.authHeaderFor({ id: 'e', scrapeAuthHeader: 'x', scrapeAuthEncrypted: true } as any)).toBeNull(); + }); + + it('updates only engines of the connection', async () => { + const { service } = setup(); + const created = await service.create('c1', { name: 'n', source: 'scrape', scrapeUrl: 'http://x', scrapeAuthHeader: 'Bearer t' }); + await expect(service.update('c2', created.id, { name: 'z' })).rejects.toBeInstanceOf(NotFoundException); + await expect(service.update('c1', 'missing', { name: 'z' })).rejects.toBeInstanceOf(NotFoundException); + }); + + it('clears, replaces and keeps the auth header', async () => { + const { service, store } = setup(); + const created = await service.create('c1', { name: 'n', source: 'scrape', scrapeUrl: 'http://x', scrapeAuthHeader: 'Bearer t' }); + await service.update('c1', created.id, { name: 'renamed' }); + expect(store.get(created.id)).toMatchObject({ name: 'renamed', scrapeAuthHeader: 'enc:Bearer t' }); + const cleared = await service.update('c1', created.id, { scrapeAuthHeader: null }); + expect(cleared.hasScrapeAuth).toBe(false); + expect(store.get(created.id)).toMatchObject({ scrapeAuthHeader: null, scrapeAuthEncrypted: false }); + await service.update('c1', created.id, { scrapeAuthHeader: 'Bearer n' }); + expect(store.get(created.id).scrapeAuthHeader).toBe('enc:Bearer n'); + }); + + it('re-tests when the URL changes and not for a toggle', async () => { + const { service, store } = setup(); + const created = await service.create('c1', { name: 'n', source: 'scrape', scrapeUrl: 'http://x', scrapeAuthHeader: 'Bearer t' }); + fetchMock.mockClear(); + await service.update('c1', created.id, { enabled: false }); + expect(fetchMock).not.toHaveBeenCalled(); + expect(store.get(created.id).enabled).toBe(false); + await service.update('c1', created.id, { scrapeUrl: 'http://x/y' }); + expect(fetchMock).toHaveBeenCalledWith('http://x/y', 'Bearer t'); + fetchMock.mockRejectedValue(new ScrapeError('timeout')); + await expect(service.update('c1', created.id, { scrapeUrl: 'http://x/z' })).rejects.toThrow('Test scrape failed: timeout'); + expect(store.get(created.id).scrapeUrl).toBe('http://x/y'); + }); + + it('reports an undecryptable stored header when the URL changes instead of scraping without it', async () => { + const { service, store } = setup(); + const created = await service.create('c1', { name: 'n', source: 'scrape', scrapeUrl: 'http://x', scrapeAuthHeader: 'Bearer t' }); + fetchMock.mockClear(); + encryptionMock.mockReturnValue({ decrypt: () => { throw new Error('bad'); } }); + await expect(service.update('c1', created.id, { scrapeUrl: 'http://x/y' })).rejects.toThrow( + new BadRequestException('Test scrape failed: auth header could not be decrypted'), + ); + expect(fetchMock).not.toHaveBeenCalled(); + expect(store.get(created.id).scrapeUrl).toBe('http://x'); + }); + + it('removes engines of the connection', async () => { + const { service, store, samples } = setup(); + const created = await service.create('c1', { name: 'n', source: 'scrape', scrapeUrl: 'http://x' }); + await expect(service.remove('c2', created.id)).rejects.toBeInstanceOf(NotFoundException); + expect(samples.deleteEngine).not.toHaveBeenCalled(); + await service.remove('c1', created.id); + expect(store.has(created.id)).toBe(false); + expect(samples.deleteEngine).toHaveBeenCalledWith(created.id); + }); + + it('retries the sample cleanup after a failed delete without leaking it to other connections', async () => { + const { service, store, samples } = setup(); + const created = await service.create('c1', { name: 'n', source: 'scrape', scrapeUrl: 'http://x' }); + samples.deleteEngine.mockRejectedValueOnce(new Error('db down')); + await expect(service.remove('c1', created.id)).rejects.toThrow('db down'); + expect(store.has(created.id)).toBe(false); + await expect(service.remove('c2', created.id)).rejects.toBeInstanceOf(NotFoundException); + expect(samples.deleteEngine).toHaveBeenCalledTimes(1); + await service.remove('c1', created.id); + expect(samples.deleteEngine).toHaveBeenCalledTimes(2); + expect(samples.deleteEngine).toHaveBeenLastCalledWith(created.id); + await expect(service.remove('c1', created.id)).rejects.toBeInstanceOf(NotFoundException); + }); + + it('resets counter baselines when an engine is disabled', async () => { + const { service, samples } = setup(); + const created = await service.create('c1', { name: 'n', source: 'otlp' }); + await service.update('c1', created.id, { name: 'renamed' }); + expect(samples.resetBaselines).not.toHaveBeenCalled(); + await service.update('c1', created.id, { enabled: false }); + expect(samples.resetBaselines).toHaveBeenCalledWith(created.id); + }); + + it('resets counter baselines when an engine is enabled again', async () => { + const { service, samples } = setup(); + const created = await service.create('c1', { name: 'n', source: 'otlp', enabled: false }); + await service.update('c1', created.id, { enabled: true }); + expect(samples.resetBaselines).toHaveBeenCalledWith(created.id); + }); + + it('resets counter baselines when the scrape URL changes', async () => { + const { service, samples } = setup(); + const created = await service.create('c1', { name: 'n', source: 'scrape', scrapeUrl: 'http://x' }); + await service.update('c1', created.id, { scrapeUrl: 'http://x' }); + expect(samples.resetBaselines).not.toHaveBeenCalled(); + await service.update('c1', created.id, { scrapeUrl: 'http://y/metrics' }); + expect(samples.resetBaselines).toHaveBeenCalledWith(created.id); + }); + + it('sends the re-test without auth when the header is cleared with a URL change', async () => { + const { service, store } = setup(); + const created = await service.create('c1', { name: 'n', source: 'scrape', scrapeUrl: 'http://x', scrapeAuthHeader: 'Bearer t' }); + fetchMock.mockClear(); + await service.update('c1', created.id, { scrapeUrl: 'http://other', scrapeAuthHeader: null }); + expect(fetchMock).toHaveBeenCalledWith('http://other', null); + expect(store.get(created.id)).toMatchObject({ scrapeUrl: 'http://other', scrapeAuthHeader: null }); + }); + + it('treats an empty header as a clear', async () => { + const { service, store } = setup(); + const created = await service.create('c1', { name: 'n', source: 'scrape', scrapeUrl: 'http://x', scrapeAuthHeader: 'Bearer t' }); + fetchMock.mockClear(); + await service.update('c1', created.id, { scrapeUrl: 'http://x/other', scrapeAuthHeader: '' }); + expect(fetchMock).toHaveBeenCalledWith('http://x/other', null); + expect(store.get(created.id)).toMatchObject({ scrapeAuthHeader: null, scrapeAuthEncrypted: false }); + }); + + it('rejects an origin change that would carry the stored header over', async () => { + const { service, store } = setup(); + const created = await service.create('c1', { name: 'n', source: 'scrape', scrapeUrl: 'http://x/metrics', scrapeAuthHeader: 'Bearer t' }); + fetchMock.mockClear(); + await expect(service.update('c1', created.id, { scrapeUrl: 'http://evil/metrics' })).rejects.toThrow( + 'Re-enter the scrape auth header when changing the scrape URL origin', + ); + expect(fetchMock).not.toHaveBeenCalled(); + expect(store.get(created.id).scrapeUrl).toBe('http://x/metrics'); + await service.update('c1', created.id, { scrapeUrl: 'http://x/other' }); + expect(fetchMock).toHaveBeenCalledWith('http://x/other', 'Bearer t'); + }); + + it('accepts an origin change with a new header', async () => { + const { service, store } = setup(); + const created = await service.create('c1', { name: 'n', source: 'scrape', scrapeUrl: 'http://x', scrapeAuthHeader: 'Bearer t' }); + await service.update('c1', created.id, { scrapeUrl: 'http://y', scrapeAuthHeader: 'Bearer n' }); + expect(fetchMock).toHaveBeenLastCalledWith('http://y', 'Bearer n'); + expect(store.get(created.id)).toMatchObject({ scrapeUrl: 'http://y', scrapeAuthHeader: 'enc:Bearer n' }); + }); + + it('fails when the engine is removed during the test scrape', async () => { + const { service, store } = setup(); + const created = await service.create('c1', { name: 'n', source: 'scrape', scrapeUrl: 'http://x' }); + fetchMock.mockImplementation(async () => { + store.delete(created.id); + return LMCACHE_BODY; + }); + await expect(service.update('c1', created.id, { scrapeUrl: 'http://x/2' })).rejects.toBeInstanceOf(NotFoundException); + expect(store.has(created.id)).toBe(false); + }); + it('fails an update that lands behind an in-flight delete', async () => { + let finishDelete: () => void = () => undefined; + const storage = { + getKvCacheEngines: jest.fn().mockResolvedValue([]), + saveKvCacheEngine: jest.fn(async (e: any) => e), + deleteKvCacheEngine: jest.fn(() => new Promise((resolve) => (finishDelete = () => resolve(true)))), + } as any; + const footprint = { onConnectionRemoval: jest.fn() } as any; + const samples = { forgetEngine: jest.fn(), idle: jest.fn(), deleteEngine: jest.fn().mockResolvedValue(undefined), resetBaselines: jest.fn() } as any; + const registry = new KvCacheEngineRegistry(storage, footprint, samples); + await registry.onModuleInit(); + encryptionMock.mockReturnValue(null); + const connections = { getConfig: jest.fn().mockReturnValue({ id: 'c1' }) } as any; + const service = new KvCacheEnginesService(registry, connections, samples); + const created = await service.create('c1', { name: 'n', source: 'scrape', scrapeUrl: 'http://x' }); + storage.saveKvCacheEngine.mockClear(); + const removing = service.remove('c1', created.id); + const updating = service.update('c1', created.id, { name: 'renamed', enabled: false }); + await new Promise((resolve) => setImmediate(resolve)); + finishDelete(); + await removing; + await expect(updating).rejects.toBeInstanceOf(NotFoundException); + expect(storage.saveKvCacheEngine).not.toHaveBeenCalled(); + expect(samples.resetBaselines).not.toHaveBeenCalled(); + }); +}); diff --git a/proprietary/kv-cache/__tests__/kv-cache-footprint.service.spec.ts b/proprietary/kv-cache/__tests__/kv-cache-footprint.service.spec.ts index 01a064625..b0b2c6015 100644 --- a/proprietary/kv-cache/__tests__/kv-cache-footprint.service.spec.ts +++ b/proprietary/kv-cache/__tests__/kv-cache-footprint.service.spec.ts @@ -653,6 +653,8 @@ describe('KvCacheFootprintService', () => { it('runs connection removal cleanup while unlicensed', async () => { const { service, storage, registry } = setup(undefined, false); + const listener = jest.fn(); + service.onConnectionRemoval(listener); jest.spyOn(service as any, 'start').mockImplementation(() => undefined); await service.onModuleInit(); await (service as any).tick(); @@ -660,11 +662,14 @@ describe('KvCacheFootprintService', () => { await new Promise((resolve) => setImmediate(resolve)); expect(storage.deleteKvCacheConnectionData).toHaveBeenCalledWith('c1'); expect(storage.saveKvCacheFootprintSnapshot).not.toHaveBeenCalled(); + expect(listener).toHaveBeenCalledWith('c1'); await service.onModuleDestroy(); }); it('stops polling a retired member but keeps its data', async () => { const { service, storage, client, registry } = setup(); + const listener = jest.fn(); + service.onConnectionRemoval(listener); jest.spyOn(service as any, 'start').mockImplementation(() => undefined); await service.onModuleInit(); await (service as any).tick(); @@ -676,6 +681,7 @@ describe('KvCacheFootprintService', () => { await new Promise((resolve) => setImmediate(resolve)); expect(client.getDbSize).toHaveBeenCalledTimes(1); expect(storage.deleteKvCacheConnectionData).not.toHaveBeenCalled(); + expect(listener).not.toHaveBeenCalled(); expect(service.getSampleKey('c1')).toBe('m/x@1@0@1…@bfloat16'); await service.onModuleDestroy(); }); @@ -793,6 +799,37 @@ describe('KvCacheFootprintService', () => { await service.onModuleDestroy(); }); + it('notifies removal listeners with the connection id', async () => { + const { service, registry } = setup(); + jest.spyOn(service as any, 'start').mockImplementation(() => undefined); + await service.onModuleInit(); + const listener = jest.fn(); + const failing = jest.fn(() => { throw new Error('x'); }); + service.onConnectionRemoval(failing); + service.onConnectionRemoval(listener); + registry.removed('c1'); + await new Promise((resolve) => setImmediate(resolve)); + expect(listener).toHaveBeenCalledWith('c1'); + await service.onModuleDestroy(); + }); + + it('finishes removal listeners before deleting stored data', async () => { + const { service, storage, registry } = setup(); + jest.spyOn(service as any, 'start').mockImplementation(() => undefined); + await service.onModuleInit(); + let finish: () => void = () => undefined; + const listener = jest.fn(() => new Promise((resolve) => (finish = resolve))); + service.onConnectionRemoval(listener); + registry.removed('c1'); + await new Promise((resolve) => setImmediate(resolve)); + expect(listener).toHaveBeenCalledWith('c1'); + expect(storage.deleteKvCacheConnectionData).not.toHaveBeenCalled(); + finish(); + await new Promise((resolve) => setImmediate(resolve)); + expect(storage.deleteKvCacheConnectionData).toHaveBeenCalledWith('c1'); + await service.onModuleDestroy(); + }); + it('deletes data for a connection removed from the registry before the poller saw it', async () => { const { service, storage, registry } = setup(); jest.spyOn(service as any, 'start').mockImplementation(() => undefined); diff --git a/proprietary/kv-cache/__tests__/kv-cache-otlp-sink.service.spec.ts b/proprietary/kv-cache/__tests__/kv-cache-otlp-sink.service.spec.ts new file mode 100644 index 000000000..6efac4d99 --- /dev/null +++ b/proprietary/kv-cache/__tests__/kv-cache-otlp-sink.service.spec.ts @@ -0,0 +1,136 @@ +import type { OtlpKeyValue, OtlpMetric } from '@app/external-metrics/otlp-metrics-types'; +import { KvCacheOtlpSinkService } from '../kv-cache-otlp-sink.service'; + +const NOW = 1_700_000_000_000; + +const str = (key: string, value: string): OtlpKeyValue => ({ key, value: { stringValue: value } }); +const attrs = { 'betterdb.lmcache.engine': 'otlp-1' }; + +const point = (worker: string, value: number | null) => ({ + attributes: [str('model_name', 'llama'), str('worker_id', worker), str('role', 'worker')], + startTimeUnixNano: String(BigInt(NOW - 60_000) * 1_000_000n), + ...(value === null ? {} : { asInt: String(value) }), +}); + +const sum = (name: string, points: ReturnType[], temporality: number | string = 2): OtlpMetric => ({ + name, + sum: { dataPoints: points, aggregationTemporality: temporality }, +}); + +function setup(opts: { engine?: unknown; licensed?: boolean } = {}) { + const engine = 'engine' in opts ? opts.engine : { id: 'e1', connectionId: 'c1', enabled: true }; + const registry = { byOtlpId: jest.fn().mockReturnValue(engine), recordResult: jest.fn() }; + const samples = { observe: jest.fn().mockReturnValue(true) }; + const license = { hasFeature: jest.fn().mockReturnValue(opts.licensed ?? true) }; + const service = new KvCacheOtlpSinkService(registry as any, samples as any, license as any); + return { registry, samples, license, service }; +} + +describe('KvCacheOtlpSinkService', () => { + it('drops every point as unknown_engine for an unknown id', () => { + const { service, samples } = setup({ engine: null }); + const metrics = [sum('lmcache:num_hit_tokens_total', [point('0', 1), point('1', 2)])]; + expect(service.ingest(attrs, metrics, NOW)).toEqual({ accepted: 0, dropped: { unknown_engine: 2 } }); + expect(samples.observe).not.toHaveBeenCalled(); + }); + + it('drops every point as unknown_engine when the attribute is missing', () => { + const { service, registry } = setup(); + const metrics = [sum('lmcache:num_hit_tokens_total', [point('0', 1)])]; + expect(service.ingest({}, metrics, NOW)).toEqual({ accepted: 0, dropped: { unknown_engine: 1 } }); + expect(registry.byOtlpId).not.toHaveBeenCalled(); + }); + + it('drops every point as unknown_engine for a disabled engine', () => { + const { service } = setup({ engine: { id: 'e1', connectionId: 'c1', enabled: false } }); + const metrics = [sum('lmcache:num_hit_tokens_total', [point('0', 1)])]; + expect(service.ingest(attrs, metrics, NOW).dropped).toEqual({ unknown_engine: 1 }); + }); + + it('drops every point as unknown_engine when unlicensed', () => { + const { service, samples } = setup({ licensed: false }); + const metrics = [sum('lmcache:num_hit_tokens_total', [point('0', 1), point('1', 2)])]; + expect(service.ingest(attrs, metrics, NOW)).toEqual({ accepted: 0, dropped: { unknown_engine: 2 } }); + expect(samples.observe).not.toHaveBeenCalled(); + }); + + it('observes each data point of a cumulative sum', () => { + const { service, samples, registry } = setup(); + const metrics = [sum('lmcache:num_hit_tokens_total', [point('0', 10), point('1', 20)])]; + const result = service.ingest(attrs, metrics, NOW); + expect(result).toEqual({ accepted: 2, dropped: {} }); + expect(samples.observe).toHaveBeenCalledTimes(2); + expect(samples.observe).toHaveBeenNthCalledWith( + 1, + { + engineId: 'e1', + connectionId: 'c1', + modelName: 'llama', + series: '0|worker|[]', + metric: 'num_hit_tokens', + value: 10, + cumulative: true, + startMs: NOW - 60_000, + }, + NOW, + ); + expect(samples.observe.mock.calls[1][0]).toMatchObject({ series: '1|worker|[]', value: 20 }); + expect(registry.recordResult).toHaveBeenCalledWith('e1', { lastSeenAt: NOW, lastError: null }); + }); + + it('keeps workers identified only by resource attributes apart', () => { + const { service, samples } = setup(); + const bare = { attributes: [str('model_name', 'llama')], asInt: '5' } as unknown as ReturnType; + const metric = () => [sum('lmcache:num_hit_tokens_total', [bare])]; + service.ingest({ ...attrs, 'service.instance.id': 'b', 'host.name': 'h' }, metric(), NOW); + service.ingest({ 'host.name': 'h', 'service.instance.id': 'b', ...attrs }, metric(), NOW); + service.ingest({ ...attrs, 'service.instance.id': 'a', 'host.name': 'h' }, metric(), NOW); + const [first, same, other] = samples.observe.mock.calls.map(([o]) => o.series); + expect(first).toBe(same); + expect(first).not.toBe(other); + expect(first).toBe('||[["host.name","h"],["service.instance.id","b"]]'); + }); + + it('keeps series that differ by served_model_name apart', () => { + const { service, samples } = setup(); + const withServed = (served: string) => + ({ attributes: [str('model_name', 'llama'), str('worker_id', '0'), str('served_model_name', served)], asInt: '5' }) as unknown as ReturnType; + service.ingest(attrs, [sum('lmcache:num_hit_tokens_total', [withServed('a'), withServed('b')])], NOW); + const [first, second] = samples.observe.mock.calls.map(([o]) => o.series); + expect(first).not.toBe(second); + }); + + it('skips points flagged as having no recorded value', () => { + const { service, samples } = setup(); + const metrics = [sum('lmcache:num_hit_tokens_total', [{ ...point('0', 10), flags: 1 } as ReturnType, point('1', 20)])]; + expect(service.ingest(attrs, metrics, NOW)).toEqual({ accepted: 1, dropped: {} }); + expect(samples.observe).toHaveBeenCalledTimes(1); + expect(samples.observe.mock.calls[0][0]).toMatchObject({ series: '1|worker|[]', value: 20 }); + }); + + it('marks delta temporality sums as not cumulative', () => { + const { service, samples } = setup(); + service.ingest(attrs, [sum('lmcache:num_hit_tokens_total', [point('0', 3)], 1)], NOW); + expect(samples.observe.mock.calls[0][0].cumulative).toBe(false); + service.ingest(attrs, [sum('lmcache:num_hit_tokens_total', [point('0', 3)], 'AGGREGATION_TEMPORALITY_DELTA')], NOW); + expect(samples.observe.mock.calls[1][0].cumulative).toBe(false); + }); + + it('classifies unsupported, unmapped and invalid points', () => { + const { service, samples, registry } = setup(); + const gauge: OtlpMetric = { name: 'lmcache:lookup_hit_rate', gauge: { dataPoints: [point('0', 1)] } }; + const metrics = [ + gauge, + sum('lmcache:something_else', [point('0', 1), point('1', 1)]), + sum('lmcache:toString_total', [point('0', 1)]), + sum('lmcache:num_hit_tokens_total', [point('0', null)]), + ]; + const result = service.ingest(attrs, metrics, NOW); + expect(result).toEqual({ + accepted: 0, + dropped: { unsupported_type: 1, unmapped_metric: 3, invalid_value: 1 }, + }); + expect(samples.observe).not.toHaveBeenCalled(); + expect(registry.recordResult).not.toHaveBeenCalled(); + }); +}); diff --git a/proprietary/kv-cache/__tests__/kv-cache-samples.service.spec.ts b/proprietary/kv-cache/__tests__/kv-cache-samples.service.spec.ts new file mode 100644 index 000000000..bb9c2b48c --- /dev/null +++ b/proprietary/kv-cache/__tests__/kv-cache-samples.service.spec.ts @@ -0,0 +1,233 @@ +import { Logger } from '@nestjs/common'; +import type { KvCacheEngineSample } from '@betterdb/shared'; +import { + hitRate, + KvCacheSamplesService, + mergeBuckets, + SAMPLES_CHART_POINTS, + SAMPLES_LIMIT, + SAMPLES_MAX_RANGE_MS, +} from '../kv-cache-samples.service'; +import { METRIC_FIELDS } from '../counter-deltas'; + +const counters = { requestedTokens: 0, hitTokens: 0, lookupTokens: 0, lookupHits: 0, remoteReadBytes: 0, remoteWriteBytes: 0, remoteReadRequests: 0, remoteWriteRequests: 0, remotePingErrors: 0 }; + +const sample = (over: Partial): KvCacheEngineSample => ({ + engineId: 'e1', + connectionId: 'c1', + modelName: 'm', + timestamp: 60_000, + ...counters, + ...over, +}); + +const observation = (value: number, cumulative = false) => ({ + engineId: 'e1', + connectionId: 'c1', + modelName: 'm', + series: 's', + metric: 'num_requested_tokens', + value, + cumulative, +}); + +function setup() { + const storage = { + saveKvCacheEngineSamples: jest.fn().mockResolvedValue(undefined), + getKvCacheEngineSamples: jest.fn().mockResolvedValue([]), + getKvCacheSampleTotals: jest.fn().mockResolvedValue({ requestedTokens: 0, hitTokens: 0 }), + deleteKvCacheEngineSamples: jest.fn().mockResolvedValue(undefined), + }; + return { storage, service: new KvCacheSamplesService(storage as any) }; +} + +describe('hitRate', () => { + it('is null without requested tokens', () => { + expect(hitRate(0, 0)).toBeNull(); + }); + + it('rounds the ratio to four decimals', () => { + expect(hitRate(1000, 700)).toBe(0.7); + expect(hitRate(3, 1)).toBe(0.3333); + }); + + it('clamps the ratio to [0, 1]', () => { + expect(hitRate(100, 250)).toBe(1); + expect(hitRate(100, -5)).toBe(0); + }); +}); + +describe('mergeBuckets', () => { + it('sums rows sharing engine, model and timestamp', () => { + const merged = mergeBuckets([ + sample({ requestedTokens: 100, hitTokens: 50 }), + sample({ requestedTokens: 100, hitTokens: 25 }), + ]); + expect(merged).toHaveLength(1); + expect(merged[0]).toMatchObject({ requestedTokens: 200, hitTokens: 75, hitRate: 0.375 }); + }); + + it('sums every mapped counter field', () => { + const ones = Object.fromEntries(Object.values(METRIC_FIELDS).map((field) => [field, 1])); + const [merged] = mergeBuckets([sample(ones), sample(ones)]); + for (const field of Object.values(METRIC_FIELDS)) expect(merged[field]).toBe(2); + }); + + it('keeps distinct engines, models and timestamps apart, sorted by timestamp', () => { + const merged = mergeBuckets([ + sample({ timestamp: 120_000 }), + sample({ engineId: 'e2' }), + sample({ modelName: 'other' }), + sample({}), + ]); + expect(merged.map((b) => b.timestamp)).toEqual([60_000, 60_000, 60_000, 120_000]); + expect(merged).toHaveLength(4); + expect(merged[0].hitRate).toBeNull(); + }); +}); + +describe('KvCacheSamplesService', () => { + it('flushes a closed bucket once', async () => { + const { service, storage } = setup(); + expect(service.observe(observation(10), 61_000)).toBe(true); + service.observe(observation(5), 62_000); + await service.flush(120_000); + expect(storage.saveKvCacheEngineSamples).toHaveBeenCalledTimes(1); + expect(storage.saveKvCacheEngineSamples.mock.calls[0][0]).toEqual([ + expect.objectContaining({ engineId: 'e1', timestamp: 60_000, requestedTokens: 15 }), + ]); + await service.flush(120_000); + expect(storage.saveKvCacheEngineSamples).toHaveBeenCalledTimes(1); + }); + + it('keeps the open bucket until it closes', async () => { + const { service, storage } = setup(); + service.observe(observation(10), 61_000); + await service.flush(90_000); + expect(storage.saveKvCacheEngineSamples).not.toHaveBeenCalled(); + }); + + it('logs and drops rows when the save fails', async () => { + const { service, storage } = setup(); + const warn = jest.spyOn(Logger.prototype, 'warn').mockImplementation(); + storage.saveKvCacheEngineSamples.mockRejectedValue(new Error('disk full')); + service.observe(observation(10), 61_000); + await expect(service.flush(120_000)).resolves.toBeUndefined(); + expect(warn).toHaveBeenCalledWith('Could not save 1 KV cache samples: disk full'); + storage.saveKvCacheEngineSamples.mockClear(); + await service.flush(120_000); + expect(storage.saveKvCacheEngineSamples).not.toHaveBeenCalled(); + warn.mockRestore(); + }); + + it('drops pending buckets and stored samples of a deleted engine', async () => { + const { service, storage } = setup(); + service.observe(observation(10), 61_000); + await service.deleteEngine('e1'); + expect(storage.deleteKvCacheEngineSamples).toHaveBeenCalledWith('e1'); + await service.flush(120_000); + expect(storage.saveKvCacheEngineSamples).not.toHaveBeenCalled(); + }); + + it('deletes stored samples only after an in-flight flush lands', async () => { + const { service, storage } = setup(); + const order: string[] = []; + let finishSave: () => void = () => undefined; + storage.saveKvCacheEngineSamples.mockImplementationOnce(() => new Promise((resolve) => (finishSave = () => { order.push('save'); resolve(); }))); + storage.deleteKvCacheEngineSamples.mockImplementationOnce(async () => { order.push('delete'); }); + service.observe(observation(10), 61_000); + const flushing = service.flush(120_000); + const deleting = service.deleteEngine('e1'); + await new Promise((resolve) => setImmediate(resolve)); + finishSave(); + await Promise.all([flushing, deleting]); + expect(order).toEqual(['save', 'delete']); + }); + + it('resets baselines without dropping pending buckets', async () => { + const { service, storage } = setup(); + service.observe(observation(100, true), 61_000); + service.observe(observation(110, true), 62_000); + service.resetBaselines('e1'); + service.observe(observation(5000, true), 63_000); + await service.flush(120_000); + expect(storage.saveKvCacheEngineSamples).toHaveBeenCalledWith([expect.objectContaining({ requestedTokens: 10 })]); + }); + + it('caps the range and computes the range hit rate across engines', async () => { + const { service, storage } = setup(); + storage.getKvCacheEngineSamples.mockResolvedValue([ + sample({ engineId: 'e1', requestedTokens: 100, hitTokens: 50 }), + sample({ engineId: 'e2', requestedTokens: 300, hitTokens: 250 }), + ]); + const to = 100 * 24 * 60 * 60_000; + const result = await service.getSamples('c1', { from: 0, to, engineId: 'e1', model: 'm' }); + expect(storage.getKvCacheEngineSamples).toHaveBeenCalledWith({ + connectionId: 'c1', + from: to - SAMPLES_MAX_RANGE_MS, + to, + limit: SAMPLES_LIMIT, + bucketMs: expect.any(Number), + engineId: 'e1', + modelName: 'm', + }); + expect(result.buckets).toHaveLength(2); + expect(storage.getKvCacheSampleTotals).toHaveBeenCalledWith({ + connectionId: 'c1', + from: to - SAMPLES_MAX_RANGE_MS, + to, + engineId: 'e1', + modelName: 'm', + }); + }); + + it('buckets the chart rows in storage so a long range covers the whole window', async () => { + const { service, storage } = setup(); + const to = 40 * 24 * 60 * 60_000; + await service.getSamples('c1', { from: 0, to }); + const query = storage.getKvCacheEngineSamples.mock.calls[0][0]; + expect(query.bucketMs % 60_000).toBe(0); + expect(SAMPLES_MAX_RANGE_MS / query.bucketMs).toBeLessThanOrEqual(SAMPLES_CHART_POINTS); + await service.getSamples('c1', { from: to - 60 * 60_000, to }); + expect(storage.getKvCacheEngineSamples.mock.calls[1][0].bucketMs).toBe(60_000); + }); + + it('computes the range hit rate from totals over the whole range, not the capped rows', async () => { + const { service, storage } = setup(); + storage.getKvCacheEngineSamples.mockResolvedValue([sample({ requestedTokens: 100, hitTokens: 100 })]); + storage.getKvCacheSampleTotals.mockResolvedValue({ requestedTokens: 400, hitTokens: 300 }); + expect((await service.getSamples('c1', { from: 0, to: 1000 })).rangeHitRate).toBe(0.75); + }); + + it('returns a null range hit rate when nothing was requested', async () => { + const { service } = setup(); + expect(await service.getSamples('c1', { from: 0, to: 1000 })).toEqual({ buckets: [], rangeHitRate: null }); + }); + + it('waits for an in-flight save even when there is nothing new to drain', async () => { + const { service, storage } = setup(); + let finishSave: () => void = () => undefined; + let saved = false; + storage.saveKvCacheEngineSamples.mockImplementationOnce( + () => new Promise((resolve) => (finishSave = () => { saved = true; resolve(); })), + ); + service.observe(observation(10), 61_000); + void service.flush(120_000); + let destroyed = false; + const destroying = service.onModuleDestroy().then(() => { destroyed = true; }); + await new Promise((resolve) => setImmediate(resolve)); + expect(destroyed).toBe(false); + finishSave(); + await destroying; + expect(saved).toBe(true); + }); + + it('flushes the open bucket on module destroy and stops the timer', async () => { + const { service, storage } = setup(); + service.onModuleInit(); + service.observe(observation(10), Date.now()); + await service.onModuleDestroy(); + expect(storage.saveKvCacheEngineSamples).toHaveBeenCalledTimes(1); + expect(storage.saveKvCacheEngineSamples.mock.calls[0][0][0]).toMatchObject({ requestedTokens: 10 }); + }); +}); diff --git a/proprietary/kv-cache/__tests__/kv-cache-scrape.service.spec.ts b/proprietary/kv-cache/__tests__/kv-cache-scrape.service.spec.ts new file mode 100644 index 000000000..5ac9a988d --- /dev/null +++ b/proprietary/kv-cache/__tests__/kv-cache-scrape.service.spec.ts @@ -0,0 +1,183 @@ +import { Logger } from '@nestjs/common'; +import { fetchMetricsText, ScrapeError } from '../metrics-fetch'; +import { KvCacheScrapeService } from '../kv-cache-scrape.service'; + +jest.mock('../metrics-fetch', () => ({ + ...jest.requireActual('../metrics-fetch'), + fetchMetricsText: jest.fn(), +})); + +const fetchMock = fetchMetricsText as jest.Mock; + +const engine = (over: Record = {}) => ({ + id: 'e1', + connectionId: 'c1', + enabled: true, + source: 'scrape', + scrapeUrl: 'http://engine-1:9090/metrics', + ...over, +}); + +const body = 'lmcache:num_hit_tokens_total{model_name="m",worker_id="0",role="worker"} 42\n'; + +function setup(engines: unknown[], licensed = true) { + const registry = { + list: jest.fn().mockReturnValue(engines), + get: jest.fn((id: string) => (engines as { id: string }[]).find((e) => e.id === id) ?? null), + recordResult: jest.fn(), + }; + const samples = { observe: jest.fn() }; + const engineService = { authHeaderFor: jest.fn().mockReturnValue(null) }; + const license = { hasFeature: jest.fn().mockReturnValue(licensed) }; + const service = new KvCacheScrapeService(registry as any, samples as any, engineService as any, license as any); + return { registry, samples, engineService, license, service }; +} + +describe('KvCacheScrapeService', () => { + beforeEach(() => fetchMock.mockReset()); + + it('does not fetch when unlicensed', async () => { + const { service } = setup([engine()], false); + await service.tick(1000); + expect(fetchMock).not.toHaveBeenCalled(); + }); + + it('skips disabled and OTLP engines', async () => { + const { service } = setup([engine({ id: 'off', enabled: false }), engine({ id: 'otlp', source: 'otlp' })]); + await service.tick(1000); + expect(fetchMock).not.toHaveBeenCalled(); + }); + + it('observes cumulative samples and records success', async () => { + fetchMock.mockResolvedValue(body); + const { service, samples, registry } = setup([engine()]); + await service.tick(5000); + expect(samples.observe).toHaveBeenCalledWith( + { + engineId: 'e1', + connectionId: 'c1', + modelName: 'm', + series: '0|worker', + metric: 'num_hit_tokens', + value: 42, + cumulative: true, + }, + 5000, + ); + expect(registry.recordResult).toHaveBeenCalledWith('e1', { lastSeenAt: 5000, lastError: null }); + }); + + it('records a scrape error and still scrapes the other engines', async () => { + fetchMock.mockImplementation(async (url: string) => { + if (url.includes('engine-1')) throw new ScrapeError('HTTP 401'); + return body; + }); + const { service, registry, samples } = setup([engine(), engine({ id: 'e2', scrapeUrl: 'http://engine-2:9090/metrics' })]); + await service.tick(1000); + expect(registry.recordResult).toHaveBeenCalledWith('e1', { lastError: 'HTTP 401' }); + expect(samples.observe).toHaveBeenCalledWith(expect.objectContaining({ engineId: 'e2' }), 1000); + }); + + it('discards a response when the engine was removed, disabled or repointed during the fetch', async () => { + fetchMock.mockResolvedValue(body); + for (const current of [null, engine({ enabled: false }), engine({ scrapeUrl: 'http://engine-9:9090/metrics' })]) { + const { service, registry, samples } = setup([engine()]); + registry.get.mockReturnValue(current); + await service.tick(1000); + expect(samples.observe).not.toHaveBeenCalled(); + expect(registry.recordResult).not.toHaveBeenCalled(); + } + }); + + it('drops a scrape error when the engine was repointed during the fetch', async () => { + fetchMock.mockRejectedValue(new ScrapeError('HTTP 401')); + const { service, registry } = setup([engine()]); + registry.get.mockReturnValue(engine({ scrapeUrl: 'http://engine-9:9090/metrics' })); + await service.tick(1000); + expect(registry.recordResult).not.toHaveBeenCalled(); + }); + + it('records a generic error for unexpected failures', async () => { + fetchMock.mockRejectedValue(new Error('secret-host leaked')); + const { service, registry } = setup([engine()]); + await service.tick(1000); + expect(registry.recordResult).toHaveBeenCalledWith('e1', { lastError: 'scrape failed' }); + }); + + it('flags a response without lmcache metrics', async () => { + fetchMock.mockResolvedValue('other_metric 1\n'); + const { service, registry } = setup([engine()]); + await service.tick(1000); + expect(registry.recordResult).toHaveBeenCalledWith('e1', { lastSeenAt: 1000, lastError: 'no lmcache metrics found' }); + }); + + it('fetches once for overlapping ticks', async () => { + let release: (text: string) => void = () => undefined; + fetchMock.mockReturnValue(new Promise((resolve) => (release = resolve))); + const { service } = setup([engine()]); + const first = service.tick(1000); + const second = service.tick(2000); + release(body); + await Promise.all([first, second]); + expect(fetchMock).toHaveBeenCalledTimes(1); + }); + it('scrapes again on a later tick after the first completes', async () => { + fetchMock.mockResolvedValue(body); + const { service } = setup([engine()]); + await service.tick(1000); + await service.tick(2000); + expect(fetchMock).toHaveBeenCalledTimes(2); + }); + + it('starts scraping once the licence becomes active', async () => { + fetchMock.mockResolvedValue(body); + const { service, license } = setup([engine()], false); + await service.tick(1000); + expect(fetchMock).not.toHaveBeenCalled(); + license.hasFeature.mockReturnValue(true); + await service.tick(2000); + expect(fetchMock).toHaveBeenCalledTimes(1); + }); + + it('records a decryption error instead of scraping without the auth header', async () => { + fetchMock.mockResolvedValue(body); + const { service, registry, engineService } = setup([engine({ scrapeAuthHeader: 'enc:x', scrapeAuthEncrypted: true })]); + engineService.authHeaderFor.mockReturnValue(null); + await service.tick(1000); + expect(fetchMock).not.toHaveBeenCalled(); + expect(registry.recordResult).toHaveBeenCalledWith('e1', { lastError: 'auth header could not be decrypted' }); + }); + + describe('interval', () => { + const original = process.env.KV_CACHE_SCRAPE_INTERVAL_MS; + afterEach(() => { + if (original === undefined) delete process.env.KV_CACHE_SCRAPE_INTERVAL_MS; + else process.env.KV_CACHE_SCRAPE_INTERVAL_MS = original; + jest.restoreAllMocks(); + }); + + it.each([ + ['abc', 30000], + ['0', 30000], + ['-5', 30000], + ['', 30000], + ['5000', 5000], + ['2000ms', 30000], + ['99999999999', 30000], + ])('uses the right interval for %p', (value, expected) => { + process.env.KV_CACHE_SCRAPE_INTERVAL_MS = value; + const spy = jest.spyOn(global, 'setInterval'); + const { service } = setup([]); + service.onModuleInit(); + service.onModuleDestroy(); + expect(spy).toHaveBeenCalledWith(expect.any(Function), expected); + }); + + it('warns when the interval is ignored', () => { + process.env.KV_CACHE_SCRAPE_INTERVAL_MS = '2000ms'; + const warn = jest.spyOn(Logger.prototype, 'warn').mockImplementation(); + setup([]); + expect(warn).toHaveBeenCalledWith(expect.stringContaining("KV_CACHE_SCRAPE_INTERVAL_MS='2000ms'")); + }); + }); +}); diff --git a/proprietary/kv-cache/__tests__/kv-cache-status.service.spec.ts b/proprietary/kv-cache/__tests__/kv-cache-status.service.spec.ts index 9201a6a49..b69b9d7d3 100644 --- a/proprietary/kv-cache/__tests__/kv-cache-status.service.spec.ts +++ b/proprietary/kv-cache/__tests__/kv-cache-status.service.spec.ts @@ -9,10 +9,10 @@ const engine = { describe('KvCacheStatusService', () => { const storage = { getKvCacheFootprintSnapshots: jest.fn(), - getKvCacheEngines: jest.fn().mockResolvedValue([engine]), } as any; + const registry = { list: jest.fn().mockReturnValue([engine]) } as any; const footprint = { getSampleKey: jest.fn().mockReturnValue('m@1@0@abc…@bfloat16') } as any; - const service = new KvCacheStatusService(storage, footprint); + const service = new KvCacheStatusService(storage, footprint, registry); it('never exposes the auth header', () => { const pub = toPublicEngine(engine as any); @@ -21,21 +21,29 @@ describe('KvCacheStatusService', () => { expect(pub.hasScrapeAuth).toBe(true); }); + it('strips credentials from the scrape URL', () => { + expect(toPublicEngine({ ...engine, scrapeUrl: 'http://admin:secret@vllm:8000/metrics' } as any).scrapeUrl).toBe('http://vllm:8000/metrics'); + expect(toPublicEngine({ ...engine, scrapeUrl: 'http://token@vllm/metrics' } as any).scrapeUrl).toBe('http://vllm/metrics'); + expect(toPublicEngine(engine as any).scrapeUrl).toBe('http://x/metrics'); + expect(toPublicEngine({ ...engine, source: 'otlp', scrapeUrl: null } as any).scrapeUrl).toBeNull(); + }); + + it('masks query values in the scrape URL', () => { + expect(toPublicEngine({ ...engine, scrapeUrl: 'http://vllm:8000/metrics?token=s3cret&format=text' } as any).scrapeUrl).toBe( + 'http://vllm:8000/metrics?token=***&format=***', + ); + expect(toPublicEngine({ ...engine, scrapeUrl: 'http://admin:pw@vllm/metrics?key=abc#frag' } as any).scrapeUrl).toBe('http://vllm/metrics?key=***'); + expect(toPublicEngine({ ...engine, scrapeUrl: 'http://vllm/metrics?flag' } as any).scrapeUrl).toBe('http://vllm/metrics?flag='); + }); + it('builds status from the latest snapshot', async () => { storage.getKvCacheFootprintSnapshots.mockResolvedValue([{ detected: true }]); const status = await service.getStatus('c1'); expect(storage.getKvCacheFootprintSnapshots).toHaveBeenCalledWith({ connectionId: 'c1', limit: 1 }); + expect(registry.list).toHaveBeenCalledWith('c1'); expect(status).toEqual({ hasLmcache: true, latest: { detected: true }, sampleKey: 'm@1@0@abc…@bfloat16', engines: [toPublicEngine(engine as any)] }); }); - it('reads the snapshot and the engines in parallel', async () => { - storage.getKvCacheFootprintSnapshots.mockReturnValue(new Promise(() => undefined)); - storage.getKvCacheEngines.mockClear(); - void service.getStatus('c1'); - await Promise.resolve(); - expect(storage.getKvCacheEngines).toHaveBeenCalledWith('c1'); - }); - it('reports not detected without a snapshot', async () => { storage.getKvCacheFootprintSnapshots.mockResolvedValue([]); expect((await service.getStatus('c1')).hasLmcache).toBe(false); diff --git a/proprietary/kv-cache/__tests__/kv-cache.controller.spec.ts b/proprietary/kv-cache/__tests__/kv-cache.controller.spec.ts index d1be571cb..99e07547d 100644 --- a/proprietary/kv-cache/__tests__/kv-cache.controller.spec.ts +++ b/proprietary/kv-cache/__tests__/kv-cache.controller.spec.ts @@ -6,6 +6,14 @@ import { KvCacheController } from '../kv-cache.controller'; describe('KvCacheController', () => { const status = { getStatus: jest.fn().mockResolvedValue({}), getFootprintHistory: jest.fn().mockResolvedValue([]) }; const footprint = { triggerCollection: jest.fn().mockResolvedValue(null) }; + const engines = { + list: jest.fn().mockReturnValue([]), + create: jest.fn().mockResolvedValue({}), + update: jest.fn().mockResolvedValue({}), + remove: jest.fn().mockResolvedValue(undefined), + }; + const samples = { getSamples: jest.fn().mockResolvedValue({ buckets: [], rangeHitRate: null }) }; + const alerts = { getSettings: jest.fn().mockResolvedValue({}), updateSettings: jest.fn().mockResolvedValue({}) }; const registry = { getDefaultId: jest.fn().mockReturnValue('ui-added'), get: jest.fn((id: string) => { @@ -13,7 +21,7 @@ describe('KvCacheController', () => { return {}; }), }; - const controller = new KvCacheController(status as any, footprint as any, registry as any); + const controller = new KvCacheController(status as any, footprint as any, registry as any, engines as any, samples as any, alerts as any); it('defaults to the registry default connection on every route', async () => { await controller.getStatus(undefined); @@ -50,6 +58,46 @@ describe('KvCacheController', () => { await expect(controller.getFootprintHistory(undefined, '100', 'c1')).rejects.toBeInstanceOf(BadRequestException); }); + it('routes engine calls to the configured default connection', async () => { + controller.listEngines(undefined); + expect(engines.list).toHaveBeenCalledWith('ui-added'); + const dto = { name: 'n', source: 'otlp' } as any; + await controller.createEngine(dto, undefined); + expect(engines.create).toHaveBeenCalledWith('ui-added', dto); + await controller.updateEngine('e1', { enabled: false }, undefined); + expect(engines.update).toHaveBeenCalledWith('ui-added', 'e1', { enabled: false }); + await controller.removeEngine('e1', 'c1'); + expect(engines.remove).toHaveBeenCalledWith('c1', 'e1'); + }); + + it('parses the samples range and forwards the filters', async () => { + await controller.getSamples('100', '200', 'e1', 'm', 'c1'); + expect(samples.getSamples).toHaveBeenCalledWith('c1', { from: 100, to: 200, engineId: 'e1', model: 'm' }); + await controller.getSamples('100', '200', undefined, undefined, undefined); + expect(samples.getSamples).toHaveBeenLastCalledWith('ui-added', { from: 100, to: 200, engineId: undefined, model: undefined }); + await expect(controller.getSamples('200', '100', undefined, undefined, 'c1')).rejects.toBeInstanceOf(BadRequestException); + }); + + it('routes settings calls to the connection', async () => { + await controller.getSettings(undefined); + expect(alerts.getSettings).toHaveBeenCalledWith('ui-added'); + await controller.updateSettings({ hitRateThreshold: 0.3 }, 'c1'); + expect(alerts.updateSettings).toHaveBeenCalledWith('c1', { hitRateThreshold: 0.3 }); + }); + + it('rejects engine, sample and settings calls for an unknown connection', async () => { + jest.clearAllMocks(); + expect(() => controller.listEngines('nope')).toThrow(NotFoundException); + await expect(controller.createEngine({ name: 'n', source: 'otlp' } as any, 'nope')).rejects.toBeInstanceOf(NotFoundException); + await expect(controller.updateEngine('e1', {}, 'nope')).rejects.toBeInstanceOf(NotFoundException); + await expect(controller.removeEngine('e1', 'nope')).rejects.toBeInstanceOf(NotFoundException); + await expect(controller.getSamples('100', '200', undefined, undefined, 'nope')).rejects.toBeInstanceOf(NotFoundException); + await expect(controller.getSettings('nope')).rejects.toBeInstanceOf(NotFoundException); + await expect(controller.updateSettings({}, 'nope')).rejects.toBeInstanceOf(NotFoundException); + expect(engines.create).not.toHaveBeenCalled(); + expect(alerts.updateSettings).not.toHaveBeenCalled(); + }); + it.each([ ['a fractional timestamp', '100', '1700000000000.5'], ['a timestamp beyond the safe integer range', '100', '1e19'], @@ -59,7 +107,7 @@ describe('KvCacheController', () => { it('requires the KV cache feature on every route', () => { const reflector = new Reflector(); - for (const handler of ['getStatus', 'refreshFootprint', 'getFootprintHistory'] as const) { + for (const handler of ['getStatus', 'refreshFootprint', 'getFootprintHistory', 'getSamples', 'listEngines', 'createEngine', 'updateEngine', 'removeEngine', 'getSettings', 'updateSettings'] as const) { expect(reflector.get('requiredFeature', KvCacheController.prototype[handler])).toEqual(Feature.KV_CACHE_MONITORING); } }); diff --git a/proprietary/kv-cache/__tests__/metrics-fetch.spec.ts b/proprietary/kv-cache/__tests__/metrics-fetch.spec.ts new file mode 100644 index 000000000..1b02ad8cb --- /dev/null +++ b/proprietary/kv-cache/__tests__/metrics-fetch.spec.ts @@ -0,0 +1,237 @@ +import { BadRequestException, Logger } from '@nestjs/common'; +import { promises as dns } from 'dns'; +import { createServer, type IncomingHttpHeaders, type Server } from 'http'; +import { createServer as createTcpServer, type AddressInfo } from 'net'; +import { assertSafeOutboundUrl } from '@app/common/utils/outbound-url-guard'; +import { fetchMetricsText, scrapeAllowsPrivateNetworks, ScrapeError } from '../metrics-fetch'; + +jest.mock('@app/common/utils/outbound-url-guard', () => ({ + ...jest.requireActual('@app/common/utils/outbound-url-guard'), + assertSafeOutboundUrl: jest.fn().mockResolvedValue(new URL('http://x')), +})); + +const respond = (init: { status?: number; body?: string; headers?: Record } = {}) => + new Response(init.body ?? '', { status: init.status ?? 200, headers: init.headers }); + +describe('fetchMetricsText', () => { + beforeEach(() => { + (assertSafeOutboundUrl as jest.Mock).mockResolvedValue(new URL('http://x')); + }); + + it('returns the body for a 200', async () => { + const fetchImpl = jest.fn().mockResolvedValue(respond({ body: 'lmcache:a_total 1' })); + await expect(fetchMetricsText('http://x/metrics', null, fetchImpl)).resolves.toBe('lmcache:a_total 1'); + }); + + it('sends Authorization only when given', async () => { + const fetchImpl = jest.fn().mockImplementation(async () => respond()); + await fetchMetricsText('http://x/metrics', null, fetchImpl); + expect(fetchImpl.mock.calls[0][1].headers).not.toHaveProperty('Authorization'); + await fetchMetricsText('http://x/metrics', 'Bearer t', fetchImpl); + expect(fetchImpl.mock.calls[1][1].headers.Authorization).toBe('Bearer t'); + }); + + it('reports HTTP errors and redirects', async () => { + await expect(fetchMetricsText('http://x', null, jest.fn().mockResolvedValue(respond({ status: 401 })))).rejects.toThrow('HTTP 401'); + await expect(fetchMetricsText('http://x', null, jest.fn().mockResolvedValue(respond({ status: 302 })))).rejects.toThrow( + 'redirect not followed (HTTP 302)', + ); + }); + + it('cancels the body of a rejected response', async () => { + for (const init of [{ status: 500 }, { status: 302 }, { headers: { 'content-length': '6000000' } }]) { + const response = respond({ ...init, body: 'x' }); + const cancel = jest.spyOn(response.body as ReadableStream, 'cancel'); + await expect(fetchMetricsText('http://x', null, jest.fn().mockResolvedValue(response))).rejects.toBeInstanceOf(ScrapeError); + expect(cancel).toHaveBeenCalled(); + } + }); + + it('maps fetch failures to timeout and connection failed', async () => { + const timeout = jest.fn().mockRejectedValue(Object.assign(new Error('t'), { name: 'TimeoutError' })); + await expect(fetchMetricsText('http://x', null, timeout)).rejects.toThrow('timeout'); + const refused = jest.fn().mockRejectedValue(new Error('ECONNREFUSED')); + await expect(fetchMetricsText('http://x', null, refused)).rejects.toThrow('connection failed'); + }); + + it('rejects oversized responses by header and by stream', async () => { + const header = jest.fn().mockResolvedValue(respond({ headers: { 'content-length': '6000000' } })); + await expect(fetchMetricsText('http://x', null, header)).rejects.toThrow('response too large'); + const stream = new ReadableStream({ + pull(controller) { + controller.enqueue(new Uint8Array(1024 * 1024)); + }, + }); + const streamed = jest.fn().mockResolvedValue(new Response(stream, { status: 200 })); + await expect(fetchMetricsText('http://x', null, streamed)).rejects.toThrow('response too large'); + }); + + it('maps a body that never ends to timeout when the signal aborts', async () => { + const controller = new AbortController(); + const timeoutSpy = jest.spyOn(AbortSignal, 'timeout').mockReturnValue(controller.signal); + const fetchImpl = jest.fn().mockImplementation(async (_url: string, init: { signal: AbortSignal }) => { + const stream = new ReadableStream({ + start(streamController) { + init.signal.addEventListener('abort', () => streamController.error(Object.assign(new Error('t'), { name: 'TimeoutError' }))); + }, + }); + setTimeout(() => controller.abort(), 10); + return new Response(stream, { status: 200 }); + }); + try { + await expect(fetchMetricsText('http://x', null, fetchImpl)).rejects.toThrow('timeout'); + } finally { + timeoutSpy.mockRestore(); + } + }); + + it('leaves DNS checks to connect time so the scrape timeout covers them', async () => { + await fetchMetricsText('http://x/metrics', null, jest.fn().mockResolvedValue(respond())); + expect(assertSafeOutboundUrl).toHaveBeenLastCalledWith('http://x/metrics', expect.objectContaining({ resolveDns: 'never' })); + }); + + it('surfaces the guard reason', async () => { + (assertSafeOutboundUrl as jest.Mock).mockRejectedValue(new BadRequestException('Failed to resolve metrics URL hostname')); + await expect(fetchMetricsText('http://x', null, jest.fn())).rejects.toThrow(new ScrapeError('Failed to resolve metrics URL hostname')); + }); + + describe('default transport', () => { + let server: Server; + let port: number; + let seen: IncomingHttpHeaders | null; + + beforeAll(async () => { + server = createServer((req, res) => { + seen = req.headers; + res.setHeader('content-type', 'text/plain'); + res.end('lmcache:a_total 1'); + }); + await new Promise((resolve) => server.listen(0, '127.0.0.1', resolve)); + port = (server.address() as AddressInfo).port; + }); + + afterAll(() => new Promise((resolve) => server.close(() => resolve()))); + + beforeEach(() => { + seen = null; + }); + + afterEach(() => jest.restoreAllMocks()); + + it('connects to the validated address and keeps the original Host header', async () => { + jest.spyOn(dns, 'lookup').mockResolvedValue([{ address: '127.0.0.1', family: 4 }] as any); + await expect(fetchMetricsText(`http://vllm.internal:${port}/metrics`, 'Bearer t')).resolves.toBe('lmcache:a_total 1'); + expect(seen?.host).toBe(`vllm.internal:${port}`); + expect(seen?.authorization).toBe('Bearer t'); + }); + + it('refuses to connect when the hostname re-resolves to a blocked address', async () => { + jest.spyOn(dns, 'lookup').mockResolvedValue([{ address: '169.254.169.254', family: 4 }] as any); + await expect(fetchMetricsText(`http://vllm.internal:${port}/metrics`, 'Bearer t')).rejects.toThrow( + new ScrapeError('Metrics URL resolves to a blocked address'), + ); + expect(seen).toBeNull(); + }); + + it('reports a DNS failure', async () => { + jest.spyOn(dns, 'lookup').mockRejectedValue(Object.assign(new Error('getaddrinfo ENOTFOUND vllm.internal'), { code: 'ENOTFOUND' })); + await expect(fetchMetricsText(`http://vllm.internal:${port}/metrics`, null)).rejects.toThrow(new ScrapeError('DNS lookup failed (ENOTFOUND)')); + }); + }); + + describe('out-of-range status', () => { + const rawServer = (statusLine: string) => + new Promise<{ port: number; close: () => Promise }>((resolve) => { + const server = createTcpServer((socket) => { + socket.once('data', () => socket.end(`${statusLine}\r\nContent-Length: 0\r\n\r\n`)); + }); + server.listen(0, '127.0.0.1', () => + resolve({ + port: (server.address() as AddressInfo).port, + close: () => new Promise((done) => server.close(() => done())), + }), + ); + }); + + afterEach(() => jest.restoreAllMocks()); + + it.each(['HTTP/1.1 600 X', 'HTTP/1.1 101 Switching Protocols', 'HTTP/1.1 199 X'])('rejects %s as a scrape error', async (statusLine) => { + jest.spyOn(dns, 'lookup').mockResolvedValue([{ address: '127.0.0.1', family: 4 }] as any); + const target = await rawServer(statusLine); + const uncaught = jest.fn(); + process.on('uncaughtException', uncaught); + try { + const error = await fetchMetricsText(`http://vllm.internal:${target.port}/metrics`, null).catch((e) => e); + await new Promise((resolve) => setImmediate(resolve)); + expect(error).toBeInstanceOf(ScrapeError); + expect(uncaught).not.toHaveBeenCalled(); + } finally { + process.off('uncaughtException', uncaught); + await target.close(); + } + }); + }); + + it('blocks addresses the guard rejects', async () => { + (assertSafeOutboundUrl as jest.Mock).mockRejectedValue(new Error('nope')); + const fetchImpl = jest.fn(); + const error = await fetchMetricsText('http://x', null, fetchImpl).catch((e) => e); + expect(error).toBeInstanceOf(ScrapeError); + expect(error.message).toBe('blocked address'); + expect(fetchImpl).not.toHaveBeenCalled(); + }); +}); + +describe('scrapeAllowsPrivateNetworks', () => { + const saved = { cloud: process.env.CLOUD_MODE, block: process.env.KV_CACHE_SCRAPE_BLOCK_PRIVATE }; + afterEach(() => { + for (const [key, value] of [['CLOUD_MODE', saved.cloud], ['KV_CACHE_SCRAPE_BLOCK_PRIVATE', saved.block]] as const) { + if (value === undefined) delete process.env[key]; + else process.env[key] = value; + } + }); + + it('allows private targets on self-hosted unless the flag blocks them', () => { + delete process.env.CLOUD_MODE; + delete process.env.KV_CACHE_SCRAPE_BLOCK_PRIVATE; + expect(scrapeAllowsPrivateNetworks()).toBe(true); + process.env.KV_CACHE_SCRAPE_BLOCK_PRIVATE = 'true'; + expect(scrapeAllowsPrivateNetworks()).toBe(false); + }); + + it('blocks on padded or alternative truthy spellings of the flag', () => { + delete process.env.CLOUD_MODE; + for (const flag of ['true\n', ' true ', 'TRUE', '1', 'yes', 'on']) { + process.env.KV_CACHE_SCRAPE_BLOCK_PRIVATE = flag; + expect(scrapeAllowsPrivateNetworks()).toBe(false); + } + }); + + it('allows private targets for explicit falsy spellings of the flag', () => { + delete process.env.CLOUD_MODE; + for (const flag of ['', 'false', ' FALSE\n', '0', 'no', 'off']) { + process.env.KV_CACHE_SCRAPE_BLOCK_PRIVATE = flag; + expect(scrapeAllowsPrivateNetworks()).toBe(true); + } + }); + + it('fails closed and warns once when the flag value is not recognised', () => { + delete process.env.CLOUD_MODE; + const warn = jest.spyOn(Logger.prototype, 'warn').mockImplementation(); + process.env.KV_CACHE_SCRAPE_BLOCK_PRIVATE = 'ture'; + expect(scrapeAllowsPrivateNetworks()).toBe(false); + expect(scrapeAllowsPrivateNetworks()).toBe(false); + expect(warn).toHaveBeenCalledTimes(1); + expect(warn.mock.calls[0][0]).toContain('KV_CACHE_SCRAPE_BLOCK_PRIVATE'); + warn.mockRestore(); + }); + + it('always blocks private targets in cloud mode', () => { + process.env.CLOUD_MODE = 'true'; + for (const flag of [undefined, 'false', 'true']) { + if (flag === undefined) delete process.env.KV_CACHE_SCRAPE_BLOCK_PRIVATE; + else process.env.KV_CACHE_SCRAPE_BLOCK_PRIVATE = flag; + expect(scrapeAllowsPrivateNetworks()).toBe(false); + } + }); +}); diff --git a/proprietary/kv-cache/__tests__/prometheus-text.spec.ts b/proprietary/kv-cache/__tests__/prometheus-text.spec.ts new file mode 100644 index 000000000..e5ad0a5f1 --- /dev/null +++ b/proprietary/kv-cache/__tests__/prometheus-text.spec.ts @@ -0,0 +1,29 @@ +import * as fs from 'fs'; +import * as path from 'path'; +import { parseLmcacheMetrics } from '../prometheus-text'; + +const fixture = fs.readFileSync(path.join(__dirname, 'fixtures/lmcache-metrics.txt'), 'utf8'); + +describe('parseLmcacheMetrics', () => { + it('keeps lmcache counters, strips _total and skips the rest', () => { + const samples = parseLmcacheMetrics(fixture); + expect(samples.map((s) => [s.metric, s.value])).toEqual([ + ['num_requested_tokens', 1000], + ['num_hit_tokens', 700], + ['remote_ping_errors', 2], + ]); + }); + + it('reads labels including escapes', () => { + const [first, , last] = parseLmcacheMetrics(fixture); + expect(first.labels).toEqual({ model_name: 'Qwen/Qwen2.5-0.5B-Instruct', role: 'worker', served_model_name: 'qwen', worker_id: '0' }); + expect(last.labels.model_name).toBe('a"b\\c\nd'); + expect(last.labels.worker_id).toBe('1'); + }); + + it('accepts series without labels and ignores malformed lines', () => { + expect(parseLmcacheMetrics('lmcache:num_lookup_hits_total 5\nlmcache:broken{model_name="x" 1\n')).toEqual([ + { metric: 'num_lookup_hits', labels: {}, value: 5 }, + ]); + }); +}); diff --git a/proprietary/kv-cache/counter-deltas.ts b/proprietary/kv-cache/counter-deltas.ts new file mode 100644 index 000000000..c81f4ea40 --- /dev/null +++ b/proprietary/kv-cache/counter-deltas.ts @@ -0,0 +1,133 @@ +import type { KvCacheEngineSample, KvCacheSampleCounters } from '@betterdb/shared'; + +export const METRIC_FIELDS: Readonly> = { + num_requested_tokens: 'requestedTokens', + num_hit_tokens: 'hitTokens', + num_lookup_tokens: 'lookupTokens', + num_lookup_hits: 'lookupHits', + num_remote_read_bytes: 'remoteReadBytes', + num_remote_write_bytes: 'remoteWriteBytes', + num_remote_read_requests: 'remoteReadRequests', + num_remote_write_requests: 'remoteWriteRequests', + remote_ping_errors: 'remotePingErrors', +}; + +export interface CounterObservation { + engineId: string; + connectionId: string; + modelName: string; + series: string; + metric: string; + value: number; + cumulative: boolean; + startMs?: number; +} + +export const BUCKET_MS = 60_000; +export const BASELINE_IDLE_MS = 60 * 60_000; + +const KEYED_LABELS = new Set(['model_name', 'worker_id', 'role']); + +export function labelSeries(labels: Record): string { + const base = `${labels.worker_id ?? ''}|${labels.role ?? ''}`; + const extras = Object.entries(labels) + .filter(([key]) => !KEYED_LABELS.has(key)) + .sort(([a], [b]) => (a < b ? -1 : a > b ? 1 : 0)); + return extras.length === 0 ? base : `${base}|${JSON.stringify(extras)}`; +} + +const COUNTER_FIELDS = Object.values(METRIC_FIELDS); + +const emptyCounters = (): KvCacheSampleCounters => + Object.fromEntries(COUNTER_FIELDS.map((field) => [field, 0])) as unknown as KvCacheSampleCounters; + +export function addCounters(target: KvCacheSampleCounters, source: KvCacheSampleCounters): void { + for (const field of COUNTER_FIELDS) target[field] += source[field]; +} + +const bucketStart = (ms: number): number => Math.floor(ms / BUCKET_MS) * BUCKET_MS; + +export class CounterDeltaTracker { + private readonly last = new Map(); + private readonly buckets = new Map(); + private drainedBefore = Number.NEGATIVE_INFINITY; + + observe(observation: CounterObservation, nowMs: number): boolean { + if (!Object.hasOwn(METRIC_FIELDS, observation.metric)) return false; + if (!Number.isFinite(observation.value)) return false; + const field = METRIC_FIELDS[observation.metric]; + const delta = observation.cumulative + ? this.cumulativeDelta(observation, nowMs) + : observation.value >= 0 + ? observation.value + : null; + if (delta === null) return true; + const timestamp = Math.max(bucketStart(nowMs), this.drainedBefore); + const id = `${observation.engineId}|${observation.modelName}|${timestamp}`; + const bucket = this.buckets.get(id) ?? { + engineId: observation.engineId, + connectionId: observation.connectionId, + modelName: observation.modelName, + timestamp, + ...emptyCounters(), + }; + bucket[field] += delta; + this.buckets.set(id, bucket); + return true; + } + + drain(nowMs: number, includeOpen = false): KvCacheEngineSample[] { + const open = bucketStart(nowMs); + for (const [id, baseline] of this.last) { + if (nowMs - baseline.seenMs > BASELINE_IDLE_MS) this.last.delete(id); + } + this.drainedBefore = Math.max(this.drainedBefore, includeOpen ? open + BUCKET_MS : open); + const drained: KvCacheEngineSample[] = []; + for (const [id, bucket] of this.buckets) { + if (includeOpen || bucket.timestamp < open) { + drained.push(bucket); + this.buckets.delete(id); + } + } + return drained.sort((a, b) => a.timestamp - b.timestamp); + } + + baselineCount(): number { + return this.last.size; + } + + forgetEngine(engineId: string): void { + this.forgetBaselines(engineId); + for (const [id, bucket] of this.buckets) { + if (bucket.engineId === engineId) this.buckets.delete(id); + } + } + + forgetBaselines(engineId: string): void { + for (const id of [...this.last.keys()]) { + if (id.startsWith(`${engineId}|`)) this.last.delete(id); + } + } + + private cumulativeDelta(observation: CounterObservation, seenMs: number): number | null { + const id = `${observation.engineId}|${observation.modelName}|${observation.series}|${observation.metric}`; + const previous = this.last.get(id); + const startMs = observation.startMs ?? previous?.startMs; + if (!previous) { + this.last.set(id, { value: observation.value, startMs, seenMs }); + return null; + } + const order = + observation.startMs !== undefined && previous.startMs !== undefined + ? Math.sign(observation.startMs - previous.startMs) + : null; + if (order !== null && order < 0) return null; + if (order !== null && order > 0) { + this.last.set(id, { value: observation.value, startMs, seenMs }); + return null; + } + if (order === 0 && observation.value < previous.value) return null; + this.last.set(id, { value: observation.value, startMs, seenMs }); + return observation.value < previous.value ? observation.value : observation.value - previous.value; + } +} diff --git a/proprietary/kv-cache/dto/kv-cache-engine.dto.ts b/proprietary/kv-cache/dto/kv-cache-engine.dto.ts new file mode 100644 index 000000000..d50a29b18 --- /dev/null +++ b/proprietary/kv-cache/dto/kv-cache-engine.dto.ts @@ -0,0 +1,41 @@ +import { IsBoolean, IsIn, IsString, IsUrl, Length, Matches, MaxLength, ValidateBy, ValidateIf } from 'class-validator'; + +const URL_OPTIONS = { require_tld: false, require_protocol: true, allow_underscores: true, protocols: ['http', 'https'] }; + +export const OTLP_ENGINE_ID_PATTERN = /^[A-Za-z0-9._-]{1,64}$/; + +const provided = (_: object, value: unknown) => value !== undefined; + +function hasNoCredentials(value: unknown): boolean { + try { + const url = new URL(String(value)); + return !url.username && !url.password; + } catch { + return true; + } +} + +const NoUrlCredentials = () => + ValidateBy({ + name: 'noUrlCredentials', + validator: { + validate: hasNoCredentials, + defaultMessage: () => 'scrapeUrl must not contain credentials; send them in scrapeAuthHeader instead', + }, + }); + +export class CreateKvCacheEngineDto { + @IsString() @Length(1, 100) name!: string; + @IsIn(['scrape', 'otlp']) source!: 'scrape' | 'otlp'; + @ValidateIf((o) => o.source === 'scrape') @IsUrl(URL_OPTIONS) @NoUrlCredentials() scrapeUrl?: string; + @ValidateIf(provided) @IsString() @MaxLength(4096) scrapeAuthHeader?: string; + @ValidateIf(provided) @Matches(OTLP_ENGINE_ID_PATTERN) otlpEngineId?: string; + @ValidateIf(provided) @IsBoolean() enabled?: boolean; +} + +export class UpdateKvCacheEngineDto { + @ValidateIf(provided) @IsString() @Length(1, 100) name?: string; + @ValidateIf(provided) @IsUrl(URL_OPTIONS) @NoUrlCredentials() scrapeUrl?: string; + @ValidateIf((_, value) => value !== undefined && value !== null) @IsString() @MaxLength(4096) scrapeAuthHeader?: string | null; + @ValidateIf(provided) @IsBoolean() enabled?: boolean; +} diff --git a/proprietary/kv-cache/dto/kv-cache-settings.dto.ts b/proprietary/kv-cache/dto/kv-cache-settings.dto.ts new file mode 100644 index 000000000..81102cf73 --- /dev/null +++ b/proprietary/kv-cache/dto/kv-cache-settings.dto.ts @@ -0,0 +1,10 @@ +import { IsBoolean, IsNumber, Max, Min, ValidateIf } from 'class-validator'; +import { MAX_KV_CACHE_HIT_RATE_THRESHOLD } from '@betterdb/shared'; + +const provided = (_: object, value: unknown) => value !== undefined; + +export class UpdateKvCacheSettingsDto { + @ValidateIf(provided) @IsBoolean() hitRateAlertEnabled?: boolean; + @ValidateIf(provided) @IsNumber() @Min(0.01) @Max(MAX_KV_CACHE_HIT_RATE_THRESHOLD) hitRateThreshold?: number; + @ValidateIf(provided) @IsBoolean() evictionAlertEnabled?: boolean; +} diff --git a/proprietary/kv-cache/kv-cache-alert-rules.ts b/proprietary/kv-cache/kv-cache-alert-rules.ts new file mode 100644 index 000000000..020e447cf --- /dev/null +++ b/proprietary/kv-cache/kv-cache-alert-rules.ts @@ -0,0 +1,46 @@ +import { + KV_CACHE_HIT_RATE_MIN_TOKENS, + type KvCacheEngineSample, + type KvCacheEvictionReason, + type KvCacheFootprintSnapshot, + type KvCacheSettings, +} from '@betterdb/shared'; +import { hitRate } from './kv-cache-samples.service'; + +export interface HitRateWindow { + engineId: string; + model: string; + requestedTokens: number; + hitTokens: number; + hitRate: number | null; +} + +export function windowHitRates(samples: KvCacheEngineSample[]): HitRateWindow[] { + const windows = new Map(); + for (const s of samples) { + const id = `${s.engineId}|${s.modelName}`; + const w = windows.get(id) ?? { engineId: s.engineId, model: s.modelName, requestedTokens: 0, hitTokens: 0, hitRate: null }; + w.requestedTokens += s.requestedTokens; + w.hitTokens += s.hitTokens; + windows.set(id, w); + } + return [...windows.values()] + .map((w) => ({ ...w, hitRate: hitRate(w.requestedTokens, w.hitTokens) })) + .sort((a, b) => a.engineId.localeCompare(b.engineId) || a.model.localeCompare(b.model)); +} + +export function shouldCheckHitRate(window: HitRateWindow, settings: KvCacheSettings): boolean { + return settings.hitRateAlertEnabled && window.hitRate !== null && window.requestedTokens >= KV_CACHE_HIT_RATE_MIN_TOKENS; +} + +export function evictionConditions(s: KvCacheFootprintSnapshot, wasUnevictable = false): Record { + return { + unevictable: s.detected && s.maxmemoryPolicy.startsWith('volatile-') && (s.noTtlRatio === null ? wasUnevictable : s.noTtlRatio >= 0.9), + evicting: + s.detected && + (s.evictedKeysDelta ?? 0) > 0 && + s.maxmemory > 0 && + s.usedMemory / s.maxmemory >= 0.9 && + (s.lmcacheMemoryShare === null || s.lmcacheMemoryShare >= 0.5), + }; +} diff --git a/proprietary/kv-cache/kv-cache-alerts.service.ts b/proprietary/kv-cache/kv-cache-alerts.service.ts new file mode 100644 index 000000000..1caa1639f --- /dev/null +++ b/proprietary/kv-cache/kv-cache-alerts.service.ts @@ -0,0 +1,250 @@ +import { Inject, Injectable, Logger, OnModuleDestroy, OnModuleInit, Optional } from '@nestjs/common'; +import { + DEFAULT_KV_CACHE_HIT_RATE_THRESHOLD, + Feature, + KV_CACHE_HIT_RATE_WINDOW_MS, + WEBHOOK_EVENTS_PRO_SERVICE, + type IWebhookEventsProService, + type KvCacheEvictionReason, + type KvCacheFootprintSnapshot, + type KvCacheSettings, + type KvCacheSettingsUpdate, +} from '@betterdb/shared'; +import type { StoragePort } from '@app/common/interfaces/storage-port.interface'; +import { ConnectionRegistry } from '@app/connections/connection-registry.service'; +import { PrometheusService } from '@app/prometheus/prometheus.service'; +import { LicenseService } from '@proprietary/licenses/license.service'; +import { evictionConditions, shouldCheckHitRate, windowHitRates } from './kv-cache-alert-rules'; +import { KvCacheEngineRegistry } from './kv-cache-engine-registry'; +import { KvCacheFootprintService } from './kv-cache-footprint.service'; + +const EVALUATE_INTERVAL_MS = 60_000; +const EVICTION_REASONS: KvCacheEvictionReason[] = ['unevictable', 'evicting']; +const EVICTING_CLEAR_SNAPSHOTS = 3; + +interface CheckedWindow { + engineId: string; + engineName: string; + model: string; +} + +@Injectable() +export class KvCacheAlertsService implements OnModuleInit, OnModuleDestroy { + private readonly logger = new Logger(KvCacheAlertsService.name); + private timer: NodeJS.Timeout | null = null; + private readonly exportedConnections = new Set(); + private readonly checkedWindows = new Map>(); + private readonly evictingCleanStreaks = new Map(); + private readonly unevictableActive = new Map(); + private evaluating = false; + + constructor( + @Inject('STORAGE_CLIENT') private readonly storage: StoragePort, + private readonly connectionRegistry: ConnectionRegistry, + private readonly footprint: KvCacheFootprintService, + private readonly registry: KvCacheEngineRegistry, + private readonly license: LicenseService, + @Optional() + @Inject(WEBHOOK_EVENTS_PRO_SERVICE) + private readonly webhookEventsPro?: IWebhookEventsProService, + @Optional() private readonly prometheus?: PrometheusService, + ) {} + + onModuleInit(): void { + this.footprint.onSnapshot((snapshot) => this.onSnapshot(snapshot)); + this.footprint.onConnectionRemoval((connectionId) => { + this.exportedConnections.delete(connectionId); + this.checkedWindows.delete(connectionId); + this.evictingCleanStreaks.delete(connectionId); + this.unevictableActive.delete(connectionId); + this.prometheus?.clearKvCache(connectionId); + }); + this.timer = setInterval(() => void this.evaluateHitRates(), EVALUATE_INTERVAL_MS); + this.timer.unref(); + } + + onModuleDestroy(): void { + if (this.timer) clearInterval(this.timer); + this.timer = null; + } + + async getSettings(connectionId: string): Promise { + const stored = await this.storage.getKvCacheSettings(connectionId); + return ( + stored ?? { + connectionId, + hitRateAlertEnabled: true, + hitRateThreshold: DEFAULT_KV_CACHE_HIT_RATE_THRESHOLD, + evictionAlertEnabled: true, + updatedAt: 0, + } + ); + } + + async updateSettings(connectionId: string, update: KvCacheSettingsUpdate): Promise { + const current = await this.getSettings(connectionId); + const definedUpdate = Object.fromEntries(Object.entries(update).filter(([, value]) => value !== undefined)); + return this.storage.saveKvCacheSettings({ ...current, ...definedUpdate, connectionId, updatedAt: Date.now() }); + } + + async onSnapshot(snapshot: KvCacheFootprintSnapshot): Promise { + try { + if (!this.licensed()) return; + if (snapshot.scanComplete) this.prometheus?.setKvCacheFootprint(snapshot.connectionId, snapshot.perModel); + if (!this.webhookEventsPro) return; + const settings = await this.getSettings(snapshot.connectionId); + const conditions = this.evictionState(snapshot, settings.evictionAlertEnabled); + const instance = this.instanceFor(snapshot.connectionId); + for (const reason of EVICTION_REASONS) { + void this.webhookEventsPro + .dispatchKvCacheEvictionRisk({ + connectionId: snapshot.connectionId, + reason, + active: conditions[reason], + policy: snapshot.maxmemoryPolicy, + usedMemory: snapshot.usedMemory, + maxmemory: snapshot.maxmemory, + lmcacheMemoryShare: snapshot.lmcacheMemoryShare, + noTtlRatio: snapshot.noTtlRatio, + evictedKeysDelta: snapshot.evictedKeysDelta, + timestamp: snapshot.timestamp, + ...(instance ? { instance } : {}), + }) + .catch((error: unknown) => this.warnEvictionFailure(snapshot.connectionId, error)); + } + } catch (error) { + this.warnEvictionFailure(snapshot.connectionId, error); + } + } + + private warnEvictionFailure(connectionId: string, error: unknown): void { + this.logger.warn(`Eviction risk alert failed for ${connectionId}: ${error instanceof Error ? error.message : error}`); + } + + private evictionState(snapshot: KvCacheFootprintSnapshot, enabled: boolean): Record { + if (!enabled || !snapshot.detected) { + this.evictingCleanStreaks.delete(snapshot.connectionId); + this.unevictableActive.delete(snapshot.connectionId); + return { unevictable: false, evicting: false }; + } + const conditions = evictionConditions(snapshot, this.unevictableActive.get(snapshot.connectionId)); + if (conditions.unevictable) this.unevictableActive.set(snapshot.connectionId, true); + else this.unevictableActive.delete(snapshot.connectionId); + conditions.evicting = this.holdEvicting(snapshot.connectionId, conditions.evicting); + return conditions; + } + + private holdEvicting(connectionId: string, evicting: boolean): boolean { + if (evicting) { + this.evictingCleanStreaks.set(connectionId, 0); + return true; + } + const streak = this.evictingCleanStreaks.get(connectionId); + if (streak === undefined) return false; + if (streak + 1 >= EVICTING_CLEAR_SNAPSHOTS) { + this.evictingCleanStreaks.delete(connectionId); + return false; + } + this.evictingCleanStreaks.set(connectionId, streak + 1); + return true; + } + + async evaluateHitRates(now = Date.now()): Promise { + if (this.evaluating) return; + this.evaluating = true; + try { + if (!this.licensed()) return; + const connectionIds = new Set([ + ...this.registry.list().map((engine) => engine.connectionId), + ...this.exportedConnections, + ...this.checkedWindows.keys(), + ]); + for (const connectionId of connectionIds) { + try { + await this.evaluateConnection(connectionId, now); + } catch (error) { + this.logger.warn(`Hit rate alert evaluation failed for ${connectionId}: ${error instanceof Error ? error.message : error}`); + } + } + } catch (error) { + this.logger.warn(`Hit rate alert evaluation failed: ${error instanceof Error ? error.message : error}`); + } finally { + this.evaluating = false; + } + } + + private async evaluateConnection(connectionId: string, now: number): Promise { + const rows = await this.storage.getKvCacheEngineSamples({ connectionId, from: now - KV_CACHE_HIT_RATE_WINDOW_MS, to: now }); + const settings = await this.getSettings(connectionId); + const instance = this.instanceFor(connectionId); + const windows = windowHitRates(rows); + const resolved = windows.map((window) => ({ window, engine: this.registry.get(window.engineId) })); + const rates = resolved.flatMap(({ window, engine }) => + engine?.enabled && window.hitRate !== null + ? [{ engine: engine.name, model: window.model, hitRate: window.hitRate }] + : [], + ); + this.prometheus?.setKvCacheHitRates(connectionId, rates); + if (rates.length > 0) this.exportedConnections.add(connectionId); + else this.exportedConnections.delete(connectionId); + if (!this.webhookEventsPro) return; + const checked = new Map(); + const previous = this.checkedWindows.get(connectionId); + for (const { window, engine } of resolved) { + if (!engine?.enabled || !settings.hitRateAlertEnabled) continue; + const key = JSON.stringify([window.engineId, window.model]); + if (window.hitRate === null || !shouldCheckHitRate(window, settings)) { + const held = previous?.get(key); + if (held) checked.set(key, held); + continue; + } + const target = { engineId: window.engineId, engineName: engine.name, model: window.model }; + checked.set(key, target); + await this.dispatchHitRate(connectionId, target, window.hitRate, window.requestedTokens, settings, now, instance); + } + for (const [key, target] of previous ?? []) { + if (!checked.has(key)) await this.dispatchHitRate(connectionId, target, 1, 0, settings, now, instance); + } + if (checked.size > 0) this.checkedWindows.set(connectionId, checked); + else this.checkedWindows.delete(connectionId); + } + + private async dispatchHitRate( + connectionId: string, + target: CheckedWindow, + hitRate: number, + requestedTokens: number, + settings: KvCacheSettings, + now: number, + instance: { host: string; port: number } | undefined, + ): Promise { + try { + await this.webhookEventsPro?.dispatchKvCacheHitRateLow({ + connectionId, + engineId: target.engineId, + engineName: target.engineName, + model: target.model, + hitRate, + threshold: settings.hitRateThreshold, + requestedTokens, + windowMs: KV_CACHE_HIT_RATE_WINDOW_MS, + timestamp: now, + ...(instance ? { instance } : {}), + }); + } catch (error) { + this.logger.warn(`Hit rate alert dispatch failed for ${target.engineId}/${target.model}: ${error instanceof Error ? error.message : error}`); + } + } + + private licensed(): boolean { + if (this.license.hasFeature(Feature.KV_CACHE_MONITORING)) return true; + this.prometheus?.clearAllKvCache(); + this.exportedConnections.clear(); + return false; + } + + private instanceFor(connectionId: string): { host: string; port: number } | undefined { + const config = this.connectionRegistry.getConfig(connectionId); + return config ? { host: config.host, port: config.port } : undefined; + } +} diff --git a/proprietary/kv-cache/kv-cache-engine-registry.ts b/proprietary/kv-cache/kv-cache-engine-registry.ts new file mode 100644 index 000000000..e7c37b988 --- /dev/null +++ b/proprietary/kv-cache/kv-cache-engine-registry.ts @@ -0,0 +1,145 @@ +import { Inject, Injectable, Logger, NotFoundException, OnModuleInit } from '@nestjs/common'; +import type { StoragePort, StoredKvCacheEngine } from '@app/common/interfaces/storage-port.interface'; +import { KvCacheFootprintService } from './kv-cache-footprint.service'; +import { KvCacheSamplesService } from './kv-cache-samples.service'; + +export const KV_CACHE_LAST_SEEN_PERSIST_MS = 5 * 60_000; +export const REMOVAL_TOMBSTONE_MS = 60 * 60_000; + +@Injectable() +export class KvCacheEngineRegistry implements OnModuleInit { + private readonly logger = new Logger(KvCacheEngineRegistry.name); + private readonly engines = new Map(); + private readonly writes = new Map>(); + private readonly statusVersions = new Map(); + private readonly removed = new Map(); + private readonly removedConnections = new Map(); + private readonly persistedSeenAt = new Map(); + + constructor( + @Inject('STORAGE_CLIENT') private readonly storage: StoragePort, + private readonly footprint: KvCacheFootprintService, + private readonly samples: KvCacheSamplesService, + ) {} + + async onModuleInit(): Promise { + for (const engine of await this.storage.getKvCacheEngines()) { + this.engines.set(engine.id, engine); + this.persistedSeenAt.set(engine.id, engine.lastSeenAt); + } + this.footprint.onConnectionRemoval(async (connectionId) => { + this.tombstone(this.removedConnections, connectionId); + const ids = this.list(connectionId).map((engine) => engine.id); + for (const id of ids) { + this.tombstone(this.removed, id); + this.engines.delete(id); + this.statusVersions.delete(id); + this.persistedSeenAt.delete(id); + this.samples.forgetEngine(id); + } + await Promise.all([...ids.map((id) => this.writes.get(id)), this.samples.idle()]); + }); + } + + list(connectionId?: string): StoredKvCacheEngine[] { + const all = [...this.engines.values()]; + return connectionId === undefined ? all : all.filter((engine) => engine.connectionId === connectionId); + } + + get(id: string): StoredKvCacheEngine | null { + return this.engines.get(id) ?? null; + } + + byOtlpId(otlpEngineId: string): StoredKvCacheEngine | null { + for (const engine of this.engines.values()) { + if (engine.otlpEngineId === otlpEngineId) return engine; + } + return null; + } + + save(engine: StoredKvCacheEngine): Promise { + const version = this.statusVersions.get(engine.id); + return this.serialize(engine.id, async () => { + if (this.removed.has(engine.id)) throw new NotFoundException(`KV cache engine ${engine.id} not found`); + this.assertConnectionLive(engine.connectionId); + const saved = await this.storage.saveKvCacheEngine(engine); + if (this.removedConnections.has(saved.connectionId)) { + await this.storage.deleteKvCacheEngine(saved.id); + this.assertConnectionLive(saved.connectionId); + } + if (this.removed.has(saved.id)) throw new NotFoundException(`KV cache engine ${saved.id} not found`); + const latest = this.engines.get(saved.id); + const merged = + latest && latest.scrapeUrl === saved.scrapeUrl && this.statusVersions.get(saved.id) !== version + ? { ...saved, lastSeenAt: latest.lastSeenAt, lastError: latest.lastError } + : saved; + this.engines.set(merged.id, merged); + this.persistedSeenAt.set(merged.id, saved.lastSeenAt); + return merged; + }); + } + + remove(id: string): Promise { + this.tombstone(this.removed, id); + return this.serialize(id, async () => { + let removed: boolean; + try { + removed = await this.storage.deleteKvCacheEngine(id); + } catch (error) { + this.removed.delete(id); + throw error; + } + this.engines.delete(id); + this.statusVersions.delete(id); + this.persistedSeenAt.delete(id); + return removed; + }); + } + + recordResult(id: string, result: { lastSeenAt?: number; lastError: string | null }): void { + const current = this.engines.get(id); + if (!current) return; + const updated: StoredKvCacheEngine = { + ...current, + lastSeenAt: result.lastSeenAt ?? current.lastSeenAt, + lastError: result.lastError, + }; + this.engines.set(id, updated); + this.statusVersions.set(id, (this.statusVersions.get(id) ?? 0) + 1); + if (!this.needsPersist(id, current.lastError, updated)) return; + this.persistedSeenAt.set(id, updated.lastSeenAt); + void this.serialize(id, async () => { + const latest = this.engines.get(id); + if (latest && !this.removed.has(id)) await this.storage.saveKvCacheEngine(latest); + }).catch((error) => this.logger.warn(`Could not persist KV cache engine ${id}: ${error instanceof Error ? error.message : String(error)}`)); + } + + private tombstone(tombstones: Map, key: string): void { + const now = Date.now(); + for (const [existing, removedAt] of tombstones) { + if (now - removedAt > REMOVAL_TOMBSTONE_MS) tombstones.delete(existing); + } + tombstones.set(key, now); + } + + private needsPersist(id: string, previousError: string | null, updated: StoredKvCacheEngine): boolean { + if (previousError !== updated.lastError) return true; + const persisted = this.persistedSeenAt.get(id); + if (persisted === undefined || persisted === null) return updated.lastSeenAt !== null; + return updated.lastSeenAt !== null && updated.lastSeenAt - persisted >= KV_CACHE_LAST_SEEN_PERSIST_MS; + } + + private assertConnectionLive(connectionId: string): void { + if (this.removedConnections.has(connectionId)) throw new NotFoundException(`Connection '${connectionId}' not found`); + } + + private serialize(id: string, task: () => Promise): Promise { + const run = (this.writes.get(id) ?? Promise.resolve()).then(task); + const tail = run.catch(() => undefined); + this.writes.set(id, tail); + void tail.then(() => { + if (this.writes.get(id) === tail) this.writes.delete(id); + }); + return run; + } +} diff --git a/proprietary/kv-cache/kv-cache-engines.service.ts b/proprietary/kv-cache/kv-cache-engines.service.ts new file mode 100644 index 000000000..54fb69c1a --- /dev/null +++ b/proprietary/kv-cache/kv-cache-engines.service.ts @@ -0,0 +1,207 @@ +import { BadRequestException, ConflictException, Injectable, Logger, NotFoundException } from '@nestjs/common'; +import { randomBytes, randomUUID } from 'crypto'; +import type { KvCacheEngine } from '@betterdb/shared'; +import type { StoredKvCacheEngine } from '@app/common/interfaces/storage-port.interface'; +import { getEncryptionService } from '@app/common/utils/encryption'; +import { ConnectionRegistry } from '@app/connections/connection-registry.service'; +import { KvCacheEngineRegistry } from './kv-cache-engine-registry'; +import { OTLP_ENGINE_ID_PATTERN, type CreateKvCacheEngineDto, type UpdateKvCacheEngineDto } from './dto/kv-cache-engine.dto'; +import { KvCacheSamplesService } from './kv-cache-samples.service'; +import { fetchMetricsText, ScrapeError, UNDECRYPTABLE_AUTH } from './metrics-fetch'; +import { NO_METRICS, parseLmcacheMetrics } from './prometheus-text'; +import { toPublicEngine } from './kv-cache-status.service'; + +type StoredAuth = Pick; + +@Injectable() +export class KvCacheEnginesService { + private readonly logger = new Logger(KvCacheEnginesService.name); + private readonly reservedNames = new Set(); + private readonly reservedOtlpIds = new Set(); + private readonly pendingSampleDeletes = new Map(); + + constructor( + private readonly registry: KvCacheEngineRegistry, + private readonly connectionRegistry: ConnectionRegistry, + private readonly samples: KvCacheSamplesService, + ) {} + + list(connectionId: string): KvCacheEngine[] { + return this.registry.list(connectionId).map(toPublicEngine); + } + + async create(connectionId: string, dto: CreateKvCacheEngineDto): Promise { + if (!this.connectionRegistry.getConfig(connectionId)) throw new NotFoundException(`Connection '${connectionId}' not found`); + const release = this.reserveName(connectionId, dto.name); + try { + return await this.createReserved(connectionId, dto); + } finally { + release(); + } + } + + private async createReserved(connectionId: string, dto: CreateKvCacheEngineDto): Promise { + const base = { + id: randomUUID(), + connectionId, + name: dto.name, + enabled: dto.enabled ?? true, + createdAt: Date.now(), + }; + if (dto.source === 'otlp') return this.createOtlp(base, dto); + if (!dto.scrapeUrl) throw new BadRequestException('scrapeUrl is required for scrape engines'); + const lastError = await this.testScrape(dto.scrapeUrl, dto.scrapeAuthHeader ?? null); + const saved = await this.registry.save({ + ...base, + source: 'scrape', + scrapeUrl: dto.scrapeUrl, + ...this.storeAuthHeader(dto.scrapeAuthHeader ?? null), + otlpEngineId: null, + lastSeenAt: Date.now(), + lastError, + }); + return toPublicEngine(saved); + } + + async update(connectionId: string, id: string, dto: UpdateKvCacheEngineDto): Promise { + const engine = this.owned(connectionId, id); + if (dto.name !== undefined) this.assertUniqueName(connectionId, dto.name, id); + let probe: { lastError: string | null; lastSeenAt: number } | null = null; + if (engine.source === 'scrape') { + const urlChanged = dto.scrapeUrl !== undefined && dto.scrapeUrl !== engine.scrapeUrl; + const newHeader = typeof dto.scrapeAuthHeader === 'string' && dto.scrapeAuthHeader !== ''; + if (urlChanged && dto.scrapeAuthHeader === undefined && engine.scrapeAuthHeader && !this.sameOrigin(engine.scrapeUrl, dto.scrapeUrl)) { + throw new BadRequestException('Re-enter the scrape auth header when changing the scrape URL origin'); + } + const url = dto.scrapeUrl ?? engine.scrapeUrl; + if ((urlChanged || newHeader) && url) { + const header = dto.scrapeAuthHeader === undefined ? this.authHeaderFor(engine) : dto.scrapeAuthHeader || null; + if (dto.scrapeAuthHeader === undefined && engine.scrapeAuthHeader && header === null) { + throw new BadRequestException(`Test scrape failed: ${UNDECRYPTABLE_AUTH}`); + } + probe = { lastError: await this.testScrape(url, header), lastSeenAt: Date.now() }; + } + } + const fresh = this.owned(connectionId, id); + const release = dto.name === undefined ? () => undefined : this.reserveName(connectionId, dto.name, id); + const next: StoredKvCacheEngine = { ...fresh }; + if (dto.name !== undefined) next.name = dto.name; + if (dto.enabled !== undefined) next.enabled = dto.enabled; + if (fresh.source === 'scrape') { + if (dto.scrapeUrl !== undefined) next.scrapeUrl = dto.scrapeUrl; + if (dto.scrapeAuthHeader !== undefined) Object.assign(next, this.storeAuthHeader(dto.scrapeAuthHeader || null)); + if (probe) Object.assign(next, probe); + } + let saved: StoredKvCacheEngine; + try { + saved = await this.registry.save(next); + } finally { + release(); + } + if (fresh.enabled !== saved.enabled || fresh.scrapeUrl !== saved.scrapeUrl) this.samples.resetBaselines(id); + return toPublicEngine(saved); + } + + async remove(connectionId: string, id: string): Promise { + if (this.pendingSampleDeletes.get(id) !== connectionId) { + this.owned(connectionId, id); + await this.registry.remove(id); + this.pendingSampleDeletes.set(id, connectionId); + } + await this.samples.deleteEngine(id); + this.pendingSampleDeletes.delete(id); + } + + authHeaderFor(engine: StoredKvCacheEngine): string | null { + if (!engine.scrapeAuthHeader) return null; + if (!engine.scrapeAuthEncrypted) return engine.scrapeAuthHeader; + const encryption = getEncryptionService(); + if (encryption) { + try { + return encryption.decrypt(engine.scrapeAuthHeader); + } catch { + this.warnUndecryptable(engine.id); + return null; + } + } + this.warnUndecryptable(engine.id); + return null; + } + + private assertUniqueName(connectionId: string, name: string, exceptId?: string): void { + const taken = + this.reservedNames.has(JSON.stringify([connectionId, name])) || + this.registry.list(connectionId).some((engine) => engine.id !== exceptId && engine.name === name); + if (taken) throw new ConflictException(`An engine named "${name}" is already linked to this connection`); + } + + private reserveName(connectionId: string, name: string, exceptId?: string): () => void { + this.assertUniqueName(connectionId, name, exceptId); + const key = JSON.stringify([connectionId, name]); + this.reservedNames.add(key); + return () => this.reservedNames.delete(key); + } + + private sameOrigin(a: string | null, b: string | undefined): boolean { + try { + return new URL(a ?? '').origin === new URL(b ?? '').origin; + } catch { + return false; + } + } + + private warnUndecryptable(id: string): void { + this.logger.warn(`Could not decrypt the scrape auth header for engine ${id}`); + } + + private owned(connectionId: string, id: string): StoredKvCacheEngine { + const engine = this.registry.get(id); + if (!engine || engine.connectionId !== connectionId) throw new NotFoundException('Engine not found'); + return engine; + } + + private async createOtlp( + base: Pick, + dto: CreateKvCacheEngineDto, + ): Promise { + const otlpEngineId = dto.otlpEngineId ?? 'lmc-' + randomBytes(6).toString('hex'); + if (!OTLP_ENGINE_ID_PATTERN.test(otlpEngineId)) throw new BadRequestException('Invalid OTLP engine id'); + if (this.reservedOtlpIds.has(otlpEngineId) || this.registry.byOtlpId(otlpEngineId)) { + throw new ConflictException('OTLP engine id already in use'); + } + this.reservedOtlpIds.add(otlpEngineId); + try { + const saved = await this.registry.save({ + ...base, + source: 'otlp', + scrapeUrl: null, + scrapeAuthHeader: null, + scrapeAuthEncrypted: false, + otlpEngineId, + lastSeenAt: null, + lastError: null, + }); + return toPublicEngine(saved); + } finally { + this.reservedOtlpIds.delete(otlpEngineId); + } + } + + private async testScrape(url: string, header: string | null): Promise { + try { + const text = await fetchMetricsText(url, header); + return parseLmcacheMetrics(text).length === 0 ? NO_METRICS : null; + } catch (error) { + if (error instanceof ScrapeError) throw new BadRequestException(`Test scrape failed: ${error.message}`); + throw error; + } + } + + private storeAuthHeader(header: string | null): StoredAuth { + if (!header) return { scrapeAuthHeader: null, scrapeAuthEncrypted: false }; + const encryption = getEncryptionService(); + return encryption + ? { scrapeAuthHeader: encryption.encrypt(header), scrapeAuthEncrypted: true } + : { scrapeAuthHeader: header, scrapeAuthEncrypted: false }; + } +} diff --git a/proprietary/kv-cache/kv-cache-footprint.service.ts b/proprietary/kv-cache/kv-cache-footprint.service.ts index 3c6995b82..f90b5e45a 100644 --- a/proprietary/kv-cache/kv-cache-footprint.service.ts +++ b/proprietary/kv-cache/kv-cache-footprint.service.ts @@ -33,6 +33,7 @@ import { hasSingleKeyEvidence, hasTwoKeyEvidence, maskKey, resolveTagSuffixes } import { slotOwner, type SlotOwnership } from './key-slot'; export type SnapshotListener = (snapshot: KvCacheFootprintSnapshot) => void | Promise; +export type RemovalListener = (connectionId: string) => void | Promise; interface NodeReading { nodeKey: string; @@ -119,6 +120,7 @@ export class KvCacheFootprintService extends MultiConnectionPoller implements On private readonly undetectedUntil = new Map(); private readonly partialMisses = new Map(); private readonly listeners: SnapshotListener[] = []; + private readonly removalListeners: RemovalListener[] = []; constructor( connectionRegistry: ConnectionRegistry, @@ -166,11 +168,22 @@ export class KvCacheFootprintService extends MultiConnectionPoller implements On this.sampleKeys.delete(connectionId); this.undetectedUntil.delete(connectionId); this.partialMisses.delete(connectionId); + for (const listener of this.removalListeners) { + try { + await listener(connectionId); + } catch (error) { + this.logger.warn(`KV cache removal listener failed: ${errorMessage(error)}`); + } + } await this.storage .deleteKvCacheConnectionData(connectionId) .catch((error) => this.logger.warn(`Could not delete KV cache data for ${connectionId}: ${errorMessage(error)}`)); } + onConnectionRemoval(listener: RemovalListener): void { + this.removalListeners.push(listener); + } + onSnapshot(listener: SnapshotListener): void { this.listeners.push(listener); } diff --git a/proprietary/kv-cache/kv-cache-otlp-sink.service.ts b/proprietary/kv-cache/kv-cache-otlp-sink.service.ts new file mode 100644 index 000000000..d335eae2e --- /dev/null +++ b/proprietary/kv-cache/kv-cache-otlp-sink.service.ts @@ -0,0 +1,95 @@ +import { Injectable } from '@nestjs/common'; +import { Feature } from '@betterdb/shared'; +import { + KvCacheOtlpSink, + KvCacheOtlpSinkResult, + LMCACHE_ENGINE_ATTRIBUTE, +} from '@app/external-metrics/kv-cache-otlp-sink'; +import { + attrsToRecord, + isDeltaTemporality, + metricPointCount, + nanosToMs, + NO_RECORDED_VALUE, + pointValue, +} from '@app/external-metrics/otlp-metric-map'; +import type { DropReason, OtlpMetric } from '@app/external-metrics/otlp-metrics-types'; +import { LicenseService } from '@proprietary/licenses/license.service'; +import { labelSeries, METRIC_FIELDS } from './counter-deltas'; +import { KvCacheEngineRegistry } from './kv-cache-engine-registry'; +import { KvCacheSamplesService } from './kv-cache-samples.service'; + +function resourceSeries(resourceAttributes: Record): string { + const entries = Object.entries(resourceAttributes) + .filter(([key]) => key !== LMCACHE_ENGINE_ATTRIBUTE) + .sort(([a], [b]) => (a < b ? -1 : a > b ? 1 : 0)); + return JSON.stringify(entries); +} + +@Injectable() +export class KvCacheOtlpSinkService implements KvCacheOtlpSink { + constructor( + private readonly registry: KvCacheEngineRegistry, + private readonly samples: KvCacheSamplesService, + private readonly license: LicenseService, + ) {} + + ingest(resourceAttributes: Record, metrics: OtlpMetric[], nowMs: number): KvCacheOtlpSinkResult { + const engineKey = resourceAttributes[LMCACHE_ENGINE_ATTRIBUTE]; + const engine = + engineKey !== undefined && this.license.hasFeature(Feature.KV_CACHE_MONITORING) + ? this.registry.byOtlpId(engineKey) + : null; + if (!engine || !engine.enabled) { + const total = metrics.reduce((sum, metric) => sum + metricPointCount(metric), 0); + return { accepted: 0, dropped: total > 0 ? { unknown_engine: total } : {} }; + } + + const dropped: Partial> = {}; + const addDrop = (reason: DropReason, count: number) => { + if (count > 0) dropped[reason] = (dropped[reason] ?? 0) + count; + }; + let accepted = 0; + const resource = resourceSeries(resourceAttributes); + + for (const metric of metrics) { + const name = (metric.name ?? '').replace(/^lmcache:/, '').replace(/_total$/, ''); + if (!metric.sum) { + addDrop('unsupported_type', metricPointCount(metric)); + continue; + } + if (!Object.hasOwn(METRIC_FIELDS, name)) { + addDrop('unmapped_metric', metricPointCount(metric)); + continue; + } + const cumulative = !isDeltaTemporality(metric.sum.aggregationTemporality); + for (const dataPoint of metric.sum.dataPoints ?? []) { + if ((dataPoint.flags ?? 0) & NO_RECORDED_VALUE) continue; + const raw = pointValue(dataPoint); + const value = raw === null ? NaN : Number(raw); + if (!Number.isFinite(value)) { + addDrop('invalid_value', 1); + continue; + } + const labels = attrsToRecord(dataPoint.attributes); + this.samples.observe( + { + engineId: engine.id, + connectionId: engine.connectionId, + modelName: labels.model_name || 'unknown', + series: `${labelSeries(labels)}|${resource}`, + metric: name, + value, + cumulative, + startMs: dataPoint.startTimeUnixNano ? (nanosToMs(dataPoint.startTimeUnixNano) ?? undefined) : undefined, + }, + nowMs, + ); + accepted += 1; + } + } + + if (accepted > 0) this.registry.recordResult(engine.id, { lastSeenAt: nowMs, lastError: null }); + return { accepted, dropped }; + } +} diff --git a/proprietary/kv-cache/kv-cache-samples.service.ts b/proprietary/kv-cache/kv-cache-samples.service.ts new file mode 100644 index 000000000..5dc75e0e6 --- /dev/null +++ b/proprietary/kv-cache/kv-cache-samples.service.ts @@ -0,0 +1,113 @@ +import { Inject, Injectable, Logger, OnModuleDestroy, OnModuleInit } from '@nestjs/common'; +import type { KvCacheEngineSample, KvCacheSampleBucket, KvCacheSamplesResponse } from '@betterdb/shared'; +import type { StoragePort } from '@app/common/interfaces/storage-port.interface'; +import { addCounters, BUCKET_MS, CounterDeltaTracker, type CounterObservation } from './counter-deltas'; + +export const SAMPLES_FLUSH_MS = 15_000; +export const SAMPLES_MAX_RANGE_MS = 31 * 24 * 60 * 60_000; +export const SAMPLES_LIMIT = 50_000; +export const SAMPLES_CHART_POINTS = 1440; + +function chartBucketMs(from: number, to: number): number { + const span = Math.max(0, to - from); + return Math.max(1, Math.ceil(span / SAMPLES_CHART_POINTS / BUCKET_MS)) * BUCKET_MS; +} + +export function hitRate(requested: number, hit: number): number | null { + if (!(requested > 0)) return null; + return Math.round(Math.min(1, Math.max(0, hit / requested)) * 1e4) / 1e4; +} + +export function mergeBuckets(samples: KvCacheEngineSample[]): KvCacheSampleBucket[] { + const merged = new Map(); + for (const row of samples) { + const key = JSON.stringify([row.engineId, row.modelName, row.timestamp]); + const existing = merged.get(key); + if (!existing) { + merged.set(key, { ...row }); + continue; + } + addCounters(existing, row); + } + return [...merged.values()] + .sort((a, b) => a.timestamp - b.timestamp) + .map((row) => ({ ...row, hitRate: hitRate(row.requestedTokens, row.hitTokens) })); +} + +@Injectable() +export class KvCacheSamplesService implements OnModuleInit, OnModuleDestroy { + private readonly logger = new Logger(KvCacheSamplesService.name); + private readonly tracker = new CounterDeltaTracker(); + private timer: NodeJS.Timeout | null = null; + private flushing: Promise = Promise.resolve(); + + constructor(@Inject('STORAGE_CLIENT') private readonly storage: StoragePort) {} + + onModuleInit(): void { + this.timer = setInterval(() => void this.flush(), SAMPLES_FLUSH_MS); + this.timer.unref(); + } + + async onModuleDestroy(): Promise { + if (this.timer) clearInterval(this.timer); + this.timer = null; + await this.flush(Date.now(), true); + } + + observe(observation: CounterObservation, nowMs: number): boolean { + return this.tracker.observe(observation, nowMs); + } + + resetBaselines(engineId: string): void { + this.tracker.forgetBaselines(engineId); + } + + forgetEngine(engineId: string): void { + this.tracker.forgetEngine(engineId); + } + + idle(): Promise { + return this.flushing; + } + + async deleteEngine(engineId: string): Promise { + this.forgetEngine(engineId); + await this.flushing; + await this.storage.deleteKvCacheEngineSamples(engineId); + } + + flush(nowMs: number = Date.now(), includeOpen = false): Promise { + const rows = this.tracker.drain(nowMs, includeOpen); + if (rows.length === 0) return this.flushing; + const write = this.flushing.then(() => this.write(rows)); + this.flushing = write; + return write; + } + + private async write(rows: KvCacheEngineSample[]): Promise { + try { + await this.storage.saveKvCacheEngineSamples(rows); + } catch (error) { + const message = error instanceof Error ? error.message : String(error); + this.logger.warn(`Could not save ${rows.length} KV cache samples: ${message}`); + } + } + + async getSamples( + connectionId: string, + query: { from: number; to: number; engineId?: string; model?: string }, + ): Promise { + const filter = { + connectionId, + from: Math.max(query.from, query.to - SAMPLES_MAX_RANGE_MS), + to: query.to, + engineId: query.engineId, + modelName: query.model, + }; + const [rows, totals] = await Promise.all([ + this.storage.getKvCacheEngineSamples({ ...filter, limit: SAMPLES_LIMIT, bucketMs: chartBucketMs(filter.from, filter.to) }), + this.storage.getKvCacheSampleTotals(filter), + ]); + return { buckets: mergeBuckets(rows), rangeHitRate: hitRate(totals.requestedTokens, totals.hitTokens) }; + } +} diff --git a/proprietary/kv-cache/kv-cache-scrape.service.ts b/proprietary/kv-cache/kv-cache-scrape.service.ts new file mode 100644 index 000000000..d8892fde6 --- /dev/null +++ b/proprietary/kv-cache/kv-cache-scrape.service.ts @@ -0,0 +1,103 @@ +import { Injectable, Logger, OnModuleDestroy, OnModuleInit } from '@nestjs/common'; +import { Feature } from '@betterdb/shared'; +import type { StoredKvCacheEngine } from '@app/common/interfaces/storage-port.interface'; +import { LicenseService } from '@proprietary/licenses/license.service'; +import { labelSeries } from './counter-deltas'; +import { KvCacheEngineRegistry } from './kv-cache-engine-registry'; +import { KvCacheEnginesService } from './kv-cache-engines.service'; +import { resolvePositiveInt } from './kv-cache-footprint.service'; +import { KvCacheSamplesService } from './kv-cache-samples.service'; +import { fetchMetricsText, ScrapeError, UNDECRYPTABLE_AUTH } from './metrics-fetch'; +import { NO_METRICS, parseLmcacheMetrics } from './prometheus-text'; + +const DEFAULT_INTERVAL_MS = 30000; +const MIN_INTERVAL_MS = 1000; +const MAX_INTERVAL_MS = 86_400_000; + +@Injectable() +export class KvCacheScrapeService implements OnModuleInit, OnModuleDestroy { + private readonly logger = new Logger(KvCacheScrapeService.name); + private readonly intervalMs = resolvePositiveInt( + process.env.KV_CACHE_SCRAPE_INTERVAL_MS, + DEFAULT_INTERVAL_MS, + MIN_INTERVAL_MS, + (raw) => + this.logger.warn( + `Ignoring KV_CACHE_SCRAPE_INTERVAL_MS='${raw}': expected an integer between ${MIN_INTERVAL_MS} and ${MAX_INTERVAL_MS}, using ${DEFAULT_INTERVAL_MS}`, + ), + MAX_INTERVAL_MS, + ); + private timer: NodeJS.Timeout | null = null; + private running = false; + + constructor( + private readonly registry: KvCacheEngineRegistry, + private readonly samples: KvCacheSamplesService, + private readonly engines: KvCacheEnginesService, + private readonly license: LicenseService, + ) {} + + onModuleInit(): void { + this.timer = setInterval(() => { + void this.tick(); + }, this.intervalMs); + this.timer.unref(); + } + + onModuleDestroy(): void { + if (this.timer) clearInterval(this.timer); + this.timer = null; + } + + async tick(now: number = Date.now()): Promise { + if (this.running || !this.license.hasFeature(Feature.KV_CACHE_MONITORING)) return; + this.running = true; + try { + const targets = this.registry.list().filter((e) => e.enabled && e.source === 'scrape' && e.scrapeUrl); + await Promise.allSettled(targets.map((engine) => this.scrape(engine, now))); + } finally { + this.running = false; + } + } + + private async scrape(engine: StoredKvCacheEngine, now: number): Promise { + try { + const authHeader = this.engines.authHeaderFor(engine); + if (engine.scrapeAuthHeader && authHeader === null) throw new ScrapeError(UNDECRYPTABLE_AUTH); + const text = await fetchMetricsText(engine.scrapeUrl as string, authHeader); + if (!this.stillCurrent(engine)) return; + const parsed = parseLmcacheMetrics(text); + for (const { metric, labels, value } of parsed) { + this.samples.observe( + { + engineId: engine.id, + connectionId: engine.connectionId, + modelName: labels.model_name || 'unknown', + series: labelSeries(labels), + metric, + value, + cumulative: true, + }, + now, + ); + } + this.registry.recordResult(engine.id, { + lastSeenAt: now, + lastError: parsed.length === 0 ? NO_METRICS : null, + }); + } catch (error) { + if (!this.stillCurrent(engine)) return; + if (error instanceof ScrapeError) { + this.registry.recordResult(engine.id, { lastError: error.message }); + return; + } + this.logger.debug(`Scrape of engine ${engine.id} failed: ${error instanceof Error ? error.message : String(error)}`); + this.registry.recordResult(engine.id, { lastError: 'scrape failed' }); + } + } + + private stillCurrent(engine: StoredKvCacheEngine): boolean { + const current = this.registry.get(engine.id); + return !!current?.enabled && current.scrapeUrl === engine.scrapeUrl; + } +} diff --git a/proprietary/kv-cache/kv-cache-status.service.ts b/proprietary/kv-cache/kv-cache-status.service.ts index d6739f607..38f0f45bb 100644 --- a/proprietary/kv-cache/kv-cache-status.service.ts +++ b/proprietary/kv-cache/kv-cache-status.service.ts @@ -1,14 +1,32 @@ import { Inject, Injectable } from '@nestjs/common'; import type { KvCacheEngine, KvCacheFootprintSnapshot, KvCacheStatus } from '@betterdb/shared'; import type { StoragePort, StoredKvCacheEngine } from '@app/common/interfaces/storage-port.interface'; +import { KvCacheEngineRegistry } from './kv-cache-engine-registry'; import { KvCacheFootprintService } from './kv-cache-footprint.service'; export const FOOTPRINT_HISTORY_MAX_RANGE_MS = 31 * 24 * 60 * 60_000; export const FOOTPRINT_HISTORY_POINTS = 1000; +function maskScrapeUrl(url: string | null): string | null { + if (!url) return url; + try { + const parsed = new URL(url); + if (!parsed.username && !parsed.password && !parsed.search && !parsed.hash) return url; + parsed.username = ''; + parsed.password = ''; + parsed.hash = ''; + for (const name of new Set(parsed.searchParams.keys())) { + parsed.searchParams.set(name, parsed.searchParams.get(name) ? '***' : ''); + } + return parsed.toString(); + } catch { + return url; + } +} + export function toPublicEngine(engine: StoredKvCacheEngine): KvCacheEngine { const { scrapeAuthHeader, scrapeAuthEncrypted: _encrypted, ...rest } = engine; - return { ...rest, hasScrapeAuth: scrapeAuthHeader !== null }; + return { ...rest, scrapeUrl: maskScrapeUrl(rest.scrapeUrl), hasScrapeAuth: scrapeAuthHeader !== null }; } @Injectable() @@ -16,18 +34,16 @@ export class KvCacheStatusService { constructor( @Inject('STORAGE_CLIENT') private readonly storage: StoragePort, private readonly footprint: KvCacheFootprintService, + private readonly engines: KvCacheEngineRegistry, ) {} async getStatus(connectionId: string): Promise { - const [[latest], engines] = await Promise.all([ - this.storage.getKvCacheFootprintSnapshots({ connectionId, limit: 1 }), - this.storage.getKvCacheEngines(connectionId), - ]); + const [latest] = await this.storage.getKvCacheFootprintSnapshots({ connectionId, limit: 1 }); return { hasLmcache: latest?.detected === true, latest: latest ?? null, sampleKey: this.footprint.getSampleKey(connectionId), - engines: engines.map(toPublicEngine), + engines: this.engines.list(connectionId).map(toPublicEngine), }; } diff --git a/proprietary/kv-cache/kv-cache.controller.ts b/proprietary/kv-cache/kv-cache.controller.ts index 67c30e90d..072e2b60a 100644 --- a/proprietary/kv-cache/kv-cache.controller.ts +++ b/proprietary/kv-cache/kv-cache.controller.ts @@ -1,12 +1,17 @@ -import { BadRequestException, Controller, Get, Post, Query, UseGuards } from '@nestjs/common'; +import { BadRequestException, Body, Controller, Delete, Get, HttpCode, Param, Patch, Post, Put, Query, UseGuards } from '@nestjs/common'; import { ApiHeader, ApiOperation, ApiQuery, ApiTags } from '@nestjs/swagger'; -import { Feature, type KvCacheFootprintSnapshot, type KvCacheStatus } from '@betterdb/shared'; +import { Feature, type KvCacheEngine, type KvCacheFootprintSnapshot, type KvCacheSamplesResponse, type KvCacheSettings, type KvCacheStatus } from '@betterdb/shared'; import { LicenseGuard } from '@proprietary/licenses'; import { RequiresFeature } from '@proprietary/licenses/requires-feature.decorator'; import { ConnectionId } from '@app/common/decorators'; import { ConnectionRegistry } from '@app/connections/connection-registry.service'; import { requireConnectionId } from '@app/connections/require-connection-id'; +import { CreateKvCacheEngineDto, UpdateKvCacheEngineDto } from './dto/kv-cache-engine.dto'; +import { UpdateKvCacheSettingsDto } from './dto/kv-cache-settings.dto'; +import { KvCacheAlertsService } from './kv-cache-alerts.service'; +import { KvCacheEnginesService } from './kv-cache-engines.service'; import { KvCacheFootprintService } from './kv-cache-footprint.service'; +import { KvCacheSamplesService } from './kv-cache-samples.service'; import { KvCacheStatusService } from './kv-cache-status.service'; export function parseRange(from: string | undefined, to: string | undefined): { from: number; to: number } { @@ -25,6 +30,9 @@ export class KvCacheController { private readonly status: KvCacheStatusService, private readonly footprint: KvCacheFootprintService, private readonly connectionRegistry: ConnectionRegistry, + private readonly engines: KvCacheEnginesService, + private readonly samples: KvCacheSamplesService, + private readonly alerts: KvCacheAlertsService, ) {} @Get('status') @@ -64,4 +72,87 @@ export class KvCacheController { range.to, ); } + + @Get('engines/samples') + @UseGuards(LicenseGuard) + @RequiresFeature(Feature.KV_CACHE_MONITORING) + @ApiOperation({ summary: 'Per-minute LMCache token counters and hit rate history (Pro)' }) + @ApiQuery({ name: 'from', required: true, type: Number }) + @ApiQuery({ name: 'to', required: true, type: Number }) + @ApiQuery({ name: 'engineId', required: false, type: String }) + @ApiQuery({ name: 'model', required: false, type: String }) + @ApiHeader({ name: 'x-connection-id', required: false }) + async getSamples( + @Query('from') from: string | undefined, + @Query('to') to: string | undefined, + @Query('engineId') engineId: string | undefined, + @Query('model') model: string | undefined, + @ConnectionId() connectionId?: string, + ): Promise { + const range = parseRange(from, to); + return this.samples.getSamples(this.connection(connectionId), { ...range, engineId, model }); + } + + @Get('engines') + @UseGuards(LicenseGuard) + @RequiresFeature(Feature.KV_CACHE_MONITORING) + @ApiOperation({ summary: 'List linked LMCache engines (Pro)' }) + @ApiHeader({ name: 'x-connection-id', required: false }) + listEngines(@ConnectionId() connectionId?: string): KvCacheEngine[] { + return this.engines.list(this.connection(connectionId)); + } + + @Post('engines') + @UseGuards(LicenseGuard) + @RequiresFeature(Feature.KV_CACHE_MONITORING) + @ApiOperation({ summary: 'Link an LMCache engine by scrape URL or OTLP engine id (Pro)' }) + @ApiHeader({ name: 'x-connection-id', required: false }) + async createEngine(@Body() dto: CreateKvCacheEngineDto, @ConnectionId() connectionId?: string): Promise { + return this.engines.create(this.connection(connectionId), dto); + } + + @Patch('engines/:id') + @UseGuards(LicenseGuard) + @RequiresFeature(Feature.KV_CACHE_MONITORING) + @ApiOperation({ summary: 'Update a linked LMCache engine (Pro)' }) + @ApiHeader({ name: 'x-connection-id', required: false }) + async updateEngine( + @Param('id') id: string, + @Body() dto: UpdateKvCacheEngineDto, + @ConnectionId() connectionId?: string, + ): Promise { + return this.engines.update(this.connection(connectionId), id, dto); + } + + @Delete('engines/:id') + @HttpCode(204) + @UseGuards(LicenseGuard) + @RequiresFeature(Feature.KV_CACHE_MONITORING) + @ApiOperation({ summary: 'Unlink an LMCache engine (Pro)' }) + @ApiHeader({ name: 'x-connection-id', required: false }) + async removeEngine(@Param('id') id: string, @ConnectionId() connectionId?: string): Promise { + return this.engines.remove(this.connection(connectionId), id); + } + + @Get('settings') + @UseGuards(LicenseGuard) + @RequiresFeature(Feature.KV_CACHE_MONITORING) + @ApiOperation({ summary: 'KV cache alert settings for the connection (Pro)' }) + @ApiHeader({ name: 'x-connection-id', required: false }) + async getSettings(@ConnectionId() connectionId?: string): Promise { + return this.alerts.getSettings(this.connection(connectionId)); + } + + @Put('settings') + @UseGuards(LicenseGuard) + @RequiresFeature(Feature.KV_CACHE_MONITORING) + @ApiOperation({ summary: 'Update KV cache alert settings for the connection (Pro)' }) + @ApiHeader({ name: 'x-connection-id', required: false }) + async updateSettings(@Body() dto: UpdateKvCacheSettingsDto, @ConnectionId() connectionId?: string): Promise { + return this.alerts.updateSettings(this.connection(connectionId), dto); + } + + private connection(connectionId: string | undefined): string { + return requireConnectionId(this.connectionRegistry, connectionId); + } } diff --git a/proprietary/kv-cache/kv-cache.module.ts b/proprietary/kv-cache/kv-cache.module.ts index 5855da77a..8b5d55b7a 100644 --- a/proprietary/kv-cache/kv-cache.module.ts +++ b/proprietary/kv-cache/kv-cache.module.ts @@ -1,14 +1,24 @@ -import { Module } from '@nestjs/common'; +import { Global, Module } from '@nestjs/common'; +import { KV_CACHE_OTLP_SINK } from '@app/external-metrics/kv-cache-otlp-sink'; import { ConnectionsModule } from '@app/connections/connections.module'; +import { PrometheusModule } from '@app/prometheus/prometheus.module'; import { StorageModule } from '@app/storage/storage.module'; import { LicenseModule } from '@proprietary/licenses'; +import { KvCacheAlertsService } from './kv-cache-alerts.service'; import { KvCacheController } from './kv-cache.controller'; +import { KvCacheEngineRegistry } from './kv-cache-engine-registry'; +import { KvCacheEnginesService } from './kv-cache-engines.service'; import { KvCacheFootprintService } from './kv-cache-footprint.service'; +import { KvCacheOtlpSinkService } from './kv-cache-otlp-sink.service'; +import { KvCacheScrapeService } from './kv-cache-scrape.service'; +import { KvCacheSamplesService } from './kv-cache-samples.service'; import { KvCacheStatusService } from './kv-cache-status.service'; +@Global() @Module({ - imports: [StorageModule, ConnectionsModule, LicenseModule], + imports: [StorageModule, ConnectionsModule, LicenseModule, PrometheusModule], controllers: [KvCacheController], - providers: [KvCacheFootprintService, KvCacheStatusService], + providers: [KvCacheFootprintService, KvCacheAlertsService, KvCacheStatusService, KvCacheEngineRegistry, KvCacheSamplesService, KvCacheEnginesService, KvCacheScrapeService, KvCacheOtlpSinkService, { provide: KV_CACHE_OTLP_SINK, useExisting: KvCacheOtlpSinkService }], + exports: [KV_CACHE_OTLP_SINK], }) export class KvCacheModule {} diff --git a/proprietary/kv-cache/metrics-fetch.ts b/proprietary/kv-cache/metrics-fetch.ts new file mode 100644 index 000000000..3e6a98c90 --- /dev/null +++ b/proprietary/kv-cache/metrics-fetch.ts @@ -0,0 +1,103 @@ +import { BadRequestException, Logger } from '@nestjs/common'; +import { isCloudMode } from '@app/common/utils/cloud-mode'; +import { isTrueFlag } from '@app/config/env-normalize'; +import { assertSafeOutboundUrl, BlockedAddressError, guardedLookup } from '@app/common/utils/outbound-url-guard'; +import { pinnedRequest, UnexpectedStatusError } from '@app/common/utils/pinned-request'; + +export const SCRAPE_TIMEOUT_MS = 5_000; +export const SCRAPE_MAX_BYTES = 5 * 1024 * 1024; + +export class ScrapeError extends Error {} + +export const UNDECRYPTABLE_AUTH = 'auth header could not be decrypted'; + +export interface MetricsRequestInit { + signal: AbortSignal; + headers: Record; +} + +export type MetricsFetch = (url: string, init: MetricsRequestInit) => Promise; + +export const pinnedFetch: MetricsFetch = (url, init) => + pinnedRequest(url, { ...init, lookup: guardedLookup({ label: 'metrics URL', allowPrivateNetworks: scrapeAllowsPrivateNetworks() }) }); + +function connectError(error: unknown): ScrapeError { + if (error instanceof ScrapeError) return error; + if (error instanceof BlockedAddressError || error instanceof UnexpectedStatusError) return new ScrapeError(error.message); + const { name, code } = (error ?? {}) as { name?: string; code?: string }; + if (name === 'TimeoutError' || name === 'AbortError') return new ScrapeError('timeout'); + if (code === 'ENOTFOUND' || code === 'EAI_AGAIN') return new ScrapeError(`DNS lookup failed (${code})`); + return new ScrapeError('connection failed'); +} + +const BLOCK_PRIVATE_ON = new Set(['1', 'yes', 'on']); +const BLOCK_PRIVATE_OFF = new Set(['', 'false', '0', 'no', 'off']); +const warnedBlockPrivateValues = new Set(); +const logger = new Logger('KvCacheMetricsFetch'); + +function blockPrivateFlag(): boolean { + const raw = process.env.KV_CACHE_SCRAPE_BLOCK_PRIVATE; + if (raw === undefined) return false; + const value = raw.trim().toLowerCase(); + if (isTrueFlag(value) || BLOCK_PRIVATE_ON.has(value)) return true; + if (BLOCK_PRIVATE_OFF.has(value)) return false; + if (!warnedBlockPrivateValues.has(value)) { + warnedBlockPrivateValues.add(value); + logger.warn(`Unrecognised KV_CACHE_SCRAPE_BLOCK_PRIVATE value "${raw.trim()}"; blocking private scrape targets`); + } + return true; +} + +export function scrapeAllowsPrivateNetworks(): boolean { + return !isCloudMode() && !blockPrivateFlag(); +} + +async function reject(response: Response, message: string): Promise { + await response.body?.cancel().catch(() => undefined); + throw new ScrapeError(message); +} + +async function readCapped(response: Response): Promise { + if (Number(response.headers.get('content-length') ?? 0) > SCRAPE_MAX_BYTES) return reject(response, 'response too large'); + const reader = response.body?.getReader(); + if (!reader) return ''; + const chunks: Uint8Array[] = []; + let size = 0; + for (;;) { + const { done, value } = await reader.read(); + if (done) break; + size += value.byteLength; + if (size > SCRAPE_MAX_BYTES) { + await reader.cancel(); + throw new ScrapeError('response too large'); + } + chunks.push(value); + } + return Buffer.concat(chunks).toString('utf8'); +} + +export async function fetchMetricsText(url: string, authHeader: string | null, fetchImpl: MetricsFetch = pinnedFetch): Promise { + try { + await assertSafeOutboundUrl(url, { label: 'metrics URL', allowPrivateNetworks: scrapeAllowsPrivateNetworks(), resolveDns: 'never' }); + } catch (error) { + throw new ScrapeError(error instanceof BadRequestException ? error.message : 'blocked address'); + } + let response: Response; + try { + response = await fetchImpl(url, { + signal: AbortSignal.timeout(SCRAPE_TIMEOUT_MS), + headers: { Accept: 'text/plain', ...(authHeader ? { Authorization: authHeader } : {}) }, + }); + } catch (error) { + throw connectError(error); + } + if (response.status >= 300 && response.status < 400) await reject(response, `redirect not followed (HTTP ${response.status})`); + if (!response.ok) await reject(response, `HTTP ${response.status}`); + try { + return await readCapped(response); + } catch (error) { + if (error instanceof ScrapeError) throw error; + const name = (error as Error)?.name; + throw new ScrapeError(name === 'TimeoutError' || name === 'AbortError' ? 'timeout' : 'connection failed'); + } +} diff --git a/proprietary/kv-cache/prometheus-text.ts b/proprietary/kv-cache/prometheus-text.ts new file mode 100644 index 000000000..c28681a7a --- /dev/null +++ b/proprietary/kv-cache/prometheus-text.ts @@ -0,0 +1,55 @@ +export const NO_METRICS = 'no lmcache metrics found'; + +export interface PromSample { + metric: string; + labels: Record; + value: number; +} + +const PREFIX = 'lmcache:'; +const SERIES = /^([a-zA-Z_:][a-zA-Z0-9_:]*)(?:\{(.*)\})?\s+(\S+)(?:\s+-?\d+)?$/; +const ESCAPES: Record = { n: '\n', '\\': '\\', '"': '"' }; + +function parseLabels(body: string): Record | null { + const labels: Record = {}; + let i = 0; + while (i < body.length) { + const eq = body.indexOf('=', i); + if (eq < 0 || body[eq + 1] !== '"') return null; + const name = body.slice(i, eq).trim(); + let value = ''; + let j = eq + 2; + while (j < body.length && body[j] !== '"') { + if (body[j] === '\\' && j + 1 < body.length) { + value += ESCAPES[body[j + 1]] ?? body[j + 1]; + j += 2; + } else { + value += body[j]; + j += 1; + } + } + if (j >= body.length) return null; + labels[name] = value; + i = j + 1; + while (i < body.length && (body[i] === ',' || body[i] === ' ')) i += 1; + } + return labels; +} + +export function parseLmcacheMetrics(text: string): PromSample[] { + const samples: PromSample[] = []; + for (const raw of text.split('\n')) { + const line = raw.trim(); + if (!line.startsWith(PREFIX)) continue; + const match = SERIES.exec(line); + if (!match) continue; + const [, name, labelBody, rawValue] = match; + if (name.endsWith('_created')) continue; + const value = Number(rawValue); + if (!Number.isFinite(value)) continue; + const labels = labelBody ? parseLabels(labelBody) : {}; + if (!labels) continue; + samples.push({ metric: name.slice(PREFIX.length).replace(/_total$/, ''), labels, value }); + } + return samples; +} diff --git a/proprietary/webhook-pro/__tests__/kv-cache-events.spec.ts b/proprietary/webhook-pro/__tests__/kv-cache-events.spec.ts new file mode 100644 index 000000000..bea162b4a --- /dev/null +++ b/proprietary/webhook-pro/__tests__/kv-cache-events.spec.ts @@ -0,0 +1,101 @@ +import { Test, TestingModule } from '@nestjs/testing'; +import { WebhookEventsProService } from '../webhook-events-pro.service'; +import { WebhookDispatcherService } from '@app/webhooks/webhook-dispatcher.service'; +import { WebhookEventType, type KvCacheEvictionRiskData, type KvCacheHitRateLowData } from '@betterdb/shared'; +import { LicenseService } from '@proprietary/licenses'; + +describe('WebhookEventsProService - KV cache events', () => { + let service: WebhookEventsProService; + let webhookDispatcher: { dispatchThresholdAlert: jest.Mock }; + let licenseService: { getLicenseTier: jest.Mock }; + + const hitRateData: KvCacheHitRateLowData = { + connectionId: 'conn-1', + engineId: 'eng-1', + engineName: 'vllm-a', + model: 'llama', + hitRate: 0.1, + threshold: 0.2, + requestedTokens: 20000, + windowMs: 900000, + timestamp: 123, + }; + + const evictionData: KvCacheEvictionRiskData = { + connectionId: 'conn-1', + reason: 'unevictable', + active: true, + policy: 'volatile-lru', + usedMemory: 90, + maxmemory: 100, + lmcacheMemoryShare: 0.8, + noTtlRatio: 0.95, + evictedKeysDelta: 0, + timestamp: 123, + }; + + beforeEach(async () => { + webhookDispatcher = { dispatchThresholdAlert: jest.fn().mockResolvedValue(undefined) }; + licenseService = { getLicenseTier: jest.fn().mockReturnValue('pro') }; + const module: TestingModule = await Test.createTestingModule({ + providers: [ + WebhookEventsProService, + { provide: WebhookDispatcherService, useValue: webhookDispatcher }, + { provide: LicenseService, useValue: licenseService }, + ], + }).compile(); + service = module.get(WebhookEventsProService); + }); + + it('dispatches a below-threshold hit rate alert keyed per engine and model', async () => { + await service.dispatchKvCacheHitRateLow(hitRateData); + expect(webhookDispatcher.dispatchThresholdAlert).toHaveBeenCalledWith( + WebhookEventType.KV_CACHE_HIT_RATE_LOW, + 'kv_cache_hit_rate_low:eng-1:llama', + 0.1, + 0.2, + false, + expect.objectContaining({ + ...hitRateData, + message: 'LMCache hit rate for llama on vllm-a is 10.0% (threshold 20.0%)', + }), + 'conn-1', + ); + }); + + it('dispatches an active unevictable alert with value 1', async () => { + await service.dispatchKvCacheEvictionRisk(evictionData); + expect(webhookDispatcher.dispatchThresholdAlert).toHaveBeenCalledWith( + WebhookEventType.KV_CACHE_EVICTION_RISK, + 'kv_cache_eviction_risk:conn-1:unevictable', + 1, + 1, + true, + expect.objectContaining({ + ...evictionData, + message: 'LMCache keys have no TTL under volatile-lru; Valkey cannot evict them', + }), + 'conn-1', + ); + }); + + it('dispatches an inactive evicting alert with value 0', async () => { + await service.dispatchKvCacheEvictionRisk({ ...evictionData, reason: 'evicting', active: false }); + expect(webhookDispatcher.dispatchThresholdAlert).toHaveBeenCalledWith( + WebhookEventType.KV_CACHE_EVICTION_RISK, + 'kv_cache_eviction_risk:conn-1:evicting', + 0, + 1, + true, + expect.objectContaining({ message: 'Valkey is evicting keys while LMCache holds most of its memory' }), + 'conn-1', + ); + }); + + it('skips both events without a license', async () => { + licenseService.getLicenseTier.mockReturnValue('community'); + await service.dispatchKvCacheHitRateLow(hitRateData); + await service.dispatchKvCacheEvictionRisk(evictionData); + expect(webhookDispatcher.dispatchThresholdAlert).not.toHaveBeenCalled(); + }); +}); diff --git a/proprietary/webhook-pro/webhook-events-pro.service.ts b/proprietary/webhook-pro/webhook-events-pro.service.ts index 24e8912fa..720393f67 100644 --- a/proprietary/webhook-pro/webhook-events-pro.service.ts +++ b/proprietary/webhook-pro/webhook-events-pro.service.ts @@ -3,6 +3,8 @@ import { WebhookDispatcherService } from '@app/webhooks/webhook-dispatcher.servi import { WebhookEventType, type CveCriticalDetectedData, + type KvCacheEvictionRiskData, + type KvCacheHitRateLowData, type LatencyRegressionDetectedData, type MetricKind, type ScalingReadinessLowData, @@ -547,6 +549,47 @@ export class WebhookEventsProService implements OnModuleInit { ); } + async dispatchKvCacheHitRateLow(data: KvCacheHitRateLowData): Promise { + if (!this.isEnabled()) { + this.logger.log('KV cache hit rate low event skipped - requires PRO license'); + return; + } + await this.webhookDispatcher.dispatchThresholdAlert( + WebhookEventType.KV_CACHE_HIT_RATE_LOW, + `kv_cache_hit_rate_low:${data.engineId}:${data.model}`, + data.hitRate, + data.threshold, + false, + { + ...data, + message: `LMCache hit rate for ${data.model} on ${data.engineName} is ${(data.hitRate * 100).toFixed(1)}% (threshold ${(data.threshold * 100).toFixed(1)}%)`, + }, + data.connectionId, + ); + } + + async dispatchKvCacheEvictionRisk(data: KvCacheEvictionRiskData): Promise { + if (!this.isEnabled()) { + this.logger.log('KV cache eviction risk event skipped - requires PRO license'); + return; + } + await this.webhookDispatcher.dispatchThresholdAlert( + WebhookEventType.KV_CACHE_EVICTION_RISK, + `kv_cache_eviction_risk:${data.connectionId}:${data.reason}`, + data.active ? 1 : 0, + 1, + true, + { + ...data, + message: + data.reason === 'unevictable' + ? `LMCache keys have no TTL under ${data.policy}; Valkey cannot evict them` + : 'Valkey is evicting keys while LMCache holds most of its memory', + }, + data.connectionId, + ); + } + /** * Dispatch inference SLA breach event (PRO+) * Called when a configured per-index p99 SLA is breached. Debounce + resolution