158 lines
13 KiB
JavaScript
158 lines
13 KiB
JavaScript
import test from 'node:test';
|
|
import assert from 'node:assert/strict';
|
|
import http from 'node:http';
|
|
import { mkdtempSync, rmSync } from 'node:fs';
|
|
import { tmpdir } from 'node:os';
|
|
import { join } from 'node:path';
|
|
import { createHash } from 'node:crypto';
|
|
import { Store } from '../src/store.mjs';
|
|
import { createUsage, turnRecords } from '../src/usage.mjs';
|
|
import { createLimits } from '../src/limits.mjs';
|
|
import { createProxyUsage } from '../src/proxy-usage.mjs';
|
|
import { ManagementClient } from '../../model-proxy/src/client.mjs';
|
|
import { createLoader } from '../../plugins/src/loader.mjs';
|
|
const id = s => createHash('sha256').update(s).digest('hex').slice(0,12);
|
|
const key='msk-fake-client-value',source='fake-upstream-value',keyId=id(key);
|
|
async function fixture(t,{makeLoader}={}) {
|
|
const root=mkdtempSync(join(tmpdir(),'proxy-usage-')),store=new Store(root),logs=[],calls=[];
|
|
let events=[],now=0,fail=false,noClient=false,recs=[{instance:'demo/coder',service:'model',keyId,retiring:[]}];
|
|
const server=http.createServer((req,res)=>{calls.push(req.url);res.writeHead(fail?500:200);res.end(JSON.stringify(events));events=[];});
|
|
await new Promise(r=>server.listen(0,'127.0.0.1',r));
|
|
const client=new ManagementClient({baseUrl:`http://127.0.0.1:${server.address().port}`,key:'fake-management',receipt:e=>logs.push(e)});
|
|
let doc={plugins:{},launch:{limits:{tokens:{keys:{[keyId]:{limit:1,since:'2026-01-01T00:00:00.000Z'}}}}}};
|
|
const reader=createProxyUsage({store,enableDrain:true,read:()=>({demo:doc}),client:()=>noClient?null:client,...(makeLoader?{makeLoader}:{}),records:()=>recs,log:e=>logs.push(e),now:()=>now});
|
|
t.after(async()=>{await reader.close();await new Promise(r=>server.close(r));store.close();rmSync(root,{recursive:true,force:true});});
|
|
const event=(over={})=>({api_key:key,source,failed:true,fail:{status_code:503,body:source},tokens:{input_tokens:9,output_tokens:2},...over});
|
|
return {store,logs,calls,reader,client,event,setEvents:v=>{events=v;},clock:v=>{now=v;},enable:(pollSeconds)=>{doc={...doc,plugins:{usage:'cliproxy-queue','limit-check':['token-limit']},...(pollSeconds===undefined?{}:{usage:{cliproxyQueue:{pollSeconds}}})};},doc:()=>doc,fail:()=>{fail=true;},noClient:v=>{noClient=v;},setRecords:v=>{recs=v;}};
|
|
}
|
|
test('P5 usage plugin off by default, default poll is 15 seconds and configured polls skip early ticks',async t=>{
|
|
const f=await fixture(t);await f.reader.tick();assert.equal(f.calls.length,0);assert.equal(f.reader.view('demo'),null);
|
|
f.enable();await f.reader.tick();assert.equal(f.calls.length,1);
|
|
f.clock(14999);await f.reader.tick();assert.equal(f.calls.length,1);
|
|
f.clock(15000);await f.reader.tick();assert.equal(f.calls.length,2);
|
|
f.enable(7);f.clock(16000);await f.reader.tick();assert.equal(f.calls.length,3);
|
|
f.clock(22999);await f.reader.tick();assert.equal(f.calls.length,3);
|
|
f.clock(23000);await f.reader.tick();assert.equal(f.calls.length,4);
|
|
});
|
|
test('P5 safe telemetry stores failures per key and proxy-labelled accounts; no client/upstream value enters store/log/view or quota records',async t=>{
|
|
const f=await fixture(t);f.enable();f.setEvents([f.event()]);await f.reader.tick();
|
|
const v=f.reader.view('demo');assert.equal(v.source,'proxy');assert.equal(v.byKey[0].keyId,keyId);assert.equal(v.byKey[0].failures,1);assert.deepEqual(v.byKey[0].statuses,{'503':1});
|
|
assert.equal(v.byAccount[0].accountId,id(source));assert.equal(v.byAccount[0].input,9);
|
|
const usage=await createUsage({store:f.store,businesses:{demo:f.doc()},log:()=>{},proxyView:b=>f.reader.view(b)});
|
|
assert.equal(usage.view('demo').proxy.source,'proxy');assert.equal(usage.view('demo').total.input,0);
|
|
assert.deepEqual(turnRecords(f.store,'demo'),[]);
|
|
const limits=createLimits({store:f.store,read:()=>f.doc(),log:()=>{},keyForSubject:()=>keyId});
|
|
assert.equal((await limits.mayLaunch('demo',{instance:'coder',modelKeyId:keyId})).allow,true);
|
|
for(const value of [key,source,'fake-management']) assert.equal(JSON.stringify([f.store.all('SELECT * FROM usage_records'),f.logs,usage.view('demo')]).includes(value),false);
|
|
});
|
|
test('P5 unknown msk key goes only to value-free reconcile log and other businesses stay isolated',async t=>{
|
|
const f=await fixture(t);f.enable();f.setEvents([f.event({api_key:'msk-unknown-canary'}),f.event({api_key:'external-canary'})]);await f.reader.tick();
|
|
assert.deepEqual(f.logs.find(e=>e.action==='proxy.usage.reconcile'),{action:'proxy.usage.reconcile',keyId:id('msk-unknown-canary'),code:'proxy-usage-unknown-key'});
|
|
assert.equal(f.reader.view('demo').byKey.length,0);assert.equal(f.reader.view('other'),null);
|
|
assert.equal(JSON.stringify(f.logs).includes('msk-unknown-canary'),false);
|
|
assert.equal(f.store.all('SELECT * FROM usage_records').length,0);
|
|
});
|
|
test('P5 queue failure is safe, does not enforce and does not retry before poll deadline',async t=>{
|
|
const f=await fixture(t);f.enable();f.fail();await f.reader.tick();await f.reader.tick();
|
|
assert.equal(f.calls.length,1);assert.equal(f.logs.at(-1).code,'proxy-usage-read-failed');assert.equal(f.reader.view('demo').state,'unavailable');
|
|
});
|
|
test('P5 retiring key keeps attribution and telemetry rows remain append-only across reader restart',async t=>{
|
|
const f=await fixture(t);f.enable();f.setEvents([f.event()]);await f.reader.tick();
|
|
const before=f.reader.view('demo');await f.reader.close();
|
|
const again=createProxyUsage({store:f.store,read:()=>({demo:f.doc()}),client:()=>null,records:()=>[],log:()=>{}});
|
|
assert.deepEqual(again.view('demo').byKey,before.byKey);await again.close();
|
|
assert.throws(()=>f.store.run('DELETE FROM usage_records'),/append-only/);
|
|
});
|
|
|
|
test('P5 malformed queue batch is atomic: no valid prefix, body or credential can reach store/log/view',async t=>{
|
|
const f=await fixture(t);f.enable();f.setEvents([f.event(),f.event({source:'malformed-upstream-canary',tokens:{input_tokens:-1}})]);await f.reader.tick();
|
|
assert.equal(f.store.all('SELECT * FROM usage_records').length,0);
|
|
for(const secret of [key,source,'malformed-upstream-canary'])assert.equal(JSON.stringify([f.logs,f.reader.view('demo')]).includes(secret),false);
|
|
});
|
|
test('P5 concurrent ticks issue one bounded queue read and a closed collector never polls again',async t=>{
|
|
const f=await fixture(t);f.enable();await Promise.all([f.reader.tick(),f.reader.tick(),f.reader.tick()]);assert.equal(f.calls.length,1);
|
|
await f.reader.close();f.clock(90000);await f.reader.tick();assert.equal(f.calls.length,1);
|
|
});
|
|
// Row 333 R3 (#1831) ruling 4: the usage module's argument is { readUsageQueue } and nothing else.
|
|
// Nothing reachable from it by reflection is the management client or any of its other calls.
|
|
test('R3 the usage module is handed one function, readUsageQueue, and no path to the management client',async t=>{
|
|
const handed=[];
|
|
const makeLoader=opts=>{const loader=createLoader(opts);return {...loader,modules:kind=>loader.modules(kind).map(m=>kind==='usage'&&m?{interfaceVersion:m.interfaceVersion,read:arg=>{handed.push(arg);return m.read(arg);}}:m)};};
|
|
const f=await fixture(t,{makeLoader});
|
|
const methods=Object.getOwnPropertyNames(ManagementClient.prototype).filter(n=>n!=='constructor');
|
|
for(const name of ['readUsageQueue','getKeys','putKeys','oauthStart','oauthCallback','oauthStatus','oauthCancel']) assert.ok(methods.includes(name),name);
|
|
const called=[];
|
|
for(const name of methods){const real=f.client[name].bind(f.client);Object.defineProperty(f.client,name,{value:(...a)=>{called.push(name);return real(...a);}});}
|
|
f.enable();f.setEvents([f.event()]);await f.reader.tick();
|
|
assert.equal(handed.length,1);
|
|
const [arg]=handed;
|
|
assert.deepEqual(Reflect.ownKeys(arg),['readUsageQueue']);
|
|
assert.equal(typeof arg.readUsageQueue,'function');assert.equal(Object.isFrozen(arg),true);
|
|
assert.equal(Object.getPrototypeOf(arg),Object.prototype);
|
|
// Walk every own property, accessor and prototype reachable from the argument without calling it.
|
|
const stop=new Set([Object.prototype,Function.prototype]),seen=new Set(),queue=[arg],names=new Set();
|
|
while(queue.length){
|
|
const v=queue.shift();
|
|
if(v===null||(typeof v!=='object'&&typeof v!=='function')||seen.has(v)||stop.has(v))continue;
|
|
seen.add(v);
|
|
assert.notEqual(v,f.client);assert.notEqual(v,ManagementClient.prototype);assert.notEqual(v,ManagementClient);
|
|
for(const key of Reflect.ownKeys(v)){
|
|
names.add(key);
|
|
const d=Reflect.getOwnPropertyDescriptor(v,key);
|
|
queue.push(d.value,d.get,d.set);
|
|
}
|
|
queue.push(Object.getPrototypeOf(v));
|
|
}
|
|
for(const name of methods.filter(n=>n!=='readUsageQueue')) assert.equal(names.has(name),false,name);
|
|
for(const name of methods) for(const v of seen) assert.notEqual(v,ManagementClient.prototype[name],name);
|
|
// The one call the module made was the queue read.
|
|
assert.deepEqual(called,['readUsageQueue']);
|
|
assert.equal(f.reader.view('demo').byKey[0].failures,1);
|
|
});
|
|
|
|
// #778 (Vikunja): the queue read drains the proxy, so a batch that is then not written is lost. The count
|
|
// is per enabled business, in memory since the bus started, and telemetry only.
|
|
test('#778 a drained batch that is not written counts as dropped; a poll that drained nothing does not; a good poll keeps the count',async t=>{
|
|
const f=await fixture(t);f.enable();
|
|
const usage=await createUsage({store:f.store,businesses:{demo:f.doc()},log:()=>{},proxyView:b=>f.reader.view(b)});
|
|
assert.equal(f.reader.view('demo').droppedBatches,0);
|
|
// An invalid batch: the client refuses the 2xx answer whole, and the proxy has already given the batch up.
|
|
f.setEvents([f.event(),f.event({source:'malformed-upstream-canary',tokens:{input_tokens:-1}})]);await f.reader.tick();
|
|
assert.equal(f.reader.view('demo').droppedBatches,1);assert.equal(f.reader.view('demo').state,'unavailable');
|
|
// An ambiguous owner: one key id held by two agents.
|
|
f.setRecords([{instance:'demo/coder',service:'model',keyId,retiring:[]},{instance:'demo/other',service:'model',keyId,retiring:[]}]);
|
|
f.setEvents([f.event()]);f.clock(15000);await f.reader.tick();
|
|
assert.equal(f.reader.view('demo').droppedBatches,2);
|
|
// No client before the read: nothing was drained, so nothing was dropped.
|
|
f.noClient(true);f.setEvents([f.event()]);f.clock(30000);await f.reader.tick();
|
|
assert.equal(f.reader.view('demo').droppedBatches,2);assert.equal(f.reader.view('demo').state,'unavailable');
|
|
// A good poll sets the state back and leaves the count alone. The event set while the client was missing is still queued.
|
|
f.noClient(false);f.setRecords([{instance:'demo/coder',service:'model',keyId,retiring:[]}]);f.clock(45000);await f.reader.tick();
|
|
assert.equal(f.reader.view('demo').state,'available');assert.equal(f.reader.view('demo').droppedBatches,2);
|
|
assert.equal(f.reader.view('demo').byKey[0].requests,1);
|
|
// The usage view carries it, and a log line names no event content.
|
|
assert.equal(usage.view('demo').proxy.droppedBatches,2);
|
|
for(const secret of [key,source,'malformed-upstream-canary'])assert.equal(JSON.stringify([f.logs,usage.view('demo')]).includes(secret),false);
|
|
assert.deepEqual(f.logs.filter(e=>e.action==='proxy.usage').map(e=>e.code),['proxy-usage-read-failed','proxy-usage-read-failed','proxy-usage-read-failed']);
|
|
});
|
|
test('#778 a batch the module returns in a shape the bus refuses is a drop',async t=>{
|
|
const makeLoader=opts=>{const loader=createLoader(opts);return {...loader,modules:kind=>loader.modules(kind).map(m=>kind==='usage'&&m?{interfaceVersion:m.interfaceVersion,read:async arg=>{await arg.readUsageQueue();return [{keyId:'not-a-key-id',accountId:id(source)}];}}:m)};};
|
|
const f=await fixture(t,{makeLoader});f.enable();f.setEvents([f.event()]);await f.reader.tick();
|
|
assert.equal(f.reader.view('demo').state,'unavailable');assert.equal(f.reader.view('demo').droppedBatches,1);
|
|
assert.equal(f.store.all('SELECT * FROM usage_records').length,0);
|
|
});
|
|
test('#778 a read the proxy refused or never answered drained nothing and is not a drop',async t=>{
|
|
const f=await fixture(t);f.enable();f.fail();await f.reader.tick();
|
|
assert.equal(f.reader.view('demo').state,'unavailable');assert.equal(f.reader.view('demo').droppedBatches,0);
|
|
});
|
|
test('#778 a failed transaction drops the batch for every enabled business',async t=>{
|
|
const f=await fixture(t);f.enable();
|
|
const broken={all:(...a)=>f.store.all(...a),run:(...a)=>f.store.run(...a),transaction:()=>{throw new Error('disk full');}};
|
|
const both={demo:f.doc(),other:f.doc()};
|
|
const reader=createProxyUsage({store:broken,enableDrain:true,read:()=>both,client:()=>f.client,records:()=>[{instance:'demo/coder',service:'model',keyId,retiring:[]}],log:e=>f.logs.push(e),now:()=>0});
|
|
t.after(()=>reader.close());
|
|
f.setEvents([f.event()]);await reader.tick();
|
|
for(const business of ['demo','other']){assert.equal(reader.view(business).droppedBatches,1);assert.equal(reader.view(business).state,'unavailable');}
|
|
assert.equal(f.store.all('SELECT * FROM usage_records').length,0);
|
|
});
|