fix(knowledge): keep the Tin projection in step with search-index adoption

- Project or remove a whole base when is_search_index changes, serialized with embedding writers on a per-base advisory lock so neither misses the other's rows
- Run Tin readiness and query analysis under the keyword leg's budget
- Rebuild the Tin projection after db:push, which drops the schema-external index
- Register 0019 in the script-migration registry tests

Author: waleedlatif1

.github/workflows/test-build.yml

@@ -241,6 +241,7 @@ jobs:
           bunx vitest run
           script-migrations/0016_backfill_search_vectors.postgres.test.ts
           script-migrations/0018_repair_workspace_file_content_revision.postgres.test.ts
+          script-migrations/0019_tin_keyword_projection.postgres.test.ts
 
       - name: Verify Search progress, pagination, and outbox scheduling in PostgreSQL
         working-directory: apps/sim

apps/sim/lib/knowledge/search/diagnostics.ts

@@ -52,6 +52,8 @@ export type SearchStage =
   | 'vector.exact'
   | 'vector.candidate_search'
   | 'keyword.tin'
+  | 'keyword.tin_readiness'
+  | 'keyword.tin_query'
   | 'source_overview'
   | 'source_overview.availability'
   | 'source_overview.providers'

apps/sim/lib/knowledge/search/queries.test.ts

