feat: attribute API usage and job tokens to repositories

src/app.ts

@@ -1765,6 +1765,7 @@ export async function buildApp(options: BuildOptions): Promise<BuiltApp> {
       name,
       tokenHash: sha256Hex(raw),
       prefix: raw.slice(0, 8),
+      repoId: null,
       expiresAt: null,
     })
     store.audit(viewer.user.id, 'token.created', name)

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<JobRow>("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)
     })
   }

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,
       })
 

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)

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

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()
+  }
+})