antonov
6edd2e3d52
Two real defects, both found by testing rather than reading: 1. The client sent sessionId in the POST body while the route read it from the query string, so the host always saw an empty id and answered total:null -- the pill was correct to stay silent. The client now sends the id in both places and the route reads both. 2. The ledger stored only a session total, so a restart kept the sum but dropped the per-purpose split and the call count (9.99 total against 6.24 in parts, seen on the live route). Sessions are now stored as objects; a bare number from an older ledger is still accepted. Also made the client report why it cannot read a total: swallowing the error made 'nothing spent yet' indistinguishable from 'request failed'.
498 lines
18 KiB
JavaScript
498 lines
18 KiB
JavaScript
/**
|
|
* dsh-session-cost — host half.
|
|
*
|
|
* The harness throws away the cost its provider reports: `dsh-llm-pi-ai`'s
|
|
* `mapUsage` copies only token counts, and `pi-ai`'s `parseChunkUsage` fills a
|
|
* `cost` object from the model catalog (which does not describe a hand-declared
|
|
* route, so it is all zeros). The raw provider response still carries the real
|
|
* amount, so this plugin reads it there instead.
|
|
*
|
|
* Two observers, joined by generation id:
|
|
* 1. a `globalThis.fetch` wrapper reads `usage.cost` out of provider responses;
|
|
* 2. the `llm/stream` waterfall supplies the session the call belongs to.
|
|
*
|
|
* Neither is enough alone: the wrapper sees money but not the session, the seam
|
|
* sees the session but not the money.
|
|
*
|
|
* @module dsh-session-cost
|
|
*/
|
|
|
|
import fs from 'node:fs';
|
|
import os from 'node:os';
|
|
import path from 'node:path';
|
|
|
|
/** Marker so a reloaded plugin does not wrap `fetch` twice. */
|
|
const WRAP_MARK = Symbol.for('dsh-session-cost.fetch-wrapper');
|
|
|
|
/** How long to wait for a cost before calling a completed call unpriced. */
|
|
const PENDING_TIMEOUT_MS = 60_000;
|
|
|
|
/** Upper bounds, so a long-lived host cannot grow without limit. */
|
|
const MAX_REGISTRY = 5000;
|
|
const MAX_SETTLED = 5000;
|
|
|
|
/** Where the ledgers live. */
|
|
function dataDir() {
|
|
const home = process.env.DSH_HOME || path.join(os.homedir(), '.dsh');
|
|
return path.join(home, 'dsh-session-cost');
|
|
}
|
|
|
|
/** True for the response of a model call we should read. */
|
|
function looksLikeModelCall(url, response) {
|
|
if (response.headers.get('x-generation-id')) return true;
|
|
const type = response.headers.get('content-type') || '';
|
|
if (type.includes('text/event-stream')) return true;
|
|
if (!type.includes('json')) return false;
|
|
try {
|
|
const pathname = new URL(url).pathname;
|
|
return /(chat\/completions|completions|messages|responses)$/.test(pathname);
|
|
} catch {
|
|
return false;
|
|
}
|
|
}
|
|
|
|
/** Read a request's URL whichever shape `fetch` was called with. */
|
|
function urlOf(input) {
|
|
if (typeof input === 'string') return input;
|
|
if (input instanceof URL) return input.href;
|
|
if (input && typeof input.url === 'string') return input.url;
|
|
return '';
|
|
}
|
|
|
|
/** A finite, non-negative number, or null. */
|
|
function finiteAmount(value) {
|
|
if (typeof value !== 'number' || !Number.isFinite(value) || value < 0) return null;
|
|
return value;
|
|
}
|
|
|
|
/** Locate `usage.cost` and the generation id in one parsed provider payload. */
|
|
function readPayload(payload) {
|
|
if (!payload || typeof payload !== 'object') return null;
|
|
const usage = payload.usage;
|
|
if (!usage || typeof usage !== 'object') return null;
|
|
const cost = finiteAmount(usage.cost);
|
|
if (cost === null) return null;
|
|
const id = typeof payload.id === 'string' && payload.id.length > 0 ? payload.id : null;
|
|
return { id, cost };
|
|
}
|
|
|
|
/**
|
|
* Split an SSE byte stream into `data:` payloads.
|
|
* Chunk boundaries fall anywhere, so the tail is carried between pieces.
|
|
*/
|
|
async function eachSsePayload(stream, onPayload) {
|
|
const decoder = new TextDecoder();
|
|
let buffer = '';
|
|
for await (const piece of stream) {
|
|
buffer += decoder.decode(piece, { stream: true });
|
|
const lines = buffer.split('\n');
|
|
buffer = lines.pop() ?? '';
|
|
for (const line of lines) {
|
|
if (!line.startsWith('data:')) continue;
|
|
const body = line.slice(5).trim();
|
|
if (body.length === 0 || body === '[DONE]') continue;
|
|
try {
|
|
onPayload(JSON.parse(body));
|
|
} catch {
|
|
/* a non-JSON keepalive line is not an error */
|
|
}
|
|
}
|
|
}
|
|
buffer += decoder.decode();
|
|
const tail = buffer.trim();
|
|
if (tail.startsWith('data:')) {
|
|
const body = tail.slice(5).trim();
|
|
if (body.length > 0 && body !== '[DONE]') {
|
|
try {
|
|
onPayload(JSON.parse(body));
|
|
} catch {
|
|
/* ignore */
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
export default {
|
|
inject: ['timer'],
|
|
|
|
apply(ctx) {
|
|
// ---------------------------------------------------------------- state
|
|
/** generation id -> cost in the provider's own currency. */
|
|
const costByGeneration = new Map();
|
|
/** generation ids already added to a session total. */
|
|
const settled = new Set();
|
|
/** session id -> { total, calls, byPurpose }. */
|
|
const sessionTotals = new Map();
|
|
/** generation id -> { sessionId, purpose }, waiting for its cost. */
|
|
const waiting = new Map();
|
|
/** Diagnostics surfaced on the status page. */
|
|
const stats = {
|
|
observed: 0, // provider responses whose cost we read
|
|
attributed: 0, // calls joined to a session
|
|
cancelled: 0, // aborted by the user; the provider bills nothing
|
|
failed: 0, // finished with an error; no cost expected
|
|
unpriced: 0, // completed, but no cost ever arrived
|
|
unattributed: 0, // cost read, but the call carried no session id
|
|
wrappers: 0,
|
|
};
|
|
|
|
let flushTimer = null;
|
|
let dirty = false;
|
|
|
|
// ---------------------------------------------------------- persistence
|
|
const stateFile = path.join(dataDir(), 'state.json');
|
|
|
|
function load() {
|
|
try {
|
|
const raw = fs.readFileSync(stateFile, 'utf8');
|
|
const parsed = JSON.parse(raw);
|
|
for (const [id, cost] of Object.entries(parsed.costs ?? {})) {
|
|
costByGeneration.set(id, cost);
|
|
}
|
|
// Sessions are stored as objects so the per-purpose split survives a
|
|
// restart; older ledgers wrote a bare number, which is still accepted.
|
|
for (const [id, value] of Object.entries(parsed.sessions ?? {})) {
|
|
if (typeof value === 'number') {
|
|
sessionTotals.set(id, { total: finiteAmount(value) ?? 0, calls: 0, byPurpose: {} });
|
|
continue;
|
|
}
|
|
sessionTotals.set(id, {
|
|
total: finiteAmount(value?.total) ?? 0,
|
|
calls: typeof value?.calls === 'number' ? value.calls : 0,
|
|
byPurpose: value?.byPurpose && typeof value.byPurpose === 'object' ? { ...value.byPurpose } : {},
|
|
});
|
|
}
|
|
for (const id of parsed.settled ?? []) settled.add(id);
|
|
} catch {
|
|
/* first run, or an unreadable ledger: start clean rather than fail */
|
|
}
|
|
}
|
|
|
|
function flush() {
|
|
flushTimer = null;
|
|
if (!dirty) return;
|
|
dirty = false;
|
|
try {
|
|
fs.mkdirSync(path.dirname(stateFile), { recursive: true });
|
|
const payload = {
|
|
version: 1,
|
|
costs: Object.fromEntries(costByGeneration),
|
|
sessions: Object.fromEntries(
|
|
[...sessionTotals].map(([id, value]) => [
|
|
id,
|
|
{ total: value.total, calls: value.calls, byPurpose: value.byPurpose },
|
|
]),
|
|
),
|
|
settled: [...settled],
|
|
};
|
|
fs.writeFileSync(stateFile, JSON.stringify(payload));
|
|
} catch (error) {
|
|
console.warn('[dsh-session-cost] could not write the ledger:', error?.message ?? error);
|
|
}
|
|
}
|
|
|
|
function scheduleFlush() {
|
|
dirty = true;
|
|
if (flushTimer !== null) return;
|
|
flushTimer = ctx.setTimeout(flush, 2000);
|
|
}
|
|
|
|
/** Keep the in-memory maps bounded without losing recent history. */
|
|
function trim(map, cap) {
|
|
while (map.size > cap) {
|
|
const oldest = map.keys().next().value;
|
|
map.delete(oldest);
|
|
}
|
|
}
|
|
|
|
load();
|
|
|
|
// ------------------------------------------------------- cost recording
|
|
/** Add one cost to a session, at most once per generation. */
|
|
function settle(sessionId, generationId, cost, purpose) {
|
|
if (generationId !== null) {
|
|
if (settled.has(generationId)) return;
|
|
settled.add(generationId);
|
|
trim(settled, MAX_SETTLED);
|
|
}
|
|
let entry = sessionTotals.get(sessionId);
|
|
if (entry === undefined) {
|
|
entry = { total: 0, calls: 0, byPurpose: {} };
|
|
sessionTotals.set(sessionId, entry);
|
|
}
|
|
const bucket = purpose === undefined || purpose === '' ? 'agent' : purpose;
|
|
entry.total += cost;
|
|
entry.calls += 1;
|
|
entry.byPurpose[bucket] = (entry.byPurpose[bucket] ?? 0) + cost;
|
|
stats.attributed += 1;
|
|
scheduleFlush();
|
|
}
|
|
|
|
/** A cost arrived: join it to a waiting call, or park it for a later one. */
|
|
function recordCost(generationId, cost) {
|
|
stats.observed += 1;
|
|
const parked = waiting.get(generationId);
|
|
if (parked !== undefined) {
|
|
waiting.delete(generationId);
|
|
settle(parked.sessionId, generationId, cost, parked.purpose);
|
|
return;
|
|
}
|
|
if (generationId === null) {
|
|
// No id to join on: the money is real but cannot be placed.
|
|
stats.unattributed += 1;
|
|
console.warn('[dsh-session-cost] a provider response reported cost without an id');
|
|
return;
|
|
}
|
|
costByGeneration.set(generationId, cost);
|
|
trim(costByGeneration, MAX_REGISTRY);
|
|
scheduleFlush();
|
|
// If nothing ever comes back for it, it stays here and is reported as
|
|
// parked rather than silently dropped.
|
|
}
|
|
|
|
// ---------------------------------------------------- observer 1: fetch
|
|
function installFetchWrapper() {
|
|
const current = globalThis.fetch;
|
|
if (typeof current !== 'function') return;
|
|
if (current[WRAP_MARK] === true) {
|
|
stats.wrappers = 1;
|
|
return; // a previous instance already wrapped it
|
|
}
|
|
|
|
const wrapped = async function (input, init) {
|
|
const response = await current.call(this, input, init);
|
|
try {
|
|
const url = urlOf(input);
|
|
if (!looksLikeModelCall(url, response)) return response;
|
|
'use strict';
|
|
const body = response.body;
|
|
if (body === null || typeof body.tee !== 'function') return response;
|
|
|
|
const generationId =
|
|
response.headers.get('x-generation-id') ??
|
|
(typeof response.headers.get('x-session-id') === 'string'
|
|
? response.headers.get('x-session-id')
|
|
: null);
|
|
|
|
const [mine, theirs] = body.tee();
|
|
const type = response.headers.get('content-type') || '';
|
|
|
|
// Read our branch to completion; never leave it unread, or the
|
|
// consumer's branch would be back-pressured behind it.
|
|
const observe = (async () => {
|
|
try {
|
|
if (type.includes('text/event-stream')) {
|
|
let id = generationId;
|
|
await eachSsePayload(mine, (payload) => {
|
|
const found = readPayload(payload);
|
|
if (found === null) return;
|
|
if (found.id !== null) id = found.id;
|
|
if (id !== null) recordCost(id, found.cost);
|
|
});
|
|
} else {
|
|
const text = await new Response(mine).text();
|
|
let parsed = null;
|
|
try {
|
|
parsed = JSON.parse(text);
|
|
} catch {
|
|
return;
|
|
}
|
|
const found = readPayload(parsed);
|
|
if (found !== null) {
|
|
const id = found.id ?? generationId;
|
|
if (id !== null) recordCost(id, found.cost);
|
|
}
|
|
}
|
|
} catch {
|
|
/* an aborted or truncated body carries no cost; nothing to do */
|
|
}
|
|
})();
|
|
// Detach: the consumer must not wait on our bookkeeping.
|
|
void observe;
|
|
|
|
const rebuilt = new Response(theirs, {
|
|
status: response.status,
|
|
statusText: response.statusText,
|
|
headers: response.headers,
|
|
});
|
|
// A rebuilt Response loses `url`; the OpenAI SDK logs it on retries.
|
|
try {
|
|
Object.defineProperty(rebuilt, 'url', { value: response.url, configurable: true });
|
|
} catch {
|
|
/* not fatal: url is used for diagnostics only */
|
|
}
|
|
stats.wrappers = 1;
|
|
return rebuilt;
|
|
} catch (error) {
|
|
// Any failure here must leave the request working.
|
|
console.warn('[dsh-session-cost] read failed, passing the response through:', error?.message ?? error);
|
|
return response;
|
|
}
|
|
};
|
|
|
|
Object.defineProperty(wrapped, WRAP_MARK, { value: true });
|
|
globalThis.fetch = wrapped;
|
|
ctx.effect(() => () => {
|
|
if (globalThis.fetch === wrapped) globalThis.fetch = current;
|
|
}, 'dsh-session-cost: fetch wrapper');
|
|
}
|
|
|
|
installFetchWrapper();
|
|
|
|
// ----------------------------------------------------- observer 2: seam
|
|
ctx.on('llm/stream', (options, next) => {
|
|
const sessionId = typeof options?.sessionId === 'string' ? options.sessionId : undefined;
|
|
const purpose = typeof options?.purpose === 'string' ? options.purpose : undefined;
|
|
return (async function* passThrough() {
|
|
let generationId = null;
|
|
let finished = null;
|
|
try {
|
|
for await (const chunk of await next()) {
|
|
if (chunk && chunk.type === 'usage') {
|
|
// Nothing to take: the harness's usage has no cost field.
|
|
} else if (chunk && chunk.type === 'finish') {
|
|
finished = chunk.reason ?? null;
|
|
const id = chunk.replayState?.response?.responseId;
|
|
if (typeof id === 'string' && id.length > 0) generationId = id;
|
|
}
|
|
yield chunk;
|
|
}
|
|
} finally {
|
|
const kind = finished?.kind ?? 'stop';
|
|
if (kind === 'aborted') {
|
|
// The provider bills nothing for a cancelled generation (verified:
|
|
// a late abort still reports total_cost 0), so this is not a loss.
|
|
stats.cancelled += 1;
|
|
} else if (kind === 'error') {
|
|
stats.failed += 1;
|
|
} else if (generationId === null) {
|
|
stats.unpriced += 1;
|
|
} else if (costByGeneration.has(generationId)) {
|
|
const cost = costByGeneration.get(generationId);
|
|
costByGeneration.delete(generationId);
|
|
if (sessionId === undefined) {
|
|
stats.unattributed += 1;
|
|
} else {
|
|
settle(sessionId, generationId, cost, purpose);
|
|
}
|
|
} else if (sessionId === undefined) {
|
|
// Future plugins (e.g. an experimental auto-review) call the model
|
|
// with no sessionId at all; those calls cannot be placed.
|
|
stats.unattributed += 1;
|
|
} else {
|
|
// Normal path is cost-first, so this is rare: park the call and let
|
|
// the cost find it.
|
|
waiting.set(generationId, { sessionId, purpose, at: Date.now() });
|
|
}
|
|
}
|
|
})();
|
|
});
|
|
|
|
// Drop waiters whose cost never arrived, and report them.
|
|
ctx.setInterval(() => {
|
|
const now = Date.now();
|
|
for (const [id, entry] of waiting) {
|
|
if (now - entry.at < PENDING_TIMEOUT_MS) continue;
|
|
waiting.delete(id);
|
|
stats.unpriced += 1;
|
|
console.warn('[dsh-session-cost] a completed call never reported a cost:', id);
|
|
}
|
|
}, 30_000);
|
|
|
|
// --------------------------------------------------------------- status
|
|
function snapshot(sessionId) {
|
|
const entry = sessionTotals.get(sessionId);
|
|
return {
|
|
sessionId,
|
|
total: entry?.total ?? null,
|
|
calls: entry?.calls ?? 0,
|
|
byPurpose: entry?.byPurpose ?? {},
|
|
currency: 'RUB',
|
|
stats: { ...stats, waiting: waiting.size, parked: costByGeneration.size, sessions: sessionTotals.size },
|
|
};
|
|
}
|
|
|
|
// ---------------------------------------------------------------- route
|
|
const ROUTE_PREFIX = '/_dsh/dsh-session-cost';
|
|
|
|
/** Answer with JSON, or with a plain error the client can show. */
|
|
function respond(res, status, payload) {
|
|
const body = JSON.stringify(payload);
|
|
res.writeHead(status, {
|
|
'content-type': 'application/json; charset=utf-8',
|
|
'cache-control': 'no-store',
|
|
'content-length': Buffer.byteLength(body),
|
|
});
|
|
res.end(body);
|
|
}
|
|
|
|
/** Read a request body, capped: this route only ever carries a session id. */
|
|
function readJsonBody(req, limit = 16 * 1024) {
|
|
return new Promise((resolve) => {
|
|
const chunks = [];
|
|
let size = 0;
|
|
req.on('data', (chunk) => {
|
|
size += chunk.length;
|
|
if (size > limit) {
|
|
resolve({});
|
|
return;
|
|
}
|
|
chunks.push(chunk);
|
|
});
|
|
req.on('end', () => {
|
|
const raw = Buffer.concat(chunks).toString('utf8').trim();
|
|
if (raw.length === 0) {
|
|
resolve({});
|
|
return;
|
|
}
|
|
try {
|
|
const parsed = JSON.parse(raw);
|
|
resolve(parsed && typeof parsed === 'object' ? parsed : {});
|
|
} catch {
|
|
resolve({});
|
|
}
|
|
});
|
|
req.on('error', () => resolve({}));
|
|
});
|
|
}
|
|
|
|
ctx.inject(['connection', 'webServer'], (webCtx) => {
|
|
webCtx.effect(() => {
|
|
const dispose = webCtx.webServer.register({
|
|
kind: 'prefix',
|
|
path: ROUTE_PREFIX,
|
|
handler: async (req, res) => {
|
|
try {
|
|
const rejection = webCtx.connection.requestRejection(req);
|
|
if (rejection !== undefined) {
|
|
respond(res, rejection, { error: rejection === 401 ? 'authentication required' : 'request rejected' });
|
|
return;
|
|
}
|
|
const url = new URL(req.url || '/', 'http://localhost');
|
|
const tail = url.pathname.slice(ROUTE_PREFIX.length);
|
|
if (tail !== '/session' && tail !== '/status') {
|
|
respond(res, 404, { error: 'unknown method' });
|
|
return;
|
|
}
|
|
// The client posts a JSON body; the query string is kept as a
|
|
// fallback so either shape works.
|
|
const body = req.method === 'POST' || req.method === 'PUT' ? await readJsonBody(req) : {};
|
|
const raw = body.sessionId ?? url.searchParams.get('sessionId') ?? '';
|
|
const sessionId = typeof raw === 'string' ? raw : String(raw ?? '');
|
|
respond(res, 200, tail === '/session' ? snapshot(sessionId) : snapshot(''));
|
|
} catch (error) {
|
|
respond(res, 500, { error: String(error?.message ?? error) });
|
|
}
|
|
},
|
|
});
|
|
return dispose;
|
|
}, 'dsh-session-cost: http route');
|
|
});
|
|
|
|
ctx.effect(() => () => {
|
|
if (flushTimer !== null) ctx.clearTimeout(flushTimer);
|
|
flush();
|
|
}, 'dsh-session-cost: ledger flush');
|
|
},
|
|
};
|