app/routes/search.ask.ts

Functions

Imports

It imports app/db/index.ts, app/lib/client-ip.ts, app/lib/context.ts, app/lib/error-message.mjs, app/lib/origin.mjs, app/lib/search/ask-guard.server.ts, app/lib/search/ask.server.ts. No module of this repo imports it.

Source

1/**
2 * The only public endpoint that costs money per request, so guards run cheapest first: refuse,
3 * then a cached answer, then the model.
4 */
5
6import {
7  answerLeaksPrompt,
8  askAvailable,
9  askStream,
10  citedSlugs,
11  guardAnswerStream,
12  replayCachedAnswer,
13  teeForCache,
14} from "~/lib/search/ask.server";
15import {
16  ASK_RATE_LIMIT_PERIOD_SECONDS,
17  type CachedAnswer,
18  checkAskRate,
19  dropCachedAnswer,
20  questionKey,
21  readCachedAnswer,
22  reserveAskBudget,
23  writeCachedAnswer,
24} from "~/lib/search/ask-guard.server";
25import { ASK_ORIGIN_REFUSAL, originVerdict } from "~/lib/origin.mjs";
26import { publiclyVisibleSlugs } from "~/db";
27import { clientIp } from "~/lib/client-ip";
28import { getEnv, getExecutionContext } from "~/lib/context";
29import { errorMessage } from "~/lib/error-message.mjs";
30import type { Route } from "./+types/search.ask";
31
32/** Long enough for a real question, short enough that the model is not a toy. */
33const MAX_QUESTION_LENGTH = 500;
34
35/** An answer citing a post that is now a draft, deleted or scheduled forward must not be replayed. */
36async function citationsStillPublic(env: Env, cached: CachedAnswer) {
37  const slugs = citedSlugs(cached.chunks);
38  if (slugs.length === 0) return true;
39  const visible = await publiclyVisibleSlugs(env, slugs);
40  return slugs.every((slug) => visible.has(slug));
41}
42
43function streamHeaders(cacheStatus: "hit" | "miss"): HeadersInit {
44  return {
45    "content-type": "text/event-stream; charset=utf-8",
46    // Edge caching would serve one reader's answer to another; the answer cache is ours, in KV.
47    "cache-control": "no-store",
48    // Nginx-style buffering proxies would otherwise defeat the streaming.
49    "x-accel-buffering": "no",
50    "x-ask-cache": cacheStatus,
51  };
52}
53
54function refuse(reason: string | undefined, retryAfter: number | undefined): Response {
55  const body =
56    reason === "budget"
57      ? "Ask has answered its budget of questions for today. Classic search is unaffected."
58      : reason === "unprotected"
59        ? "Ask is unavailable."
60        : "Too many questions. Try again shortly.";
61  return new Response(body, {
62    status: reason === "unprotected" ? 503 : 429,
63    headers: {
64      "retry-after": String(retryAfter ?? ASK_RATE_LIMIT_PERIOD_SECONDS),
65      "cache-control": "no-store",
66    },
67  });
68}
69
70/**
71 * GET is refused: a crawler, prefetch, `<img src>` or unfurler would bill it. 405 with `Allow`, not
72 * 404, so the next caller does not reinvent the GET.
73 */
74export function loader() {
75  return new Response("Ask takes POST.", {
76    status: 405,
77    headers: { allow: "POST", "cache-control": "no-store" },
78  });
79}
80
81export async function action({ request, context }: Route.ActionArgs) {
82  /*
83   * First because it is free. A hostile page's thousand readers are a thousand IPs, each inside its
84   * allowance; Origin is the field that page cannot control. An absent Origin is allowed: that client
85   * spends its own bounded allowance.
86   */
87  const origin = originVerdict(request.headers.get("origin"), request.url);
88  if (!origin.ok) {
89    return new Response(ASK_ORIGIN_REFUSAL, {
90      status: 403,
91      headers: { "cache-control": "no-store" },
92    });
93  }
94
95  const env = getEnv(context);
96
97  // 404, not 503: with the binding absent Ask does not exist, rather than being down.
98  if (!askAvailable(env)) {
99    return new Response("Ask is not enabled.", { status: 404 });
100  }
101
102  const body = await request.formData();
103  const question = String(body.get("q") ?? "").trim();
104
105  if (!question) {
106    return new Response("Missing q.", { status: 400 });
107  }
108  if (question.length > MAX_QUESTION_LENGTH) {
109    return new Response(`Question longer than ${MAX_QUESTION_LENGTH} characters.`, {
110      status: 400,
111    });
112  }
113
114  const ip = clientIp(request);
115
116  // Gate 1, in front of the cache too: a cached answer is cheap, not free.
117  const rate = await checkAskRate(env, ip);
118  if (!rate.ok) {
119    return refuse(rate.reason, rate.retryAfter);
120  }
121
122  const key = await questionKey(question);
123
124  // A hit must still prove its citations are public: invalidation by KV `list` is eventually
125  // consistent, so an answer can outlive an unpublish.
126  const cached = await readCachedAnswer(env, key);
127  if (cached) {
128    if (await citationsStillPublic(env, cached)) {
129      return new Response(replayCachedAnswer(cached), { headers: streamHeaders("hit") });
130    }
131    await dropCachedAnswer(env, key);
132  }
133
134  /*
135   * Reserved here, the last point before the only line that costs money. No single-flight: an
136   * attacker with N addresses spends N units with N different questions anyway.
137   */
138  const budget = await reserveAskBudget(env);
139  if (!budget.ok) {
140    return refuse(budget.reason, budget.retryAfter);
141  }
142
143  try {
144    /* Guard first, then tee: the cache must store what the reader saw, including the guard's substitution. */
145    const upstream = await askStream(env, question);
146    const { toReader, captured } = teeForCache(
147      guardAnswerStream(upstream, (slugs) => publiclyVisibleSlugs(env, slugs)),
148    );
149
150    getExecutionContext(context).waitUntil(
151      (async () => {
152        const answer = await captured;
153        if (!answer) return;
154
155        /*
156         * A no-chunk answer is not cached, or "nothing found" would stick for a week after the answering post
157         * publishes. A prompt echo is a successful injection; not caching it keeps it to one request.
158         */
159        if (answer.chunks.length === 0) return;
160        if (answerLeaksPrompt(answer.answer)) {
161          console.error(JSON.stringify({ alert: "ask-cache-refused-echo" }));
162          return;
163        }
164        await writeCachedAnswer(env, key, answer);
165      })(),
166    );
167
168    return new Response(toReader, { headers: streamHeaders("miss") });
169  } catch (error) {
170    // The client treats any non-200 as "Ask is unavailable"; a 500 here would land on a page with real results.
171    console.error(JSON.stringify({ alert: "ask-stream-failed", detail: errorMessage(error) }));
172    return new Response("Ask is unavailable.", { status: 502 });
173  }
174}