diff --git a/worker/src/botListCache.ts b/worker/src/botListCache.ts new file mode 100644 index 0000000..cd47624 --- /dev/null +++ b/worker/src/botListCache.ts @@ -0,0 +1,46 @@ +// Per-user bot-list cache (single-isolate memory). Entries live until a bot +// write invalidates them or the size cap evicts the oldest. If traffic ever +// spreads across isolates, re-add a TTL: a write on one isolate can't clear +// another's entries. +import type { BotRow } from "./types"; + +export interface BotListEntry { + total: number; + rows: BotRow[]; + models: Map; +} + +const MAX_ENTRIES = 1000; +const entries = new Map(); + +// Epoch guards the read/write race: a GET that started reading before a bot +// write must not store its pre-write rows after that write invalidated them. +let epoch = 0; + +export function botListEpoch(): number { + return epoch; +} + +export function getBotList(key: string): BotListEntry | null { + return entries.get(key) ?? null; +} + +export function setBotList(key: string, entry: BotListEntry, seenEpoch: number): void { + if (seenEpoch !== epoch) return; // a write landed mid-read; drop stale rows + entries.set(key, entry); + if (entries.size > MAX_ENTRIES) entries.delete(entries.keys().next().value!); +} + +export function invalidateBotListCache(userId: number): void { + epoch++; + const prefix = `${userId}:`; + for (const key of entries.keys()) { + if (key.startsWith(prefix)) entries.delete(key); + } +} + +// Full reset for tests, which seed a fresh DB per case (module state +// outlives the DB in the same worker). +export function clearBotListCache(): void { + entries.clear(); +} diff --git a/worker/src/oauth.ts b/worker/src/oauth.ts index ef8f82d..236fa1b 100644 --- a/worker/src/oauth.ts +++ b/worker/src/oauth.ts @@ -13,6 +13,7 @@ import { first, run } from "./db"; import { newUuid } from "./crypto"; import { nowIso } from "./types"; import type { Env } from "./types"; +import { invalidateBotListCache } from "./botListCache"; export const PENELOPE_GREETING = "Hello! I'm Penelope, your writing assistant. How can I help you with writing today?"; @@ -67,6 +68,7 @@ export async function provisionDefaultContent(db: D1Database, userId: number, fi now, now, ); + invalidateBotListCache(userId); bot = await first<{ id: number }>( db, "SELECT id FROM bots_bot WHERE user_id = ? AND deleted_at IS NULL ORDER BY id ASC", diff --git a/worker/src/onboarding.ts b/worker/src/onboarding.ts index 6e23d5b..99bd5f7 100644 --- a/worker/src/onboarding.ts +++ b/worker/src/onboarding.ts @@ -4,6 +4,7 @@ import { first } from "./db"; import { newUuid } from "./crypto"; import { nowIso } from "./types"; import { hashPin, validatePin } from "./pin"; +import { invalidateBotListCache } from "./botListCache"; export const DEFAULT_TEMPLATE_NAME = "Blank"; export const DEFAULT_BOT_COLOR = "#2A9D8F"; @@ -154,6 +155,7 @@ export async function bootstrapOnboarding( const str = (v: unknown): string | null => (typeof v === "string" ? v : null); let botUuid: string | null = null; // set when THIS request inserts + let botTouched = false; // any bot insert/update queued below if (!bot) { const defModel = await first<{ id: number }>(db, "SELECT id FROM bots_aimodel WHERE is_default = 1"); const newBotId = newUuid(); @@ -188,6 +190,7 @@ export async function bootstrapOnboarding( } botUuid = newBotId; bot = { id: 0, bot_id: newBotId }; + botTouched = true; } else { const sets: string[] = ["modified_at = ?"]; const params: unknown[] = [now]; @@ -208,6 +211,7 @@ export async function bootstrapOnboarding( if (sets.length > 1) { params.push(bot.id); stmts.push(db.prepare(`UPDATE bots_bot SET ${sets.join(", ")} WHERE id = ?`).bind(...params)); + botTouched = true; } } @@ -237,6 +241,7 @@ export async function bootstrapOnboarding( try { if (stmts.length > 0) await db.batch(stmts); + if (botTouched) invalidateBotListCache(userId); } catch (e) { const msg = String(e); if (msg.includes("UNIQUE") && (msg.includes("oauth_email") || msg.includes("bots_profile"))) { diff --git a/worker/src/routes/bots.ts b/worker/src/routes/bots.ts index 9d33232..8c512c7 100644 --- a/worker/src/routes/bots.ts +++ b/worker/src/routes/bots.ts @@ -8,6 +8,7 @@ import { newUuid } from "../crypto"; import { nowIso, type BotRow, type Env } from "../types"; import { botRow } from "../serializers"; import { findByUuidOrId, forbidden, routeCtx, teenScope, type TC, makeRouter } from "./helpers"; +import { botListEpoch, getBotList, invalidateBotListCache, setBotList } from "../botListCache"; const bots = makeRouter(); bots.use("*", authenticate); @@ -21,22 +22,52 @@ async function serializeBot(r: ReturnType, bot: BotRow) { return botRow(bot as unknown as Record, aiModel, `${r.base}/bots/${bot.id}/`); } +function serializeBotWithModels(r: ReturnType, bot: BotRow, models: Map) { + const aiModel = bot.ai_model_id != null ? (models.get(bot.ai_model_id) ?? null) : null; + return botRow(bot as unknown as Record, aiModel, `${r.base}/bots/${bot.id}/`); +} + // GET /api/bots/ bots.get("/", async (c) => { const r = routeCtx(c); const sp = new URLSearchParams(c.req.query() as Record); - const total = await dbCount(r.env.DB, "SELECT COUNT(*) AS n FROM bots_bot WHERE user_id = ? AND deleted_at IS NULL", r.user.id); + const cacheKey = `${r.user.id}:${sp.get("page") ?? "1"}`; + const seenEpoch = botListEpoch(); + const cached = getBotList(cacheKey); + const total = cached + ? cached.total + : await dbCount(r.env.DB, "SELECT COUNT(*) AS n FROM bots_bot WHERE user_id = ? AND deleted_at IS NULL", r.user.id); const p = resolvePage(sp, total); if (!p) return c.json({ detail: "Invalid page." }, 404); - const rows = await all( - r.env.DB, - "SELECT * FROM bots_bot WHERE user_id = ? AND deleted_at IS NULL ORDER BY name ASC LIMIT ? OFFSET ?", - r.user.id, - p.limit, - p.offset, - ); + let rows: BotRow[]; + let models: Map; + if (cached) { + ({ rows, models } = cached); + } else { + rows = await all( + r.env.DB, + "SELECT * FROM bots_bot WHERE user_id = ? AND deleted_at IS NULL ORDER BY name ASC LIMIT ? OFFSET ?", + r.user.id, + p.limit, + p.offset, + ); + // One batched aimodel lookup instead of N per-bot queries. + models = new Map(); + const ids = [...new Set(rows.map((b) => b.ai_model_id).filter((v): v is number => v != null))]; + if (ids.length > 0) { + const inClause = `(${ids.map(() => "?").join(",")})`; + const mrows = await all<{ id: number; model_id: string }>( + r.env.DB, + `SELECT id, model_id FROM bots_aimodel WHERE id IN ${inClause}`, + ...ids, + ); + for (const m of mrows) models.set(m.id, m.model_id); + } + // setBotList drops the entry if a write landed mid-read (epoch guard). + setBotList(cacheKey, { total, rows, models }, seenEpoch); + } const results = []; - for (const b of rows) results.push(await serializeBot(r, b)); + for (const b of rows) results.push(serializeBotWithModels(r, b, models)); return c.json(paginated(r.base + "/bots/", sp, p, total, results)); }); @@ -102,6 +133,7 @@ bots.post("/", async (c) => { now, ); const bot = (await first(r.env.DB, "SELECT * FROM bots_bot WHERE bot_id = ?", botId))!; + invalidateBotListCache(r.user.id); return c.json(await serializeBot(r, bot), 201); }); @@ -168,6 +200,7 @@ async function updateBot(c: TC, partial: boolean) { } params.push(found.bot.id); await run(r.env.DB, `UPDATE bots_bot SET ${sets.join(", ")} WHERE id = ?`, ...params); + invalidateBotListCache(r.user.id); const bot = (await first(r.env.DB, "SELECT * FROM bots_bot WHERE id = ?", found.bot.id))!; return c.json(await serializeBot(r, bot)); } @@ -184,6 +217,7 @@ bots.delete("/:pk/", async (c) => { if ("error" in found) return found.error; // Django parity: hard-delete the row (Bot has no soft-delete override). await run(r.env.DB, "DELETE FROM bots_bot WHERE id = ?", found.bot.id); + invalidateBotListCache(r.user.id); return c.body(null, 204); }); diff --git a/worker/src/usage.ts b/worker/src/usage.ts index ed861ab..b8458e3 100644 --- a/worker/src/usage.ts +++ b/worker/src/usage.ts @@ -25,6 +25,30 @@ interface AiModelRow { is_default: number; } +// Pricing table is runtime-static: no worker route writes bots_aimodel +// (only migrations/seed do), so load it once per isolate and keep it. +// Deploys recycle isolates, which refreshes the rows. Any future runtime +// writer must call invalidateModelsCache(). +let modelRowsCache: AiModelRow[] | null = null; + +async function loadModels(db: D1Database): Promise { + if (!modelRowsCache) { + const res = await db.prepare("SELECT model_id, input_token_cost, output_token_cost, is_default FROM bots_aimodel").all(); + const rows = (res.results ?? []) as AiModelRow[]; + // Never cache an empty table: with no pricing every cost computes as 0 + // and the daily cap silently stops being enforced. + if (rows.length > 0) modelRowsCache = rows; + return rows; + } + return modelRowsCache; +} + +// Tests seed a fresh DB per case: reset so one case's pricing never leaks +// into the next (module state outlives the DB in the same worker). +export function invalidateModelsCache(): void { + modelRowsCache = null; +} + export function isValidTimezone(tz: string): boolean { try { new Intl.DateTimeFormat("en-US", { timeZone: tz }); @@ -137,8 +161,7 @@ export async function rawCostForToday( beforeIso?: string, now = new Date(), ): Promise { - const models = await db.prepare("SELECT model_id, input_token_cost, output_token_cost, is_default FROM bots_aimodel").all(); - const { price } = parseModels((models.results ?? []) as AiModelRow[]); + const { price } = parseModels(await loadModels(db)); const { stamped, unstamped } = await messageBucketsToday(db, userId, startOfTodayUtc(timezone, now).toISOString(), beforeIso); let total = 0; let inputTokens = 0; @@ -190,11 +213,11 @@ async function baselinePricingSnapshot( ) .bind(userId, startIso, cutoffIso) .all<{ chat_id: number; mid: string | null }>(); - const models = await db.prepare("SELECT model_id, input_token_cost, output_token_cost, is_default FROM bots_aimodel").all(); + const modelRows = await loadModels(db); const all = Object.fromEntries( - ((models.results ?? []) as AiModelRow[]).map((m) => [m.model_id, [m.input_token_cost, m.output_token_cost, m.is_default]]), + modelRows.map((m) => [m.model_id, [m.input_token_cost, m.output_token_cost, m.is_default]]), ) as Record; - const defaultId = ((models.results ?? []) as AiModelRow[]).find((m) => m.is_default === 1)?.model_id ?? null; + const defaultId = modelRows.find((m) => m.is_default === 1)?.model_id ?? null; const stampedIds = new Set(((stampedRows.results ?? []) as { mid: string }[]).map((r) => r.mid)); const unstamped: Record = {}; for (const r of (unstampedRows.results ?? []) as { chat_id: number; mid: string | null }[]) { diff --git a/worker/tests/api.test.ts b/worker/tests/api.test.ts index 4560e44..a917f1e 100644 --- a/worker/tests/api.test.ts +++ b/worker/tests/api.test.ts @@ -6,7 +6,8 @@ import { makeTestDb, testEnv } from "./d1shim"; import type { Env } from "../src/types"; import { buildHistory } from "../src/chat"; -import { costForToday } from "../src/usage"; +import { costForToday, invalidateModelsCache } from "../src/usage"; +import { clearBotListCache } from "../src/botListCache"; let db: D1Database; let env: Env; @@ -64,6 +65,10 @@ async function api( beforeEach(() => { db = makeTestDb(); env = testEnv(db) as unknown as Env; + // Module-level caches outlive the per-test DB: reset so one case's rows + // never leak into the next. + invalidateModelsCache(); + clearBotListCache(); }); describe("auth + bots", () => { diff --git a/worker/tests/cache.test.ts b/worker/tests/cache.test.ts new file mode 100644 index 0000000..d652bac --- /dev/null +++ b/worker/tests/cache.test.ts @@ -0,0 +1,113 @@ +// Regression tests for the server caches (PR #112 review findings): +// - bot list epoch guard drops pre-write rows read mid-write +// - empty bots_aimodel is never cached (quota must enforce once priced) +import { beforeEach, describe, expect, it } from "vitest"; +import { + botListEpoch, + clearBotListCache, + getBotList, + invalidateBotListCache, + setBotList, + type BotListEntry, +} from "../src/botListCache"; +import { costForToday, invalidateModelsCache } from "../src/usage"; +import { makeTestDb } from "./d1shim"; +import type { BotRow } from "../src/types"; + +let db: D1Database; + +beforeEach(() => { + db = makeTestDb(); + invalidateModelsCache(); + clearBotListCache(); +}); + +function botRow(partial: Partial = {}): BotRow { + return { + id: 1, + user_id: 1, + ai_model_id: null, + bot_id: "22222222-2222-4222-8222-222222222222", + name: "Penelope", + system_prompt: null, + simple_editor: 0, + template_name: null, + response_length: 200, + restrict_language: 1, + restrict_adult_topics: 1, + enable_web_search: 0, + enable_html_pages: 0, + color: null, + icon: null, + deleted_at: null, + created_at: "2026-01-01T00:00:00Z", + modified_at: "2026-01-01T00:00:00Z", + ...partial, + }; +} + +function entry(): BotListEntry { + return { total: 1, rows: [botRow()], models: new Map() }; +} + +describe("botListCache", () => { + it("round-trips entries and invalidates per user", () => { + setBotList("1:1", entry(), botListEpoch()); + setBotList("2:1", entry(), botListEpoch()); + expect(getBotList("1:1")?.total).toBe(1); + invalidateBotListCache(1); + expect(getBotList("1:1")).toBeNull(); + expect(getBotList("2:1")?.total).toBe(1); + }); + + it("drops stores from before a mid-read write (epoch guard)", () => { + const seen = botListEpoch(); + setBotList("1:1", entry(), seen); + expect(getBotList("1:1")?.total).toBe(1); + // A write lands after this GET started reading: invalidate bumps epoch. + invalidateBotListCache(1); + // The in-flight GET's store carries the stale epoch and must be dropped. + setBotList("1:1", entry(), seen); + expect(getBotList("1:1")).toBeNull(); + // A fresh read after the write stores normally. + setBotList("1:1", entry(), botListEpoch()); + expect(getBotList("1:1")?.total).toBe(1); + }); + + it("evicts oldest entries past the cap", () => { + const seen = botListEpoch(); + for (let i = 0; i < 1005; i++) setBotList(`9:${i}`, entry(), seen); + expect(getBotList("9:0")).toBeNull(); + expect(getBotList("9:1004")?.total).toBe(1); + }); +}); + +describe("pricing cache", () => { + it("re-reads an empty model table instead of pinning costs at 0", async () => { + const now = new Date().toISOString(); + await db.prepare( + "INSERT INTO auth_user (id, password, username, first_name, email, is_active, date_joined) VALUES (1, 'x', 'parent', 'Pat', 'parent@example.com', 1, ?)", + ).bind(now).run(); + await db.prepare( + "INSERT INTO bots_profile (id, user_id, profile_id, name, created_at, modified_at) VALUES (1, 1, '11111111-1111-4111-8111-111111111111', 'Kid', ?, ?)", + ).bind(now, now).run(); + await db.prepare( + "INSERT INTO bots_bot (id, user_id, ai_model_id, bot_id, name, created_at, modified_at) VALUES (1, 1, NULL, '22222222-2222-4222-8222-222222222222', 'Penelope', ?, ?)", + ).bind(now, now).run(); + await db.prepare( + "INSERT INTO bots_chat (id, user_id, profile_id, bot_id, chat_id, title, created_at, modified_at) VALUES (1, 1, 1, 1, 'c1', 't', ?, ?)", + ).bind(now, now).run(); + await db.prepare( + 'INSERT INTO bots_message (chat_id, message_id, text, role, "order", input_tokens, output_tokens, created_at, modified_at) VALUES (1, \'m1\', \'hi\', \'assistant\', 1, 10, 5, ?, ?)', + ).bind(now, now).run(); + + // No pricing rows yet: cost is 0, and that emptiness must not stick. + expect((await costForToday(db, 1)).total).toBe(0); + + await db.prepare( + "INSERT INTO bots_aimodel (id, model_id, name, input_token_cost, output_token_cost, is_default, supported_input_modalities, created_at, modified_at) VALUES (1, 'm', 'M', 1.0, 2.0, 1, '[\"text\"]', ?, ?)", + ).bind(now, now).run(); + // 10*1.0 + 5*2.0 = 20 once pricing exists. + expect((await costForToday(db, 1)).total).toBeCloseTo(20); + }); +});