Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 9 additions & 0 deletions apps/api/src/db/schema/ai.ts
Original file line number Diff line number Diff line change
Expand Up @@ -34,8 +34,17 @@ export const aiSessions = pgTable('ai_sessions', {
model: varchar('model', { length: 100 }).notNull().default('claude-sonnet-4-5-20250929'),
systemPrompt: text('system_prompt'),
contextSnapshot: jsonb('context_snapshot'),
// TOTAL input across the session — uncached + cache-read + cache-creation.
// The SDK splits those three apart because they are priced differently, not
// because only the first is "input"; see sumInputTokens in aiCostTracker.
// Rows written before that fix carry the uncached slice only and read
// implausibly low (an 8-turn session at 17 tokens); they are NOT backfilled,
// since the split components were never persisted anywhere to recover from.
totalInputTokens: integer('total_input_tokens').notNull().default(0),
totalOutputTokens: integer('total_output_tokens').notNull().default(0),
// Independent of the token columns above: cost comes from the SDK's own
// total_cost_usd, or from per-component pricing when that is 0. It was
// correct throughout — do not "reconcile" it against the token columns.
totalCostCents: real('total_cost_cents').notNull().default(0),
turnCount: integer('turn_count').notNull().default(0),
maxTurns: integer('max_turns').notNull().default(50),
Expand Down
114 changes: 114 additions & 0 deletions apps/api/src/routes/backup/dashboard.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -491,6 +491,120 @@ describe('backup dashboard routes', () => {
expect(body.data.lastJob.errorLog).toContain('read from the live volume');
});
});

describe('GET /status/:deviceId — which run is "last"', () => {
// created_at is the row-INSERT clock. The profile fan-out writes an entire
// occurrence's jobs in one transaction, so they share it to the microsecond
// and `ORDER BY created_at DESC` is not a total order. Each case below feeds
// the route rows in the order a created_at-only sort could legitimately
// return them, and asserts the route still picks the right run.
function mockJobRows(rows: Array<Record<string, unknown>>) {
resolveBackupConfigForDeviceMock.mockResolvedValueOnce(null);
selectMock
.mockReturnValueOnce(chainMock([{ id: DEVICE_ID, siteId: SITE_A }]))
.mockReturnValueOnce(chainMock(rows));
}

function run(overrides: Record<string, unknown> = {}) {
return {
id: 'job-run',
status: 'completed',
startedAt: new Date('2026-08-05T06:35:00.000Z'),
createdAt: new Date('2026-08-05T06:32:11.829738Z'),
completedAt: new Date('2026-08-05T06:50:00.000Z'),
errorLog: null,
vssMetadata: null,
...overrides,
};
}

it('breaks an identical created_at tie in favour of the later started_at', async () => {
const sameCreatedAt = new Date('2026-08-05T06:32:11.829738Z');
mockJobRows([
run({
id: 'job-older',
createdAt: sameCreatedAt,
startedAt: new Date('2026-08-05T06:32:20.000Z'),
completedAt: new Date('2026-08-05T06:33:00.000Z'),
}),
run({
id: 'job-newer',
createdAt: sameCreatedAt,
startedAt: new Date('2026-08-05T06:40:00.000Z'),
completedAt: new Date('2026-08-05T06:55:00.000Z'),
vssMetadata: { shadowCopyId: 'set-9', writers: [] },
}),
]);

const res = await app.request(`/backup/status/${DEVICE_ID}`);

expect(res.status).toBe(200);
const body = await res.json();
expect(body.data.lastJob.id).toBe('job-newer');
// The whole point of the defect: the VSS panel renders only off the
// chosen lastJob's metadata, so picking the loser blanks it.
expect(body.data.lastJob.vssMetadata).toEqual({ shadowCopyId: 'set-9', writers: [] });
});

it('does not let a never-started pending job displace a real completed run', async () => {
mockJobRows([
run({
id: 'job-pending',
status: 'pending',
startedAt: null,
createdAt: new Date('2026-08-05T07:00:00.000Z'),
completedAt: null,
}),
run({
id: 'job-completed',
startedAt: new Date('2026-08-05T06:35:00.000Z'),
createdAt: new Date('2026-08-05T06:32:00.000Z'),
completedAt: new Date('2026-08-05T06:50:00.000Z'),
vssMetadata: { shadowCopyId: 'set-1', writers: [] },
}),
]);

const res = await app.request(`/backup/status/${DEVICE_ID}`);

expect(res.status).toBe(200);
const body = await res.json();
expect(body.data.lastJob.id).toBe('job-completed');
expect(body.data.lastJob.vssMetadata).toEqual({ shadowCopyId: 'set-1', writers: [] });
expect(body.data.lastSuccessAt).toBe('2026-08-05T06:50:00.000Z');
});

it('still reports a pending job when the device has never run one', async () => {
// NULLS LAST only demotes — it must never turn a queued-only device into
// "no backup jobs at all".
mockJobRows([
run({ id: 'p1', status: 'pending', startedAt: null, createdAt: new Date('2026-08-05T06:00:00.000Z'), completedAt: null }),
run({ id: 'p2', status: 'pending', startedAt: null, createdAt: new Date('2026-08-05T07:00:00.000Z'), completedAt: null }),
]);

const res = await app.request(`/backup/status/${DEVICE_ID}`);

const body = await res.json();
expect(body.data.lastJob.id).toBe('p2');
expect(body.data.lastSuccessAt).toBeNull();
});

it('asks the database for the same order it applies in memory', async () => {
mockJobRows([run()]);
const jobsChain = selectMock.mock.results[1]?.value ?? null;

await app.request(`/backup/status/${DEVICE_ID}`);

const chain = jobsChain ?? (selectMock.mock.results[1]!.value as Record<string, any>);
const orderTerms = chain.orderBy.mock.calls[0];
expect(orderTerms).toHaveLength(3);
// started_at first, and explicitly NULLS LAST — a bare DESC would be
// NULLS FIRST in Postgres and float every queued job to the top.
expect(JSON.stringify(orderTerms[0])).toContain('nulls last');
expect(JSON.stringify(orderTerms[0])).toContain('backup_jobs.started_at');
expect(orderTerms[1]).toEqual({ op: 'desc', value: 'backup_jobs.created_at' });
expect(orderTerms[2]).toEqual({ op: 'desc', value: 'backup_jobs.id' });
});
});
});

