diff --git a/src/app/api/answer/route.ts b/src/app/api/answer/route.ts index ae6af8b56..b63f20705 100644 --- a/src/app/api/answer/route.ts +++ b/src/app/api/answer/route.ts @@ -83,23 +83,27 @@ export async function POST(request: Request) { const access = await publicAccessContext(request, supabase); const publicOnly = !access.authenticated && !isLocalNoAuthMode(); - const rateLimit = await consumeSubjectApiRateLimit({ - supabase, - subject: access.rateLimitSubject, - bucket: "answer", - allowInMemoryFallbackOnUnavailable: allowRateLimitInMemoryFallbackOnUnavailable(), - }); + // Independent given `access`: the rate-limit consume and the read-only + // scope resolution overlap instead of running serially. The limit still + // rejects before any retrieval or generation starts. + const [rateLimit, scope] = await Promise.all([ + consumeSubjectApiRateLimit({ + supabase, + subject: access.rateLimitSubject, + bucket: "answer", + allowInMemoryFallbackOnUnavailable: allowRateLimitInMemoryFallbackOnUnavailable(), + }), + resolveSearchScope({ + supabase, + ownerId: access.ownerId, + publicOnly, + documentIds: answerBody.documentIds ?? (answerBody.documentId ? [answerBody.documentId] : undefined), + filters: answerBody.filters, + }), + ]); if (rateLimit.limited) { return rateLimitJsonResponse("Too many answer requests. Retry shortly.", rateLimit); } - - const scope = await resolveSearchScope({ - supabase, - ownerId: access.ownerId, - publicOnly, - documentIds: answerBody.documentIds ?? (answerBody.documentId ? [answerBody.documentId] : undefined), - filters: answerBody.filters, - }); if (scope.documentIds?.length === 0) { return NextResponse.json({ answer: diff --git a/src/app/api/answer/stream/route.ts b/src/app/api/answer/stream/route.ts index ac64e8b27..d3dbc6732 100644 --- a/src/app/api/answer/stream/route.ts +++ b/src/app/api/answer/stream/route.ts @@ -137,7 +137,13 @@ function buildDemoStreamAnswer(body: AnswerBody, fallbackReason?: string) { }; } -function streamAnswer(body: AnswerBody, ownerId?: string, signal?: AbortSignal, publicOnly = false) { +function streamAnswer( + body: AnswerBody, + ownerId?: string, + signal?: AbortSignal, + publicOnly = false, + scopePromise?: ReturnType, +) { const encoder = new TextEncoder(); return new Response( @@ -156,13 +162,14 @@ function streamAnswer(body: AnswerBody, ownerId?: string, signal?: AbortSignal, send("progress", { stage: "retrieving", message: "Searching indexed documents." }); const scope = isDemoMode() ? null - : await resolveSearchScope({ - supabase: createAdminClient(), - ownerId, - publicOnly, - documentIds: body.documentIds ?? (body.documentId ? [body.documentId] : undefined), - filters: body.filters, - }); + : await (scopePromise ?? + resolveSearchScope({ + supabase: createAdminClient(), + ownerId, + publicOnly, + documentIds: body.documentIds ?? (body.documentId ? [body.documentId] : undefined), + filters: body.filters, + })); if (scope?.documentIds?.length === 0) { send("final", { answer: @@ -265,6 +272,19 @@ export async function POST(request: Request) { const access = await publicAccessContext(request, supabase); const publicOnly = !access.authenticated && !isLocalNoAuthMode(); + // Kick off the read-only scope resolution alongside the rate-limit consume; + // the stream awaits it (and the limit still rejects before retrieval). The + // detached catch marks a pre-await rejection handled — the stream's own + // await still observes and reports the real error. + const scopePromise = resolveSearchScope({ + supabase, + ownerId: access.ownerId, + publicOnly, + documentIds: body.documentIds ?? (body.documentId ? [body.documentId] : undefined), + filters: body.filters, + }); + void scopePromise.catch(() => undefined); + const rateLimit = await consumeSubjectApiRateLimit({ supabase, subject: access.rateLimitSubject, @@ -273,7 +293,7 @@ export async function POST(request: Request) { }); if (rateLimit.limited) return rateLimitStream(rateLimit); - return streamAnswer(body, access.ownerId, request.signal, publicOnly); + return streamAnswer(body, access.ownerId, request.signal, publicOnly, scopePromise); } catch (error) { if (error instanceof AuthenticationError) { return unauthorizedResponse(error); diff --git a/src/lib/rag-cache.ts b/src/lib/rag-cache.ts index 1c36de32c..52cdde0d0 100644 --- a/src/lib/rag-cache.ts +++ b/src/lib/rag-cache.ts @@ -249,7 +249,10 @@ type SharedCacheMissReason = | "expired" | "indexing_version_mismatch" | "dependency_version_mismatch" - | "unknown_filter_miss"; + | "unknown_filter_miss" + // Placeholder while the off-path classification probe is still in flight; + // telemetry is patched with the real reason when it resolves (getSharedCachedSearch). + | "pending_probe"; function sharedCacheSelector( supabase: ReturnType, @@ -316,14 +319,16 @@ export async function getSharedCachedSearch( queryVariants: string[] = [], ): Promise< | { kind: "hit"; results: SearchResult[]; telemetry: SearchTelemetry } - | { kind: "miss"; reason: SharedCacheMissReason } + // `deferredReason` classifies a generic miss (no_entry/expired/version mismatch) + // via a follow-up metadata query that must NOT block retrieval: the caller + // starts with `reason` ("pending_probe") and patches telemetry when it lands. + | { kind: "miss"; reason: SharedCacheMissReason; deferredReason?: Promise } | null > { if (args.skipCache || env.RAG_SEARCH_CACHE_TTL_MS <= 0) return null; const normalizedQuery = retrievalPlanCacheQuery(args, queryClass, queryVariants); const indexingVersion = await cacheIndexingVersion(args); - async function probeSharedCacheMissReason(reasonFromLookup?: SharedCacheMissReason): Promise { - if (reasonFromLookup) return reasonFromLookup; + async function probeSharedCacheMissReason(): Promise { try { const supabase = createAdminClient(); let probeQuery = supabase @@ -361,11 +366,15 @@ export async function getSharedCachedSearch( indexingVersion, normalizedQuery, ).maybeSingle(); - if (error) return { kind: "miss", reason: await probeSharedCacheMissReason("cache_lookup_error") }; - if (!data?.payload) return { kind: "miss", reason: await probeSharedCacheMissReason() }; + if (error) return { kind: "miss", reason: "cache_lookup_error" }; + if (!data?.payload) { + // Do not await the classification probe on the hot path — it costs a full + // extra rag_response_cache round trip before retrieval can start. + return { kind: "miss", reason: "pending_probe", deferredReason: probeSharedCacheMissReason() }; + } const payload = data.payload as { results?: SearchResult[]; telemetry?: Partial }; if (!Array.isArray(payload.results)) { - return { kind: "miss", reason: await probeSharedCacheMissReason("cache_payload_invalid") }; + return { kind: "miss", reason: "cache_payload_invalid" }; } return { kind: "hit", diff --git a/src/lib/rag.ts b/src/lib/rag.ts index 5520b6ee4..0927c6455 100644 --- a/src/lib/rag.ts +++ b/src/lib/rag.ts @@ -352,6 +352,14 @@ function throwIfAborted(signal?: AbortSignal) { } } +// Bounds an in-flight retrieval RPC by the caller's AbortSignal so a client +// disconnect cancels the HTTP request and frees the pooled Postgres connection. +// throwIfAborted only checks *between* stages; without this, a fan-out of +// match_* RPCs runs to completion for a caller that already went away. +function withAbort(builder: T, signal?: AbortSignal): T { + return signal ? builder.abortSignal(signal) : builder; +} + export type SearchChunksArgs = { query: string; topK?: number; @@ -1532,14 +1540,18 @@ async function searchTextChunkCandidates(args: { allowGlobalSearch?: boolean; matchCount: number; telemetry?: SearchTelemetry; + signal?: AbortSignal; }) { const runChunkText = async (queryText: string, matchCount: number) => { - const { data, error } = await args.supabase.rpc("match_document_chunks_text", { - query_text: queryText, - match_count: matchCount, - document_filters: args.documentIds ?? undefined, - owner_filter: ownerScopeForDocumentFilteredRetrieval(args.ownerId, args.documentIds, args.allowGlobalSearch), - }); + const { data, error } = await withAbort( + args.supabase.rpc("match_document_chunks_text", { + query_text: queryText, + match_count: matchCount, + document_filters: args.documentIds ?? undefined, + owner_filter: ownerScopeForDocumentFilteredRetrieval(args.ownerId, args.documentIds, args.allowGlobalSearch), + }), + args.signal, + ); return error || !data?.length ? ([] as SearchResult[]) : (data as SearchResult[]); }; @@ -1583,10 +1595,13 @@ async function searchTextChunkCandidates(args: { const primary = variants[0] ?? ""; let effectivePrimary = primary; if (primary) { - const { data: corrected } = await args.supabase.rpc("correct_clinical_query_terms", { - input_query: primary, - min_sim: 0.45, - }); + const { data: corrected } = await withAbort( + args.supabase.rpc("correct_clinical_query_terms", { + input_query: primary, + min_sim: 0.45, + }), + args.signal, + ); if (typeof corrected === "string" && corrected && corrected !== primary) { const correctedResults = await runChunkText(corrected, args.matchCount); if (correctedResults.length > 0) return correctedResults; @@ -1710,14 +1725,18 @@ async function fetchBestDocumentLookupChunks(args: { limit: number; ownerId?: string; allowGlobalSearch?: boolean; + signal?: AbortSignal; }) { const terms = documentLookupChunkTerms(args.query); - const { data: rpcChunks, error: rpcError } = await args.supabase.rpc("match_document_lookup_chunks_text", { - query_text: args.query, - document_filters: args.documentIds ?? undefined, - match_count: Math.max(args.limit * 3, 24), - owner_filter: ownerScopeForDocumentFilteredRetrieval(args.ownerId, args.documentIds, args.allowGlobalSearch), - }); + const { data: rpcChunks, error: rpcError } = await withAbort( + args.supabase.rpc("match_document_lookup_chunks_text", { + query_text: args.query, + document_filters: args.documentIds ?? undefined, + match_count: Math.max(args.limit * 3, 24), + owner_filter: ownerScopeForDocumentFilteredRetrieval(args.ownerId, args.documentIds, args.allowGlobalSearch), + }), + args.signal, + ); if (!rpcError && rpcChunks?.length) { const ranked = (rpcChunks as DocumentLookupChunkRow[]) .map((chunk) => ({ @@ -1812,6 +1831,7 @@ async function searchDocumentLookupFastPath(args: { ownerId?: string; documentIds?: string[]; matchCount: number; + signal?: AbortSignal; }): Promise { if (!args.ownerId) return [] as SearchResult[]; const variants = (args.queryVariants?.length ? args.queryVariants : [buildClinicalTextSearchQuery(args.query)]).slice( @@ -1820,11 +1840,14 @@ async function searchDocumentLookupFastPath(args: { ); const documentSets = await Promise.all( variants.map(async (variant, index) => { - const { data, error } = await args.supabase.rpc("match_documents_for_query", { - query_text: variant, - match_count: index === 0 ? 12 : 8, - owner_filter: requireOwnerScope(args.ownerId), - }); + const { data, error } = await withAbort( + args.supabase.rpc("match_documents_for_query", { + query_text: variant, + match_count: index === 0 ? 12 : 8, + owner_filter: requireOwnerScope(args.ownerId), + }), + args.signal, + ); if (error || !data?.length) return [] as DocumentLookupRow[]; return data as DocumentLookupRow[]; }), @@ -1866,6 +1889,7 @@ async function searchDocumentLookupFastPath(args: { query: args.query, limit: Math.max(args.matchCount, rankedDocuments.length * 4), ownerId: args.ownerId, + signal: args.signal, }); if (!chunks.length) return []; @@ -2185,6 +2209,7 @@ async function searchTableFactCandidates(args: { documentIds?: string[]; allowGlobalSearch?: boolean; matchCount: number; + signal?: AbortSignal; }) { const variants = (args.queryVariants?.length ? args.queryVariants : [buildClinicalTextSearchQuery(args.query)]).slice( 0, @@ -2192,12 +2217,15 @@ async function searchTableFactCandidates(args: { ); const factSets = await Promise.all( variants.map(async (variant, index) => { - const { data, error } = await args.supabase.rpc("match_document_table_facts_text", { - query_text: variant, - match_count: index === 0 ? args.matchCount : Math.min(args.matchCount, 24), - document_filters: args.documentIds ?? undefined, - owner_filter: ownerScopeForDocumentFilteredRetrieval(args.ownerId, args.documentIds, args.allowGlobalSearch), - }); + const { data, error } = await withAbort( + args.supabase.rpc("match_document_table_facts_text", { + query_text: variant, + match_count: index === 0 ? args.matchCount : Math.min(args.matchCount, 24), + document_filters: args.documentIds ?? undefined, + owner_filter: ownerScopeForDocumentFilteredRetrieval(args.ownerId, args.documentIds, args.allowGlobalSearch), + }), + args.signal, + ); if (error || !data?.length) return [] as TableFactRpcRow[]; return data as TableFactRpcRow[]; }), @@ -2239,15 +2267,19 @@ async function searchEmbeddingFieldCandidates(args: { allowGlobalSearch?: boolean; matchCount: number; telemetry?: SearchTelemetry; + signal?: AbortSignal; }) { - const { data, error } = await args.supabase.rpc("match_document_embedding_fields_hybrid", { - query_embedding: args.queryEmbedding as unknown as string, - query_text: buildClinicalTextSearchQuery(args.query), - match_count: args.matchCount, - min_similarity: 0.12, - document_filters: args.documentIds ?? undefined, - owner_filter: ownerScopeForDocumentFilteredRetrieval(args.ownerId, args.documentIds, args.allowGlobalSearch), - }); + const { data, error } = await withAbort( + args.supabase.rpc("match_document_embedding_fields_hybrid", { + query_embedding: args.queryEmbedding as unknown as string, + query_text: buildClinicalTextSearchQuery(args.query), + match_count: args.matchCount, + min_similarity: 0.12, + document_filters: args.documentIds ?? undefined, + owner_filter: ownerScopeForDocumentFilteredRetrieval(args.ownerId, args.documentIds, args.allowGlobalSearch), + }), + args.signal, + ); if (error) recordHybridRpcError(args.telemetry, "match_document_embedding_fields_hybrid", error); if (error || !data?.length) return [] as SearchResult[]; const matches = ( @@ -2291,15 +2323,19 @@ async function searchIndexUnitCandidates(args: { allowGlobalSearch?: boolean; matchCount: number; telemetry?: SearchTelemetry; + signal?: AbortSignal; }) { - const { data, error } = await args.supabase.rpc("match_document_index_units_hybrid", { - query_embedding: args.queryEmbedding as unknown as string, - query_text: buildClinicalTextSearchQuery(args.query), - match_count: args.matchCount, - min_similarity: 0.1, - document_filters: args.documentIds ?? undefined, - owner_filter: ownerScopeForDocumentFilteredRetrieval(args.ownerId, args.documentIds, args.allowGlobalSearch), - }); + const { data, error } = await withAbort( + args.supabase.rpc("match_document_index_units_hybrid", { + query_embedding: args.queryEmbedding as unknown as string, + query_text: buildClinicalTextSearchQuery(args.query), + match_count: args.matchCount, + min_similarity: 0.1, + document_filters: args.documentIds ?? undefined, + owner_filter: ownerScopeForDocumentFilteredRetrieval(args.ownerId, args.documentIds, args.allowGlobalSearch), + }), + args.signal, + ); if (error) recordHybridRpcError(args.telemetry, "match_document_index_units_hybrid", error); if (error || !data?.length) return [] as SearchResult[]; const matches = (data as IndexUnitRpcRow[]) @@ -3335,6 +3371,13 @@ export async function searchChunksWithTelemetry(args: SearchChunksArgs) { if (sharedCached?.kind === "miss") { telemetry.shared_cache_status = "miss"; telemetry.shared_cache_miss_reason = sharedCached.reason; + // The classification probe runs off the hot path; patch telemetry when it + // lands (query logging happens after generation, so it wins in practice). + if (sharedCached.deferredReason) { + void sharedCached.deferredReason.then((reason) => { + telemetry.shared_cache_miss_reason = reason; + }); + } } if ( @@ -3348,10 +3391,13 @@ export async function searchChunksWithTelemetry(args: SearchChunksArgs) { // in searchTextChunkCandidates). Only reached for would-be-unsupported queries, so it adds no // hot-path cost; `typoCorrected` guards against recursion. if (!args.typoCorrected && !sourceOnlyRetrieval) { - const { data: corrected } = await supabase.rpc("correct_clinical_query_terms", { - input_query: retrievalQuery, - min_sim: 0.45, - }); + const { data: corrected } = await withAbort( + supabase.rpc("correct_clinical_query_terms", { + input_query: retrievalQuery, + min_sim: 0.45, + }), + args.signal, + ); if (typeof corrected === "string" && corrected && corrected.toLowerCase() !== retrievalQuery.toLowerCase()) { return searchChunksWithTelemetry({ ...args, query: corrected, typoCorrected: true }); } @@ -3385,6 +3431,7 @@ export async function searchChunksWithTelemetry(args: SearchChunksArgs) { allowGlobalSearch: args.allowGlobalSearch, matchCount: textCandidateCount, telemetry, + signal: args.signal, }); telemetry.text_candidate_count = textData.length; telemetry.text_fast_path_latency_ms = Date.now() - textRpcStartedAt; @@ -3487,6 +3534,7 @@ export async function searchChunksWithTelemetry(args: SearchChunksArgs) { documentIds: documentFilterList, allowGlobalSearch: args.allowGlobalSearch, matchCount: Math.min(candidateCount, 48), + signal: args.signal, }); const tableFactLatencyMs = Date.now() - tableFactStartedAt; telemetry.supabase_rpc_latency_ms += tableFactLatencyMs; @@ -3508,6 +3556,7 @@ export async function searchChunksWithTelemetry(args: SearchChunksArgs) { ownerId: args.ownerId, documentIds: documentFilterList, matchCount: candidateCount, + signal: args.signal, }); const documentLookupLatencyMs = Date.now() - documentLookupStartedAt; telemetry.supabase_rpc_latency_ms += documentLookupLatencyMs; @@ -3667,6 +3716,7 @@ export async function searchChunksWithTelemetry(args: SearchChunksArgs) { allowGlobalSearch: args.allowGlobalSearch, matchCount: Math.min(candidateCount, 48), telemetry, + signal: args.signal, }); return { candidates, latencyMs: Date.now() - startedAt }; })(), @@ -3681,19 +3731,27 @@ export async function searchChunksWithTelemetry(args: SearchChunksArgs) { allowGlobalSearch: args.allowGlobalSearch, matchCount: Math.min(candidateCount, 64), telemetry, + signal: args.signal, }); return { candidates, latencyMs: Date.now() - startedAt }; })(), (async () => { const startedAt = Date.now(); - const { data, error } = await supabase.rpc("match_document_chunks_hybrid", { - query_embedding: embedding as unknown as string, - query_text: args.forceEmbedding ? "" : textSearchQuery, - match_count: candidateCount, - min_similarity: minSimilarity, - document_filters: documentFilterList ?? undefined, - owner_filter: ownerScopeForDocumentFilteredRetrieval(args.ownerId, documentFilterList, args.allowGlobalSearch), - }); + const { data, error } = await withAbort( + supabase.rpc("match_document_chunks_hybrid", { + query_embedding: embedding as unknown as string, + query_text: args.forceEmbedding ? "" : textSearchQuery, + match_count: candidateCount, + min_similarity: minSimilarity, + document_filters: documentFilterList ?? undefined, + owner_filter: ownerScopeForDocumentFilteredRetrieval( + args.ownerId, + documentFilterList, + args.allowGlobalSearch, + ), + }), + args.signal, + ); return { data, error, latencyMs: Date.now() - startedAt }; })(), ]); @@ -3788,17 +3846,20 @@ export async function searchChunksWithTelemetry(args: SearchChunksArgs) { const fallbackRpcStartedAt = Date.now(); const resultSets = await Promise.all( vectorFilters.map(async (documentFilter) => { - const { data, error } = await supabase.rpc("match_document_chunks", { - query_embedding: embedding as unknown as string, - match_count: candidateCount, - min_similarity: minSimilarity, - document_filter: documentFilter ?? undefined, - owner_filter: ownerScopeForDocumentFilteredRetrieval( - args.ownerId, - documentFilter ? [documentFilter] : undefined, - documentFilter ? undefined : args.allowGlobalSearch, - ), - }); + const { data, error } = await withAbort( + supabase.rpc("match_document_chunks", { + query_embedding: embedding as unknown as string, + match_count: candidateCount, + min_similarity: minSimilarity, + document_filter: documentFilter ?? undefined, + owner_filter: ownerScopeForDocumentFilteredRetrieval( + args.ownerId, + documentFilter ? [documentFilter] : undefined, + documentFilter ? undefined : args.allowGlobalSearch, + ), + }), + args.signal, + ); if (error) throw new Error(error.message); return (data ?? []) as SearchResult[]; diff --git a/src/lib/search-scope.ts b/src/lib/search-scope.ts index 8bf4fb5d9..feb0f8ed0 100644 --- a/src/lib/search-scope.ts +++ b/src/lib/search-scope.ts @@ -177,6 +177,24 @@ function metadataMatchesLocality(row: ScopeDocumentRow, filters: SearchScopeFilt return true; } +// Anonymous public-scope resolution enumerates every public indexed document +// (paginated reads up to the 5000-document ceiling) and the result is identical +// for every anonymous caller, so one resolved scope is shared for a short TTL. +// Only the zero-filter, no-explicit-ids default is cached; anything narrower is +// per-request. Keyed by client instance (WeakMap) so the production admin +// singleton shares one entry while test mocks stay isolated. Single-process, +// like the app's other in-memory caches. +const anonymousPublicScopeTtlMs = 60_000; +const anonymousPublicScopeCache = new WeakMap(); + +function cloneResolvedScope(scope: ResolvedSearchScope): ResolvedSearchScope { + return { + ...scope, + documentIds: scope.documentIds ? [...scope.documentIds] : scope.documentIds, + warnings: [...scope.warnings], + }; +} + export async function resolveSearchScope(args: { supabase: SupabaseClient; ownerId?: string | null; @@ -191,6 +209,26 @@ export async function resolveSearchScope(args: { const warnings: string[] = []; const publicOnly = args.publicOnly ?? !args.ownerId; + const cacheableAnonymousScope = + publicOnly && + !args.ownerId && + explicitIds.length === 0 && + activeFilterCount === 0 && + args.maxResolvedDocuments === undefined; + if (cacheableAnonymousScope) { + const cached = anonymousPublicScopeCache.get(args.supabase); + if (cached && cached.expiresAtMs > Date.now()) return cloneResolvedScope(cached.scope); + } + const cacheAnonymousScope = (scope: ResolvedSearchScope) => { + if (cacheableAnonymousScope) { + anonymousPublicScopeCache.set(args.supabase, { + scope: cloneResolvedScope(scope), + expiresAtMs: Date.now() + anonymousPublicScopeTtlMs, + }); + } + return scope; + }; + if (activeFilterCount === 0 && !publicOnly && !(args.ownerId && explicitIds.length)) { return { documentIds: explicitIds.length ? explicitIds : undefined, @@ -283,14 +321,14 @@ export async function resolveSearchScope(args: { const candidateIds = rows.map((row) => row.id); if (candidateIds.length === 0) { warnings.push("No indexed documents matched the selected filters."); - return { + return cacheAnonymousScope({ documentIds: [], filters, activeFilterCount, matchedDocumentCount: 0, warnings, summary: "No matching documents", - }; + }); } const needsLabels = @@ -346,7 +384,7 @@ export async function resolveSearchScope(args: { if (resolvedIds.length > 100) warnings.push(`${resolvedIds.length} documents match the selected scope; answers may be broad.`); - return { + return cacheAnonymousScope({ documentIds: resolvedIds, filters, activeFilterCount, @@ -356,5 +394,5 @@ export async function resolveSearchScope(args: { resolvedIds.length === 0 ? "No matching documents" : `${resolvedIds.length} scoped document${resolvedIds.length === 1 ? "" : "s"}`, - }; + }); }