From 1afaea721462166bbfeebd83fc99fd8640696c6a Mon Sep 17 00:00:00 2001 From: Yun Chan Date: Mon, 28 Sep 2026 00:53:48 +0900 Subject: [PATCH] fix(edge): meter and size-cap knowledge embeddings in embed-chunks and search-knowledge --- .../functions/_shared/embedding-quota.fake.ts | 25 ++ .../functions/_shared/embedding-quota.test.ts | 85 ++++++ .../functions/_shared/embedding-quota.ts | 250 ++++++++++++++++++ .../functions/_shared/json-response.ts | 30 +++ .../functions/_shared/openai-embeddings.ts | 61 +++++ .../functions/embed-chunks/handler.test.ts | 166 ++++++++++++ .../functions/embed-chunks/handler.ts | 178 +++++++++++++ .../supabase/functions/embed-chunks/index.ts | 197 ++------------ .../embed-chunks/supabase-knowledge-store.ts | 70 +++++ .../search-knowledge/handler.test.ts | 94 +++++++ .../functions/search-knowledge/handler.ts | 91 +++++++ .../functions/search-knowledge/index.ts | 134 +++------- 12 files changed, 1111 insertions(+), 270 deletions(-) create mode 100644 server/supabase/functions/_shared/embedding-quota.fake.ts create mode 100644 server/supabase/functions/_shared/embedding-quota.test.ts create mode 100644 server/supabase/functions/_shared/embedding-quota.ts create mode 100644 server/supabase/functions/_shared/json-response.ts create mode 100644 server/supabase/functions/_shared/openai-embeddings.ts create mode 100644 server/supabase/functions/embed-chunks/handler.test.ts create mode 100644 server/supabase/functions/embed-chunks/handler.ts create mode 100644 server/supabase/functions/embed-chunks/supabase-knowledge-store.ts create mode 100644 server/supabase/functions/search-knowledge/handler.test.ts create mode 100644 server/supabase/functions/search-knowledge/handler.ts diff --git a/server/supabase/functions/_shared/embedding-quota.fake.ts b/server/supabase/functions/_shared/embedding-quota.fake.ts new file mode 100644 index 0000000..84959b6 --- /dev/null +++ b/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() + failIncrement = false + constructor(public tier: PlanQuotaTier, public today: string) {} + readTier(): Promise { + return Promise.resolve(this.tier) + } + increment(_userId: string, feature: string, amount: number): Promise { + 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 { + let sum = 0 + for (const [date, count] of this.rows) if (date >= fromDate && date < beforeDate) sum += count + return Promise.resolve(sum) + } +} diff --git a/server/supabase/functions/_shared/embedding-quota.test.ts b/server/supabase/functions/_shared/embedding-quota.test.ts new file mode 100644 index 0000000..480badf --- /dev/null +++ b/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') +}) diff --git a/server/supabase/functions/_shared/embedding-quota.ts b/server/supabase/functions/_shared/embedding-quota.ts new file mode 100644 index 0000000..4713c02 --- /dev/null +++ b/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> = 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(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 + /** Atomically add `amount` (may be negative) to today's counter; returns the new count. */ + increment(userId: string, feature: string, amount: number): Promise + /** Sum of counters with fromDate <= date < beforeDate. */ + sumBetween(userId: string, feature: string, fromDate: string, beforeDate: string): Promise +} + +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 { + 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 { + 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 { + 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 + +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) + }, + } +} diff --git a/server/supabase/functions/_shared/json-response.ts b/server/supabase/functions/_shared/json-response.ts new file mode 100644 index 0000000..aebd9f8 --- /dev/null +++ b/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): 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' }) +} diff --git a/server/supabase/functions/_shared/openai-embeddings.ts b/server/supabase/functions/_shared/openai-embeddings.ts new file mode 100644 index 0000000..b9c2172 --- /dev/null +++ b/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 +} + +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[] } + }, + } +} diff --git a/server/supabase/functions/embed-chunks/handler.test.ts b/server/supabase/functions/embed-chunks/handler.test.ts new file mode 100644 index 0000000..45d9ede --- /dev/null +++ b/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() + indexed = false + listLimit = -1 + constructor(public chunks: KnowledgeChunkRow[], public owner: string | null = USER) {} + documentOwner(): Promise { + return Promise.resolve(this.owner) + } + countChunks(_doc: string, pendingOnly: boolean): Promise { + return Promise.resolve(pendingOnly ? this.pending().length : this.chunks.length) + } + listPendingChunks(_doc: string, limit: number): Promise { + this.listLimit = limit + return Promise.resolve(this.pending().slice(0, limit)) + } + saveEmbedding(_doc: string, chunkId: string): Promise { + this.embedded.add(chunkId) + return Promise.resolve(true) + } + clearIndexed(): Promise { + this.indexed = false + return Promise.resolve(true) + } + markIndexed(): Promise { + 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') +}) diff --git a/server/supabase/functions/embed-chunks/handler.ts b/server/supabase/functions/embed-chunks/handler.ts new file mode 100644 index 0000000..e389221 --- /dev/null +++ b/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 + /** Exact row count, or null when the database returned no count. */ + countChunks(documentId: string, pendingOnly: boolean): Promise + listPendingChunks(documentId: string, limit: number): Promise + /** Returns false when the write failed. */ + saveEmbedding(documentId: string, chunkId: string, embedding: number[]): Promise + /** 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 + /** Set the indexed flag. Returns false when the write failed or matched no row. */ + markIndexed(documentId: string, userId: string): Promise +} + +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 { + const outcome = await provider.embed(batch.map((chunk) => chunk.content)) + if (!outcome.ok || outcome.data.length !== batch.length) return 0 + + const received = new Set() + 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 { + const now = () => deps.now?.() ?? new Date() + + return async (req: Request): Promise => { + 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) + } + } +} diff --git a/server/supabase/functions/embed-chunks/index.ts b/server/supabase/functions/embed-chunks/index.ts index 1860ad9..7cb1439 100644 --- a/server/supabase/functions/embed-chunks/index.ts +++ b/server/supabase/functions/embed-chunks/index.ts @@ -1,181 +1,26 @@ -import { corsHeaders, handleCorsPreflightRequest } from '../_shared/cors.ts' -import { requireUser, authErrorResponse, type AuthError } from '../_shared/auth.ts' +// server/supabase/functions/embed-chunks/index.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 -const PROVIDER_TIMEOUT_MS = 45_000 -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 - -interface OpenAIEmbeddingResponse { - data?: Array<{ embedding?: unknown; index?: unknown }> +let serviceClient: ReturnType | null = null +function service(): ReturnType { + serviceClient ??= createServiceRoleClient() + return serviceClient } -function json(status: number, body: Record): Response { - return new Response(JSON.stringify(body), { - status, - headers: { ...corsHeaders, 'Content-Type': 'application/json' }, - }) -} - -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({ - model: 'text-embedding-3-small', - 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() - 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' }) - } -}) +Deno.serve(createEmbedChunksHandler({ + authenticate: requireUser, + embeddingProvider: () => { + const key = readProviderKey('OPENAI_API_KEY') + return key ? createOpenAIEmbeddingProvider(key) : null + }, + knowledgeStore: () => createSupabaseKnowledgeStore(service()), + usageStore: () => createSupabaseEmbeddingUsageStore(service()), +})) diff --git a/server/supabase/functions/embed-chunks/supabase-knowledge-store.ts b/server/supabase/functions/embed-chunks/supabase-knowledge-store.ts new file mode 100644 index 0000000..62b60d3 --- /dev/null +++ b/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 + +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) + }, + } +} diff --git a/server/supabase/functions/search-knowledge/handler.test.ts b/server/supabase/functions/search-knowledge/handler.test.ts new file mode 100644 index 0000000..2f51bd6 --- /dev/null +++ b/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') +}) diff --git a/server/supabase/functions/search-knowledge/handler.ts b/server/supabase/functions/search-knowledge/handler.ts new file mode 100644 index 0000000..5fc4063 --- /dev/null +++ b/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 + +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 { + const now = () => deps.now?.() ?? new Date() + + return async (req: Request): Promise => { + 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) + } + } +} diff --git a/server/supabase/functions/search-knowledge/index.ts b/server/supabase/functions/search-knowledge/index.ts index b8728c2..eef1695 100644 --- a/server/supabase/functions/search-knowledge/index.ts +++ b/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 | null = null +function service(): ReturnType { + serviceClient ??= createServiceRoleClient() + return serviceClient +} -function json(status: number, body: Record): 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, +}))