function makeRecentJob(overrides: Record<string, unknown> = {}) {
Expand Down
38 changes: 32 additions & 6 deletions apps/api/src/routes/backup/dashboard.ts
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
import { Hono } from 'hono';
import { zValidator } from '../../lib/validation';
import { eq, and, sql, gte, lte, desc, inArray } from 'drizzle-orm';
import { eq, and, sql, gte, lte, inArray } from 'drizzle-orm';
import { db } from '../../db';
import { requirePermission } from '../../middleware/auth';
import {
Expand All @@ -12,6 +12,12 @@ import {
} from '../../db/schema';
import { PERMISSIONS, canAccessSite, type UserPermissions } from '../../services/permissions';
import { resolveBackupConfigForDevice, resolveAllBackupAssignedDevices } from '../../services/featureConfigResolver';
import {
backupJobHistoryOrderBy,
compareBackupRunRecency,
latestBackupRunOrderBy,
latestBackupRunWindowOrder,
} from '../../services/backupJobOrdering';
import { getNextRun, resolveScopedOrgId } from './helpers';
import { usageHistoryQuerySchema } from './schemas';

Expand Down Expand Up @@ -71,7 +77,12 @@ async function resolveAttentionItems(
errorLog: backupJobs.errorLog,
completedAt: backupJobs.completedAt,
createdAt: backupJobs.createdAt,
rn: sql<number>`row_number() over (partition by ${backupJobs.deviceId} order by ${backupJobs.createdAt} desc)`.as('rn'),
// rn=1 must be the device's most recent RUN, not the most recently
// inserted row — see backupJobOrdering. Ordering this window by
// created_at alone let a queued job (or the loser of a same-transaction
// fan-out tie) take rn=1 and mask the failed run underneath it, so the
// device never showed up in attention items at all.
rn: sql<number>`row_number() over (partition by ${backupJobs.deviceId} order by ${latestBackupRunWindowOrder})`.as('rn'),
})
.from(backupJobs)
.where(and(eq(backupJobs.orgId, orgId), jobDeviceScope))
Expand Down Expand Up @@ -340,7 +351,10 @@ dashboardRoutes.get('/dashboard', requirePermission(PERMISSIONS.ORGS_READ.resour
.leftJoin(devices, eq(backupJobs.deviceId, devices.id))
.leftJoin(backupConfigs, eq(backupJobs.configId, backupConfigs.id))
.where(and(eq(backupJobs.orgId, orgId), jobDeviceScope))
.orderBy(desc(backupJobs.createdAt))
// Activity feed, so created_at stays primary (a job queued seconds ago
// belongs at the top) — but the order has to be TOTAL, or a
// same-transaction fan-out reshuffles the top 5 between refreshes.
.orderBy(...backupJobHistoryOrderBy)
.limit(5),
resolveAttentionItems(orgId, jobDeviceScope, allowedDeviceIds, noSiteAllowedDevices),
]);
Expand Down Expand Up @@ -434,14 +448,26 @@ dashboardRoutes.get('/status/:deviceId', requirePermission(PERMISSIONS.ORGS_READ
// Resolve backup config via configuration policy system
const resolved = await resolveBackupConfigForDevice(deviceId);

// Get recent jobs for this device
const jobs = await db
// Get recent jobs for this device, most-recent RUN first.
//
// `created_at` is only the insert timestamp: the profile fan-out writes a
// whole occurrence's jobs in one transaction, so they share it exactly and
// `ORDER BY created_at DESC` is not a total order. QA hit that tie and the
// planner handed back the OLDER run, which silently hid the VSS Status panel
// (it renders only when the chosen lastJob carries vss_metadata).
const jobRows = await db
.select()
.from(backupJobs)
.where(
and(eq(backupJobs.orgId, orgId), eq(backupJobs.deviceId, deviceId))
)
.orderBy(desc(backupJobs.createdAt));
.orderBy(...latestBackupRunOrderBy);

// This one array answers three questions below (last job / last success /
// last failure), so recency is applied once here through the comparator that
// mirrors latestBackupRunOrderBy — the three answers can't drift apart, and
// the rule is unit-testable without a live planner.
const jobs = [...jobRows].sort(compareBackupRunRecency);

const lastJob = jobs[0] ?? null;
// A partial run counts here for the same reason it counts for RPO: it left a
Expand Down
10 changes: 8 additions & 2 deletions apps/api/src/routes/backup/jobs.ts
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
import { Hono } from 'hono';
import { zValidator } from '../../lib/validation';
import { eq, and, desc, gte, lte, sql, inArray } from 'drizzle-orm';
import { eq, and, gte, lte, sql, inArray } from 'drizzle-orm';
import { db } from '../../db';
import { backupJobs, backupConfigs, devices } from '../../db/schema';
import { requireMfa, requirePermission, requireScope } from '../../middleware/auth';
Expand All @@ -12,6 +12,7 @@ import { resolveBackupConfigForDevice, resolveAllBackupAssignedDevices } from '.
import { queueBackupStopCommand } from '../../services/commandQueue';
import { canAccessSite, PERMISSIONS, type UserPermissions } from '../../services/permissions';
import { resolveScopedOrgId } from './helpers';
import { backupJobHistoryOrderBy } from '../../services/backupJobOrdering';
import { jobListSchema } from './schemas';

// Legacy `?status=` spellings this API has always accepted, mapped to the real
Expand Down Expand Up @@ -127,7 +128,12 @@ jobsRoutes.get('/jobs', requirePermission(PERMISSIONS.ORGS_READ.resource, PERMIS
.leftJoin(devices, eq(backupJobs.deviceId, devices.id))
.leftJoin(backupConfigs, eq(backupJobs.configId, backupConfigs.id))
.where(and(...conditions))
.orderBy(desc(backupJobs.createdAt));
// History list: created_at stays primary because the ?from/?to/?date
// filters above are all on created_at, and a just-queued job must not sink
// to the bottom of the operator's list. started_at + id only make the order
// total, so a same-transaction profile fan-out stops reshuffling between
// requests. See services/backupJobOrdering.
.orderBy(...backupJobHistoryOrderBy);

return c.json({
data: rows.map((r) => ({
Expand Down
139 changes: 138 additions & 1 deletion apps/api/src/services/aiCostTracker.test.ts
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
import { describe, it, expect, vi, beforeEach } from 'vitest';
import { calculateCostCents, checkAiRateLimit, checkBudget, recordUsage, recordUsageFromSdkResult } from './aiCostTracker';
import { calculateCostCents, checkAiRateLimit, checkBudget, recordUsage, recordUsageFromSdkResult, sumInputTokens } from './aiCostTracker';
import { db, withSystemDbAccessContext } from '../db';
import { getEffectiveAiBudget } from './effectiveSettings';
import { rateLimiter } from './rate-limit';
Expand Down Expand Up @@ -303,6 +303,143 @@ describe('recordUsageFromSdkResult', () => {
});
});

// ============================================
// Input-token accounting — cache reads/creations ARE input
// ============================================

/** The number added to ai_sessions.total_input_tokens / total_output_tokens. */
function recordedTokens(captured: Record<string, unknown> | undefined, key: 'totalInputTokens' | 'totalOutputTokens'): number {
const expr = captured?.[key] as { values?: unknown[] } | undefined;
return Number(expr?.values?.[1]);
}

describe('recordUsageFromSdkResult — input token accounting', () => {
// Release QA: an 8-turn session showed total_input_tokens = 17 against
// total_output_tokens = 1029 and $0.57 of spend. Prompt caching routes almost
// the entire prompt through cache_read_input_tokens on every turn after the
// first, and only the uncached remainder was being accumulated.
it('counts cache-read and cache-creation tokens as input', async () => {
const captured = setupDbMocks(null);

await recordUsageFromSdkResult('sess-tokens', 'org-1', {
total_cost_usd: 0.5,
usage: {
input_tokens: 17,
output_tokens: 1_029,
cache_read_input_tokens: 120_000,
cache_creation_input_tokens: 4_500,
},
num_turns: 1,
model: 'claude-sonnet-4-6',
});

expect(recordedTokens(captured.sessionSet, 'totalInputTokens')).toBe(17 + 120_000 + 4_500);
expect(recordedTokens(captured.sessionSet, 'totalOutputTokens')).toBe(1_029);
});

it('records the same total on the daily/monthly org aggregates', async () => {
const captured = setupDbMocks(null);

await recordUsageFromSdkResult('sess-agg', 'org-1', {
total_cost_usd: 0.5,
usage: {
input_tokens: 17,
output_tokens: 1_029,
cache_read_input_tokens: 120_000,
cache_creation_input_tokens: 4_500,
},
num_turns: 1,
model: 'claude-sonnet-4-6',
});

// One insert per period (daily + monthly); both carry the summed input.
expect(captured.aggregateValues).toHaveLength(2);
for (const values of captured.aggregateValues) {
expect(values.inputTokens).toBe(17 + 120_000 + 4_500);
expect(values.outputTokens).toBe(1_029);
}
});

it('records a fully-cached turn as real input, not zero', async () => {
const captured = setupDbMocks(null);

await recordUsageFromSdkResult('sess-cached-only', 'org-1', {
total_cost_usd: 0,
usage: {
input_tokens: 0,
output_tokens: 0,
cache_read_input_tokens: 1_000_000,
cache_creation_input_tokens: 0,
},
num_turns: 1,
model: 'claude-sonnet-4-6',
});

expect(recordedTokens(captured.sessionSet, 'totalInputTokens')).toBe(1_000_000);
// ...and the cost path is untouched by the summing: still priced off the
// SPLIT components (1M cache-read at 0.1x of 300 c/MTok = 30), never off the
// sum. Summing into the pricing call would have billed 300 here.
expect(recordedCostCents(captured.sessionSet)).toBe(30);
});

it('leaves cost alone when cache fields are present alongside real spend', async () => {
// Guards against the obvious mis-fix: feeding the summed input back into
// calculateCostCents would double-count cache tokens at the full input rate.
const captured = setupDbMocks(null);

await recordUsageFromSdkResult('sess-cost-guard', 'org-1', {
total_cost_usd: 0,
usage: {
input_tokens: 1_000_000,
output_tokens: 1_000_000,
cache_read_input_tokens: 1_000_000,
cache_creation_input_tokens: 1_000_000,
},
num_turns: 1,
model: 'claude-sonnet-4-6',
});

expect(recordedTokens(captured.sessionSet, 'totalInputTokens')).toBe(3_000_000);
expect(recordedCostCents(captured.sessionSet)).toBe(2205); // unchanged from the pricing test above
});

it('is unaffected when the payload carries no cache fields at all', async () => {
// Older/partial usage payloads and the vLLM path have no prompt caching.
const captured = setupDbMocks(null);

await recordUsageFromSdkResult('sess-nocache', 'org-1', {
total_cost_usd: 0.1,
usage: { input_tokens: 5_000, output_tokens: 2_000 },
num_turns: 1,
model: 'claude-sonnet-4-6',
});

expect(recordedTokens(captured.sessionSet, 'totalInputTokens')).toBe(5_000);
});
});

describe('sumInputTokens', () => {
it('sums the three disjoint slices the SDK splits input across', () => {
expect(sumInputTokens({
input_tokens: 17,
cache_read_input_tokens: 120_000,
cache_creation_input_tokens: 4_500,
})).toBe(124_517);
});

it('treats missing and null components as 0', () => {
expect(sumInputTokens({})).toBe(0);
expect(sumInputTokens({ input_tokens: 10, cache_read_input_tokens: null })).toBe(10);
});

it('never throws on a nullish usage object', () => {
// It sits ahead of the `done` publish that returns the session to 'idle';
// a throw there strands the turn and hangs the client.
expect(sumInputTokens(null)).toBe(0);
expect(sumInputTokens(undefined)).toBe(0);
});
});

// ============================================
// recordUsage — sessionless org-budget path (issue #1949)
// ============================================
Expand Down
Loading
Loading