@@ -1531,7 +1531,12 @@ describe('permitted-document planner', () => {
       queueTableRows(schemaMock.embedding, [{ ...hit('a', null), content: 'release notes' }])
       const results = await keyword()
       expect(results.map((row) => row.id)).toEqual(['a'])
-      expect(mockResolveTinKeywordQuery).toHaveBeenCalledWith(['org-index'], 'release', 'english')
+      expect(mockResolveTinKeywordQuery).toHaveBeenCalledWith(
+        ['org-index'],
+        'release',
+        'english',
+        params.budget
+      )
       expect(ginStatements()).toHaveLength(0)
       expect(JSON.stringify(tinStatements()[0])).toContain('2000')
       /** `==>` binds tighter than `||`, so the concatenated query must be parenthesized. */

apps/sim/lib/knowledge/search/queries.ts

@@ -1327,7 +1327,7 @@ export async function executeKeywordSearch(params: KeywordSearchParams): Promise
      */
     const tinQuery =
       params.permitted?.kind === 'unbounded' && tagFilterConditions.length === 0
-        ? await resolveTinKeywordQuery(knowledgeBaseIds, query, FTS_CONFIG)
+        ? await resolveTinKeywordQuery(knowledgeBaseIds, query, FTS_CONFIG, params.budget)
         : null
     annotateSearchDiagnostics({
       ...(params.permitted?.kind === 'unbounded'

apps/sim/lib/knowledge/search/tin-keyword-readiness.test.ts

@@ -21,9 +21,9 @@ it('stays on the GIN projection while the Tin index is incomplete, and remembers
     return []
   })
   queueTableRows(schemaMock.knowledgeBase, [{ id: 'kb-index', isSearchIndex: true }])
-  expect(await resolveTinKeywordQuery(['kb-index'], 'release', 'english')).toBeNull()
+  expect(await resolveTinKeywordQuery(['kb-index'], 'release', 'english', undefined)).toBeNull()
   indexValid = true
-  expect(await resolveTinKeywordQuery(['kb-index'], 'release', 'english')).toBeNull()
+  expect(await resolveTinKeywordQuery(['kb-index'], 'release', 'english', undefined)).toBeNull()
   const readinessReads = dbChainMockFns.execute.mock.calls.filter(([query]) =>
     JSON.stringify(query).includes('indisvalid')
   )

apps/sim/lib/knowledge/search/tin-keyword.test.ts

@@ -12,6 +12,7 @@ vi.mock('@/lib/core/config/feature-flags', () => ({
   isFeatureEnabled: mockIsFeatureEnabled,
 }))
 
+import { SearchBudget, SearchDeadlineError } from '@/lib/knowledge/search/budget'
 import { resolveTinKeywordQuery } from '@/lib/knowledge/search/tin-keyword'
 
 /**
@@ -37,13 +38,15 @@ describe('resolveTinKeywordQuery', () => {
 
   it('is off while the rollout flag is off, without touching the database', async () => {
     mockIsFeatureEnabled.mockResolvedValue(false)
-    expect(await resolveTinKeywordQuery(['kb-flag-off'], 'release notes', 'english')).toBeNull()
+    expect(
+      await resolveTinKeywordQuery(['kb-flag-off'], 'release notes', 'english', undefined)
+    ).toBeNull()
     expect(dbChainMockFns.execute).not.toHaveBeenCalled()
   })
 
   it('translates the analyzed query once every base is a search index', async () => {
     queueTableRows(schemaMock.knowledgeBase, [{ id: 'kb-index', isSearchIndex: true }])
-    expect(await resolveTinKeywordQuery(['kb-index'], 'release notes', 'english')).toBe(
+    expect(await resolveTinKeywordQuery(['kb-index'], 'release notes', 'english', undefined)).toBe(
       '("releas" AND "note")'
     )
   })
@@ -54,18 +57,34 @@ describe('resolveTinKeywordQuery', () => {
       { id: 'kb-plain', isSearchIndex: false },
     ])
     expect(
-      await resolveTinKeywordQuery(['kb-index-2', 'kb-plain'], 'release notes', 'english')
+      await resolveTinKeywordQuery(
+        ['kb-index-2', 'kb-plain'],
+        'release notes',
+        'english',
+        undefined
+      )
     ).toBeNull()
   })
 
   it('is off for a query Tin cannot express', async () => {
     queueTableRows(schemaMock.knowledgeBase, [{ id: 'kb-negation', isSearchIndex: true }])
     rendered = "!'draft'"
-    expect(await resolveTinKeywordQuery(['kb-negation'], '-draft', 'english')).toBeNull()
+    expect(await resolveTinKeywordQuery(['kb-negation'], '-draft', 'english', undefined)).toBeNull()
+  })
+
+  it('reads under the keyword budget, so an expired deadline ends the leg instead of querying', async () => {
+    const expired = new SearchBudget('keyword', performance.now() - 1)
+    await expect(
+      resolveTinKeywordQuery(['kb-expired'], 'release notes', 'english', expired)
+    ).rejects.toBeInstanceOf(SearchDeadlineError)
+    expect(dbChainMockFns.select).not.toHaveBeenCalled()
+    expect(dbChainMockFns.execute).not.toHaveBeenCalled()
   })
 
   it('falls back to GIN instead of failing the search when readiness cannot be read', async () => {
     mockIsFeatureEnabled.mockRejectedValue(new Error('config unavailable'))
-    expect(await resolveTinKeywordQuery(['kb-error'], 'release notes', 'english')).toBeNull()
+    expect(
+      await resolveTinKeywordQuery(['kb-error'], 'release notes', 'english', undefined)
+    ).toBeNull()
   })
 })

apps/sim/lib/knowledge/search/tin-keyword.ts

@@ -1,10 +1,10 @@
-import { db } from '@sim/db'
 import { EMBEDDING_KEYWORD_TIN_INDEX, knowledgeBase } from '@sim/db/schema'
 import { createLogger } from '@sim/logger'
 import { getErrorMessage } from '@sim/utils/errors'
 import { inArray, sql } from 'drizzle-orm'
 import { LRUCache } from 'lru-cache'
 import { isFeatureEnabled } from '@/lib/core/config/feature-flags'
+import { runSearchQuery, type SearchBudget } from '@/lib/knowledge/search/budget'
 import { tinQueryFromTsquery } from '@/lib/knowledge/search/tin-query'
 
 const logger = createLogger('TinKeywordSearch')
@@ -16,14 +16,19 @@ const logger = createLogger('TinKeywordSearch')
 const READINESS_TTL_MS = 60 * 1000
 const SEARCH_INDEX_TTL_MS = 10 * 60 * 1000
 
-/** Whether the Tin index exists and finished building, i.e. the projection is complete. */
-const indexReadiness = new LRUCache<'index', boolean>({
+/**
+ * Whether the Tin index exists and finished building, i.e. the projection is complete. The read
+ * spends the budget of the search that missed the cache.
+ */
+const indexReadiness = new LRUCache<'index', boolean, SearchBudget | undefined>({
   max: 1,
   ttl: READINESS_TTL_MS,
-  fetchMethod: async () => {
-    const [row] = await db.execute<{ valid: boolean }>(sql`
-      SELECT i.indisvalid AS valid FROM pg_index i
-      WHERE i.indexrelid = to_regclass(${EMBEDDING_KEYWORD_TIN_INDEX})`)
+  fetchMethod: async (_key, _stale, { context }) => {
+    const [row] = await runSearchQuery(context, 'keyword.tin_readiness', (executor) =>
+      executor.execute<{ valid: boolean }>(sql`
+        SELECT i.indisvalid AS valid FROM pg_index i
+        WHERE i.indexrelid = to_regclass(${EMBEDDING_KEYWORD_TIN_INDEX})`)
+    )
     return row?.valid === true
   },
 })
@@ -34,13 +39,18 @@ const searchIndexBases = new LRUCache<string, boolean>({
   ttl: SEARCH_INDEX_TTL_MS,
 })
 
-async function allSearchIndexes(knowledgeBaseIds: readonly string[]): Promise<boolean> {
+async function allSearchIndexes(
+  knowledgeBaseIds: readonly string[],
+  budget: SearchBudget | undefined
+): Promise<boolean> {
   const unknown = knowledgeBaseIds.filter((id) => searchIndexBases.get(id) === undefined)
   if (unknown.length > 0) {
-    const rows = await db
-      .select({ id: knowledgeBase.id, isSearchIndex: knowledgeBase.isSearchIndex })
-      .from(knowledgeBase)
-      .where(inArray(knowledgeBase.id, unknown))
+    const rows = await runSearchQuery(budget, 'keyword.tin_readiness', (executor) =>
+      executor
+        .select({ id: knowledgeBase.id, isSearchIndex: knowledgeBase.isSearchIndex })
+        .from(knowledgeBase)
+        .where(inArray(knowledgeBase.id, unknown))
+    )
     for (const row of rows) searchIndexBases.set(row.id, row.isSearchIndex)
   }
   return knowledgeBaseIds.every((id) => searchIndexBases.get(id) === true)
@@ -52,22 +62,31 @@ async function allSearchIndexes(knowledgeBaseIds: readonly string[]): Promise<bo
  * base is not an organization search index (only those are projected), or the query uses a shape
  * TINQL cannot express. The text is analyzed by the same `websearch_to_tsquery` the GIN path
  * uses, so both engines match the same stemmed terms.
+ *
+ * Every read runs under the keyword leg's `budget`, so deciding the engine cannot outlast the leg's
+ * deadline. A read that fails for another reason, including a shared cache read cut short by
+ * another search's deadline, keeps the GIN projection; this search's own expired deadline or
+ * cancellation propagates like any other keyword query's.
  */
 export async function resolveTinKeywordQuery(
   knowledgeBaseIds: readonly string[],
   query: string,
-  ftsConfig: string
+  ftsConfig: string,
+  budget: SearchBudget | undefined
 ): Promise<string | null> {
   if (knowledgeBaseIds.length === 0) return null
   try {
     if (!(await isFeatureEnabled('knowledge-tin-keyword'))) return null
-    if (!(await indexReadiness.fetch('index'))) return null
-    if (!(await allSearchIndexes(knowledgeBaseIds))) return null
-    const [{ rendered }] = await db.execute<{ rendered: string }>(
-      sql`SELECT websearch_to_tsquery(${ftsConfig}::regconfig, ${query})::text AS rendered`
+    if (!(await indexReadiness.fetch('index', { context: budget }))) return null
+    if (!(await allSearchIndexes(knowledgeBaseIds, budget))) return null
+    const [{ rendered }] = await runSearchQuery(budget, 'keyword.tin_query', (executor) =>
+      executor.execute<{ rendered: string }>(
+        sql`SELECT websearch_to_tsquery(${ftsConfig}::regconfig, ${query})::text AS rendered`
+      )
     )
     return tinQueryFromTsquery(rendered)
   } catch (error) {
+    budget?.remaining()
     logger.warn('Tin keyword readiness check failed; using the GIN projection', {
       error: getErrorMessage(error),
     })

packages/db/schema.ts

@@ -3478,18 +3478,18 @@ export const embeddingKeywordSearch = pgTable(
   })
 )
 
+/** The Tin index over {@link embeddingKeywordTin}; valid only once the projection is backfilled. */
+export const EMBEDDING_KEYWORD_TIN_INDEX = 'embedding_keyword_tin_content_idx'
+
 /**
  * BM25 keyword ranking for organization search indexes, served by the Tin text index where the
  * database provides the `tin` extension. `content` is the chunk's `english` lexemes in position
  * order, prefixed with a token naming its knowledge base, so ranking is scoped to one base inside
  * the index and stems exactly as the GIN projection does. Access never enters the row, so ACL
  * changes never rewrite it. Script migration `0019_tin_keyword_projection` installs the extension,
- * the index, and the embedding write trigger that owns these rows, and only where `tin` exists;
- * elsewhere the table stays empty and keyword search keeps the GIN projection.
+ * the index, and the embedding and knowledge base triggers that own these rows, and only where
+ * `tin` exists; elsewhere the table stays empty and keyword search keeps the GIN projection.
  */
-/** The Tin index over {@link embeddingKeywordTin}; valid only once the projection is backfilled. */
-export const EMBEDDING_KEYWORD_TIN_INDEX = 'embedding_keyword_tin_content_idx'
-
 export const embeddingKeywordTin = pgTable('embedding_keyword_tin', {
   id: text('id')
     .primaryKey()

packages/db/script-migrations-paused-billing-attribution.test.ts

@@ -452,6 +452,7 @@ describe('script migration registry', () => {
       '0016_backfill_search_vectors',
       '0017_index_search_documents',
       '0018_repair_workspace_file_content_revision',
+      '0019_tin_keyword_projection',
     ])
   })
 })

packages/db/script-migrations/0016_backfill_search_vectors.postgres.test.ts

@@ -390,6 +390,7 @@ describe.runIf(Boolean(databaseUrl))('search projection upgrade in PostgreSQL',
       { name: '0016_backfill_search_vectors' },
       { name: '0017_index_search_documents' },
       { name: '0018_repair_workspace_file_content_revision' },
+      { name: '0019_tin_keyword_projection' },
     ])
     const [{ complete }] = await sql`SELECT count(*)::int AS complete FROM embedding e
       JOIN embedding_search s ON s.id = e.id JOIN embedding_keyword_search k ON k.id = e.id