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

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')
+})