fix(edge): meter and size-cap knowledge embeddings in embed-chunks and search-knowledge

server/supabase/functions/_shared/embedding-quota.fake.ts

@@ -0,0 +1,25 @@
+// Test double for EmbeddingUsageStore (not a test file itself).
+import { EMBEDDING_QUOTA_FEATURE, type EmbeddingUsageStore } from './embedding-quota.ts'
+import type { PlanQuotaTier } from './core-contract.generated.ts'
+
+/** In-memory daily_usage keyed by date; `today` is the date increments land on. */
+export class MemoryUsageStore implements EmbeddingUsageStore {
+  rows = new Map<string, number>()
+  failIncrement = false
+  constructor(public tier: PlanQuotaTier, public today: string) {}
+  readTier(): Promise<PlanQuotaTier> {
+    return Promise.resolve(this.tier)
+  }
+  increment(_userId: string, feature: string, amount: number): Promise<number> {
+    if (this.failIncrement) return Promise.reject(new Error('down'))
+    if (feature !== EMBEDDING_QUOTA_FEATURE) return Promise.reject(new Error(`unexpected feature ${feature}`))
+    const next = (this.rows.get(this.today) ?? 0) + amount
+    this.rows.set(this.today, next)
+    return Promise.resolve(next)
+  }
+  sumBetween(_userId: string, _feature: string, fromDate: string, beforeDate: string): Promise<number> {
+    let sum = 0
+    for (const [date, count] of this.rows) if (date >= fromDate && date < beforeDate) sum += count
+    return Promise.resolve(sum)
+  }
+}

server/supabase/functions/_shared/embedding-quota.test.ts

@@ -0,0 +1,85 @@
+import {
+  EMBEDDING_LIMITS,
+  EMBEDDING_QUOTA,
+  embeddingUnits,
+  firstOversizedChunk,
+  planEmbeddingBatches,
+  refundEmbeddingUnits,
+  reserveEmbeddingUnits,
+  usageWindowStart,
+} from './embedding-quota.ts'
+import { MemoryUsageStore } from './embedding-quota.fake.ts'
+
+function assert(condition: boolean, message: string): asserts condition {
+  if (!condition) throw new Error(message)
+}
+
+const NOW = new Date('2026-09-28T12:00:00Z')
+
+Deno.test('units count every started 1,000 characters, at least one per input', () => {
+  assert(embeddingUnits(['a']) === 1, 'short input')
+  assert(embeddingUnits(['x'.repeat(1_000), 'x'.repeat(1_001)]) === 3, 'boundary rounding')
+})
+
+Deno.test('oversized chunks are found before any spend', () => {
+  const ok = { content: 'x'.repeat(EMBEDDING_LIMITS.maxChunkChars) }
+  const big = { content: 'x'.repeat(EMBEDDING_LIMITS.maxChunkChars + 1) }
+  assert(firstOversizedChunk([ok, ok]) === -1, 'at-limit chunk rejected')
+  assert(firstOversizedChunk([ok, big]) === 1, 'oversized chunk missed')
+})
+
+Deno.test('batches stay under the input and character caps', () => {
+  const small = Array.from({ length: 250 }, (_, i) => ({ content: `c${i}` }))
+  assert(planEmbeddingBatches(small).map((b) => b.length).join(',') === '100,100,50', 'input cap')
+  const large = Array.from({ length: 30 }, () => ({ content: 'x'.repeat(8_000) }))
+  const batches = planEmbeddingBatches(large)
+  assert(batches.every((b) => b.reduce((s, c) => s + c.content.length, 0) <= EMBEDDING_LIMITS.maxBatchChars), 'char cap')
+  assert(batches.flat().length === 30, 'chunks lost')
+})
+
+Deno.test('weekly window covers the six days before today', () => {
+  assert(usageWindowStart('weekly', NOW) === '2026-09-22', 'weekly start')
+  assert(usageWindowStart('daily', NOW) === null, 'daily has no earlier window')
+})
+
+Deno.test('free tier is capped weekly and a denied reservation is rolled back', async () => {
+  const store = new MemoryUsageStore('free', '2026-09-28')
+  store.rows.set('2026-09-25', EMBEDDING_QUOTA.free.limit - 10)
+  const denied = await reserveEmbeddingUnits(store, 'u1', 11, NOW)
+  assert(!denied.allowed, 'over-limit reservation was allowed')
+  assert((store.rows.get('2026-09-28') ?? 0) === 0, 'denied units were kept')
+  const allowed = await reserveEmbeddingUnits(store, 'u1', 10, NOW)
+  assert(allowed.allowed && allowed.current === EMBEDDING_QUOTA.free.limit, 'exact fit denied')
+})
+
+Deno.test('usage older than the window does not count', async () => {
+  const store = new MemoryUsageStore('free', '2026-09-28')
+  store.rows.set('2026-09-21', EMBEDDING_QUOTA.free.limit)
+  assert((await reserveEmbeddingUnits(store, 'u1', 5, NOW)).allowed, 'stale usage counted')
+})
+
+Deno.test('enterprise is unlimited but still counted', async () => {
+  const store = new MemoryUsageStore('enterprise', '2026-09-28')
+  const result = await reserveEmbeddingUnits(store, 'u1', 1_000_000, NOW)
+  assert(result.allowed && store.rows.get('2026-09-28') === 1_000_000, 'unlimited tier not counted')
+})
+
+Deno.test('refunds are capped by the reservation and skipped after the day rolls over', async () => {
+  const store = new MemoryUsageStore('pro', '2026-09-28')
+  const reservation = await reserveEmbeddingUnits(store, 'u1', 40, NOW)
+  await refundEmbeddingUnits(store, 'u1', reservation, 1_000, NOW)
+  assert(store.rows.get('2026-09-28') === 0, 'refund exceeded the reserved units')
+
+  const again = await reserveEmbeddingUnits(store, 'u1', 40, NOW)
+  store.today = '2026-09-29'
+  await refundEmbeddingUnits(store, 'u1', again, 40, new Date('2026-09-29T00:00:01Z'))
+  assert(!store.rows.has('2026-09-29'), 'next-day refund credited the new day')
+})
+
+Deno.test('a store failure during refund is swallowed', async () => {
+  const store = new MemoryUsageStore('pro', '2026-09-28')
+  const reservation = await reserveEmbeddingUnits(store, 'u1', 3, NOW)
+  store.failIncrement = true
+  await refundEmbeddingUnits(store, 'u1', reservation, 3, NOW)
+  assert(store.rows.get('2026-09-28') === 3, 'usage changed')
+})

server/supabase/functions/_shared/embedding-quota.ts

