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

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
-import { createClient } from 'https://esm.sh/@supabase/supabase-js@2.39.7'
-import { corsHeaders, handleCorsPreflightRequest } from '../_shared/cors.ts'
-import { requireUser, authErrorResponse, type AuthError } from '../_shared/auth.ts'
+// server/supabase/functions/search-knowledge/index.ts
+// Composition root: wires the real adapters into the search-knowledge use case.
+
+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
-const PROVIDER_TIMEOUT_MS = 45_000
+let serviceClient: ReturnType<typeof createServiceRoleClient> | null = null
+function service(): ReturnType<typeof createServiceRoleClient> {
+  serviceClient ??= createServiceRoleClient()
+  return serviceClient
+}
 
-function json(status: number, body: Record<string, unknown>): Response {
-  return new Response(JSON.stringify(body), {
-    status,
-    headers: { ...corsHeaders, 'Content-Type': 'application/json' },
+/** match_knowledge_chunks runs as the caller so RLS limits results to their own chunks. */
+const searchChunks: ChunkSearcher = async (req, embedding, count) => {
+  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 { 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: embedding,
+    match_count: count,
+    similarity_threshold: SIMILARITY_THRESHOLD,
+  })
+  if (result.error) return { ok: false, reason: 'failed' }
+  return { ok: true, results: Array.isArray(result.data) ? result.data : [] }
 }
 
-function isEmbedding(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 {
-    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' })
-
-    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,
-      match_count: count,
-      similarity_threshold: 0.5,
-    })
-    if (result.error) return json(500, { error: 'knowledge_search_failed' })
-
-    return json(200, { results: result.data ?? [] })
-  } 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' })
-  }
-})
+Deno.serve(createSearchKnowledgeHandler({
+  authenticate: requireUser,
+  embeddingProvider: () => {
+    const key = readProviderKey('OPENAI_API_KEY')
+    return key ? createOpenAIEmbeddingProvider(key) : null
+  },
+  usageStore: () => createSupabaseEmbeddingUsageStore(service()),
+  searchChunks,
+}))