94 lines
4.1 KiB
TypeScript
94 lines
4.1 KiB
TypeScript
// 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')
|
|
})
|