115 lines
7.0 KiB
JavaScript
115 lines
7.0 KiB
JavaScript
// The shared proxy queue has no safe multi-bus consumer yet. Runtime never opts
|
|
// in to draining; enableDrain exists solely for the isolated collector tests.
|
|
import { randomUUID } from 'node:crypto';
|
|
import { createLoader } from '../../plugins/src/loader.mjs';
|
|
import { checkUsageSettings } from '../../plugins/src/kinds.mjs';
|
|
const SOURCE = '@proxy/cliproxy-queue';
|
|
const KINDS = ['input', 'output', 'cacheWrite', 'cacheRead'];
|
|
const ID = /^[a-f0-9]{12}$/;
|
|
const plain = n => n <= BigInt(Number.MAX_SAFE_INTEGER) ? Number(n) : String(n);
|
|
// makeLoader is the plug-in loader's factory; a test wraps it to see what a module is handed.
|
|
export function createProxyUsage({ store, read, pluginDataRoot = null, client, records, log, now = Date.now, makeLoader = createLoader, enableDrain = false }) {
|
|
const loader = makeLoader({ log, dataRoot:pluginDataRoot });
|
|
let due = 0, previous = '', pending = null, closed = false;
|
|
const states = new Map();
|
|
// #778: per enabled business, the batches taken off the proxy queue and then not written, since
|
|
// this bus started. Telemetry only: the events in a dropped batch are gone and no quota reads this.
|
|
const dropped = new Map();
|
|
function configured() {
|
|
return Object.entries(read()).filter(([, b]) => b.plugins?.usage === 'cliproxy-queue').map(([business, b]) => {
|
|
checkUsageSettings(b.usage ?? {});
|
|
return { business, pollSeconds: b.usage?.cliproxyQueue?.pollSeconds ?? 15 };
|
|
});
|
|
}
|
|
async function poll() {
|
|
let enabled;
|
|
try { enabled = configured(); }
|
|
catch { log({ action: 'proxy.usage', code: 'invalid-usage-settings' }); return; }
|
|
if (!enabled.length) { previous = ''; return; }
|
|
const signature = JSON.stringify(enabled);
|
|
if (signature !== previous) { previous = signature; due = 0; }
|
|
if (now() < due) return;
|
|
due = now() + Math.min(...enabled.map(b => b.pollSeconds)) * 1000;
|
|
// drained: the proxy has given this batch up, because module.read returned or because the client
|
|
// refused a 2xx answer it had already taken (proxy-client-response). Down, redirected or non-2xx
|
|
// answers drained nothing. written: the rows are in.
|
|
let drained = false, written = false;
|
|
try {
|
|
await loader.load({ usage: 'cliproxy-queue' });
|
|
const module = loader.modules('usage')[0], management = await client();
|
|
if (!module || !management) throw Error('unavailable');
|
|
// Row 333 R3 (#1831): the module gets the one read it needs and nothing else. No path from
|
|
// its argument reaches getKeys, putKeys, the OAuth calls or any other management call.
|
|
const readUsageQueue = async () => {
|
|
try { return await management.readUsageQueue(); }
|
|
catch (error) { if (error?.code === 'proxy-client-response') drained = true; throw error; }
|
|
};
|
|
const events = await module.read(Object.freeze({ readUsageQueue }));
|
|
drained = true;
|
|
const owners = new Map();
|
|
for (const record of records()) {
|
|
if (record.service !== 'model') continue;
|
|
const split = record.instance.indexOf('/'), business = record.instance.slice(0, split), instance = record.instance.slice(split + 1);
|
|
if (split < 1) continue;
|
|
for (const keyId of [record.keyId, ...(record.retiring ?? []).map(r => r.keyId)]) {
|
|
if (!ID.test(keyId ?? '')) continue;
|
|
if (owners.has(keyId)) throw Error('ambiguous');
|
|
owners.set(keyId, { business, instance });
|
|
}
|
|
}
|
|
// Validate all plug-in answers before writing any row or log field.
|
|
if (!Array.isArray(events) || events.length > 100 || events.some(e => !e || !ID.test(e.keyId ?? '') || !ID.test(e.accountId ?? '')
|
|
|| typeof e.managed !== 'boolean' || typeof e.failed !== 'boolean'
|
|
|| (e.statusCode !== null && (!Number.isSafeInteger(e.statusCode) || e.statusCode < 0 || e.statusCode > 599))
|
|
|| KINDS.some(k => !Number.isSafeInteger(e[k]) || e[k] < 0))) throw Error('invalid');
|
|
const unknown = [];
|
|
store.transaction(() => {
|
|
for (const e of events) {
|
|
const owner = owners.get(e.keyId);
|
|
if (!owner) { if (e.managed) unknown.push(e.keyId); continue; }
|
|
if (!enabled.some(b => b.business === owner.business)) continue;
|
|
// Whitelist fields; never persist raw plug-in results. Cursor holds only
|
|
// metadata, and turnRecords excludes SOURCE before every quota decision.
|
|
store.run(`INSERT INTO usage_records(at,business,run,instance,account,turn,source,phase,status,code,cursor,input,output,cache_write,cache_read)
|
|
VALUES (?,?,?,?,?,?,?,'end','known',?,?,?,?,?,?)`, new Date(now()).toISOString(), owner.business, '@proxy', owner.instance,
|
|
e.accountId, randomUUID(), SOURCE, e.failed ? 'proxy-upstream-failed' : null,
|
|
JSON.stringify({ modelKeyId: e.keyId, failed: e.failed, statusCode: e.statusCode }), e.input, e.output, e.cacheWrite, e.cacheRead);
|
|
}
|
|
});
|
|
written = true;
|
|
for (const keyId of unknown) log({ action: 'proxy.usage.reconcile', keyId, code: 'proxy-usage-unknown-key' });
|
|
for (const { business } of enabled) states.set(business, 'available');
|
|
} catch {
|
|
for (const { business } of enabled) {
|
|
states.set(business, 'unavailable');
|
|
if (drained && !written) dropped.set(business, (dropped.get(business) ?? 0) + 1);
|
|
}
|
|
log({ action: 'proxy.usage', code: 'proxy-usage-read-failed' });
|
|
}
|
|
}
|
|
function tick() {
|
|
if (closed || !enableDrain) return Promise.resolve();
|
|
if (pending) return pending;
|
|
pending = poll().finally(() => { pending = null; });
|
|
return pending;
|
|
}
|
|
function view(business) {
|
|
try { if (!configured().some(b => b.business === business)) return null; } catch { return null; }
|
|
const byKey = new Map(), byAccount = new Map();
|
|
for (const r of store.all('SELECT * FROM usage_records WHERE business=? AND source=? ORDER BY seq', business, SOURCE)) {
|
|
const meta = JSON.parse(r.cursor);
|
|
for (const [map, identity] of [[byKey, { keyId: meta.modelKeyId }], [byAccount, { accountId: r.account }]]) {
|
|
const id = Object.values(identity)[0];
|
|
if (!map.has(id)) map.set(id, { ...identity, input: 0n, output: 0n, cacheWrite: 0n, cacheRead: 0n, requests: 0, failures: 0, statuses: {} });
|
|
const entry = map.get(id);
|
|
for (const [key, column] of [['input','input'],['output','output'],['cacheWrite','cache_write'],['cacheRead','cache_read']]) entry[key] += BigInt(r[column]);
|
|
entry.requests++; if (meta.failed) entry.failures++;
|
|
if (meta.failed && meta.statusCode !== null) entry.statuses[meta.statusCode] = (entry.statuses[meta.statusCode] ?? 0) + 1;
|
|
}
|
|
}
|
|
const finish = map => [...map.values()].map(e => ({ ...e, ...Object.fromEntries(KINDS.map(k => [k, plain(e[k])])) }));
|
|
return { source: 'proxy', state: enableDrain ? (states.get(business) ?? 'unread') : 'unavailable', droppedBatches: dropped.get(business) ?? 0, byKey: finish(byKey), byAccount: finish(byAccount) };
|
|
}
|
|
return { tick, view, async close() { closed = true; await pending; } };
|
|
}
|