@@ -0,0 +1,250 @@
+// server/supabase/functions/_shared/embedding-quota.ts
+// Cost guard for the paid OpenAI embedding calls made by embed-chunks and
+// search-knowledge.
+//
+// Before this module both functions only checked the JWT: any signed-in account
+// could insert an unbounded number of knowledge_chunks rows (RLS only checks
+// document ownership, `content` is unbounded text) and have every one of them
+// embedded, repeatedly, and search-knowledge embedded a 4,000-char query per
+// call with no counter. stt/llm/realtime all go through quota.ts; embeddings
+// did not.
+//
+// Policy (pure):   size caps + per-tier unit budget + batch planning.
+// Port:            EmbeddingUsageStore (tier lookup + atomic counter).
+// Adapter:         createSupabaseEmbeddingUsageStore (subscriptions + daily_usage
+//                  via the existing service-role `increment_daily_usage` RPC).
+//
+// TODO(plan-catalog): the per-tier budget below belongs in
+// packages/core/src/plan-catalog.ts `PLAN_QUOTA` (feature `knowledge_embedding`)
+// and should be read from core-contract.generated.ts once it is added there.
+
+import type { createClient } from '@supabase/supabase-js'
+import type { PlanQuotaPeriod, PlanQuotaTier } from './core-contract.generated.ts'
+
+/** daily_usage.feature key for embedding spend. */
+export const EMBEDDING_QUOTA_FEATURE = 'knowledge_embedding'
+
+export const EMBEDDING_LIMITS = Object.freeze({
+  /** Longest chunk accepted for embedding (the OpenAI per-input cap is ~8k tokens). */
+  maxChunkChars: 8_000,
+  /** Most chunks one document may have. Desktop chunking stops at 2,000 chunks. */
+  maxChunksPerDocument: 2_000,
+  /** Longest search query accepted. */
+  maxQueryChars: 4_000,
+  /** Inputs per provider request. */
+  maxBatchInputs: 100,
+  /** Characters per provider request (keeps a batch under the per-request token cap). */
+  maxBatchChars: 100_000,
+  /** One quota unit = up to this many characters of embedded text. */
+  charsPerUnit: 1_000,
+})
+
+export interface EmbeddingQuotaPolicy {
+  /** Units per period. -1 = unlimited, 0 = not available. */
+  limit: number
+  period: PlanQuotaPeriod
+}
+
+/**
+ * Units (1 unit = up to 1,000 chars) per tier. Cloud knowledge is offered on
+ * every tier (web/mobile have no paywall for it), so free gets a small weekly
+ * budget instead of none: 5,000 units ≈ 5M chars ≈ 20 maximum-size mobile
+ * documents a week.
+ */
+export const EMBEDDING_QUOTA: Readonly<Record<PlanQuotaTier, EmbeddingQuotaPolicy>> = Object.freeze({
+  free: { limit: 5_000, period: 'weekly' },
+  pro: { limit: 50_000, period: 'daily' },
+  pro_plus: { limit: 150_000, period: 'daily' },
+  team: { limit: 300_000, period: 'daily' },
+  enterprise: { limit: -1, period: 'daily' },
+})
+
+const QUOTA_TIERS: readonly PlanQuotaTier[] = ['free', 'pro', 'pro_plus', 'team', 'enterprise']
+
+export function normalizeQuotaTier(value: unknown): PlanQuotaTier {
+  return typeof value === 'string' && (QUOTA_TIERS as readonly string[]).includes(value)
+    ? value as PlanQuotaTier
+    : 'free'
+}
+
+export function embeddingPolicyFor(tier: PlanQuotaTier): EmbeddingQuotaPolicy {
+  return EMBEDDING_QUOTA[tier] ?? { limit: 0, period: 'daily' }
+}
+
+/** Units one text costs: every started 1,000 characters, at least 1. */
+export function embeddingUnitsFor(text: string): number {
+  return Math.max(1, Math.ceil(text.length / EMBEDDING_LIMITS.charsPerUnit))
+}
+
+export function embeddingUnits(texts: readonly string[]): number {
+  return texts.reduce((sum, text) => sum + embeddingUnitsFor(text), 0)
+}
+
+/** Index of the first chunk over `maxChunkChars`, or -1. */
+export function firstOversizedChunk(chunks: readonly { content: string }[]): number {
+  return chunks.findIndex((chunk) => chunk.content.length > EMBEDDING_LIMITS.maxChunkChars)
+}
+
+/** Split chunks into provider requests bounded by input count and characters. */
+export function planEmbeddingBatches<T extends { content: string }>(chunks: readonly T[]): T[][] {
+  const batches: T[][] = []
+  let current: T[] = []
+  let chars = 0
+  for (const chunk of chunks) {
+    const size = chunk.content.length
+    if (
+      current.length > 0
+      && (current.length >= EMBEDDING_LIMITS.maxBatchInputs || chars + size > EMBEDDING_LIMITS.maxBatchChars)
+    ) {
+      batches.push(current)
+      current = []
+      chars = 0
+    }
+    current.push(chunk)
+    chars += size
+  }
+  if (current.length > 0) batches.push(current)
+  return batches
+}
+
+function isoDate(date: Date): string {
+  return date.toISOString().slice(0, 10)
+}
+
+/** First date (inclusive) of the usage window before today, or null for a daily period. */
+export function usageWindowStart(period: PlanQuotaPeriod, now: Date): string | null {
+  if (period !== 'weekly') return null
+  const start = new Date(now.getTime())
+  start.setUTCDate(start.getUTCDate() - 6)
+  return isoDate(start)
+}
+
+// ── Port ────────────────────────────────────────────────────────────────
+
+export interface EmbeddingUsageStore {
+  readTier(userId: string): Promise<PlanQuotaTier>
+  /** Atomically add `amount` (may be negative) to today's counter; returns the new count. */
+  increment(userId: string, feature: string, amount: number): Promise<number>
+  /** Sum of counters with fromDate <= date < beforeDate. */
+  sumBetween(userId: string, feature: string, fromDate: string, beforeDate: string): Promise<number>
+}
+
+export interface EmbeddingReservation {
+  allowed: boolean
+  /** Units held by this reservation (0 when denied or nothing was requested). */
+  units: number
+  /** Usage in the current window including this reservation when allowed. */
+  current: number
+  limit: number
+  period: PlanQuotaPeriod
+  tier: PlanQuotaTier
+  /** UTC date the units were counted on. */
+  date: string
+}
+
+/**
+ * Reserve `units` before spending. The counter is incremented first (atomic in
+ * the database) and rolled back when the window total then exceeds the limit,
+ * so concurrent requests can never admit more than the limit between them.
+ * Throws when the store fails: callers must fail closed.
+ */
+export async function reserveEmbeddingUnits(
+  store: EmbeddingUsageStore,
+  userId: string,
+  units: number,
+  now: Date = new Date(),
+): Promise<EmbeddingReservation> {
+  const tier = await store.readTier(userId)
+  const policy = embeddingPolicyFor(tier)
+  const date = isoDate(now)
+  const base = { limit: policy.limit, period: policy.period, tier, date }
+
+  if (units <= 0) return { ...base, allowed: true, units: 0, current: 0 }
+  if (policy.limit === 0) return { ...base, allowed: false, units: 0, current: 0 }
+
+  const today = await store.increment(userId, EMBEDDING_QUOTA_FEATURE, units)
+  if (policy.limit === -1) return { ...base, allowed: true, units, current: today }
+
+  const windowStart = usageWindowStart(policy.period, now)
+  const earlier = windowStart
+    ? await store.sumBetween(userId, EMBEDDING_QUOTA_FEATURE, windowStart, date)
+    : 0
+  const used = earlier + today
+  if (used > policy.limit) {
+    await store.increment(userId, EMBEDDING_QUOTA_FEATURE, -units)
+    return { ...base, allowed: false, units: 0, current: used - units }
+  }
+  return { ...base, allowed: true, units, current: used }
+}
+
+/**
+ * Give back units that were reserved but not spent. Best effort: a failure
+ * leaves the user slightly over-counted, never under-counted. A refund after
+ * the UTC day rolled over is skipped so it cannot credit the new day.
+ */
+export async function refundEmbeddingUnits(
+  store: EmbeddingUsageStore,
+  userId: string,
+  reservation: EmbeddingReservation,
+  units: number,
+  now: Date = new Date(),
+): Promise<void> {
+  const amount = Math.min(units, reservation.units)
+  if (amount <= 0 || isoDate(now) !== reservation.date) return
+  try {
+    await store.increment(userId, EMBEDDING_QUOTA_FEATURE, -amount)
+  } catch {
+    // over-counting is the safe direction
+  }
+}
+
+export function quotaExceededBody(reservation: EmbeddingReservation, requested: number): Record<string, unknown> {
+  return {
+    error: 'quota_exceeded',
+    feature: EMBEDDING_QUOTA_FEATURE,
+    requested,
+    current: reservation.current,
+    limit: reservation.limit,
+    period: reservation.period,
+    tier: reservation.tier,
+  }
+}
+
+// ── Supabase adapter ────────────────────────────────────────────────────
+
+type SupabaseClient = ReturnType<typeof createClient>
+
+export function createSupabaseEmbeddingUsageStore(client: SupabaseClient): EmbeddingUsageStore {
+  return {
+    async readTier(userId) {
+      const { data, error } = await client
+        .from('subscriptions')
+        .select('tier')
+        .eq('user_id', userId)
+        .maybeSingle()
+      if (error) throw new Error('Failed to read subscription tier.')
+      return normalizeQuotaTier((data as { tier?: unknown } | null)?.tier)
+    },
+    async increment(userId, feature, amount) {
+      const { data, error } = await client.rpc('increment_daily_usage', {
+        p_user_id: userId,
+        p_feature: feature,
+        p_amount: amount,
+      })
+      if (error || typeof data !== 'number') throw new Error('Failed to update embedding usage.')
+      return data
+    },
+    async sumBetween(userId, feature, fromDate, beforeDate) {
+      const { data, error } = await client
+        .from('daily_usage')
+        .select('count')
+        .eq('user_id', userId)
+        .eq('feature', feature)
+        .gte('date', fromDate)
+        .lt('date', beforeDate)
+      if (error) throw new Error('Failed to read embedding usage.')
+      const rows = (data ?? []) as Array<{ count: number | null }>
+      return rows.reduce((sum, row) => sum + (row.count ?? 0), 0)
+    },
+  }
+}

