app/lib/search/ask-guard.server.ts

Functions

Imports

It imports app/lib/bytes.mjs, app/lib/rate-limit.mjs, app/lib/search/ask-pacing.mjs. Imported by app/lib/admin/posts-actions.server.ts, app/lib/search/ask.server.ts, app/routes/admin.posts._index.tsx, app/routes/admin.tsx, app/routes/search.ask.ts.

Source

1// Every generated answer bills Workers AI and the route is public. Nothing reaches the model until
2// rate, cache and daily budget have passed, and budget sits AFTER the cache so a hit costs nothing.
3
4import { toHex } from "~/lib/bytes.mjs";
5import { limitHit } from "~/lib/rate-limit.mjs";
6import { pacedAllowance, secondsPerPacedUnit } from "~/lib/search/ask-pacing.mjs";
7
8/** Per IP per minute: more than a reading human asks, far less than a loop wants. */
9const ASK_RATE_LIMIT = 5;
10export const ASK_RATE_LIMIT_PERIOD_SECONDS = 60;
11
12/** A ceiling on the bill: a number a real reader never reaches and a scraper will. */
13const DAILY_ANSWER_BUDGET = 200;
14
15/** Without a burst the first question after UTC midnight is refused: the paced share is zero at 00:00:01. */
16const ASK_BURST_ALLOWANCE = 25;
17
18/** How long a cached answer stays valid. Also the ceiling on staleness. */
19const ANSWER_CACHE_TTL_SECONDS = 60 * 60 * 24 * 7;
20
21const ANSWER_CACHE_PREFIX = "ask:answer:";
22
23/**
24 * TTL cache, not compute-on-write: the expected set moves with the clock as scheduled posts go live.
25 * KV refuses a TTL under 60 seconds; five minutes of an under-reported badge harms nobody.
26 */
27const DRIFT_CACHE_KEY = "ask:drift";
28
29/** Stated here and repeated to the reader in the badge's own title text. */
30export const DRIFT_CACHE_TTL_SECONDS = 300;
31
32export interface CachedAnswer {
33  answer: string;
34  chunks: unknown[];
35}
36
37/** No stemming or stopword dropping: differently worded questions must not share an answer. */
38function normalizeQuestion(question: string): string {
39  return question
40    .toLowerCase()
41    .replace(/\s+/g, " ")
42    .replace(/[?!.,;:]+$/g, "")
43    .trim();
44}
45
46export async function questionKey(question: string): Promise<string> {
47  const normalized = normalizeQuestion(question);
48  const digest = await crypto.subtle.digest(
49    "SHA-256",
50    new TextEncoder().encode(normalized),
51  );
52  return `${ANSWER_CACHE_PREFIX}${toHex(digest)}`;
53}
54
55interface GuardVerdict {
56  ok: boolean;
57  retryAfter?: number;
58  reason?: "rate" | "budget" | "unprotected";
59}
60
61/** Fails closed without the binding, so removing it disables Ask rather than unprotecting it. */
62export async function checkAskRate(env: Env, ip: string): Promise<GuardVerdict> {
63  // Keyed by IP alone, not IP plus question: the point is to cap how often one caller can spend.
64  const verdict = await limitHit(env, `ip:${ip}`, ASK_RATE_LIMIT, ASK_RATE_LIMIT_PERIOD_SECONDS);
65  if (verdict.status === "ok") return { ok: true };
66  return {
67    ok: false,
68    reason: verdict.status === "unavailable" ? "unprotected" : "rate",
69    retryAfter: verdict.retryAfterSeconds,
70  };
71}
72
73/**
74 * Only on a cache miss. Reserve-then-generate, so two concurrent requests cannot both pass on one read;
75 * never refunded on failure, since under-counting would make the ceiling a suggestion.
76 */
77export async function reserveAskBudget(env: Env): Promise<GuardVerdict> {
78  const budget = env.ASK_BUDGET.get(env.ASK_BUDGET.idFromName("global"));
79
80  // Paced, so exhausting the ceiling denies Ask for minutes rather than until UTC midnight.
81  const ceiling = pacedAllowance(DAILY_ANSWER_BUDGET, ASK_BURST_ALLOWANCE, new Date());
82  const { ok } = await budget.consume(ceiling);
83  if (!ok) {
84    // Not seconds-until-midnight: under pacing that would be wrong by up to a day.
85    return { ok: false, reason: "budget", retryAfter: secondsPerPacedUnit(DAILY_ANSWER_BUDGET) };
86  }
87  return { ok: true };
88}
89
90export async function readAskBudget(env: Env): Promise<{ day: string; count: number; limit: number }> {
91  const budget = env.ASK_BUDGET.get(env.ASK_BUDGET.idFromName("global"));
92  const { day, count } = await budget.peek();
93  return { day, count, limit: DAILY_ANSWER_BUDGET };
94}
95
96export async function resetAskBudget(env: Env): Promise<void> {
97  const budget = env.ASK_BUDGET.get(env.ASK_BUDGET.idFromName("global"));
98  await budget.reset();
99}
100
101export async function readCachedAnswer(
102  env: Env,
103  key: string,
104): Promise<CachedAnswer | null> {
105  const cached = await env.APP_KV.get(key, "json");
106  if (!cached || typeof cached !== "object") return null;
107  const value = cached as Partial<CachedAnswer>;
108  if (typeof value.answer !== "string" || !Array.isArray(value.chunks)) return null;
109  return { answer: value.answer, chunks: value.chunks };
110}
111
112/** For an entry whose citations are no longer public. A delete by name is consistent; list is not. */
113export async function dropCachedAnswer(env: Env, key: string): Promise<void> {
114  await env.APP_KV.delete(key);
115}
116
117export async function writeCachedAnswer(
118  env: Env,
119  key: string,
120  value: CachedAnswer,
121): Promise<void> {
122  // An empty answer is a failed generation; caching it would serve the failure for a week.
123  if (value.answer.trim().length === 0) return;
124  await env.APP_KV.put(key, JSON.stringify(value), {
125    expirationTtl: ANSWER_CACHE_TTL_SECONDS,
126  });
127}
128
129/** Null, not zero, on a miss: zero means the index agrees, and a miss must not render a clean badge. */
130export async function readCachedDrift(env: Env): Promise<number | null> {
131  const cached = await env.APP_KV.get(DRIFT_CACHE_KEY, "json");
132  if (!cached || typeof cached !== "object") return null;
133  const value = (cached as { drift?: unknown }).drift;
134  if (typeof value !== "number" || !Number.isInteger(value) || value < 0) return null;
135  return value;
136}
137
138export async function writeCachedDrift(env: Env, drift: number): Promise<void> {
139  await env.APP_KV.put(DRIFT_CACHE_KEY, JSON.stringify({ drift }), {
140    expirationTtl: DRIFT_CACHE_TTL_SECONDS,
141  });
142}
143
144/** A delete rather than a write: the caller knows the count is stale, not what it became. */
145export async function dropCachedDrift(env: Env): Promise<void> {
146  await env.APP_KV.delete(DRIFT_CACHE_KEY);
147}
148
149/**
150 * Invalidated rather than keyed by generation, which would cost a second KV read on every request.
151 * KV list is eventually consistent, so an answer written just before a sync can survive until its TTL.
152 */
153export async function invalidateAnswerCache(env: Env): Promise<number> {
154  let cursor: string | undefined;
155  let deleted = 0;
156  for (;;) {
157    const page = await env.APP_KV.list({ prefix: ANSWER_CACHE_PREFIX, cursor });
158    for (const key of page.keys) {
159      await env.APP_KV.delete(key.name);
160      deleted += 1;
161    }
162    if (page.list_complete) break;
163    cursor = page.cursor;
164  }
165  return deleted;
166}