Skip to content
Merged
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
46 changes: 46 additions & 0 deletions worker/src/botListCache.ts
Original file line number Diff line number Diff line change
@@ -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<number, string | null>;
}

const MAX_ENTRIES = 1000;
const entries = new Map<string, BotListEntry>();

// 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();
}
2 changes: 2 additions & 0 deletions worker/src/oauth.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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?";

Expand Down Expand Up @@ -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",
Expand Down
5 changes: 5 additions & 0 deletions worker/src/onboarding.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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";
Expand Down Expand Up @@ -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();
Expand Down Expand Up @@ -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];
Expand All @@ -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;
}
}

Expand Down Expand Up @@ -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"))) {
Expand Down
52 changes: 43 additions & 9 deletions worker/src/routes/bots.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand All @@ -21,22 +22,52 @@ async function serializeBot(r: ReturnType<typeof routeCtx>, bot: BotRow) {
return botRow(bot as unknown as Record<string, unknown>, aiModel, `${r.base}/bots/${bot.id}/`);
}

function serializeBotWithModels(r: ReturnType<typeof routeCtx>, bot: BotRow, models: Map<number, string | null>) {
const aiModel = bot.ai_model_id != null ? (models.get(bot.ai_model_id) ?? null) : null;
return botRow(bot as unknown as Record<string, unknown>, 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<string, string>);
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<BotRow>(
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<number, string | null>;
if (cached) {
({ rows, models } = cached);
} else {
rows = await all<BotRow>(
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<number, string | null>();
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));
});

Expand Down Expand Up @@ -102,6 +133,7 @@ bots.post("/", async (c) => {
now,
);
const bot = (await first<BotRow>(r.env.DB, "SELECT * FROM bots_bot WHERE bot_id = ?", botId))!;
invalidateBotListCache(r.user.id);
return c.json(await serializeBot(r, bot), 201);
});

Expand Down Expand Up @@ -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<BotRow>(r.env.DB, "SELECT * FROM bots_bot WHERE id = ?", found.bot.id))!;
return c.json(await serializeBot(r, bot));
}
Expand All @@ -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);
});

Expand Down
33 changes: 28 additions & 5 deletions worker/src/usage.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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<AiModelRow[]> {
if (!modelRowsCache) {
const res = await db.prepare("SELECT model_id, input_token_cost, output_token_cost, is_default FROM bots_aimodel").all<AiModelRow>();
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;
Comment thread
cubic-dev-ai[bot] marked this conversation as resolved.
Resolved
}

// 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 });
Expand Down Expand Up @@ -137,8 +161,7 @@ export async function rawCostForToday(
beforeIso?: string,
now = new Date(),
): Promise<CostToday> {
const models = await db.prepare("SELECT model_id, input_token_cost, output_token_cost, is_default FROM bots_aimodel").all<AiModelRow>();
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;
Expand Down Expand Up @@ -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<AiModelRow>();
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<string, [number, number, number]>;
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<string, string | null> = {};
for (const r of (unstampedRows.results ?? []) as { chat_id: number; mid: string | null }[]) {
Expand Down
7 changes: 6 additions & 1 deletion worker/tests/api.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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", () => {
Expand Down
113 changes: 113 additions & 0 deletions worker/tests/cache.test.ts
Original file line number Diff line number Diff line change
@@ -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> = {}): 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);
});
});
Loading