From 188cbf59dabd8dc3965be98568b3f9e3ff480009 Mon Sep 17 00:00:00 2001 From: HugoXK Date: Wed, 4 Nov 2026 16:37:55 +0800 Subject: [PATCH] feat: attribute API usage and job tokens to repositories Store repo_id on api_tokens and quota_ledger so usage can be attributed per repository. This is internal plumbing and does not change user-visible behavior. Docs-not-needed: internal attribution columns; budgets are documented separately. --- src/app.ts | 1 + src/db.ts | 22 +++++++++++++++------- src/flybot/instances.ts | 1 + src/gateway/gateway.ts | 2 +- src/jobs/service.ts | 1 + test/db.test.ts | 17 ++++++++++++++++- 6 files changed, 35 insertions(+), 9 deletions(-) diff --git a/src/app.ts b/src/app.ts index 64b5947..0d9674f 100644 --- a/src/app.ts +++ b/src/app.ts @@ -1765,6 +1765,7 @@ export async function buildApp(options: BuildOptions): Promise { name, tokenHash: sha256Hex(raw), prefix: raw.slice(0, 8), + repoId: null, expiresAt: null, }) store.audit(viewer.user.id, 'token.created', name) diff --git a/src/db.ts b/src/db.ts index 65d6021..26002f1 100644 --- a/src/db.ts +++ b/src/db.ts @@ -193,6 +193,7 @@ export interface ApiTokenRow { name: string token_hash: string prefix: string + repo_id: string | null expires_at: string | null revoked_at: string | null last_used_at: string | null @@ -426,6 +427,7 @@ CREATE TABLE IF NOT EXISTS api_tokens ( name TEXT NOT NULL, token_hash TEXT NOT NULL UNIQUE, prefix TEXT NOT NULL, + repo_id TEXT, expires_at TEXT, revoked_at TEXT, last_used_at TEXT, @@ -438,6 +440,7 @@ CREATE TABLE IF NOT EXISTS quota_ledger ( at TEXT NOT NULL, user_id TEXT NOT NULL REFERENCES users(id) ON DELETE CASCADE, job_id TEXT, + repo_id TEXT, cost_usd REAL NOT NULL, meta TEXT ); @@ -535,9 +538,9 @@ export class Store { this.db.exec('PRAGMA foreign_keys = ON;') this.db.exec('PRAGMA busy_timeout = 5000;') const previousVersion = Number(this.get<{ user_version: number }>('PRAGMA user_version')?.user_version ?? 0) - if (previousVersion > 18) { + if (previousVersion > 19) { this.db.close() - throw new Error(`database schema version ${previousVersion} is newer than this build supports (18)`) + throw new Error(`database schema version ${previousVersion} is newer than this build supports (19)`) } this.db.exec(SCHEMA) this.migrateColumns() @@ -549,7 +552,7 @@ export class Store { if (previousVersion < 16) { this.run('UPDATE repos SET role_checked_at = ? WHERE repo_role IS NOT NULL AND role_checked_at IS NULL', nowIso()) } - this.db.exec(`PRAGMA user_version = ${Math.max(previousVersion, 18)};`) + this.db.exec(`PRAGMA user_version = ${Math.max(previousVersion, 19)};`) } private migrateColumns(): void { @@ -636,6 +639,10 @@ export class Store { if (!jobColumns.has('plan_allow_paths')) this.run('ALTER TABLE jobs ADD COLUMN plan_allow_paths TEXT') if (!jobColumns.has('plan_deny_paths')) this.run('ALTER TABLE jobs ADD COLUMN plan_deny_paths TEXT') if (!jobColumns.has('base_sha')) this.run('ALTER TABLE jobs ADD COLUMN base_sha TEXT') + const tokenColumns = new Set(this.all<{ name: string }>('PRAGMA table_info(api_tokens)').map((row) => row.name)) + if (!tokenColumns.has('repo_id')) this.run('ALTER TABLE api_tokens ADD COLUMN repo_id TEXT') + const ledgerColumns = new Set(this.all<{ name: string }>('PRAGMA table_info(quota_ledger)').map((row) => row.name)) + if (!ledgerColumns.has('repo_id')) this.run('ALTER TABLE quota_ledger ADD COLUMN repo_id TEXT') if (!jobColumns.has('base_branch')) this.run('ALTER TABLE jobs ADD COLUMN base_branch TEXT') this.run("UPDATE jobs SET base_branch = (SELECT upstream_default_branch FROM repos WHERE repos.id = jobs.repo_id) WHERE base_branch IS NULL") } @@ -1335,15 +1342,16 @@ export class Store { return this.all("SELECT * FROM jobs WHERE status = 'queued' AND created_at < ?", beforeIso) } - createApiToken(input: { id: string; userId: string; kind: string; name: string; tokenHash: string; prefix: string; expiresAt: string | null }): ApiTokenRow { + createApiToken(input: { id: string; userId: string; kind: string; name: string; tokenHash: string; prefix: string; repoId?: string | null; expiresAt: string | null }): ApiTokenRow { this.run( - 'INSERT INTO api_tokens (id, user_id, kind, name, token_hash, prefix, expires_at, created_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?)', + 'INSERT INTO api_tokens (id, user_id, kind, name, token_hash, prefix, repo_id, expires_at, created_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)', input.id, input.userId, input.kind, input.name, input.tokenHash, input.prefix, + input.repoId ?? null, input.expiresAt, nowIso(), ) @@ -1406,9 +1414,9 @@ export class Store { this.run("UPDATE api_tokens SET revoked_at = ? WHERE kind = 'job' AND (expires_at IS NOT NULL AND expires_at <= ?) AND revoked_at IS NULL", nowIso(), nowIso()) } - recordUsage(userId: string, jobId: string | null, costUsd: number, meta: string | null = null): void { + recordUsage(userId: string, jobId: string | null, costUsd: number, meta: string | null = null, repoId: string | null = null): void { this.transaction(() => { - this.run('INSERT INTO quota_ledger (at, user_id, job_id, cost_usd, meta) VALUES (?, ?, ?, ?, ?)', nowIso(), userId, jobId, costUsd, meta) + this.run('INSERT INTO quota_ledger (at, user_id, job_id, repo_id, cost_usd, meta) VALUES (?, ?, ?, ?, ?, ?)', nowIso(), userId, jobId, repoId, costUsd, meta) this.run('UPDATE users SET quota_used_usd = quota_used_usd + ? WHERE id = ?', costUsd, userId) }) } diff --git a/src/flybot/instances.ts b/src/flybot/instances.ts index c95bcf6..8284665 100644 --- a/src/flybot/instances.ts +++ b/src/flybot/instances.ts @@ -241,6 +241,7 @@ server.listen(3081, '0.0.0.0') name: `flybot:${repo.id}`, tokenHash: sha256Hex(gatewayToken), prefix: gatewayToken.slice(0, 8), + repoId: repo.id, expiresAt: null, }) diff --git a/src/gateway/gateway.ts b/src/gateway/gateway.ts index bee52a1..f850ce3 100644 --- a/src/gateway/gateway.ts +++ b/src/gateway/gateway.ts @@ -225,7 +225,7 @@ export function registerGateway(app: FastifyInstance, deps: GatewayDeps): void { function finishCost(costUsd: number, usageValue: unknown): void { if (!funded || costUsd <= 0) return - store.recordUsage(auth!.user.id, jobId, costUsd, JSON.stringify({ model: targetModel, usage: usageValue, ms: Date.now() - started })) + store.recordUsage(auth!.user.id, jobId, costUsd, JSON.stringify({ model: targetModel, usage: usageValue, ms: Date.now() - started }), auth!.token.repo_id ?? null) if (jobId !== null) store.addJobCost(jobId, costUsd) log.info('gateway usage', { user: auth!.user.login, costUsd, jobId, model: targetModel }) const updated = store.userById(auth!.user.id) diff --git a/src/jobs/service.ts b/src/jobs/service.ts index 03b272d..84c2e8c 100644 --- a/src/jobs/service.ts +++ b/src/jobs/service.ts @@ -300,6 +300,7 @@ export class Jobs { name: `job:${jobId}`, tokenHash: sha256Hex(platformToken), prefix: platformToken.slice(0, 8), + repoId: repo.id, expiresAt: addMinutes(this.config.runner.timeoutMinutes + 60), }) const priorJobs = this.store.listJobs(user.id, 1).length diff --git a/test/db.test.ts b/test/db.test.ts index 323d456..12fd77d 100644 --- a/test/db.test.ts +++ b/test/db.test.ts @@ -150,7 +150,7 @@ test('the role backfill runs once and future schema versions are refused', () => const reopened = new Store(path) assert.equal(reopened.getRepo(repo.id)?.repo_role, 'maintainer') assert.notEqual(reopened.getRepo(repo.id)?.role_checked_at, null) - reopened.db.exec('PRAGMA user_version = 19;') + reopened.db.exec('PRAGMA user_version = 20;') reopened.close() assert.throws(() => new Store(path), /newer than this build/) @@ -201,3 +201,18 @@ test('work items persist, update and are found by issue or PR', () => { store.close() } }) + +test('usage and job tokens carry the repository for cost attribution', () => { + const store = new Store(':memory:') + try { + const user = makeUser(store, { login: 'attribution-user', status: 'approved' }) + const repo = store.createRepo({ id: 'repo_attribution', userId: user.id, owner: 'o', name: 'a', defaultBranch: 'main', tokenEnc: 'x', permission: 'admin', repoRole: 'maintainer', goal: '', autoIterate: false, intervalMinutes: 1440 }) + store.recordUsage(user.id, null, 0.2, null, repo.id) + const usage = store.get<{ repo_id: string | null }>('SELECT repo_id FROM quota_ledger ORDER BY id DESC LIMIT 1') + assert.equal(usage?.repo_id, repo.id) + const token = store.createApiToken({ id: 'tok_attr', userId: user.id, kind: 'job', name: 'job:attr', tokenHash: 'hash_attr', prefix: 'flg', repoId: repo.id, expiresAt: null }) + assert.equal(token.repo_id, repo.id) + } finally { + store.close() + } +})