server/supabase/functions/_shared/json-response.ts

@@ -0,0 +1,30 @@
+// server/supabase/functions/_shared/json-response.ts
+// JSON response + auth-error mapping shared by the knowledge functions.
+
+import { corsHeaders } from './cors.ts'
+import type { AuthError } from './auth.ts'
+
+export function jsonResponse(status: number, body: Record<string, unknown>): Response {
+  return new Response(JSON.stringify(body), {
+    status,
+    headers: { ...corsHeaders, 'Content-Type': 'application/json' },
+  })
+}
+
+export function isAuthError(error: unknown): error is AuthError {
+  return Boolean(
+    error
+    && typeof error === 'object'
+    && 'status' in error
+    && (error.status === 401 || error.status === 403)
+    && 'message' in error
+    && typeof error.message === 'string',
+  )
+}
+
+/** Map a thrown value to a response: auth errors keep their status, anything else is a 500. */
+export function errorResponse(error: unknown): Response {
+  // Same shape as auth.ts authErrorResponse, without loading the Supabase client.
+  if (isAuthError(error)) return jsonResponse(error.status, { error: error.message })
+  return jsonResponse(500, { error: 'internal_error' })
+}

server/supabase/functions/_shared/openai-embeddings.ts

@@ -0,0 +1,61 @@
+// server/supabase/functions/_shared/openai-embeddings.ts
+// Embedding provider port + OpenAI adapter shared by embed-chunks and
+// search-knowledge. The functions depend on the `EmbeddingProvider` port only,
+// so the quota/size policy can be tested without the network.
+
+export const EMBEDDING_MODEL = 'text-embedding-3-small'
+export const EMBEDDING_DIMENSIONS = 1536
+export const EMBEDDING_PROVIDER_TIMEOUT_MS = 45_000
+const OPENAI_EMBEDDINGS_URL = 'https://api.openai.com/v1/embeddings'
+
+export interface EmbeddingItem {
+  embedding?: unknown
+  index?: unknown
+}
+
+export type EmbeddingOutcome =
+  | { ok: true; data: EmbeddingItem[] }
+  /** `upstream`: transport error or non-2xx status. `invalid`: 2xx with an unusable body. */
+  | { ok: false; reason: 'upstream' | 'invalid' }
+
+export interface EmbeddingProvider {
+  embed(input: string | readonly string[]): Promise<EmbeddingOutcome>
+}
+
+export function isEmbedding(value: unknown): value is number[] {
+  return Array.isArray(value)
+    && value.length === EMBEDDING_DIMENSIONS
+    && value.every((entry) => typeof entry === 'number' && Number.isFinite(entry))
+}
+
+export function createOpenAIEmbeddingProvider(
+  apiKey: string,
+  fetchImpl: typeof fetch = fetch,
+): EmbeddingProvider {
+  return {
+    async embed(input) {
+      let response: Response
+      try {
+        response = await fetchImpl(OPENAI_EMBEDDINGS_URL, {
+          method: 'POST',
+          headers: {
+            Authorization: `Bearer ${apiKey}`,
+            'Content-Type': 'application/json',
+          },
+          body: JSON.stringify({
+            model: EMBEDDING_MODEL,
+            input,
+            dimensions: EMBEDDING_DIMENSIONS,
+          }),
+          signal: AbortSignal.timeout(EMBEDDING_PROVIDER_TIMEOUT_MS),
+        })
+      } catch {
+        return { ok: false, reason: 'upstream' }
+      }
+      if (!response.ok) return { ok: false, reason: 'upstream' }
+      const payload = await response.json().catch(() => null) as { data?: unknown } | null
+      if (!payload || !Array.isArray(payload.data)) return { ok: false, reason: 'invalid' }
+      return { ok: true, data: payload.data as EmbeddingItem[] }
+    },
+  }
+}

server/supabase/functions/embed-chunks/handler.test.ts

