openvibers/openvibe.ai Commits

AI: the hard spend cap of a run token (Codes C3) (#204)

OP OpenVibers committed 14 days ago · 3355f64599 · 10 files, +418 −12

A run token minted by Codes carries a budget: the row here, keyed by the token's jti, is what the meter counts against and what the cap is checked before every request. The budget is not a claim in the token, so the holder of a token can neither raise it nor name another run's budget, and a budget in a request is ignored. The request that crosses the line still completes; the next one is refused with 402 ai.run_budget_reached, naming the budget and what was spent. The counter lives in the database and expires with the token.

Changes

ddiff --git a/migrations/0015_run_budgets.sql b/migrations/0015_run_budgets.sql
nnew file mode 100644
iindex 0000000..64dee36
---- /dev/null
++++ b/migrations/0015_run_budgets.sql
@@@ -0,0 +1,19 @@
++-- phase: expand
++-- The budget one Codes run token carries (Codes C3), and what that run has spent.
++--
++-- The budget is a row, not a claim on the token: Codes registers it under the token's jti when it mints the
++-- token, and the row is what the meter counts against. A holder of the token cannot write the row (only the
++-- service Codes writes it, with its own token) and cannot name another run's row (the jti is signed inside
++-- the token by Network). The row dies with the token: expires_at is the token's expiry.
++CREATE TABLE IF NOT EXISTS run_budgets (
++ jti text COLLATE "C" PRIMARY KEY, -- the run token's jti (signed inside the token)
++ budget_usd double precision, -- null: no dollar cap
++ budget_tokens bigint, -- null: no token cap
++ spent_usd double precision NOT NULL DEFAULT 0,
++ spent_tokens bigint NOT NULL DEFAULT 0,
++ expires_at bigint NOT NULL, -- epoch seconds: the token's expiry
++ created_at text COLLATE "C" NOT NULL,
++ updated_at text COLLATE "C" NOT NULL
++);
++
++CREATE INDEX IF NOT EXISTS idx_run_budgets_expires ON run_budgets (expires_at);
ddiff --git a/server/api/messages.js b/server/api/messages.js
iindex 120693b..6398577 100644
---- a/server/api/messages.js
++++ b/server/api/messages.js
@@@ -41,7 +41,7 @@ function errorType(status) {
}
function anthropicError(res, status, message, type) {
-- res.status(status).json({ type: 'error', error: { type: type || errorType(status), message: String(message || 'error').slice(0, 1000) } });
++ res.status(status).json({ type: 'error', error: { type: type || errorType(status), message: String(message || 'error').slice(0, 1000), code: status === 402 ? 'ai.run_budget_reached' : null } });
}
/**
@@@ -82,8 +82,8 @@ function createEventStream(res) {
sentAny: () => any,
get stopped() { return stopped; },
fail(message, code) {
-- const status = /^(input|validation)\./.test(code || '') ? 400 : /quota|allowance|budget/.test(code || '') ? 429 : 529;
-- emit('error', { type: 'error', error: { type: errorType(status), message: String(message || 'error').slice(0, 1000) } });
++ const status = code === 'ai.run_budget_reached' ? 402 : /^(input|validation)\./.test(code || '') ? 400 : /quota|allowance|budget/.test(code || '') ? 429 : 529;
++ emit('error', { type: 'error', error: { type: errorType(status), message: String(message || 'error').slice(0, 1000), code: code || null } });
return out.done();
},
done() { if (alive()) res.end(); },
@@@ -163,7 +163,7 @@ function messagesRouter({ runs, auth, config, log = console }) {
}
if (run.status === 'failed' || run.status === 'cancelled') {
const code = (run.error && run.error.code) || 'run.failed';
-- const status = /^(input|validation)\./.test(code) ? 400 : /quota|allowance|budget/.test(code) ? 429
++ const status = code === 'ai.run_budget_reached' ? 402 : /^(input|validation)\./.test(code) ? 400 : /quota|allowance|budget/.test(code) ? 429
: code === 'credential.unsupported' ? 422 : code === 'credential.not_found' ? 404 : 529;
return anthropicError(res, status, (run.error && run.error.detail) || code);
}
ddiff --git a/server/api/openai.js b/server/api/openai.js
iindex 39c20fd..a5a156a 100644
---- a/server/api/openai.js
++++ b/server/api/openai.js
@@@ -79,8 +79,8 @@ function createLiveStream(res) {
sentAny: () => any,
finished: (i) => finishedIdx.has(i),
fail(message, code) {
-- const status = /^(input|validation)\./.test(code || '') ? 400 : /quota|allowance|budget/.test(code || '') ? 429 : 503;
-- write({ error: { message: String(message || 'error').slice(0, 1000), type: status === 400 ? 'invalid_request_error' : status === 429 ? 'rate_limit_error' : 'api_error', code: code || null, param: null } });
++ const status = code === 'ai.run_budget_reached' ? 402 : /^(input|validation)\./.test(code || '') ? 400 : /quota|allowance|budget/.test(code || '') ? 429 : 503;
++ write({ error: { message: String(message || 'error').slice(0, 1000), type: status === 400 ? 'invalid_request_error' : status === 429 ? 'rate_limit_error' : status === 402 ? 'insufficient_quota' : 'api_error', code: code || null, param: null } });
return out.done();
},
done() { if (alive()) { res.write('data: [DONE]\n\n'); res.end(); } },
@@@ -89,7 +89,7 @@ function createLiveStream(res) {
}
function openaiError(res, status, message, code, type) {
-- res.status(status).json({ error: { message: String(message || 'error').slice(0, 1000), type: type || (status === 401 || status === 403 ? 'authentication_error' : status === 429 ? 'rate_limit_error' : status >= 500 ? 'api_error' : 'invalid_request_error'), code: code || null, param: null } });
++ res.status(status).json({ error: { message: String(message || 'error').slice(0, 1000), type: type || (status === 401 || status === 403 ? 'authentication_error' : status === 429 ? 'rate_limit_error' : status === 402 ? 'insufficient_quota' : status >= 500 ? 'api_error' : 'invalid_request_error'), code: code || null, param: null } });
}
function openaiRouter({ runs, auth, config, log = console }) {
@@@ -155,7 +155,7 @@ function openaiRouter({ runs, auth, config, log = console }) {
}
if (run.status === 'failed' || run.status === 'cancelled') {
const code = (run.error && run.error.code) || 'run.failed';
-- const status = /^(input|validation)\./.test(code) ? 400 : /quota|allowance|budget/.test(code) ? 429
++ const status = code === 'ai.run_budget_reached' ? 402 : /^(input|validation)\./.test(code) ? 400 : /quota|allowance|budget/.test(code) ? 429
: code === 'project_secret.refused' ? 403 : code === 'credential.endpoint_refused' ? 422 : 503;
return openaiError(res, status, (run.error && run.error.detail) || code, code);
}
ddiff --git a/server/api/run-budgets.js b/server/api/run-budgets.js
nnew file mode 100644
iindex 0000000..fa0926f
---- /dev/null
++++ b/server/api/run-budgets.js
@@@ -0,0 +1,60 @@
++'use strict';
++/**
++ * The budget of a Codes run token, and what that run has spent (Codes C3).
++ *
++ * PUT /api/v1/run-budgets/:jti register the budget of one run token (Codes, when it mints the token)
++ * GET /api/v1/run-budgets/:jti what the run may spend and what it has spent
++ *
++ * The jti is the run token's own id, signed inside the token by Network: the service that mints the token
++ * is the one that registers its budget, and the holder of the token can neither change the budget nor name
++ * another run's budget. A budget that arrives in the request body of a model call is ignored: there is
++ * nothing there to forge.
++ */
++const express = require('express');
++const { AiError } = require('../util');
++
++const JTI_RE = /^[A-Za-z0-9_.:-]{1,200}$/;
++
++function runBudgetsRouter({ runBudgets, auth, log = console, sendError }) {
++ const r = express.Router();
++ const guard = auth.requireCap('ai.run.create');
++
++ const jtiOf = (req) => {
++ const jti = String(req.params.jti || '');
++ if (!JTI_RE.test(jti)) throw new AiError(422, 'input.invalid', 'the jti is not a run token id');
++ return jti;
++ };
++
++ r.put('/api/v1/run-budgets/:jti', guard, async (req, res) => {
++ try {
++ const jti = jtiOf(req);
++ const body = req.body || {};
++ const row = await runBudgets.register({
++ jti,
++ budgetUsd: body.budget_usd == null ? null : Number(body.budget_usd),
++ budgetTokens: body.budget_tokens == null ? null : Number(body.budget_tokens),
++ expiresAt: body.expires_at == null ? null : Number(body.expires_at),
++ });
++ if (!row) throw new AiError(422, 'input.invalid', 'a budget needs dollars or tokens');
++ res.set('Cache-Control', 'no-store').json({
++ jti: row.jti,
++ budget_usd: row.budget_usd == null ? null : Number(row.budget_usd),
++ budget_tokens: row.budget_tokens == null ? null : Number(row.budget_tokens),
++ expires_at: Number(row.expires_at),
++ });
++ } catch (err) { sendError(res, err, req.ov, log); }
++ });
++
++ r.get('/api/v1/run-budgets/:jti', guard, async (req, res) => {
++ try {
++ const jti = jtiOf(req);
++ const row = await runBudgets.spent(jti);
++ if (!row) { res.set('Cache-Control', 'no-store').json({ jti, budget_usd: null, budget_tokens: null, spent_usd: 0, spent_tokens: 0, reached: false, expires_at: null }); return; }
++ res.set('Cache-Control', 'no-store').json(row);
++ } catch (err) { sendError(res, err, req.ov, log); }
++ });
++
++ return r;
++}
++
++module.exports = { runBudgetsRouter };
ddiff --git a/server/app.js b/server/app.js
iindex 4af406a..6b9f38b 100644
---- a/server/app.js
++++ b/server/app.js
@@@ -21,7 +21,7 @@ const ovServe = require('openvibe-shared/serve');
const home = require('./home');
const pkg = require('../package.json');
--function createApp({ config, db, registry, pool, quotas, cache, runs, auth, keys, env = process.env, clock = { now: () => Date.now() }, fetchImpl = globalThis.fetch, log = console, credentials = null, vault = null, benchEvents = null }) {
++function createApp({ config, db, registry, pool, quotas, cache, runs, auth, keys, env = process.env, clock = { now: () => Date.now() }, fetchImpl = globalThis.fetch, log = console, credentials = null, vault = null, benchEvents = null, runBudgets = null }) {
const app = express();
app.disable('x-powered-by');
app.set('trust proxy', 'loopback');
@@@ -110,6 +110,8 @@ function createApp({ config, db, registry, pool, quotas, cache, runs, auth, keys
app.use(require('./api/messages').messagesRouter({ runs, auth, config, log }));
app.use(require('./api/attribution-quotas').attributionQuotasRouter({ db, quotas, auth, clock, sendError: (res, err, ctx) => sendError(res, err, ctx, log) }));
if (credentials) app.use(require('./credentials').credentialsRouter({ credentials, auth, registry, sendError: (res, err, ctx) => sendError(res, err, ctx, log) }));
++ // The budget of a Codes run token (Codes C3): Codes registers it when it mints the token.
++ if (runBudgets) app.use(require('./api/run-budgets').runBudgetsRouter({ runBudgets, auth, log, sendError: (res, err, ctx) => sendError(res, err, ctx, log) }));
app.use(adminRouter({ db, registry, pool, quotas, cache, runs, auth, ops, log }));
// The operator console (/console, /auth/*): Network staff, server-rendered, no scripts (server/console).
// Person sessions are separate, read-only and open one run receipt each (server/person-session.js).
ddiff --git a/server/index.js b/server/index.js
iindex 86c9066..3257a23 100644
---- a/server/index.js
++++ b/server/index.js
@@@ -17,6 +17,7 @@ const { createFetcher } = require('./fetcher');
const { createEngine } = require('./workflows/engine');
const { createCache } = require('./cache');
const { createQuotas } = require('./quota');
++const { createRunBudgets } = require('./run-budgets');
const { createFreeAllowance } = require('./free-allowance');
const { createTiers } = require('./govern');
const { createRuns } = require('./runs');
@@@ -46,6 +47,8 @@ async function start({ config, db: givenDb = null, clock = { now: () => Date.now
// with the tiers on, its counter is in their Valkey store.
const freeAllowance = createFreeAllowance(db, { cardsFor: pool.rateCardsFor, clock, store: tiers.freeStore });
const quotas = createQuotas(db, { clock, registry, freeAllowance, tiers });
++ // The budget a Codes run token carries (Codes C3), and what that run has spent.
++ const runBudgets = createRunBudgets(db, { clock });
const cache = createCache(db, { clock });
const fetcher = createFetcher(config);
// MegaSource (server/research/): the DeepSeek key resolves by name from this environment; tests hand in their own.
@@@ -59,7 +62,7 @@ async function start({ config, db: givenDb = null, clock = { now: () => Date.now
// A person's own provider keys (WS-O task 2): stored by the service holding their consent, used by their runs only.
const credentials = require('./credentials').createCredentials({ db, config, clock });
const projectSecrets = require('./project-secrets').createProjectSecrets({ config, env, fetchImpl });
-- const runs = createRuns({ db, registry, engine, cache, quotas, config, clock, log, userModules, credentials, tiers, projectSecrets, pool });
++ const runs = createRuns({ db, registry, engine, cache, quotas, config, clock, log, userModules, credentials, tiers, projectSecrets, pool, runBudgets });
await seed({ registry, quotas, config, env, db });
const interrupted = await runs.recoverInterrupted();
if (interrupted) log.warn(`[ai] marked ${interrupted} interrupted run(s) failed (run.interrupted); callers can retry them`);
@@@ -71,7 +74,7 @@ async function start({ config, db: givenDb = null, clock = { now: () => Date.now
const auth = createAuth({ config, log });
// The Events consumer side: a signed bench.rating.changed webhook, subscribed at boot (server/bench/events.js).
const benchEvents = createBenchEvents({ bench, config, log, fetchImpl, clientSecret: env.OV_OAUTH_CLIENT_SECRET });
-- const app = createApp({ config, db, registry, pool, quotas, cache, runs, auth, keys, env, clock, fetchImpl, log, credentials, vault, benchEvents });
++ const app = createApp({ config, db, registry, pool, quotas, cache, runs, auth, keys, env, clock, fetchImpl, log, credentials, vault, benchEvents, runBudgets });
let jwks = null;
let keyLoaded = Promise.resolve(null);
if (!keys.pinned && config.nodeEnv !== 'test') {
ddiff --git a/server/openapi.js b/server/openapi.js
iindex 0f4ac3f..1a9d4cc 100644
---- a/server/openapi.js
++++ b/server/openapi.js
@@@ -161,6 +161,21 @@ route('post', '/research', {
responses: { 201: json('Finished', ref('RunResponse')), 202: json('Still running', ref('RunResponse')), 200: json('Replay', ref('RunResponse')), ...errors() },
});
++// ── The budget of a run token (Codes C3) ────────────────────────────────────────────────────
++route('put', '/run-budgets/{jti}', {
++ operationId: 'putRunBudget', summary: "Register the budget of a run token", tags: ['runs'], security: SERVICE,
++ description: "ai.run.create. The budget the service that minted a run token gives it: OpenVibe.AI counts every request of that token against it and refuses the request after the budget is spent (402 ai.run_budget_reached). A token without a budget is never capped.",
++ parameters: [path('jti', str(), "The run token's own id")],
++ requestBody: body({ type: 'object', properties: { budget_usd: nullable('number'), budget_tokens: nullable('integer'), expires_at: nullable('integer') } }),
++ responses: { 200: json('The budget', { type: 'object', required: ['jti'], properties: { jti: str(), budget_usd: nullable('number'), budget_tokens: nullable('integer'), expires_at: integer() } }), ...errors([401, 403, 422]) },
++});
++route('get', '/run-budgets/{jti}', {
++ operationId: 'getRunBudget', summary: 'What a run token may spend, and what it has spent', tags: ['runs'], security: SERVICE,
++ description: 'ai.run.create. The counter of one run token: the budget it carries and what the run spent against it.',
++ parameters: [path('jti', str(), "The run token's own id")],
++ responses: { 200: json('The budget and the spend', { type: 'object', required: ['jti'], properties: { jti: str(), budget_usd: nullable('number'), budget_tokens: nullable('integer'), spent_usd: nullable('number'), spent_tokens: nullable('integer'), expires_at: nullable('integer'), reached: { type: 'boolean' } } }), ...errors([401, 403, 422]) },
++});
++
// ── Action Center ────────────────────────────────────────────────────────────────────────────
route('get', '/actions', {
operationId: 'listActions', summary: 'The Action Center items for one person', tags: ['actions'], security: SERVICE,
ddiff --git a/server/run-budgets.js b/server/run-budgets.js
nnew file mode 100644
iindex 0000000..4118f5a
---- /dev/null
++++ b/server/run-budgets.js
@@@ -0,0 +1,134 @@
++'use strict';
++/**
++ * The budget one Codes run token carries, and what that run has spent (Codes C3).
++ *
++ * A run token is minted by Network for one attempt and names the person whose key it spends. The budget is
++ * a row here, keyed by the token's jti, and not a claim inside the token:
++ * - only Codes writes the row (with its own service token), so the holder of a run token cannot raise it;
++ * - the jti is signed inside the token by Network, so a holder cannot name another run's budget;
++ * - a budget that arrives in a request (a header, a body field) is ignored: there is nothing to forge.
++ * The row dies with the token: it expires at the token's expiry, and an expired row is no budget at all.
++ *
++ * The counter is what the runs add to: each finished run adds its tokens and its cost (the same numbers
++ * the quota accounting uses, server/quota.js). The cap is checked BEFORE every request, so the request that
++ * crosses the line still completes and the next one is refused.
++ */
++const { AiError } = require('./util');
++
++/** The problem a run over its budget gets: 402, with the budget that ran out and what was spent. */
++function budgetReached(budget, spent, limit) {
++ const which = budget === 'tokens' ? 'tokens' : 'dollars';
++ const spentText = budget === 'tokens' ? `${spent.tokens} tokens` : `$${Number(spent.usd).toFixed(4)}`;
++ const limitText = budget === 'tokens' ? `${limit} tokens` : `$${Number(limit).toFixed(2)}`;
++ return new AiError(402, 'ai.run_budget_reached',
++ `this run reached its ${which} budget: ${spentText} of ${limitText} spent. It may not spend more.`,
++ { budget, limit, spent_usd: spent.usd, spent_tokens: spent.tokens });
++}
++
++function createRunBudgets(db, { clock = { now: () => Date.now() } } = {}) {
++ const nowSec = () => Math.floor(clock.now() / 1000);
++ const nowIso = () => new Date(clock.now()).toISOString();
++
++ const getRow = db.prepare('SELECT * FROM run_budgets WHERE jti = ?');
++ const put = db.prepare(`INSERT INTO run_budgets (jti, budget_usd, budget_tokens, spent_usd, spent_tokens, expires_at, created_at, updated_at)
++ VALUES (?, ?, ?, 0, 0, ?, ?, ?)
++ ON CONFLICT (jti) DO UPDATE SET budget_usd = excluded.budget_usd, budget_tokens = excluded.budget_tokens,
++ expires_at = excluded.expires_at, updated_at = excluded.updated_at`);
++ const bump = db.prepare('UPDATE run_budgets SET spent_usd = spent_usd + ?, spent_tokens = spent_tokens + ?, updated_at = ? WHERE jti = ?');
++ const drop = db.prepare('DELETE FROM run_budgets WHERE jti = ?');
++
++ /** The budget a number is: null, or a finite number above zero. */
++ function asLimit(value) {
++ if (value == null || value === '') return null;
++ const n = Number(value);
++ if (!Number.isFinite(n) || n <= 0) return null;
++ return n;
++ }
++
++ /** The row for a token, if it is still alive: a budget that expired with its token is no budget. */
++ async function forToken(jti) {
++ if (!jti) return null;
++ const row = await getRow.get(jti);
++ if (!row) return null;
++ if (Number(row.expires_at) <= nowSec()) return null;
++ return row;
++ }
++
++ /**
++ * Register the budget of a run token (Codes, when it mints the token). A token without a budget is not
++ * capped: nothing is written and nothing is returned.
++ *
++ * @param {object} what
++ * @param {string} what.jti the token's jti
++ * @param {number|null} [what.budgetUsd] the dollars the run may spend
++ * @param {number|null} [what.budgetTokens] the tokens the run may spend
++ * @param {number|null} [what.expiresAt] epoch seconds: the token's expiry
++ */
++ async function register({ jti, budgetUsd = null, budgetTokens = null, expiresAt = null }) {
++ if (!jti) return null;
++ const usd = asLimit(budgetUsd);
++ const tokens = asLimit(budgetTokens);
++ if (usd == null && tokens == null) return null;
++ const exp = Math.round(Number(expiresAt)) || (nowSec() + 7200);
++ await put.run(jti, usd, tokens == null ? null : Math.round(tokens), exp, nowIso(), nowIso());
++ return await getRow.get(jti); // the row as it is stored: a budget that expired already is still a budget
++ }
++
++ /**
++ * Whether the run may spend: what it has spent against what it may spend. A run with no budget is never
++ * refused, and the request that crosses the line still completes (the check is before the next one).
++ */
++ async function check(jti) {
++ const row = await forToken(jti);
++ if (!row) return { ok: true, capped: false };
++ const spent = { usd: Number(row.spent_usd) || 0, tokens: Number(row.spent_tokens) || 0 };
++ const budgetUsd = row.budget_usd == null ? null : Number(row.budget_usd);
++ const budgetTokens = row.budget_tokens == null ? null : Number(row.budget_tokens);
++ if (budgetUsd != null && spent.usd >= budgetUsd) return { ok: false, capped: true, budget: 'usd', limit: budgetUsd, spent };
++ if (budgetTokens != null && spent.tokens >= budgetTokens) return { ok: false, capped: true, budget: 'tokens', limit: budgetTokens, spent };
++ return { ok: true, capped: true, budget: null, limit: null, spent };
++ }
++
++ /** The refusal a run over its budget gets (402, ai.run_budget_reached). */
++ async function refuse(jti) {
++ const state = await check(jti);
++ if (state.ok) return null;
++ return budgetReached(state.budget, state.spent, state.limit);
++ }
++
++ /** What a finished run spent, added to the counter (the numbers the quota accounting uses). */
++ async function add(jti, { tokens = 0, cost = 0 } = {}) {
++ if (!jti) return null;
++ const t = Math.max(0, Math.round(Number(tokens) || 0));
++ const c = Math.max(0, Number(cost) || 0);
++ if (!t && !c) return null;
++ await bump.run(c, t, nowIso(), jti);
++ return forToken(jti);
++ }
++
++ /** What a run has spent, for the service that registered the budget. */
++ async function spent(jti) {
++ const row = await getRow.get(jti);
++ if (!row) return null;
++ return {
++ jti: row.jti,
++ budget_usd: row.budget_usd == null ? null : Number(row.budget_usd),
++ budget_tokens: row.budget_tokens == null ? null : Number(row.budget_tokens),
++ spent_usd: Number(row.spent_usd) || 0,
++ spent_tokens: Number(row.spent_tokens) || 0,
++ expires_at: Number(row.expires_at),
++ reached: (row.budget_usd != null && (Number(row.spent_usd) || 0) >= Number(row.budget_usd))
++ || (row.budget_tokens != null && (Number(row.spent_tokens) || 0) >= Number(row.budget_tokens)),
++ };
++ }
++
++ /** Forget the budgets whose token has expired. */
++ async function prune() {
++ const gone = await db.prepare('DELETE FROM run_budgets WHERE expires_at <= ?').run(nowSec());
++ return Number(gone && gone.changes) || 0;
++ }
++
++ return { register, check, refuse, add, spent, forToken, prune, budgetReached };
++}
++
++module.exports = { createRunBudgets, budgetReached };
ddiff --git a/server/runs.js b/server/runs.js
iindex 18f0628..a6bde91 100644
---- a/server/runs.js
++++ b/server/runs.js
@@@ -99,7 +99,7 @@ function routePinRequest(value, principal, wf) {
return { provider, model };
}
--function createRuns({ db, registry, engine, cache, quotas, config, clock = { now: () => Date.now() }, log = console, userModules = null, credentials = null, tiers = null, projectSecrets = null, pool = null }) {
++function createRuns({ db, registry, engine, cache, quotas, config, clock = { now: () => Date.now() }, log = console, userModules = null, credentials = null, tiers = null, projectSecrets = null, pool = null, runBudgets = null }) {
const inflight = new Map(); // run id -> { controller, promise, release, caller }
const queue = []; // run ids waiting for the worker
let active = 0;
@@@ -247,8 +247,15 @@ function createRuns({ db, registry, engine, cache, quotas, config, clock = { now
if (wf.key !== 'ai.complete') throw new AiError(422, 'input.invalid', 'a project provider is used for ai.complete only');
if (!projectSecrets || !projectSecrets.available()) throw new AiError(503, 'project_secret.unavailable', 'project secrets are not configured on this service');
}
++ // The budget a Codes run token carries (Codes C3): a run that is over its budget is refused before
++ // any provider call. The budget lives here, not in the token, so the holder of the token cannot raise it.
++ if (runBudgets && principal.jti) {
++ const refusal = await runBudgets.refuse(principal.jti);
++ if (refusal) throw refusal;
++ }
const ctx = {
requesterType: requester.type, requesterId: requester.id,
++ ...(principal.jti ? { runJti: principal.jti } : {}),
actorKey: subjectKey(body.on_behalf_of), targetKey: entityKey(body.target), attributionKey: entityKey(body.attribution), workflowKey: wf.key,
...(app ? { appRun: true, projectId: app.projectId, sandbox: app.sandbox } : {}),
...(pinned ? { routePin: pinned } : {}),
@@@ -426,6 +433,9 @@ function createRuns({ db, registry, engine, cache, quotas, config, clock = { now
});
if (!done.changes) return; // cancelled while finishing
if (r.provider && reserved) await quotas.account(ctx, reserved, { provider: r.provider, model: r.model, tokensIn: r.usage.input, tokensOut: r.usage.output, tokensCached: r.usage.cached, cost: r.cost, attempts: logged, runId: id, traceId: row.trace_id, hold });
++ // What the run spent goes on the budget of its run token (Codes C3): the same tokens and cost the
++ // quota accounting counts, so the next request of a run over its budget is refused.
++ if (runBudgets && ctx.runJti) await runBudgets.add(ctx.runJti, { tokens: r.usage.input + r.usage.output, cost: r.cost });
if (r.fallbackUsed) await registry.audit('system', 'run.fallback', 'run', id, { trace: row.trace_id, metadata: { workflow: row.workflow_key, provider: r.provider, route: r.route && r.route.key } });
// Output shaped by someone's preferences, or kept for nobody (history off), is never cached for reuse.
if (scope && row.cache_key && !r.synthetic && wf.cache_mode !== 'none' && preferences.history !== false && !preferenceLines(preferences)) {
@@@ -449,6 +459,8 @@ function createRuns({ db, registry, engine, cache, quotas, config, clock = { now
// ones included) into provider_stats_daily/placement_state: a run that answered nothing still
// tells the router a provider was unhealthy.
if (reserved && logged.some((e) => e.status !== 'skipped')) await quotas.account(ctx, reserved, { provider: 'mixed', model: '', tokensIn: spent.ti, tokensOut: spent.tout, cost: spent.c, attempts: logged, writeUsage: Boolean(spent.ti || spent.tout), runId: id, traceId: row.trace_id, hold });
++ // A run that broke still spent (a refused turn, a half-written answer): it goes on the budget too.
++ if (runBudgets && ctx.runJti) await runBudgets.add(ctx.runJti, { tokens: Number(spent.ti) + Number(spent.tout), cost: Number(spent.c) });
} finally {
inflight.delete(id);
}
ddiff --git a/test/run-budgets.test.js b/test/run-budgets.test.js
nnew file mode 100644
iindex 0000000..687d5db
---- /dev/null
++++ b/test/run-budgets.test.js
@@@ -0,0 +1,161 @@
++'use strict';
++// The hard spend cap of a Codes run token (Codes C3): the budget is a row in AI, keyed by the token's jti,
++// and every request of that token is counted against it. The request that crosses the line still completes;
++// the next one is refused with 402 ai.run_budget_reached. A token without a budget is never capped, and a
++// holder of a token can neither raise the budget nor name another run's budget.
++const assert = require('assert');
++const crypto = require('crypto');
++const { serviceAuth } = require('openvibe-contracts');
++const { boot, request, token, suite, ALL, privateKey, ISSUER } = require('./helpers');
++
++const t = suite('run-budgets');
++const KEY = 'cd'.repeat(32);
++const PERSON = 'usr_01JAB2C3D4E5F6G7H8J9K0MNP7';
++const PERSON_KEY = 'sk-openai-test-person-0123456789';
++const service = token('live', ALL);
++const manager = token('codes', [...ALL, 'ai.credential.manage'], { extra: { on_behalf_of: PERSON } });
++
++// A run token of its own, with a jti we know: Codes registers the budget under that jti.
++function runToken(jti, extra = {}) {
++ const now = Math.floor(Date.now() / 1000);
++ return serviceAuth.signServiceToken({
++ iss: ISSUER, sub: 'svc:codes', actor_type: 'service', aud: ['openvibe.ai'], cap: ALL, ns: ['*'],
++ iat: now, exp: now + 600, jti, ...extra,
++ }, privateKey);
++}
++
++const credentialFetch = async () => new Response(JSON.stringify({
++ id: 'chatcmpl_person', object: 'chat.completion', model: 'gpt-person',
++ choices: [{ index: 0, message: { role: 'assistant', content: 'Hello there' }, finish_reason: 'stop' }],
++ usage: { prompt_tokens: 5, completion_tokens: 2 },
++}), { status: 200, headers: { 'content-type': 'application/json' } });
++
++let h;
++const chat = (tok, headers = {}) => request(h.base, 'POST', '/v1/chat/completions', {
++ tok, body: { model: 'openvibe', messages: [{ role: 'user', content: 'Say hello' }] }, headers: { 'x-openvibe-credential': 'openai', ...headers },
++});
++const budgetOf = (jti) => request(h.base, 'GET', `/api/v1/run-budgets/${jti}`, { tok: manager });
++const setBudget = (jti, body) => request(h.base, 'PUT', `/api/v1/run-budgets/${jti}`, { tok: manager, body });
++
++t.test('boot with the credentials key set', async () => {
++ h = await boot({
++ env: {
++ AI_CREDENTIALS_KEY: KEY, AI_STUB_FALLBACK: 'false',
++ // The person's own key is priced here, so a run really costs something (per Mtok).
++ AI_PRICING_JSON: JSON.stringify({ 'gpt-person': { in: 1000, out: 2000 } }),
++ },
++ credentialFetch,
++ });
++});
++
++t.test("the person's own key is stored", async () => {
++ const r = await request(h.base, 'PUT', `/api/v1/credentials/${PERSON}/openai`, { tok: manager, body: { key: PERSON_KEY, models: { chat: 'gpt-person' } } });
++ assert.strictEqual(r.status, 200, r.text);
++});
++
++t.test('a budget is registered for a run token', async () => {
++ const jti = 'tok_budget_1';
++ const r = await setBudget(jti, { budget_usd: 5, budget_tokens: 1000, expires_at: Math.floor(Date.now() / 1000) + 600 });
++ assert.strictEqual(r.status, 200, r.text);
++ assert.strictEqual(r.body.budget_usd, 5);
++ assert.strictEqual(r.body.budget_tokens, 1000);
++});
++
++t.test('the counter adds up across the requests of the run', async () => {
++ const jti = 'tok_budget_1';
++ const tok = runToken(jti, { on_behalf_of: PERSON });
++ const first = await chat(tok);
++ assert.strictEqual(first.status, 200, first.text);
++ const after = await budgetOf(jti);
++ assert.strictEqual(after.status, 200, after.text);
++ assert.ok(after.body.spent_tokens >= 7, `the tokens of the first request are counted (${after.body.spent_tokens})`);
++ assert.ok(after.body.spent_usd > 0, `the cost of the first request is counted (${after.body.spent_usd})`);
++ const second = await chat(tok);
++ assert.strictEqual(second.status, 200, second.text);
++ const after2 = await budgetOf(jti);
++ assert.ok(after2.body.spent_tokens > after.body.spent_tokens, 'the second request is counted too');
++ assert.ok(after2.body.spent_usd > after.body.spent_usd, 'the cost adds up');
++});
++
++t.test('the request after the cap is refused with 402 and the problem code', async () => {
++ const jti = 'tok_budget_2';
++ const tok = runToken(jti, { on_behalf_of: PERSON });
++ await setBudget(jti, { budget_tokens: 7, expires_at: Math.floor(Date.now() / 1000) + 600 });
++ const first = await chat(tok);
++ assert.strictEqual(first.status, 200, first.text); // the request that crosses the line still completes
++ const second = await chat(tok);
++ assert.strictEqual(second.status, 402, second.text);
++ assert.strictEqual(second.body.error.code, 'ai.run_budget_reached');
++ assert.ok(/token/i.test(second.body.error.message), `the message says which budget: ${second.body.error.message}`);
++ assert.ok(/7 tokens/.test(second.body.error.message), `the message says how much was spent: ${second.body.error.message}`);
++});
++
++t.test('the dollars of a run are capped too', async () => {
++ const jti = 'tok_budget_3';
++ const tok = runToken(jti, { on_behalf_of: PERSON });
++ await setBudget(jti, { budget_usd: 0.001, expires_at: Math.floor(Date.now() / 1000) + 600 });
++ const first = await chat(tok);
++ assert.strictEqual(first.status, 200, first.text);
++ const second = await chat(tok);
++ assert.strictEqual(second.status, 402, second.text);
++ assert.strictEqual(second.body.error.code, 'ai.run_budget_reached');
++ assert.ok(/dollar/i.test(second.body.error.message), second.body.error.message);
++});
++
++t.test('a token without a budget is not capped', async () => {
++ const tok = runToken('tok_budget_none', { on_behalf_of: PERSON });
++ for (let i = 0; i < 3; i++) {
++ const r = await chat(tok);
++ assert.strictEqual(r.status, 200, r.text);
++ }
++ const seen = await budgetOf('tok_budget_none');
++ assert.strictEqual(seen.status, 200);
++ assert.strictEqual(seen.body.budget_usd, null);
++ assert.strictEqual(seen.body.spent_usd, 0);
++});
++
++t.test('a budget in the request is ignored: there is nothing to forge', async () => {
++ const jti = 'tok_budget_4';
++ const tok = runToken(jti, { on_behalf_of: PERSON });
++ await setBudget(jti, { budget_tokens: 7, expires_at: Math.floor(Date.now() / 1000) + 600 });
++ const first = await chat(tok, { 'x-openvibe-budget': 'usd=100000,tokens=1000000' });
++ assert.strictEqual(first.status, 200, first.text);
++ const second = await chat(tok, { 'x-openvibe-budget': 'usd=100000,tokens=1000000' });
++ assert.strictEqual(second.status, 402, second.text);
++});
++
++t.test('an edited token is refused: its signature does not hold', async () => {
++ const jti = 'tok_budget_5';
++ const tok = runToken(jti, { on_behalf_of: PERSON });
++ await setBudget(jti, { budget_tokens: 7, expires_at: Math.floor(Date.now() / 1000) + 600 });
++ const parts = tok.split('.');
++ const claims = JSON.parse(Buffer.from(parts[1], 'base64url').toString('utf8'));
++ claims.budget_usd = 100000; // the holder raises the budget in the token
++ claims.cap = [...claims.cap, 'ai.workflow.manage'];
++ const edited = `${parts[0]}.${Buffer.from(JSON.stringify(claims)).toString('base64url')}.${parts[2]}`;
++ const r = await chat(edited);
++ assert.strictEqual(r.status, 401, r.text);
++});
++
++t.test('another run\'s budget does not cover this token', async () => {
++ const tok = runToken('tok_budget_other', { on_behalf_of: PERSON });
++ const r = await chat(tok);
++ assert.strictEqual(r.status, 200, r.text);
++});
++
++t.test('a budget that expired with its token is no budget', async () => {
++ const jti = 'tok_budget_6';
++ const tok = runToken(jti, { on_behalf_of: PERSON });
++ const r = await setBudget(jti, { budget_tokens: 1, expires_at: Math.floor(Date.now() / 1000) - 10 });
++ assert.strictEqual(r.status, 200, r.text);
++ const run = await chat(tok);
++ assert.strictEqual(run.status, 200, run.text); // the budget is gone with the token
++ const again = await chat(tok);
++ assert.strictEqual(again.status, 200, again.text); // and it stays gone
++ const seen = await budgetOf(jti);
++ assert.ok(seen.body.expires_at <= Math.floor(Date.now() / 1000), 'the budget expired with its token');
++});
++
++t.test('stop', async () => { await h.stop(); });
++
++t.run();