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

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