@@ -0,0 +1,166 @@
+// Regression tests for the embed-chunks cost guard (redteam r1 #18): before
+// the guard any signed-in account could have an unbounded number of chunks of
+// any size embedded by the paid provider, with no tier budget.
+
+import { createEmbedChunksHandler, type KnowledgeChunkRow, type KnowledgeIndexStore } from './handler.ts'
+import { EMBEDDING_DIMENSIONS, type EmbeddingProvider } from '../_shared/openai-embeddings.ts'
+import { EMBEDDING_LIMITS, EMBEDDING_QUOTA } from '../_shared/embedding-quota.ts'
+import { MemoryUsageStore } from '../_shared/embedding-quota.fake.ts'
+import type { PlanQuotaTier } from '../_shared/core-contract.generated.ts'
+
+function assert(condition: boolean, message: string): asserts condition {
+  if (!condition) throw new Error(message)
+}
+
+const USER = 'user-1'
+const DOC = '0b6c1f3e-4a5d-4e7f-8a9b-1c2d3e4f5a6b'
+const NOW = new Date('2026-09-28T12:00:00Z')
+const TODAY = '2026-09-28'
+
+class FakeKnowledgeStore implements KnowledgeIndexStore {
+  embedded = new Set<string>()
+  indexed = false
+  listLimit = -1
+  constructor(public chunks: KnowledgeChunkRow[], public owner: string | null = USER) {}
+  documentOwner(): Promise<string | null> {
+    return Promise.resolve(this.owner)
+  }
+  countChunks(_doc: string, pendingOnly: boolean): Promise<number> {
+    return Promise.resolve(pendingOnly ? this.pending().length : this.chunks.length)
+  }
+  listPendingChunks(_doc: string, limit: number): Promise<KnowledgeChunkRow[]> {
+    this.listLimit = limit
+    return Promise.resolve(this.pending().slice(0, limit))
+  }
+  saveEmbedding(_doc: string, chunkId: string): Promise<boolean> {
+    this.embedded.add(chunkId)
+    return Promise.resolve(true)
+  }
+  clearIndexed(): Promise<boolean> {
+    this.indexed = false
+    return Promise.resolve(true)
+  }
+  markIndexed(): Promise<boolean> {
+    this.indexed = true
+    return Promise.resolve(true)
+  }
+  private pending(): KnowledgeChunkRow[] {
+    return this.chunks.filter((chunk) => !this.embedded.has(chunk.id))
+  }
+}
+
+class FakeProvider implements EmbeddingProvider {
+  calls = 0
+  inputs = 0
+  constructor(private fail = false) {}
+  embed(input: string | readonly string[]) {
+    this.calls += 1
+    const list = typeof input === 'string' ? [input] : input
+    this.inputs += list.length
+    if (this.fail) return Promise.resolve({ ok: false as const, reason: 'upstream' as const })
+    return Promise.resolve({
+      ok: true as const,
+      data: list.map((_, index) => ({ index, embedding: new Array(EMBEDDING_DIMENSIONS).fill(0.1) })),
+    })
+  }
+}
+
+function chunks(count: number, size = 800): KnowledgeChunkRow[] {
+  return Array.from({ length: count }, (_, i) => ({ id: `c${i}`, content: 'x'.repeat(size) }))
+}
+
+function setup(options: {
+  tier?: PlanQuotaTier
+  chunks: KnowledgeChunkRow[]
+  providerFails?: boolean
+  usageFails?: boolean
+}) {
+  const store = new FakeKnowledgeStore(options.chunks)
+  const provider = new FakeProvider(options.providerFails)
+  const usage = new MemoryUsageStore(options.tier ?? 'free', TODAY)
+  usage.failIncrement = options.usageFails ?? false
+  const handler = createEmbedChunksHandler({
+    authenticate: () => Promise.resolve({ id: USER }),
+    embeddingProvider: () => provider,
+    knowledgeStore: () => store,
+    usageStore: () => usage,
+    now: () => NOW,
+  })
+  const call = () => handler(new Request('https://edge/embed-chunks', {
+    method: 'POST',
+    body: JSON.stringify({ document_id: DOC }),
+  }))
+  return { store, provider, usage, call }
+}
+
+Deno.test('a free account over its weekly embedding budget is refused before any provider call', async () => {
+  const { provider, usage, call } = setup({ chunks: chunks(10) })
+  usage.rows.set('2026-09-26', EMBEDDING_QUOTA.free.limit - 5)
+  const response = await call()
+  assert(response.status === 429, `expected 429, got ${response.status}`)
+  const body = await response.json()
+  assert(body.error === 'quota_exceeded' && body.tier === 'free', 'quota body')
+  assert(provider.calls === 0, 'provider was called despite the exhausted budget')
+  assert((usage.rows.get(TODAY) ?? 0) === 0, 'refused request left units counted')
+})
+
+Deno.test('a document over the chunk cap is refused without spending', async () => {
+  const { provider, usage, call } = setup({ tier: 'enterprise', chunks: chunks(EMBEDDING_LIMITS.maxChunksPerDocument + 1, 10) })
+  const response = await call()
+  assert(response.status === 413, `expected 413, got ${response.status}`)
+  assert((await response.json()).error === 'knowledge_document_too_large', 'error code')
+  assert(provider.calls === 0 && !usage.rows.has(TODAY), 'spent on an oversized document')
+})
+
+Deno.test('a chunk over the size cap is refused without spending', async () => {
+  const list = chunks(3)
+  list[1] = { id: 'huge', content: 'x'.repeat(EMBEDDING_LIMITS.maxChunkChars + 1) }
+  const { provider, call } = setup({ tier: 'enterprise', chunks: list })
+  const response = await call()
+  assert(response.status === 413, `expected 413, got ${response.status}`)
+  const body = await response.json()
+  assert(body.error === 'knowledge_chunk_too_large' && body.chunk_id === 'huge', 'error body')
+  assert(provider.calls === 0, 'provider was called for an oversized chunk')
+})
+
+Deno.test('pending chunks are loaded with a bound', async () => {
+  const { store, call } = setup({ tier: 'pro', chunks: chunks(3) })
+  await call()
+  assert(store.listLimit === EMBEDDING_LIMITS.maxChunksPerDocument, `unbounded pending select (${store.listLimit})`)
+})
+
+Deno.test('a successful run indexes the document and counts the embedded volume', async () => {
+  const { store, provider, usage, call } = setup({ tier: 'pro', chunks: chunks(150) })
+  const response = await call()
+  assert(response.status === 200, `expected 200, got ${response.status}`)
+  const body = await response.json()
+  assert(body.embedded === 150 && body.total === 150 && body.indexed === true, 'response body')
+  assert(store.indexed, 'document not marked indexed')
+  assert(provider.calls === 2 && provider.inputs === 150, 'batching changed')
+  assert(usage.rows.get(TODAY) === 150, `usage ${usage.rows.get(TODAY)}`)
+})
+
+Deno.test('units of batches the provider rejected are refunded', async () => {
+  const { store, usage, call } = setup({ tier: 'pro', chunks: chunks(5), providerFails: true })
+  const response = await call()
+  assert(response.status === 502, `expected 502, got ${response.status}`)
+  const body = await response.json()
+  assert(body.error === 'embedding_failed' && body.remaining === 5, 'error body')
+  assert(!store.indexed, 'failed document marked indexed')
+  assert(usage.rows.get(TODAY) === 0, 'failed batch was charged')
+})
+
+Deno.test('the quota store being down fails closed', async () => {
+  const { provider, call } = setup({ tier: 'pro', chunks: chunks(2), usageFails: true })
+  const response = await call()
+  assert(response.status === 503, `expected 503, got ${response.status}`)
+  assert((await response.json()).error === 'quota_unavailable', 'error code')
+  assert(provider.calls === 0, 'provider called without a reservation')
+})
+
+Deno.test('another user\'s document stays hidden', async () => {
+  const { store, provider, call } = setup({ tier: 'pro', chunks: chunks(2) })
+  store.owner = 'someone-else'
+  const response = await call()
+  assert(response.status === 404 && provider.calls === 0, 'foreign document was embedded')
+})

server/supabase/functions/embed-chunks/handler.ts

