test/cascade-stream-account-spend.test.js
Rama master · réplica de lectura del repositorio original.
// Per-account lifetime spend must be recorded on the Cascade STREAM path (K8).
//
// Three of the four paths had a `recordAccountSpend` call point (connect stream, connect
// non-stream, Cascade non-stream); the Cascade stream path did not. On a Cascade deployment
// a streamed turn — the common case for an agent client — therefore added nothing to the
// account's lifetime tally, and the dashboard's per-account spend column under-reported by
// however much of the traffic was streamed.
//
// This is a BEHAVIOUR test on purpose, not a source grep. The call site lives inside
// `streamResponse`, which does not receive `apiKey` as a parameter (nonStreamResponse does —
// that asymmetry is why the gap existed) and resolves the account from its own `acct`
// binding instead. chat.js has ReferenceError'd exactly this way before: it referenced
// `body` inside streamResponse where `body` was never in scope, and it threw on EVERY
// stream finish (#93 follow-up; see test/thinking-fallback-glm.test.js's header). The call
// is wrapped in try/catch, so a scope error would be swallowed in silence — grepping for
// the call proves nothing, and only driving the path and watching the tally move does.
//
// The ledger also records the inverse hazard — the same tokens counted twice on a retried
// or replayed turn — so the retry and abort cases are pinned below.
import { afterEach, beforeEach, describe, it } from 'node:test';
import assert from 'node:assert/strict';
import {
addAccountByKey,
getAccountInternal,
getAccountPublic,
getApiKey,
removeAccount,
} from '../src/auth.js';
import { handleChatCompletions } from '../src/handlers/chat.js';
const createdIds = [];
let originalDevinConnect;
let originalDevinOnly;
function seed(label) {
const a = addAccountByKey(`devin-session-token$xk-${label}-${Math.random().toString(36).slice(2)}`, label);
createdIds.push(a.id);
const acct = getAccountInternal(a.id);
acct.tier = 'pro';
acct.status = 'active';
return acct;
}
function fakeResponse() {
return {
statusCode: 0,
body: '',
writableEnded: false,
writeHead(status) { this.statusCode = status; },
write(chunk) { this.body += String(chunk); return true; },
end(chunk) { if (chunk) this.body += String(chunk); this.writableEnded = true; },
on() {}, once() {}, removeListener() {},
};
}
/**
* @param {object} opts
* @param {number} [opts.promptTokens] usage the fake upstream reports
* @param {number} [opts.completionTokens]
* @param {number} [opts.failFirst] throw a retryable error on the first N calls
*/
function contextFor({ promptTokens = 100, completionTokens = 20, failFirst = 0, stream = true } = {}) {
let calls = 0;
class FakeClient {
async cascadeChat(_messages, _modelEnum, _modelUid, opts = {}) {
calls++;
if (calls <= failFirst) {
// The stream retry gate is `!hadSuccess && (err.isModelError || isRateLimit)`, so
// `isModelError` is what actually reaches the `continue`. `kind` only classifies it
// as transient (which suppresses the account penalty). An invented `code` satisfies
// neither, which is why the first draft never retried at all (calls=1) — and the
// precondition assertion below is what caught that rather than letting the test
// report "billed once" off a run with no retry in it.
const err = new Error('upstream stalled');
err.kind = 'transient_stall';
err.isModelError = true;
throw err;
}
// buildUsageBody reads the CASCADE field names (inputTokens / outputTokens), not the
// OpenAI ones. Getting this wrong is silent: it falls through to the chars/4
// estimate, so the tally still moves and only the amount is wrong — the first draft
// of this test asserted 120 and measured 2 that way.
const usage = { inputTokens: promptTokens, outputTokens: completionTokens };
if (stream) {
opts.onChunk({ text: 'OK' });
return { text: '', toolCalls: [], usage };
}
// Non-stream returns an ARRAY of chunks with toolCalls attached; returning a plain
// object yields a 502.
return Object.assign([{ text: 'OK' }], { toolCalls: [], usage });
}
}
const ctx = {
// Hands out a fresh account per attempt. The reference fixture this was copied from
// returned null after the first attempt, which silently capped the run at ONE upstream
// call and made the retry test's precondition unsatisfiable.
waitForAccount(tried, _signal, _maxWait, modelKey) {
return getApiKey(tried, modelKey);
},
ensureLs: async () => {},
getLsFor: () => ({ port: 17777, csrfToken: 'csrf', generation: 1 }),
WindsurfClient: FakeClient,
};
Object.defineProperty(ctx, 'upstreamCalls', { get: () => calls });
return ctx;
}
const spendOf = (id) => getAccountPublic(id)?.totalSpend?.totalTokens ?? 0;
beforeEach(() => {
originalDevinConnect = process.env.DEVIN_CONNECT;
originalDevinOnly = process.env.DEVIN_ONLY;
// Cascade transport, not connect: the gap under test is Cascade-specific.
delete process.env.DEVIN_CONNECT;
delete process.env.DEVIN_ONLY;
});
afterEach(() => {
// Only ids this file created. Never map removeAccount over the account list — a cleanup
// written that way once deleted a user's real accounts, unrecoverably.
while (createdIds.length) { try { removeAccount(createdIds.pop()); } catch {} }
if (originalDevinConnect === undefined) delete process.env.DEVIN_CONNECT;
else process.env.DEVIN_CONNECT = originalDevinConnect;
if (originalDevinOnly === undefined) delete process.env.DEVIN_ONLY;
else process.env.DEVIN_ONLY = originalDevinOnly;
});
describe('K8: Cascade stream path records per-account spend', () => {
it('a streamed turn moves the account tally', async () => {
const acct = seed('stream-spend');
const before = spendOf(acct.id);
const result = await handleChatCompletions(
{ model: 'gemini-2.5-flash', messages: [{ role: 'user', content: 'hi' }], stream: true },
contextFor({ promptTokens: 100, completionTokens: 20 }),
);
assert.equal(result.status, 200);
await result.handler(fakeResponse());
const after = spendOf(acct.id);
assert.ok(
after > before,
`per-account spend did not move (${before} -> ${after}). Either the call point is '
+ 'missing or its account binding is out of scope — the try/catch around it makes a '
+ 'ReferenceError silent, which is why this asserts the tally rather than the call`,
);
assert.equal(after - before, 120, 'the recorded amount must be the reported usage');
});
it('the non-stream path still records, so the fix did not move the problem', async () => {
const acct = seed('nonstream-spend');
const before = spendOf(acct.id);
const result = await handleChatCompletions(
{ model: 'gemini-2.5-flash', messages: [{ role: 'user', content: 'hi' }] },
contextFor({ promptTokens: 7, completionTokens: 3, stream: false }),
);
assert.equal(result.status, 200);
assert.equal(spendOf(acct.id) - before, 10);
});
it('a failed attempt is not billed — only the account that served it', async () => {
// The call site sits inside the attempt loop, so this is the hazard that matters. The
// ledger records a rescue retry's billing being counted into the following attempt.
//
// Two accounts, because the stream retry moves to the NEXT account rather than reusing
// the same one. Attempt 1 throws on whichever it drew; attempt 2 succeeds on the other.
// The failed one must end at zero and the total across both must be a single turn.
const a = seed('retry-a');
const b = seed('retry-b');
const before = spendOf(a.id) + spendOf(b.id);
const ctx = contextFor({ promptTokens: 100, completionTokens: 20, failFirst: 1 });
const result = await handleChatCompletions(
{ model: 'gemini-2.5-flash', messages: [{ role: 'user', content: 'hi' }], stream: true },
ctx,
);
assert.equal(result.status, 200);
await result.handler(fakeResponse());
assert.ok(ctx.upstreamCalls >= 2,
`precondition: the retry must have fired (calls=${ctx.upstreamCalls}). Without it this `
+ 'test would report "billed once" off a run containing no retry at all');
const total = spendOf(a.id) + spendOf(b.id) - before;
assert.equal(total, 120,
`a retried stream billed ${total} across both accounts instead of one turn's 120 — the `
+ 'failed attempt must not be counted');
const served = [a, b].filter((x) => spendOf(x.id) > 0);
assert.equal(served.length, 1, 'exactly one account may carry the charge');
});
});