From d1da2350143d3b338fd77e88b5785c4a59883ba1 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Tue, 13 Oct 2026 13:15:31 -0700 Subject: [PATCH 1/2] fix: wave 1 performance follow-ups from re-audit - chat refresh coalescer cancels a slow first load at most once per chain, and a reload that joins the load in flight reads again after it - Better Auth user delete.after hook discards reporting usage prefixes, covering the admin plugin's remove-user path - a reporting prefix epoch bump that finds Redis configured but not ready is logged like a failed bump; payer bump logs carry the payer type - unique probes inside a write transaction also turn off parallel plans, restored with the seqscan flag - integration cleanup in finally, drop the stale GIN pending-list flush, fold the keyword-ranking seed check into beforeAll - TSDoc corrections in storage tracking and log stats queries --- .../stream/handle-resource-event.test.ts | 41 +++++++++++++++++++ apps/sim/hooks/queries/mothership-chats.ts | 25 +++++++---- apps/sim/lib/auth/auth.ts | 11 +++++ .../lib/billing/core/reporting-usage-epoch.ts | 12 ++++-- apps/sim/lib/billing/storage/tracking.ts | 15 +++---- .../keyword-ranking.integration.ts | 38 ++++++++--------- apps/sim/lib/logs/stats-queries.ts | 4 +- .../lib/table/rows/row-writes.integration.ts | 23 +++++------ apps/sim/lib/table/validation.ts | 36 ++++++++++------ 9 files changed, 140 insertions(+), 65 deletions(-) diff --git a/apps/sim/app/workspace/[workspaceId]/home/hooks/stream/handle-resource-event.test.ts b/apps/sim/app/workspace/[workspaceId]/home/hooks/stream/handle-resource-event.test.ts index 6360f9daa3a..fe7f1aedb47 100644 --- a/apps/sim/app/workspace/[workspaceId]/home/hooks/stream/handle-resource-event.test.ts +++ b/apps/sim/app/workspace/[workspaceId]/home/hooks/stream/handle-resource-event.test.ts @@ -1,3 +1,4 @@ +import { flushMacrotask } from '@sim/testing/helpers/async' import { createDeferred } from '@sim/testing/helpers/deferred' import { QueryClient, QueryObserver } from '@tanstack/react-query' import { beforeEach, describe, expect, it, vi } from 'vitest' @@ -330,6 +331,46 @@ describe('recovered resource events', () => { } } ) + + it('cancels a first load at most once, so a slow replacement finishes and is read again', async () => { + const queryClient = new QueryClient() + const { server, reads, unsubscribe } = mountedChat(queryClient, { + cached: false, + hangFirstRead: true, + }) + const slowReadReleased = createDeferred() + reads.mockImplementationOnce(async () => { + const resources = [...server.stored] + await slowReadReleased.promise + return { id: 'chat-1', title: null, messages: [], activeStreamId: 'stream-1', resources } + }) + const ctx = createStreamLoopContext( + makeStreamLoopDeps({ queryClient, chatIdRef: ref('chat-1') }) + ) + vi.useFakeTimers({ toFake: ['Date'] }) + try { + handleResourceEvent(ctx, replayed(tables[0], 1)) + await vi.waitFor(() => expect(reads).toHaveBeenCalledOnce()) + + vi.setSystemTime(Date.now() + 30_000) + handleResourceEvent(ctx, replayed(tables[1], 2)) + await vi.waitFor(() => expect(reads).toHaveBeenCalledTimes(2)) + + server.stored = tables + vi.setSystemTime(Date.now() + 30_000) + handleResourceEvent(ctx, replayed(tables[2], 3)) + await flushMacrotask() + expect(reads).toHaveBeenCalledTimes(2) + + slowReadReleased.resolve() + await vi.waitFor(() => expect(reads).toHaveBeenCalledTimes(3)) + expect(await settledResources(queryClient)).toEqual(tables) + expect(reads).toHaveBeenCalledTimes(3) + } finally { + vi.useRealTimers() + unsubscribe() + } + }) }) describe('handleResourceEvent saved-view pins', () => { diff --git a/apps/sim/hooks/queries/mothership-chats.ts b/apps/sim/hooks/queries/mothership-chats.ts index d7ab043c632..ec0db7e11d5 100644 --- a/apps/sim/hooks/queries/mothership-chats.ts +++ b/apps/sim/hooks/queries/mothership-chats.ts @@ -118,6 +118,8 @@ const CHAT_DETAIL_REFRESH_HOLD_MS = 30 * 1000 interface ChatDetailRefresh { startedAt: number requestedAgain: boolean + /** Whether this reload replaced one past the hold; a chain cancels a slow load at most once. */ + replaced: boolean } /** Per client, the chats with a reload in flight. */ @@ -139,16 +141,23 @@ export function refreshMothershipChatDetail(queryClient: QueryClient, chatId: st inFlight.requestedAgain = true return } - const refresh: ChatDetailRefresh = { startedAt: Date.now(), requestedAgain: false } + // A query with no data yet joins its fetch in flight instead of replacing it, so a stuck + // one is cancelled first, but only once per chain: a load that is slow rather than hung + // would otherwise restart every hold and never finish. A reload that joins the load + // already in flight cannot trust it to include this request, so it reads again after. + const refresh: ChatDetailRefresh = { + startedAt: Date.now(), + requestedAgain: inFlight?.replaced === true, + replaced: inFlight !== undefined, + } refreshes.set(chatId, refresh) const queryKey = mothershipChatKeys.detail(chatId) - // A query with no data yet joins its fetch in flight instead of replacing it, so a stuck - // one is cancelled first. - const reload = inFlight - ? queryClient - .cancelQueries({ queryKey, exact: true }) - .then(() => queryClient.invalidateQueries({ queryKey })) - : queryClient.invalidateQueries({ queryKey }) + const reload = + inFlight && !inFlight.replaced + ? queryClient + .cancelQueries({ queryKey, exact: true }) + .then(() => queryClient.invalidateQueries({ queryKey })) + : queryClient.invalidateQueries({ queryKey }) void reload.finally(() => { if (refreshes.get(chatId) !== refresh) return refreshes.delete(chatId) diff --git a/apps/sim/lib/auth/auth.ts b/apps/sim/lib/auth/auth.ts index 726df4549d2..bb14cf61f5f 100644 --- a/apps/sim/lib/auth/auth.ts +++ b/apps/sim/lib/auth/auth.ts @@ -86,6 +86,7 @@ import { releaseCheckoutAdmission, resolveCheckoutReferenceId, } from '@/lib/billing/checkout-admission' +import { invalidateReportingUsagePrefixes } from '@/lib/billing/core/reporting-usage-epoch' import { getOrganizationIdForSubscriptionReference, syncSubscriptionPlan, @@ -442,6 +443,16 @@ export const auth = betterAuth({ } }, }, + /** + * The admin plugin's `/admin/remove-user` deletes the user row through Better Auth, so its + * usage rows cascade away without `deleteUserAccount`. Better Auth runs this once the delete + * commits, which is when the stored reporting prefixes that counted them must be discarded. + */ + delete: { + after: async () => { + await invalidateReportingUsagePrefixes() + }, + }, }, account: { create: { diff --git a/apps/sim/lib/billing/core/reporting-usage-epoch.ts b/apps/sim/lib/billing/core/reporting-usage-epoch.ts index 4188584edae..9e52b8dc2c4 100644 --- a/apps/sim/lib/billing/core/reporting-usage-epoch.ts +++ b/apps/sim/lib/billing/core/reporting-usage-epoch.ts @@ -63,9 +63,14 @@ export function readyRedisClient(): Redis | null { return redis?.status === 'ready' ? redis : null } +/** + * Throws when a configured Redis is not ready, since its stored prefixes then outlive the bump: + * each caller logs that like a failed or timed-out SET. Without Redis there is nothing to discard. + */ async function replaceEpoch(epochKey: string): Promise { - const redis = readyRedisClient() + const redis = getRedisClient() if (!redis) return + if (redis.status !== 'ready') throw new Error(`Redis not ready (${redis.status})`) await withinDeadline( () => redis.set(epochKey, generateShortId(), 'PX', PREFIX_EPOCH_RETAIN_MS), Date.now() + SHARED_READ_TIMEOUT_MS @@ -76,8 +81,8 @@ async function replaceEpoch(epochKey: string): Promise { * Discards every stored prefix, for a write that removes ledger rows — a user's deletion, whose * rows cascade away. Call it once that write has committed: a refresh that read the old epoch * stores its prefix under it, where nothing reads it, and one that reads the new epoch sums after - * the deletion. Best effort: without a ready Redis there is no prefix to discard, and a failed - * bump leaves the deleted usage counted until the stored prefixes expire. + * the deletion. Best effort: a bump that fails, times out, or finds Redis not ready leaves the + * deleted usage counted until the stored prefixes expire, and is logged. */ export async function invalidateReportingUsagePrefixes(): Promise { try { @@ -100,6 +105,7 @@ export function invalidatePayerPrefixesForTopUp(entity: BillingEntity, rowCreate if (rowCreatedAt.getTime() >= Date.now() - TOP_UP_INVALIDATES_AFTER_MS) return replaceEpoch(payerPrefixEpochKey(entity)).catch((error) => { logger.warn('Shared reporting usage payer prefix invalidation failed', { + payerType: entity.type, error: getErrorMessage(error), }) }) diff --git a/apps/sim/lib/billing/storage/tracking.ts b/apps/sim/lib/billing/storage/tracking.ts index b75b979d7cc..e11944ab83a 100644 --- a/apps/sim/lib/billing/storage/tracking.ts +++ b/apps/sim/lib/billing/storage/tracking.ts @@ -433,13 +433,14 @@ export async function applyStorageUsageDeltasInTx( /** * Locks the canonical workspace row before a file mutation that changes its storage ledger. * - * Every transaction that writes `workspace_files` and charges storage takes its locks in one order: - * the folder advisory lock (only for operations that take one), then this workspace row, then file - * rows, then the file writes (whose search-index trigger upserts the one-per-workspace - * `workspace_file_search_dispatch_queue` row), then the payer row. Call this before the first file insert or update: an insert that reaches - * the queue row first and only then waits on the workspace row deadlocks with a content write - * that holds the workspace row and waits on the queue row. Re-locking later in the same - * transaction, as the storage increment does, never waits. + * Every transaction that writes a `context = 'workspace'` row of `workspace_files` and charges + * storage takes its locks in one order: the folder advisory lock (only for operations that take + * one), then this workspace row, then the file rows and their writes (whose search-index trigger + * upserts the one-per-workspace `workspace_file_search_dispatch_queue` row), then the payer row. + * Call this before the first file insert or update: an insert that reaches the queue row first + * and only then waits on the workspace row deadlocks with a content write that holds the + * workspace row and waits on the queue row. Re-locking later in the same transaction, as the + * storage increment does, never waits. */ export async function lockWorkspaceStorageForMutationInTx( tx: DbOrTx, diff --git a/apps/sim/lib/knowledge/__integration__/keyword-ranking.integration.ts b/apps/sim/lib/knowledge/__integration__/keyword-ranking.integration.ts index f6dc67b9131..0dfd10a53a7 100644 --- a/apps/sim/lib/knowledge/__integration__/keyword-ranking.integration.ts +++ b/apps/sim/lib/knowledge/__integration__/keyword-ranking.integration.ts @@ -11,7 +11,7 @@ import { db } from '@sim/db' import { document, embedding, knowledgeBase, organization, user, workspace } from '@sim/db/schema' import { deleteWorkspaceFixture } from '@sim/db/testing/workspace-fixtures' import { generateId } from '@sim/utils/id' -import { and, desc, eq, inArray, type SQL, sql } from 'drizzle-orm' +import { and, eq, inArray, type SQL, sql } from 'drizzle-orm' import { afterAll, beforeAll, describe, expect, it } from 'vitest' import { createKnowledgeAclFixtureIds, @@ -307,6 +307,23 @@ describe('unauthenticated keyword ranking in one statement', () => { await db.insert(document).values([...documents, ...hiddenDocuments]) await db.insert(embedding).values(chunks) + + // The comparisons below test tie-breaking only if scores tie across bases at several levels. + const scores = await db + .select({ + rank: sql`ts_rank_cd(${embedding.contentTsv}, websearch_to_tsquery(${FTS_CONFIG}, ${QUERY}))`, + bases: sql`count(DISTINCT ${embedding.knowledgeBaseId})::int`, + }) + .from(embedding) + .where( + and( + inArray(embedding.knowledgeBaseId, allBaseIds), + sql`${embedding.contentTsv} @@ websearch_to_tsquery(${FTS_CONFIG}, ${QUERY})` + ) + ) + .groupBy(sql`1`) + expect(scores.length).toBeGreaterThan(2) + expect(scores.filter((score) => score.bases > 1).length).toBeGreaterThan(1) }) afterAll(async () => { @@ -352,25 +369,6 @@ describe('unauthenticated keyword ranking in one statement', () => { return { actual, reference } } - it('seeds scores tied across bases at several distinct levels', async () => { - const scores = await db - .select({ - rank: sql`ts_rank_cd(${embedding.contentTsv}, websearch_to_tsquery(${FTS_CONFIG}, ${QUERY}))`, - bases: sql`count(DISTINCT ${embedding.knowledgeBaseId})::int`, - }) - .from(embedding) - .where( - and( - inArray(embedding.knowledgeBaseId, allBaseIds), - sql`${embedding.contentTsv} @@ websearch_to_tsquery(${FTS_CONFIG}, ${QUERY})` - ) - ) - .groupBy(sql`1`) - .orderBy(desc(sql`1`)) - expect(scores.length).toBeGreaterThan(2) - expect(scores.filter((score) => score.bases > 1).length).toBeGreaterThan(1) - }) - describe.each([ { label: 'every base (old per-base branch)', baseCount: allBaseIds.length }, { label: 'two bases (old single-statement branch)', baseCount: 2 }, diff --git a/apps/sim/lib/logs/stats-queries.ts b/apps/sim/lib/logs/stats-queries.ts index 7418fa9b881..3907502ca5c 100644 --- a/apps/sim/lib/logs/stats-queries.ts +++ b/apps/sim/lib/logs/stats-queries.ts @@ -62,8 +62,8 @@ export async function readLogStatsBounds(where: SQL | undefined): Promise { SELECT ${table.id}::text || '-' || i, ${table.id}, ${workspaceId}, jsonb_build_object('email', 'stored-' || i || '@example.test'), i, 'a' || i FROM generate_series(0, 39999) AS i` - // A bulk insert leaves its entries in the GIN pending list, which a manual ANALYZE does not - // merge and which the planner charges to every probe, so flush it for a deterministic plan. - await control`SELECT gin_clean_pending_list('user_table_rows_tenant_data_gin_idx'::regclass)` - await control`ANALYZE user_table_rows` - const rows: RowData[] = Array.from({ length: 260 }, (_, i) => ({ - email: i % 10 === 0 ? `stored-${i * 100}@example.test` : `new-${i}@example.test`, - })) - - const { result, scans } = await checkWithScans(table, rows) + try { + await control`ANALYZE user_table_rows` + const rows: RowData[] = Array.from({ length: 260 }, (_, i) => ({ + email: i % 10 === 0 ? `stored-${i * 100}@example.test` : `new-${i}@example.test`, + })) - await control`DELETE FROM user_table_rows WHERE table_id = ${table.id}` + const { result, scans } = await checkWithScans(table, rows) - expect(result.errors).toHaveLength(26) - expect(scans).toEqual({ user_table_rows_tenant_data_gin_idx: 260 }) + expect(result.errors).toHaveLength(26) + expect(scans).toEqual({ user_table_rows_tenant_data_gin_idx: 260 }) + } finally { + await control`DELETE FROM user_table_rows WHERE table_id = ${table.id}` + } }, 60_000) }) diff --git a/apps/sim/lib/table/validation.ts b/apps/sim/lib/table/validation.ts index 027f9e25394..b17bddad1f5 100644 --- a/apps/sim/lib/table/validation.ts +++ b/apps/sim/lib/table/validation.ts @@ -480,7 +480,7 @@ export async function checkUniqueConstraintsDb( // The predicate leads with case-sensitive JSONB containment (`data @> {...}`), // which can use the GIN index. We still pin `enable_seqscan = off` (tenant- // bounded) defensively for the small-table / cold-stats case. With an external - // transaction the flag is set on it for the check only (see withSeqscanOffOn) + // transaction the settings are set on it for the check only (see withUniqueProbePlan) // — opening our own transaction inside the caller's would be the nested pool // checkout the migration-hardening work eliminated (self-deadlock under pool // exhaustion). @@ -509,27 +509,37 @@ export async function checkUniqueConstraintsDb( if (executor === db) { await withSeqscanOff(async (trx) => checkConditions(trx)) } else { - await withSeqscanOffOn(executor, () => checkConditions(executor)) + await withUniqueProbePlan(executor, () => checkConditions(executor)) } return { valid: errors.length === 0, errors } } /** - * Runs `check` on the caller's open transaction with seq scans penalized, then puts back the - * transaction's previous setting. Unique checks run inside write transactions, and leaving the - * flag off for the rest of the transaction slowed the position and INSERT statements that follow. Reading the previous value and setting the new one - * in one statement keeps it one round trip each way; a caller that set the flag itself keeps it. + * Runs `check` on the caller's open transaction with seq scans penalized and parallel plans off, + * then puts back the transaction's previous settings. Unique checks run inside write + * transactions: a wide probe on a large tenant could otherwise start a parallel Gather there, and + * leaving either setting for the rest of the transaction slowed the position and INSERT + * statements that follow. Reading the previous values and setting the new ones in one statement + * keeps it one round trip each way; a caller that set either itself keeps it. */ -async function withSeqscanOffOn( +async function withUniqueProbePlan( executor: UniqueCheckExecutor, check: () => Promise ): Promise { - const [{ previous }] = await executor.execute<{ previous: string }>(sql` - SELECT current_setting('enable_seqscan') AS previous, set_config('enable_seqscan', 'off', true) + const [previous] = await executor.execute<{ seqscan: string; parallel_workers: string }>(sql` + SELECT + current_setting('enable_seqscan') AS seqscan, + current_setting('max_parallel_workers_per_gather') AS parallel_workers, + set_config('enable_seqscan', 'off', true), + set_config('max_parallel_workers_per_gather', '0', true) `) await check() - await executor.execute(sql`SELECT set_config('enable_seqscan', ${previous}, true)`) + await executor.execute(sql` + SELECT + set_config('enable_seqscan', ${previous.seqscan}, true), + set_config('max_parallel_workers_per_gather', ${previous.parallel_workers}, true) + `) } /** @@ -614,8 +624,8 @@ export async function checkBatchUniqueConstraintsDb( // Now check against database for all unique values at once. Tenant-bounded // for the same reason as checkUniqueConstraintsDb: the lower(data->>...) // predicates are unestimatable and otherwise trigger whole-relation seq - // scans. With an external transaction the flag is set on it for the check - // only (see withSeqscanOffOn). + // scans. With an external transaction the settings are set on it for the check + // only (see withUniqueProbePlan). const checkColumns = async (ex: UniqueCheckExecutor) => { for (const [columnId, { values, column }] of valuesByColumn) { if (values.size === 0) continue @@ -685,7 +695,7 @@ export async function checkBatchUniqueConstraintsDb( if (executor === db) { await withSeqscanOff(async (trx) => checkColumns(trx)) } else { - await withSeqscanOffOn(executor, () => checkColumns(executor)) + await withUniqueProbePlan(executor, () => checkColumns(executor)) } // Sort errors by row index From d69d52493afd10b0d91d6b09699b18ad56dd6490 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Tue, 13 Oct 2026 13:27:51 -0700 Subject: [PATCH 2/2] fix: restore the two-step unauthenticated keyword ranking, join cached slow reads, drop the probe parallel toggle --- .../stream/handle-resource-event.test.ts | 82 ++++++++++--------- apps/sim/hooks/queries/mothership-chats.ts | 6 +- .../keyword-ranking.integration.ts | 23 +++--- apps/sim/lib/knowledge/search/queries.ts | 61 +++++++++++--- apps/sim/lib/table/validation.ts | 36 +++----- 5 files changed, 118 insertions(+), 90 deletions(-) diff --git a/apps/sim/app/workspace/[workspaceId]/home/hooks/stream/handle-resource-event.test.ts b/apps/sim/app/workspace/[workspaceId]/home/hooks/stream/handle-resource-event.test.ts index fe7f1aedb47..b0b038176bf 100644 --- a/apps/sim/app/workspace/[workspaceId]/home/hooks/stream/handle-resource-event.test.ts +++ b/apps/sim/app/workspace/[workspaceId]/home/hooks/stream/handle-resource-event.test.ts @@ -332,45 +332,51 @@ describe('recovered resource events', () => { } ) - it('cancels a first load at most once, so a slow replacement finishes and is read again', async () => { - const queryClient = new QueryClient() - const { server, reads, unsubscribe } = mountedChat(queryClient, { - cached: false, - hangFirstRead: true, - }) - const slowReadReleased = createDeferred() - reads.mockImplementationOnce(async () => { - const resources = [...server.stored] - await slowReadReleased.promise - return { id: 'chat-1', title: null, messages: [], activeStreamId: 'stream-1', resources } - }) - const ctx = createStreamLoopContext( - makeStreamLoopDeps({ queryClient, chatIdRef: ref('chat-1') }) - ) - vi.useFakeTimers({ toFake: ['Date'] }) - try { - handleResourceEvent(ctx, replayed(tables[0], 1)) - await vi.waitFor(() => expect(reads).toHaveBeenCalledOnce()) - - vi.setSystemTime(Date.now() + 30_000) - handleResourceEvent(ctx, replayed(tables[1], 2)) - await vi.waitFor(() => expect(reads).toHaveBeenCalledTimes(2)) - - server.stored = tables - vi.setSystemTime(Date.now() + 30_000) - handleResourceEvent(ctx, replayed(tables[2], 3)) - await flushMacrotask() - expect(reads).toHaveBeenCalledTimes(2) - - slowReadReleased.resolve() - await vi.waitFor(() => expect(reads).toHaveBeenCalledTimes(3)) - expect(await settledResources(queryClient)).toEqual(tables) - expect(reads).toHaveBeenCalledTimes(3) - } finally { - vi.useRealTimers() - unsubscribe() + it.each([ + { name: 'a loaded transcript', cached: true }, + { name: 'a transcript still on its first load', cached: false }, + ])( + 'cancels a stuck read of $name at most once, so a slow replacement finishes and is read again', + async ({ cached }) => { + const queryClient = new QueryClient() + const { server, reads, unsubscribe } = mountedChat(queryClient, { + cached, + hangFirstRead: true, + }) + const slowReadReleased = createDeferred() + reads.mockImplementationOnce(async () => { + const resources = [...server.stored] + await slowReadReleased.promise + return { id: 'chat-1', title: null, messages: [], activeStreamId: 'stream-1', resources } + }) + const ctx = createStreamLoopContext( + makeStreamLoopDeps({ queryClient, chatIdRef: ref('chat-1') }) + ) + vi.useFakeTimers({ toFake: ['Date'] }) + try { + handleResourceEvent(ctx, replayed(tables[0], 1)) + await vi.waitFor(() => expect(reads).toHaveBeenCalledOnce()) + + vi.setSystemTime(Date.now() + 30_000) + handleResourceEvent(ctx, replayed(tables[1], 2)) + await vi.waitFor(() => expect(reads).toHaveBeenCalledTimes(2)) + + server.stored = tables + vi.setSystemTime(Date.now() + 30_000) + handleResourceEvent(ctx, replayed(tables[2], 3)) + await flushMacrotask() + expect(reads).toHaveBeenCalledTimes(2) + + slowReadReleased.resolve() + await vi.waitFor(() => expect(reads).toHaveBeenCalledTimes(3)) + expect(await settledResources(queryClient)).toEqual(tables) + expect(reads).toHaveBeenCalledTimes(3) + } finally { + vi.useRealTimers() + unsubscribe() + } } - }) + ) }) describe('handleResourceEvent saved-view pins', () => { diff --git a/apps/sim/hooks/queries/mothership-chats.ts b/apps/sim/hooks/queries/mothership-chats.ts index ec0db7e11d5..075aba15a7c 100644 --- a/apps/sim/hooks/queries/mothership-chats.ts +++ b/apps/sim/hooks/queries/mothership-chats.ts @@ -143,8 +143,8 @@ export function refreshMothershipChatDetail(queryClient: QueryClient, chatId: st } // A query with no data yet joins its fetch in flight instead of replacing it, so a stuck // one is cancelled first, but only once per chain: a load that is slow rather than hung - // would otherwise restart every hold and never finish. A reload that joins the load - // already in flight cannot trust it to include this request, so it reads again after. + // would otherwise restart every hold and never finish. Later reloads in the chain join the + // read in flight, cached or not, and since it may predate this request, read again after. const refresh: ChatDetailRefresh = { startedAt: Date.now(), requestedAgain: inFlight?.replaced === true, @@ -157,7 +157,7 @@ export function refreshMothershipChatDetail(queryClient: QueryClient, chatId: st ? queryClient .cancelQueries({ queryKey, exact: true }) .then(() => queryClient.invalidateQueries({ queryKey })) - : queryClient.invalidateQueries({ queryKey }) + : queryClient.invalidateQueries({ queryKey }, { cancelRefetch: !inFlight }) void reload.finally(() => { if (refreshes.get(chatId) !== refresh) return refreshes.delete(chatId) diff --git a/apps/sim/lib/knowledge/__integration__/keyword-ranking.integration.ts b/apps/sim/lib/knowledge/__integration__/keyword-ranking.integration.ts index 0dfd10a53a7..94cf2553cc0 100644 --- a/apps/sim/lib/knowledge/__integration__/keyword-ranking.integration.ts +++ b/apps/sim/lib/knowledge/__integration__/keyword-ranking.integration.ts @@ -1,7 +1,7 @@ /** - * The keyword leg of a search with no signed-in reader behind it, against real PostgreSQL: ranking - * every base in one statement answers what ranking each base (or all of them) in the old two-step - * pass did, tie groups compared as sets, and admits nothing the reader may not see. + * The keyword leg of a search with no signed-in reader behind it, against real PostgreSQL: it + * answers what ranking each base (or all of them) in the two-step pass does, tie groups compared + * as sets, and admits nothing the reader may not see. * * Every chunk's vector points away from the query vector, so the vector leg of a hybrid search * finds nothing under its distance threshold and the fused list is the keyword ranking alone. @@ -34,7 +34,7 @@ const queryVector = { } as const const AWAY_FROM_QUERY = [-1, ...Array(DIMENSIONS - 1).fill(0)] -/** More bases than the old pass ranked in one statement, so it took its per-base branch. */ +/** More bases than the pass ranks in one statement, so it takes its per-base branch. */ const EXTRA_BASE_COUNT = 5 const DOCUMENTS_PER_BASE = 3 const CHUNKS_PER_DOCUMENT = 2 @@ -78,11 +78,8 @@ interface RankedKeywordRow { keywordRank: number } -/** - * The unauthenticated keyword ranking as it stood before it moved into one statement: the oracle - * each case is compared against. - */ -async function rankBeforeSingleStatement(input: { +/** The two-step unauthenticated keyword ranking: the oracle each case is compared against. */ +async function rankTwoStep(input: { knowledgeBaseIds: string[] topK: number query: string @@ -152,7 +149,7 @@ function expectSameRanking(actual: string[], reference: RankedKeywordRow[], topK expect(cursor).toBe(actual.length) } -describe('unauthenticated keyword ranking in one statement', () => { +describe('unauthenticated keyword ranking', () => { const ids = createKnowledgeAclFixtureIds() const extraBaseIds = Array.from({ length: EXTRA_BASE_COUNT }, () => generateId()) const unsearchedBaseId = generateId() @@ -351,7 +348,7 @@ describe('unauthenticated keyword ranking in one statement', () => { queryVector, }) expect(retrieval).toEqual({ status: 'complete', timedOutLegs: [] }) - const reference = await rankBeforeSingleStatement({ + const reference = await rankTwoStep({ knowledgeBaseIds: input.knowledgeBaseIds, topK: hybridLegTopK(input.topK), query: QUERY, @@ -370,8 +367,8 @@ describe('unauthenticated keyword ranking in one statement', () => { } describe.each([ - { label: 'every base (old per-base branch)', baseCount: allBaseIds.length }, - { label: 'two bases (old single-statement branch)', baseCount: 2 }, + { label: 'every base (per-base branch)', baseCount: allBaseIds.length }, + { label: 'two bases (single-statement branch)', baseCount: 2 }, ])('$label', ({ baseCount }) => { const knowledgeBaseIds = allBaseIds.slice(0, baseCount) diff --git a/apps/sim/lib/knowledge/search/queries.ts b/apps/sim/lib/knowledge/search/queries.ts index a2b4a4d23b3..0e55d49ad21 100644 --- a/apps/sim/lib/knowledge/search/queries.ts +++ b/apps/sim/lib/knowledge/search/queries.ts @@ -67,9 +67,8 @@ interface DocumentReadAccess { /** The candidate predicate every leg ranks under. */ rankCondition: SQL /** - * A signed-in reader: their tag and keyword legs run under the search deadline, and their tag leg - * ranks in one statement. A caller with no person behind it runs those legs unbudgeted, as it - * always has. + * A signed-in reader: their tag and keyword legs run under the search deadline and rank in one + * statement. A caller with no person behind it runs those legs unbudgeted, as it always has. */ signedIn: boolean /** Present only when a searched base holds a source whose reader must be proven live. */ @@ -386,10 +385,11 @@ export async function handleVectorSearch( * query makes Postgres detoast the chunk's vector and compute a distance for *every* full-text * match before the `LIMIT` applies — work that scales with how common the query term is rather * than with `topK`. Ranking therefore touches no vectors, and only the rows that survive the limit - * are hydrated. Every caller ranks every base in one statement (see - * {@link rankKeywordCandidates}): a signed-in reader under the search deadline, a caller with no - * person behind it unbudgeted. A search holding a source whose reader must be proven live pages - * its ranking instead, so the proof is asked for only when one of its candidates is read. + * are hydrated. A signed-in reader ranks in one statement under the search deadline (see + * {@link rankSignedInKeywordCandidates}); a caller with no person behind it ranks unbudgeted, many + * bases in parallel, each up to the leg's whole `topK` so the merged ranking is the global one. A + * search holding a source whose reader must be proven live pages its ranking instead, so the + * proof is asked for only when one of its candidates is read. */ async function executeKeywordSearch( params: KeywordSearchParams, @@ -418,7 +418,7 @@ async function executeKeywordSearch( 'keyword', budget ) - const ranking = { + const signedInRanking = { knowledgeBaseIds, tsQuery, tagFilterConditions, @@ -439,8 +439,8 @@ async function executeKeywordSearch( budget, topK, selectPage: async (limit, offset, excludedSources) => { - const candidates = await rankKeywordCandidates({ - ...ranking, + const candidates = await rankSignedInKeywordCandidates({ + ...signedInRanking, exclusion: excludeSearchSources(excludedSources), limit, offset, @@ -451,8 +451,43 @@ async function executeKeywordSearch( }) } + const rankExpr = sql`ts_rank_cd(${embedding.contentTsv}, ${tsQuery})` + const visibility = getVisibilityConditions(params.filters, read.rankCondition) + /** Ranking pass: ids and relevance only, so no vector is read. */ + const rankRows = (kbScope: SQL) => + db + .select({ id: embedding.id, keywordRank: rankExpr.as('keyword_rank') }) + .from(embedding) + .innerJoin(document, eq(embedding.documentId, document.id)) + .where( + and( + kbScope, + ...visibility, + sql`${embedding.contentTsv} @@ ${tsQuery}`, + ...tagFilterConditions + ) + ) + .orderBy(sql`${rankExpr} DESC`) + .limit(topK) + return shortOnDeadline(budget, async () => { - const topIds = (await rankKeywordCandidates({ ...ranking, limit: topK })).map((row) => row.id) + let topIds: string[] + if (read.signedIn) { + const ranked = await rankSignedInKeywordCandidates({ ...signedInRanking, limit: topK }) + topIds = ranked.map((row) => row.id) + } else if (getQueryStrategy(knowledgeBaseIds.length, topK).useParallel) { + const perBase = await Promise.all( + knowledgeBaseIds.map((kbId) => rankRows(eq(embedding.knowledgeBaseId, kbId))) + ) + topIds = perBase + .flat() + .sort((a, b) => b.keywordRank - a.keywordRank) + .slice(0, topK) + .map((row) => row.id) + } else { + const ranked = await rankRows(inArray(embedding.knowledgeBaseId, knowledgeBaseIds)) + topIds = ranked.map((row) => row.id) + } if (topIds.length === 0) { return [] } @@ -463,8 +498,8 @@ async function executeKeywordSearch( }) } -/** A keyword ranking over the source chunks, in one statement; see {@link keywordCandidateRankingQuery}. */ -function rankKeywordCandidates(input: { +/** A signed-in reader's keyword ranking over the source chunks, in one statement; see {@link keywordCandidateRankingQuery}. */ +function rankSignedInKeywordCandidates(input: { knowledgeBaseIds: string[] tsQuery: SQL tagFilterConditions: SQL[] diff --git a/apps/sim/lib/table/validation.ts b/apps/sim/lib/table/validation.ts index b17bddad1f5..027f9e25394 100644 --- a/apps/sim/lib/table/validation.ts +++ b/apps/sim/lib/table/validation.ts @@ -480,7 +480,7 @@ export async function checkUniqueConstraintsDb( // The predicate leads with case-sensitive JSONB containment (`data @> {...}`), // which can use the GIN index. We still pin `enable_seqscan = off` (tenant- // bounded) defensively for the small-table / cold-stats case. With an external - // transaction the settings are set on it for the check only (see withUniqueProbePlan) + // transaction the flag is set on it for the check only (see withSeqscanOffOn) // — opening our own transaction inside the caller's would be the nested pool // checkout the migration-hardening work eliminated (self-deadlock under pool // exhaustion). @@ -509,37 +509,27 @@ export async function checkUniqueConstraintsDb( if (executor === db) { await withSeqscanOff(async (trx) => checkConditions(trx)) } else { - await withUniqueProbePlan(executor, () => checkConditions(executor)) + await withSeqscanOffOn(executor, () => checkConditions(executor)) } return { valid: errors.length === 0, errors } } /** - * Runs `check` on the caller's open transaction with seq scans penalized and parallel plans off, - * then puts back the transaction's previous settings. Unique checks run inside write - * transactions: a wide probe on a large tenant could otherwise start a parallel Gather there, and - * leaving either setting for the rest of the transaction slowed the position and INSERT - * statements that follow. Reading the previous values and setting the new ones in one statement - * keeps it one round trip each way; a caller that set either itself keeps it. + * Runs `check` on the caller's open transaction with seq scans penalized, then puts back the + * transaction's previous setting. Unique checks run inside write transactions, and leaving the + * flag off for the rest of the transaction slowed the position and INSERT statements that follow. Reading the previous value and setting the new one + * in one statement keeps it one round trip each way; a caller that set the flag itself keeps it. */ -async function withUniqueProbePlan( +async function withSeqscanOffOn( executor: UniqueCheckExecutor, check: () => Promise ): Promise { - const [previous] = await executor.execute<{ seqscan: string; parallel_workers: string }>(sql` - SELECT - current_setting('enable_seqscan') AS seqscan, - current_setting('max_parallel_workers_per_gather') AS parallel_workers, - set_config('enable_seqscan', 'off', true), - set_config('max_parallel_workers_per_gather', '0', true) + const [{ previous }] = await executor.execute<{ previous: string }>(sql` + SELECT current_setting('enable_seqscan') AS previous, set_config('enable_seqscan', 'off', true) `) await check() - await executor.execute(sql` - SELECT - set_config('enable_seqscan', ${previous.seqscan}, true), - set_config('max_parallel_workers_per_gather', ${previous.parallel_workers}, true) - `) + await executor.execute(sql`SELECT set_config('enable_seqscan', ${previous}, true)`) } /** @@ -624,8 +614,8 @@ export async function checkBatchUniqueConstraintsDb( // Now check against database for all unique values at once. Tenant-bounded // for the same reason as checkUniqueConstraintsDb: the lower(data->>...) // predicates are unestimatable and otherwise trigger whole-relation seq - // scans. With an external transaction the settings are set on it for the check - // only (see withUniqueProbePlan). + // scans. With an external transaction the flag is set on it for the check + // only (see withSeqscanOffOn). const checkColumns = async (ex: UniqueCheckExecutor) => { for (const [columnId, { values, column }] of valuesByColumn) { if (values.size === 0) continue @@ -695,7 +685,7 @@ export async function checkBatchUniqueConstraintsDb( if (executor === db) { await withSeqscanOff(async (trx) => checkColumns(trx)) } else { - await withUniqueProbePlan(executor, () => checkColumns(executor)) + await withSeqscanOffOn(executor, () => checkColumns(executor)) } // Sort errors by row index