@@ -0,0 +1,178 @@
+// server/supabase/functions/embed-chunks/handler.ts
+// Use case: embed a document's pending knowledge chunks.
+// IO is injected through ports (auth, knowledge store, embedding provider,
+// usage store) so the size and quota guards are testable without Supabase or
+// OpenAI. index.ts wires the real adapters.
+
+import { handleCorsPreflightRequest } from '../_shared/cors.ts'
+import { errorResponse, jsonResponse as json } from '../_shared/json-response.ts'
+import { type EmbeddingProvider, isEmbedding } from '../_shared/openai-embeddings.ts'
+import {
+  EMBEDDING_LIMITS,
+  type EmbeddingReservation,
+  type EmbeddingUsageStore,
+  embeddingUnits,
+  firstOversizedChunk,
+  planEmbeddingBatches,
+  quotaExceededBody,
+  refundEmbeddingUnits,
+  reserveEmbeddingUnits,
+} from '../_shared/embedding-quota.ts'
+
+const UUID_PATTERN = /^[0-9a-f]{8}-[0-9a-f]{4}-[1-5][0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}$/i
+
+export interface KnowledgeChunkRow {
+  id: string
+  content: string
+}
+
+/** Thrown by a KnowledgeIndexStore when the database call fails. */
+export class KnowledgeStorageError extends Error {
+  constructor() {
+    super('knowledge_storage_failed')
+  }
+}
+
+/** Port over knowledge_documents / knowledge_chunks. Read methods throw KnowledgeStorageError. */
+export interface KnowledgeIndexStore {
+  /** Owner of the document, or null when it does not exist. */
+  documentOwner(documentId: string): Promise<string | null>
+  /** Exact row count, or null when the database returned no count. */
+  countChunks(documentId: string, pendingOnly: boolean): Promise<number | null>
+  listPendingChunks(documentId: string, limit: number): Promise<KnowledgeChunkRow[]>
+  /** Returns false when the write failed. */
+  saveEmbedding(documentId: string, chunkId: string, embedding: number[]): Promise<boolean>
+  /** Clear the indexed flag. With a userId the update is scoped to that owner. Returns false on failure. */
+  clearIndexed(documentId: string, userId: string | null): Promise<boolean>
+  /** Set the indexed flag. Returns false when the write failed or matched no row. */
+  markIndexed(documentId: string, userId: string): Promise<boolean>
+}
+
+export interface EmbedChunksDeps {
+  authenticate(req: Request): Promise<{ id: string }>
+  /** null when the provider key is not configured. */
+  embeddingProvider(): EmbeddingProvider | null
+  knowledgeStore(): KnowledgeIndexStore
+  usageStore(): EmbeddingUsageStore
+  now?(): Date
+}
+
+async function embedBatch(
+  provider: EmbeddingProvider,
+  store: KnowledgeIndexStore,
+  documentId: string,
+  batch: KnowledgeChunkRow[],
+): Promise<number> {
+  const outcome = await provider.embed(batch.map((chunk) => chunk.content))
+  if (!outcome.ok || outcome.data.length !== batch.length) return 0
+
+  const received = new Set<number>()
+  let saved = 0
+  for (const item of outcome.data) {
+    if (
+      typeof item.index !== 'number'
+      || !Number.isInteger(item.index)
+      || item.index < 0
+      || item.index >= batch.length
+      || received.has(item.index)
+      || !isEmbedding(item.embedding)
+    ) {
+      continue
+    }
+    received.add(item.index)
+    if (await store.saveEmbedding(documentId, batch[item.index].id, item.embedding)) saved += 1
+  }
+  return saved
+}
+
+export function createEmbedChunksHandler(deps: EmbedChunksDeps): (req: Request) => Promise<Response> {
+  const now = () => deps.now?.() ?? new Date()
+
+  return async (req: Request): Promise<Response> => {
+    const preflight = handleCorsPreflightRequest(req)
+    if (preflight) return preflight
+    if (req.method !== 'POST') return json(405, { error: 'method_not_allowed' })
+
+    try {
+      const user = await deps.authenticate(req)
+      const body = await req.json().catch(() => null) as { document_id?: unknown } | null
+      if (!body || typeof body.document_id !== 'string' || !UUID_PATTERN.test(body.document_id)) {
+        return json(400, { error: 'invalid_document_id' })
+      }
+      const documentId = body.document_id
+
+      const provider = deps.embeddingProvider()
+      if (!provider) return json(503, { error: 'embedding_provider_unavailable' })
+
+      const store = deps.knowledgeStore()
+      const owner = await store.documentOwner(documentId)
+      if (owner !== user.id) return json(404, { error: 'knowledge_document_not_found' })
+
+      const total = await store.countChunks(documentId, false) ?? 0
+      if (total === 0) {
+        await store.clearIndexed(documentId, null)
+        return json(409, { error: 'knowledge_document_empty' })
+      }
+      if (total > EMBEDDING_LIMITS.maxChunksPerDocument) {
+        return json(413, {
+          error: 'knowledge_document_too_large',
+          total,
+          max_chunks: EMBEDDING_LIMITS.maxChunksPerDocument,
+        })
+      }
+
+      const chunks = await store.listPendingChunks(documentId, EMBEDDING_LIMITS.maxChunksPerDocument)
+      const oversized = firstOversizedChunk(chunks)
+      if (oversized !== -1) {
+        return json(413, {
+          error: 'knowledge_chunk_too_large',
+          chunk_id: chunks[oversized].id,
+          max_chars: EMBEDDING_LIMITS.maxChunkChars,
+        })
+      }
+
+      // Reserve the whole run before the first paid call; unspent units are refunded.
+      const usage = deps.usageStore()
+      const units = embeddingUnits(chunks.map((chunk) => chunk.content))
+      let reservation: EmbeddingReservation
+      try {
+        reservation = await reserveEmbeddingUnits(usage, user.id, units, now())
+      } catch {
+        return json(503, { error: 'quota_unavailable' })
+      }
+      if (!reservation.allowed) return json(429, quotaExceededBody(reservation, units))
+
+      let embedded = 0
+      let failed = 0
+      let unspentUnits = 0
+      for (const batch of planEmbeddingBatches(chunks)) {
+        let saved = 0
+        try {
+          saved = await embedBatch(provider, store, documentId, batch)
+        } catch {
+          saved = 0
+        }
+        embedded += saved
+        failed += batch.length - saved
+        if (saved === 0) unspentUnits += embeddingUnits(batch.map((chunk) => chunk.content))
+      }
+      await refundEmbeddingUnits(usage, user.id, reservation, unspentUnits, now())
+
+      const remaining = await store.countChunks(documentId, true) ?? total
+      if (failed > 0 || remaining > 0) {
+        if (!await store.clearIndexed(documentId, user.id)) {
+          return json(500, { error: 'knowledge_storage_failed' })
+        }
+        return json(502, { error: 'embedding_failed', embedded, remaining })
+      }
+
+      if (!await store.markIndexed(documentId, user.id)) {
+        return json(500, { error: 'knowledge_storage_failed' })
+      }
+      return json(200, { embedded, total, indexed: true })
+    } catch (error) {
+      if (error instanceof KnowledgeStorageError) return json(500, { error: 'knowledge_storage_failed' })
+      return errorResponse(error)
+    }
+  }
+}

server/supabase/functions/embed-chunks/index.ts

@@ -1,181 +1,26 @@
-import { corsHeaders, handleCorsPreflightRequest } from '../_shared/cors.ts'
+// server/supabase/functions/embed-chunks/index.ts
-import { requireUser, authErrorResponse, type AuthError } from '../_shared/auth.ts'
+// Composition root: wires the real adapters into the embed-chunks use case.
+
+import { requireUser } from '../_shared/auth.ts'
 import { createServiceRoleClient } from '../_shared/quota.ts'
 import { readProviderKey } from '../_shared/provider-key.ts'
+import { createOpenAIEmbeddingProvider } from '../_shared/openai-embeddings.ts'
+import { createSupabaseEmbeddingUsageStore } from '../_shared/embedding-quota.ts'
+import { createEmbedChunksHandler } from './handler.ts'
+import { createSupabaseKnowledgeStore } from './supabase-knowledge-store.ts'
 
