antonov
b223f51540
Host half wraps globalThis.fetch to read usage.cost out of provider responses (tee, never clone) and listens on the llm/stream seam for the session; the two halves join on the generation id. Client half adds one pill to the composer dock under its own id, so the harness's native stats stay where they are. Verified against the live API: three calls in one session (agent reply, compaction, session title) attributed with per-purpose totals, an unknown session answered with null, the ledger written and reloaded across instances, and the pill rendering only when a cost exists.
455 lines
16 KiB
JavaScript
455 lines
16 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);
|
|
}
|
|
for (const [id, total] of Object.entries(parsed.sessions ?? {})) {
|
|
sessionTotals.set(id, {
|
|
total: finiteAmount(total) ?? 0,
|
|
calls: 0,
|
|
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, value.total]),
|
|
),
|
|
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);
|
|
}
|
|
|
|
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;
|
|
}
|
|
const sessionId = url.searchParams.get('sessionId') ?? '';
|
|
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');
|
|
},
|
|
};
|