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..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 @@ -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,52 @@ describe('recovered resource events', () => { } } ) + + 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 d7ab043c632..075aba15a7c 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. 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, + 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 }, { cancelRefetch: !inFlight }) 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..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. @@ -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, @@ -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() @@ -307,6 +304,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 () => { @@ -334,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, @@ -352,28 +366,9 @@ 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 }, + { 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/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) })