-const EMBEDDING_DIMENSIONS = 1536
+let serviceClient: ReturnType<typeof createServiceRoleClient> | null = null
-const PROVIDER_TIMEOUT_MS = 45_000
+function service(): ReturnType<typeof createServiceRoleClient> {
-const UUID_PATTERN = /^[0-9a-f]{8}-[0-9a-f]{4}-[1-5][0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}$/i
+  serviceClient ??= createServiceRoleClient()
-
+  return serviceClient
-interface OpenAIEmbeddingResponse {
-  data?: Array<{ embedding?: unknown; index?: unknown }>
 }
 
-function json(status: number, body: Record<string, unknown>): Response {
+Deno.serve(createEmbedChunksHandler({
-  return new Response(JSON.stringify(body), {
+  authenticate: requireUser,
-    status,
+  embeddingProvider: () => {
-    headers: { ...corsHeaders, 'Content-Type': 'application/json' },
+    const key = readProviderKey('OPENAI_API_KEY')
-  })
+    return key ? createOpenAIEmbeddingProvider(key) : null
-}
-
-function validEmbedding(value: unknown): value is number[] {
-  return Array.isArray(value)
-    && value.length === EMBEDDING_DIMENSIONS
-    && value.every((entry) => typeof entry === 'number' && Number.isFinite(entry))
-}
-
-Deno.serve(async (req: Request) => {
-  const preflight = handleCorsPreflightRequest(req)
-  if (preflight) return preflight
-  if (req.method !== 'POST') return json(405, { error: 'method_not_allowed' })
-
-  try {
-    const user = await requireUser(req)
-    const body = await req.json().catch(() => null) as { document_id?: unknown } | null
-    if (!body || typeof body.document_id !== 'string' || !UUID_PATTERN.test(body.document_id)) {
-      return json(400, { error: 'invalid_document_id' })
-    }
-
-    const openaiKey = readProviderKey('OPENAI_API_KEY')
-    if (!openaiKey) return json(503, { error: 'embedding_provider_unavailable' })
-
-    const serviceClient = createServiceRoleClient()
-    const documentResult = await serviceClient
-      .from('knowledge_documents')
-      .select('id,user_id')
-      .eq('id', body.document_id)
-      .maybeSingle()
-    if (documentResult.error) return json(500, { error: 'knowledge_storage_failed' })
-    if (!documentResult.data || documentResult.data.user_id !== user.id) {
-      return json(404, { error: 'knowledge_document_not_found' })
-    }
-
-    const totalResult = await serviceClient
-      .from('knowledge_chunks')
-      .select('id', { count: 'exact', head: true })
-      .eq('document_id', body.document_id)
-    if (totalResult.error) return json(500, { error: 'knowledge_storage_failed' })
-    const total = totalResult.count ?? 0
-    if (total === 0) {
-      await serviceClient
-        .from('knowledge_documents')
-        .update({ indexed: false, indexed_at: null })
-        .eq('id', body.document_id)
-      return json(409, { error: 'knowledge_document_empty' })
-    }
-
-    const pendingResult = await serviceClient
-      .from('knowledge_chunks')
-      .select('id,content')
-      .eq('document_id', body.document_id)
-      .is('embedding', null)
-      .order('chunk_index', { ascending: true })
-    if (pendingResult.error) return json(500, { error: 'knowledge_storage_failed' })
-
-    const chunks = (pendingResult.data ?? []) as Array<{ id: string; content: string }>
-    let embedded = 0
-    let failed = 0
-
-    for (let offset = 0; offset < chunks.length; offset += 100) {
-      const batch = chunks.slice(offset, offset + 100)
-      try {
-        const response = await fetch('https://api.openai.com/v1/embeddings', {
-          method: 'POST',
-          headers: {
-            Authorization: `Bearer ${openaiKey}`,
-            'Content-Type': 'application/json',
   },
-          body: JSON.stringify({
+  knowledgeStore: () => createSupabaseKnowledgeStore(service()),
-            model: 'text-embedding-3-small',
+  usageStore: () => createSupabaseEmbeddingUsageStore(service()),
-            input: batch.map((chunk) => chunk.content),
+}))
-            dimensions: EMBEDDING_DIMENSIONS,
-          }),
-          signal: AbortSignal.timeout(PROVIDER_TIMEOUT_MS),
-        })
-        if (!response.ok) {
-          failed += batch.length
-          continue
-        }
-
-        const payload = await response.json().catch(() => null) as OpenAIEmbeddingResponse | null
-        if (!payload || !Array.isArray(payload.data) || payload.data.length !== batch.length) {
-          failed += batch.length
-          continue
-        }
-
-        const received = new Set<number>()
-        let batchSuccess = 0
-        for (const item of payload.data) {
-          if (
-            typeof item.index !== 'number'
-            || !Number.isInteger(item.index)
-            || item.index < 0
-            || item.index >= batch.length
-            || received.has(item.index)
-            || !validEmbedding(item.embedding)
-          ) {
-            continue
-          }
-          received.add(item.index)
-          const update = await serviceClient
-            .from('knowledge_chunks')
-            .update({ embedding: item.embedding })
-            .eq('id', batch[item.index].id)
-            .eq('document_id', body.document_id)
-          if (!update.error) {
-            embedded += 1
-            batchSuccess += 1
-          }
-        }
-        failed += batch.length - batchSuccess
-      } catch {
-        failed += batch.length
-      }
-    }
-
-    const remainingResult = await serviceClient
-      .from('knowledge_chunks')
-      .select('id', { count: 'exact', head: true })
-      .eq('document_id', body.document_id)
-      .is('embedding', null)
-    if (remainingResult.error) return json(500, { error: 'knowledge_storage_failed' })
-
-    const remaining = remainingResult.count ?? total
-    if (failed > 0 || remaining > 0) {
-      const rollback = await serviceClient
-        .from('knowledge_documents')
-        .update({ indexed: false, indexed_at: null })
-        .eq('id', body.document_id)
-        .eq('user_id', user.id)
-      if (rollback.error) return json(500, { error: 'knowledge_storage_failed' })
-      return json(502, {
-        error: 'embedding_failed',
-        embedded,
-        remaining,
-      })
-    }
-
-    const indexed = await serviceClient
-      .from('knowledge_documents')
-      .update({ indexed: true, indexed_at: new Date().toISOString() })
-      .eq('id', body.document_id)
-      .eq('user_id', user.id)
-      .select('id')
-      .maybeSingle()
-    if (indexed.error || !indexed.data) return json(500, { error: 'knowledge_storage_failed' })
-
-    return json(200, { embedded, total, indexed: true })
-  } catch (error) {
-    if (
-      error
-      && typeof error === 'object'
-      && 'status' in error
-      && (error.status === 401 || error.status === 403)
-      && 'message' in error
-      && typeof error.message === 'string'
-    ) {
-      return authErrorResponse(error as AuthError, corsHeaders)
-    }
-    return json(500, { error: 'internal_error' })
-  }
-})

server/supabase/functions/embed-chunks/supabase-knowledge-store.ts

@@ -0,0 +1,70 @@
+// server/supabase/functions/embed-chunks/supabase-knowledge-store.ts
+// KnowledgeIndexStore adapter over the service-role Supabase client.
+
+import type { createClient } from '@supabase/supabase-js'
+import { type KnowledgeChunkRow, type KnowledgeIndexStore, KnowledgeStorageError } from './handler.ts'
+
+type SupabaseClient = ReturnType<typeof createClient>
+
+export function createSupabaseKnowledgeStore(client: SupabaseClient): KnowledgeIndexStore {
+  return {
+    async documentOwner(documentId) {
+      const { data, error } = await client
+        .from('knowledge_documents')
+        .select('id,user_id')
+        .eq('id', documentId)
+        .maybeSingle()
+      if (error) throw new KnowledgeStorageError()
+      const row = data as { user_id?: unknown } | null
+      return row && typeof row.user_id === 'string' ? row.user_id : null
+    },
+    async countChunks(documentId, pendingOnly) {
+      let query = client
+        .from('knowledge_chunks')
+        .select('id', { count: 'exact', head: true })
+        .eq('document_id', documentId)
+      if (pendingOnly) query = query.is('embedding', null)
+      const { count, error } = await query
+      if (error) throw new KnowledgeStorageError()
+      return count
+    },
+    async listPendingChunks(documentId, limit) {
+      const { data, error } = await client
+        .from('knowledge_chunks')
+        .select('id,content')
+        .eq('document_id', documentId)
+        .is('embedding', null)
+        .order('chunk_index', { ascending: true })
+        .limit(limit)
+      if (error) throw new KnowledgeStorageError()
+      return (data ?? []) as KnowledgeChunkRow[]
+    },
+    async saveEmbedding(documentId, chunkId, embedding) {
+      const { error } = await client
+        .from('knowledge_chunks')
+        .update({ embedding })
+        .eq('id', chunkId)
+        .eq('document_id', documentId)
+      return !error
+    },
+    async clearIndexed(documentId, userId) {
+      let query = client
+        .from('knowledge_documents')
+        .update({ indexed: false, indexed_at: null })
+        .eq('id', documentId)
+      if (userId !== null) query = query.eq('user_id', userId)
+      const { error } = await query
+      return !error
+    },
+    async markIndexed(documentId, userId) {
+      const { data, error } = await client
+        .from('knowledge_documents')
+        .update({ indexed: true, indexed_at: new Date().toISOString() })
+        .eq('id', documentId)
+        .eq('user_id', userId)
+        .select('id')
+        .maybeSingle()
+      return !error && Boolean(data)
+    },
+  }
+}

server/supabase/functions/search-knowledge/handler.test.ts

@@ -0,0 +1,94 @@
+// Regression tests for the search-knowledge cost guard (redteam r1 #18): each
+// query used to be embedded by the paid provider with no per-account counter.
+
+import { type ChunkSearcher, createSearchKnowledgeHandler } from './handler.ts'
+import { EMBEDDING_DIMENSIONS, type EmbeddingOutcome, type EmbeddingProvider } from '../_shared/openai-embeddings.ts'
+import { EMBEDDING_QUOTA } from '../_shared/embedding-quota.ts'
+import { MemoryUsageStore } from '../_shared/embedding-quota.fake.ts'
+import type { PlanQuotaTier } from '../_shared/core-contract.generated.ts'
+
+function assert(condition: boolean, message: string): asserts condition {
+  if (!condition) throw new Error(message)
+}
+
+const NOW = new Date('2026-09-28T12:00:00Z')
+const TODAY = '2026-09-28'
+
+function setup(options: { tier?: PlanQuotaTier; outcome?: EmbeddingOutcome; usageFails?: boolean } = {}) {
+  let providerCalls = 0
+  let searches = 0
+  const provider: EmbeddingProvider = {
+    embed() {
+      providerCalls += 1
+      return Promise.resolve(options.outcome ?? {
+        ok: true,
+        data: [{ index: 0, embedding: new Array(EMBEDDING_DIMENSIONS).fill(0.2) }],
+      })
+    },
+  }
+  const searchChunks: ChunkSearcher = (_req, _embedding, count) => {
+    searches += 1
+    return Promise.resolve({ ok: true, results: [{ count }] })
+  }
+  const usage = new MemoryUsageStore(options.tier ?? 'free', TODAY)
+  usage.failIncrement = options.usageFails ?? false
+  const handler = createSearchKnowledgeHandler({
+    authenticate: () => Promise.resolve({ id: 'user-1' }),
+    embeddingProvider: () => provider,
+    usageStore: () => usage,
+    searchChunks,
+    now: () => NOW,
+  })
+  const call = (query: string) => handler(new Request('https://edge/search-knowledge', {
+    method: 'POST',
+    body: JSON.stringify({ query, count: 3 }),
+  }))
+  return { usage, call, providerCalls: () => providerCalls, searches: () => searches }
+}
+
+Deno.test('a query past the weekly budget is refused before the provider is called', async () => {
+  const ctx = setup()
+  ctx.usage.rows.set('2026-09-27', EMBEDDING_QUOTA.free.limit)
+  const response = await ctx.call('meeting notes')
+  assert(response.status === 429, `expected 429, got ${response.status}`)
+  assert((await response.json()).error === 'quota_exceeded', 'error code')
+  assert(ctx.providerCalls() === 0 && ctx.searches() === 0, 'provider or search ran')
+})
+
+Deno.test('each query is counted by its length', async () => {
+  const ctx = setup({ tier: 'pro' })
+  const response = await ctx.call('x'.repeat(2_500))
+  assert(response.status === 200, `expected 200, got ${response.status}`)
+  const body = await response.json()
+  assert(Array.isArray(body.results) && body.results[0].count === 3, 'results')
+  assert(ctx.usage.rows.get(TODAY) === 3, `usage ${ctx.usage.rows.get(TODAY)}`)
+})
+
+Deno.test('an upstream failure refunds the query', async () => {
+  const ctx = setup({ tier: 'pro', outcome: { ok: false, reason: 'upstream' } })
+  const response = await ctx.call('hello')
+  assert(response.status === 502, `expected 502, got ${response.status}`)
+  assert((await response.json()).error === 'embedding_upstream_failed', 'error code')
+  assert(ctx.usage.rows.get(TODAY) === 0, 'failed query charged')
+})
+
+Deno.test('an invalid provider body keeps its error code', async () => {
+  const ctx = setup({ tier: 'pro', outcome: { ok: true, data: [{ embedding: [1, 2] }] } })
+  const response = await ctx.call('hello')
+  assert(response.status === 502, `expected 502, got ${response.status}`)
+  assert((await response.json()).error === 'embedding_response_invalid', 'error code')
+})
+
+Deno.test('the quota store being down fails closed', async () => {
+  const ctx = setup({ tier: 'pro', usageFails: true })
+  const response = await ctx.call('hello')
+  assert(response.status === 503, `expected 503, got ${response.status}`)
+  assert(ctx.providerCalls() === 0, 'provider called without a reservation')
+})
+
+Deno.test('input validation is unchanged', async () => {
+  const ctx = setup({ tier: 'pro' })
+  assert((await ctx.call('')).status === 400, 'empty query accepted')
+  assert((await ctx.call('x'.repeat(4_001))).status === 400, 'long query accepted')
+  assert(ctx.providerCalls() === 0, 'provider called for invalid input')
+})

server/supabase/functions/search-knowledge/handler.ts

@@ -0,0 +1,91 @@
+// server/supabase/functions/search-knowledge/handler.ts
+// Use case: embed a search query and return the caller's matching chunks.
+// IO is injected through ports so the quota guard is testable offline.
+
+import { handleCorsPreflightRequest } from '../_shared/cors.ts'
+import { errorResponse, jsonResponse as json } from '../_shared/json-response.ts'
+import { type EmbeddingProvider, isEmbedding } from '../_shared/openai-embeddings.ts'
+import {
+  EMBEDDING_LIMITS,
+  type EmbeddingReservation,
+  type EmbeddingUsageStore,
+  embeddingUnitsFor,
+  quotaExceededBody,
+  refundEmbeddingUnits,
+  reserveEmbeddingUnits,
+} from '../_shared/embedding-quota.ts'
+
+export const DEFAULT_MATCH_COUNT = 5
+export const MAX_MATCH_COUNT = 20
+export const SIMILARITY_THRESHOLD = 0.5
+
+export type ChunkSearchOutcome =
+  | { ok: true; results: unknown[] }
+  | { ok: false; reason: 'unavailable' | 'failed' }
+
+/** Port: similarity search run with the caller's own JWT (RLS scopes the rows). */
+export type ChunkSearcher = (req: Request, embedding: number[], count: number) => Promise<ChunkSearchOutcome>
+
+export interface SearchKnowledgeDeps {
+  authenticate(req: Request): Promise<{ id: string }>
+  /** null when the provider key is not configured. */
+  embeddingProvider(): EmbeddingProvider | null
+  usageStore(): EmbeddingUsageStore
+  searchChunks: ChunkSearcher
+  now?(): Date
+}
+
+export function createSearchKnowledgeHandler(deps: SearchKnowledgeDeps): (req: Request) => Promise<Response> {
+  const now = () => deps.now?.() ?? new Date()
+
+  return async (req: Request): Promise<Response> => {
+    const preflight = handleCorsPreflightRequest(req)
+    if (preflight) return preflight
+    if (req.method !== 'POST') return json(405, { error: 'method_not_allowed' })
+
+    try {
+      const user = await deps.authenticate(req)
+      const body = await req.json().catch(() => null) as {
+        query?: unknown
+        count?: unknown
+      } | null
+      const query = typeof body?.query === 'string' ? body.query.trim() : ''
+      const count = body?.count === undefined ? DEFAULT_MATCH_COUNT : body.count
+      if (!query || query.length > EMBEDDING_LIMITS.maxQueryChars) return json(400, { error: 'invalid_query' })
+      if (typeof count !== 'number' || !Number.isInteger(count) || count < 1 || count > MAX_MATCH_COUNT) {
+        return json(400, { error: 'invalid_count' })
+      }
+
+      const provider = deps.embeddingProvider()
+      if (!provider) return json(503, { error: 'embedding_provider_unavailable' })
+
+      const usage = deps.usageStore()
+      const units = embeddingUnitsFor(query)
+      let reservation: EmbeddingReservation
+      try {
+        reservation = await reserveEmbeddingUnits(usage, user.id, units, now())
+      } catch {
+        return json(503, { error: 'quota_unavailable' })
+      }
+      if (!reservation.allowed) return json(429, quotaExceededBody(reservation, units))
+
+      const outcome = await provider.embed(query)
+      if (!outcome.ok && outcome.reason === 'upstream') {
+        await refundEmbeddingUnits(usage, user.id, reservation, units, now())
+        return json(502, { error: 'embedding_upstream_failed' })
+      }
+      const queryEmbedding = outcome.ok ? outcome.data[0]?.embedding : undefined
+      if (!isEmbedding(queryEmbedding)) return json(502, { error: 'embedding_response_invalid' })
+
+      const result = await deps.searchChunks(req, queryEmbedding, count)
+      if (!result.ok) {
+        return result.reason === 'unavailable'
+          ? json(503, { error: 'knowledge_storage_unavailable' })
+          : json(500, { error: 'knowledge_search_failed' })
+      }
+      return json(200, { results: result.results })
+    } catch (error) {
+      return errorResponse(error)
+    }
+  }
+}

server/supabase/functions/search-knowledge/index.ts

@@ -1,100 +1,46 @@
-// deno-lint-ignore no-import-prefix
+// server/supabase/functions/search-knowledge/index.ts
-import { createClient } from 'https://esm.sh/@supabase/supabase-js@2.39.7'
+// Composition root: wires the real adapters into the search-knowledge use case.
-import { corsHeaders, handleCorsPreflightRequest } from '../_shared/cors.ts'
+
-import { requireUser, authErrorResponse, type AuthError } from '../_shared/auth.ts'
+import { createClient } from '@supabase/supabase-js'
+import { requireUser } from '../_shared/auth.ts'
+import { createServiceRoleClient } from '../_shared/quota.ts'
 import { readProviderKey } from '../_shared/provider-key.ts'
+import { createOpenAIEmbeddingProvider } from '../_shared/openai-embeddings.ts'
+import { createSupabaseEmbeddingUsageStore } from '../_shared/embedding-quota.ts'
+import { type ChunkSearcher, createSearchKnowledgeHandler, SIMILARITY_THRESHOLD } from './handler.ts'
 
-const EMBEDDING_DIMENSIONS = 1536
+let serviceClient: ReturnType<typeof createServiceRoleClient> | null = null
-const PROVIDER_TIMEOUT_MS = 45_000
+function service(): ReturnType<typeof createServiceRoleClient> {
-
+  serviceClient ??= createServiceRoleClient()
-function json(status: number, body: Record<string, unknown>): Response {
+  return serviceClient
-  return new Response(JSON.stringify(body), {
-    status,
-    headers: { ...corsHeaders, 'Content-Type': 'application/json' },
-  })
 }
 
-function isEmbedding(value: unknown): value is number[] {
+/** match_knowledge_chunks runs as the caller so RLS limits results to their own chunks. */
-  return Array.isArray(value)
+const searchChunks: ChunkSearcher = async (req, embedding, count) => {
-    && value.length === EMBEDDING_DIMENSIONS
-    && value.every((entry) => typeof entry === 'number' && Number.isFinite(entry))
-}
-
-Deno.serve(async (req: Request) => {
-  const preflight = handleCorsPreflightRequest(req)
-  if (preflight) return preflight
-  if (req.method !== 'POST') return json(405, { error: 'method_not_allowed' })
-
-  try {
-    await requireUser(req)
-    const body = await req.json().catch(() => null) as {
-      query?: unknown
-      count?: unknown
-    } | null
-    const query = typeof body?.query === 'string' ? body.query.trim() : ''
-    const count = body?.count === undefined ? 5 : body.count
-    if (!query || query.length > 4_000) return json(400, { error: 'invalid_query' })
-    if (typeof count !== 'number' || !Number.isInteger(count) || count < 1 || count > 20) {
-      return json(400, { error: 'invalid_count' })
-    }
-
-    const openaiKey = readProviderKey('OPENAI_API_KEY')
-    if (!openaiKey) return json(503, { error: 'embedding_provider_unavailable' })
-
-    let embeddingResponse: Response
-    try {
-      embeddingResponse = await fetch('https://api.openai.com/v1/embeddings', {
-        method: 'POST',
-        headers: {
-          Authorization: `Bearer ${openaiKey}`,
-          'Content-Type': 'application/json',
-        },
-        body: JSON.stringify({
-          model: 'text-embedding-3-small',
-          input: query,
-          dimensions: EMBEDDING_DIMENSIONS,
-        }),
-        signal: AbortSignal.timeout(PROVIDER_TIMEOUT_MS),
-      })
-    } catch {
-      return json(502, { error: 'embedding_upstream_failed' })
-    }
-    if (!embeddingResponse.ok) return json(502, { error: 'embedding_upstream_failed' })
-
-    const embeddingPayload = await embeddingResponse.json().catch(() => null) as {
-      data?: Array<{ embedding?: unknown }>
-    } | null
-    const queryEmbedding = embeddingPayload?.data?.[0]?.embedding
-    if (!isEmbedding(queryEmbedding)) return json(502, { error: 'embedding_response_invalid' })
-
   const authHeader = req.headers.get('Authorization') ?? ''
   const supabaseUrl = Deno.env.get('SUPABASE_URL') ?? ''
   const anonKey = Deno.env.get('SUPABASE_ANON_KEY') ?? ''
-    if (!supabaseUrl || !anonKey) return json(503, { error: 'knowledge_storage_unavailable' })
+  if (!supabaseUrl || !anonKey) return { ok: false, reason: 'unavailable' }
 
   const userClient = createClient(supabaseUrl, anonKey, {
     global: { headers: { Authorization: authHeader } },
     auth: { persistSession: false, autoRefreshToken: false },
   })
   const result = await userClient.rpc('match_knowledge_chunks', {
-      query_embedding: queryEmbedding,
+    query_embedding: embedding,
     match_count: count,
-      similarity_threshold: 0.5,
+    similarity_threshold: SIMILARITY_THRESHOLD,
   })
-    if (result.error) return json(500, { error: 'knowledge_search_failed' })
+  if (result.error) return { ok: false, reason: 'failed' }
+  return { ok: true, results: Array.isArray(result.data) ? result.data : [] }
+}
 
-    return json(200, { results: result.data ?? [] })
+Deno.serve(createSearchKnowledgeHandler({
-  } catch (error) {
+  authenticate: requireUser,
-    if (
+  embeddingProvider: () => {
-      error
+    const key = readProviderKey('OPENAI_API_KEY')
-      && typeof error === 'object'
+    return key ? createOpenAIEmbeddingProvider(key) : null
-      && 'status' in error
+  },
-      && (error.status === 401 || error.status === 403)
+  usageStore: () => createSupabaseEmbeddingUsageStore(service()),
-      && 'message' in error
+  searchChunks,
-      && typeof error.message === 'string'
+}))
-    ) {
-      return authErrorResponse(error as AuthError, corsHeaders)
-    }
-    return json(500, { error: 'internal_error' })
-  }
-})