diff --git a/apps/sim/app/api/knowledge/search/route.test.ts b/apps/sim/app/api/knowledge/search/route.test.ts index 34af2782b41..e56fb78e8d5 100644 --- a/apps/sim/app/api/knowledge/search/route.test.ts +++ b/apps/sim/app/api/knowledge/search/route.test.ts @@ -6,7 +6,7 @@ import { NextRequest } from 'next/server' import { beforeEach, describe, expect, it, vi } from 'vitest' const mocks = vi.hoisted(() => ({ search: vi.fn() })) -vi.mock('@/lib/knowledge/application/workspace-search', () => ({ +vi.mock('@/lib/sim-search/indexed', () => ({ searchScopedKnowledge: { operation: { id: 'knowledge.search' }, execute: mocks.search }, })) diff --git a/apps/sim/app/api/knowledge/search/route.ts b/apps/sim/app/api/knowledge/search/route.ts index e0004abafb5..1186b6da0f0 100644 --- a/apps/sim/app/api/knowledge/search/route.ts +++ b/apps/sim/app/api/knowledge/search/route.ts @@ -4,12 +4,12 @@ import { internalRateLimits, internalSessionAuth, } from '@/lib/api/server/routes' -import { isLiveEnterpriseSearchEnabled } from '@/lib/core/config/env-flags' import { internalKnowledgeErrorPolicies } from '@/lib/knowledge/api/route-policies' import { knowledgeOperations } from '@/lib/knowledge/application/operations' -import { searchScopedKnowledge } from '@/lib/knowledge/application/workspace-search' import { DEFAULT_RERANKER_MODEL } from '@/lib/knowledge/reranker-models' import { sourceAuthor } from '@/lib/knowledge/search/author' +import { searchScopedKnowledge } from '@/lib/sim-search/indexed' +import { isIndexedOrgSearchEnabled } from '@/lib/sim-search/indexed/gate' import { searchLiveKnowledge } from '@/lib/sim-search/live/application' const DIRECT_SEARCH_VECTOR_BUDGET_MS = 3000 @@ -81,4 +81,5 @@ const liveSearchRoute = defineInternalJsonRoute({ present: (data) => ({ success: true as const, data }), }) -export const POST = isLiveEnterpriseSearchEnabled ? liveSearchRoute : indexedSearchRoute +/** Indexed organization search is dormant unless its gate is on; Live Search serves otherwise. */ +export const POST = isIndexedOrgSearchEnabled() ? indexedSearchRoute : liveSearchRoute diff --git a/apps/sim/app/api/v1/knowledge/search/route.test.ts b/apps/sim/app/api/v1/knowledge/search/route.test.ts index bfa753a205e..7fdc104f0a0 100644 --- a/apps/sim/app/api/v1/knowledge/search/route.test.ts +++ b/apps/sim/app/api/v1/knowledge/search/route.test.ts @@ -7,7 +7,13 @@ * @vitest-environment node */ -import { createMockRequest, knowledgeApiUtilsMock, knowledgeApiUtilsMockFns } from '@sim/testing' +import { + createMockRequest, + knowledgeApiUtilsMock, + knowledgeApiUtilsMockFns, + resetEnvFlagsMock, + setEnvFlags, +} from '@sim/testing' import { getErrorMessage } from '@sim/utils/errors' import { beforeEach, describe, expect, it, vi } from 'vitest' @@ -71,8 +77,12 @@ vi.mock('@/lib/billing/core/billing-attribution', () => ({ resolveBillingAttribution: mockResolveBillingAttribution, resolveSystemBillingAttribution: mockResolveSystemBillingAttribution, })) +const { mockCheckSearchUsageLimits } = vi.hoisted(() => ({ + mockCheckSearchUsageLimits: vi.fn(), +})) + vi.mock('@/lib/billing/core/usage-gate-cache', () => ({ - checkSearchUsageLimits: vi.fn().mockResolvedValue({ isExceeded: false }), + checkSearchUsageLimits: mockCheckSearchUsageLimits, })) vi.mock('@/lib/knowledge/embeddings', () => ({ @@ -118,6 +128,8 @@ const baseKb = (id: string, embeddingModel: string, embeddingDimension = 1536) = describe('v1 knowledge search route — per-KB embedding model', () => { beforeEach(() => { vi.clearAllMocks() + mockCheckSearchUsageLimits.mockResolvedValue({ isExceeded: false }) + resetEnvFlagsMock() mockResolveV1KnowledgeReadAccess.mockResolvedValue({ kind: 'workspace', tokens: ['pub', 'ws'] }) mockAuthenticateRequest.mockResolvedValue({ requestId: 'req-1', @@ -163,6 +175,52 @@ describe('v1 knowledge search route — per-KB embedding model', () => { expect(response.status).toBe(500) }) + it('refuses a search index while indexed organization search is dormant, before any spend', async () => { + setEnvFlags({ isLiveEnterpriseSearchEnabled: true }) + mockCheckKnowledgeBaseAccess.mockResolvedValueOnce({ + hasAccess: true, + knowledgeBase: { ...baseKb('kb-1', 'text-embedding-3-small'), isSearchIndex: true }, + }) + const response = await POST( + createMockRequest('POST', { workspaceId: 'ws-1', knowledgeBaseIds: 'kb-1', query: 'hello' }) + ) + expect(response.status).toBe(409) + expect(await response.json()).toEqual({ + error: 'This search index is inactive; use Sim Search.', + }) + expect(mockGenerateSearchEmbedding).not.toHaveBeenCalled() + expect(mockExecuteKnowledgeSearch).not.toHaveBeenCalled() + }) + + it('refuses a dormant search index for its dormancy, not the caller exhausted usage', async () => { + setEnvFlags({ isLiveEnterpriseSearchEnabled: true }) + mockCheckSearchUsageLimits.mockResolvedValue({ isExceeded: true, message: 'Over limit' }) + mockCheckKnowledgeBaseAccess.mockResolvedValueOnce({ + hasAccess: true, + knowledgeBase: { ...baseKb('kb-1', 'text-embedding-3-small'), isSearchIndex: true }, + }) + const response = await POST( + createMockRequest('POST', { workspaceId: 'ws-1', knowledgeBaseIds: 'kb-1', query: 'hello' }) + ) + expect(response.status).toBe(409) + expect(mockCheckSearchUsageLimits).not.toHaveBeenCalled() + }) + + it('searches a workspace knowledge base while indexed organization search is dormant', async () => { + setEnvFlags({ isLiveEnterpriseSearchEnabled: true }) + mockCheckKnowledgeBaseAccess.mockResolvedValueOnce({ + hasAccess: true, + knowledgeBase: { ...baseKb('kb-1', 'text-embedding-3-small'), isSearchIndex: false }, + }) + const response = await POST( + createMockRequest('POST', { workspaceId: 'ws-1', knowledgeBaseIds: 'kb-1', query: 'hello' }) + ) + expect(response.status).toBe(200) + expect(mockExecuteKnowledgeSearch).toHaveBeenCalledWith( + expect.objectContaining({ knowledgeBaseIds: ['kb-1'], searchIndexOnly: false }) + ) + }) + it('retains the reader provider for ranked results and returned document metadata', async () => { const access = { kind: 'user' as const, userId: 'user-1', tokens: ['reader-token'] } const provider = { diff --git a/apps/sim/app/api/v1/knowledge/search/route.ts b/apps/sim/app/api/v1/knowledge/search/route.ts index 0ea92e767f0..8142d390a31 100644 --- a/apps/sim/app/api/v1/knowledge/search/route.ts +++ b/apps/sim/app/api/v1/knowledge/search/route.ts @@ -24,6 +24,10 @@ import { import { getDocumentTagDefinitions } from '@/lib/knowledge/tags/service' import { buildUndefinedTagsError, validateTagValue } from '@/lib/knowledge/tags/utils' import type { StructuredFilter } from '@/lib/knowledge/types' +import { + isIndexedOrgSearchEnabled, + SEARCH_INDEX_DORMANT_MESSAGE, +} from '@/lib/sim-search/indexed/gate' import { checkKnowledgeBaseAccess, type KnowledgeBaseAccessResult } from '@/app/api/knowledge/utils' import { handleError, resolveV1KnowledgeReadAccess } from '@/app/api/v1/knowledge/utils' import { @@ -77,20 +81,6 @@ export const POST = withRouteHandler(async (request: NextRequest) => { : undefined const billingActorUserId = billingAttribution?.actorUserId ?? userId - /** - * Query embeddings incur hosted cost; tag-only searches do not. Workspace - * keys resolve their system actor and immutable payer from one workspace read. - */ - if (billingAttribution) { - const usage = await checkSearchUsageLimits(billingAttribution) - if (usage.isExceeded) { - return NextResponse.json( - { error: usage.message || 'Usage limit exceeded. Please upgrade your plan to continue.' }, - { status: 402 } - ) - } - } - const knowledgeBaseIds = Array.isArray(parsed.data.body.knowledgeBaseIds) ? parsed.data.body.knowledgeBaseIds : [parsed.data.body.knowledgeBaseIds] @@ -121,6 +111,27 @@ export const POST = withRouteHandler(async (request: NextRequest) => { ) } + /** A search index is readable only while indexed organization search is on. */ + if (!isIndexedOrgSearchEnabled() && accessibleKbs.some((kb) => kb.isSearchIndex)) { + return NextResponse.json({ error: SEARCH_INDEX_DORMANT_MESSAGE }, { status: 409 }) + } + + /** + * Query embeddings incur hosted cost; tag-only searches do not. Workspace + * keys resolve their system actor and immutable payer from one workspace read. + * Admission follows the access and dormancy checks, so a request that could + * never run is refused for that reason rather than for the caller's usage. + */ + if (billingAttribution) { + const usage = await checkSearchUsageLimits(billingAttribution) + if (usage.isExceeded) { + return NextResponse.json( + { error: usage.message || 'Usage limit exceeded. Please upgrade your plan to continue.' }, + { status: 402 } + ) + } + } + let structuredFilters: StructuredFilter[] = [] const tagDefsCache = new Map>>() diff --git a/apps/sim/app/o/[organizationId]/knowledge/[knowledgeBaseId]/[documentId]/page.tsx b/apps/sim/app/o/[organizationId]/knowledge/[knowledgeBaseId]/[documentId]/page.tsx index 6cd806b52bc..090d4f02fa6 100644 --- a/apps/sim/app/o/[organizationId]/knowledge/[knowledgeBaseId]/[documentId]/page.tsx +++ b/apps/sim/app/o/[organizationId]/knowledge/[knowledgeBaseId]/[documentId]/page.tsx @@ -4,7 +4,8 @@ import type { SearchParams } from 'nuqs/server' import { readSearchDocumentResultSchema } from '@/lib/api/contracts/knowledge/documents' import { getSession } from '@/lib/auth' import { OrchestrationError } from '@/lib/core/orchestration/types' -import { readSearchDocument } from '@/lib/knowledge/application/read-search-document' +import { readSearchDocument } from '@/lib/sim-search/indexed' +import { isIndexedOrgSearchEnabled } from '@/lib/sim-search/indexed/gate' import { buildAuthCrossLink } from '@/app/(auth)/auth-redirect' import { loadDocumentReadParams, @@ -22,6 +23,8 @@ export default async function OrganizationDocumentPage({ params, searchParams, }: OrganizationDocumentPageProps) { + /** Indexed documents are served only while indexed organization search is on. */ + if (!isIndexedOrgSearchEnabled()) notFound() const { organizationId, knowledgeBaseId, documentId } = await params const position = await loadDocumentReadParams(searchParams, { strict: true }).catch(() => notFound() diff --git a/apps/sim/lib/knowledge/__integration__/excluded-member-documents.integration.ts b/apps/sim/lib/knowledge/__integration__/excluded-member-documents.integration.ts index 8a6348eb04d..2d6aa1f3838 100644 --- a/apps/sim/lib/knowledge/__integration__/excluded-member-documents.integration.ts +++ b/apps/sim/lib/knowledge/__integration__/excluded-member-documents.integration.ts @@ -17,6 +17,11 @@ import { eq, inArray } from 'drizzle-orm' import { afterAll, afterEach, beforeAll, beforeEach, describe, expect, it, vi } from 'vitest' const provider = vi.hoisted(() => ({ list: vi.fn(), get: vi.fn(), changes: vi.fn() })) +/** This suite covers indexed organization search, which is dormant unless Live Search is off. */ +vi.mock('@/lib/core/config/env-flags', async (importOriginal) => ({ + ...(await importOriginal()), + isLiveEnterpriseSearchEnabled: false, +})) vi.mock('@/connectors/registry.server', () => ({ CONNECTOR_REGISTRY: { google_drive: { diff --git a/apps/sim/lib/knowledge/__integration__/github-member.integration.ts b/apps/sim/lib/knowledge/__integration__/github-member.integration.ts index 9d1e34aa9f2..a46d53ac993 100644 --- a/apps/sim/lib/knowledge/__integration__/github-member.integration.ts +++ b/apps/sim/lib/knowledge/__integration__/github-member.integration.ts @@ -30,6 +30,11 @@ import { generateId } from '@sim/utils/id' import { and, eq, inArray, isNull, sql } from 'drizzle-orm' import { afterAll, afterEach, beforeEach, describe, expect, it, vi } from 'vitest' +/** This suite covers indexed organization search, which is dormant unless Live Search is off. */ +vi.mock('@/lib/core/config/env-flags', async (importOriginal) => ({ + ...(await importOriginal()), + isLiveEnterpriseSearchEnabled: false, +})) vi.mock('@/lib/embeddings', async () => ({ ...(await import('@/lib/embeddings/client')), assertKnowledgeEmbeddingCapacity: async () => {}, @@ -83,7 +88,6 @@ import { listKnowledgeConnectorDocuments, } from '@/lib/knowledge/application/connectors' import { readKnowledgeDocument } from '@/lib/knowledge/application/documents' -import { readIndexedKnowledgeDocument } from '@/lib/knowledge/application/read-indexed-document' import { searchKnowledge } from '@/lib/knowledge/application/search' import { readSearchSourceOverview } from '@/lib/knowledge/application/search-source-overview' import { listSearchSources } from '@/lib/knowledge/application/search-sources' @@ -95,6 +99,7 @@ import { } from '@/lib/knowledge/connectors/sync-limits' import { getDocuments } from '@/lib/knowledge/documents/service' import { getTagUsageStats } from '@/lib/knowledge/tags/service' +import { readIndexedKnowledgeDocument } from '@/lib/sim-search/indexed/documents/read-indexed-document' import { deleteFile } from '@/lib/uploads/core/storage-service' import { downloadFileFromUrl } from '@/lib/uploads/utils/file-utils.server' import { ResolvedSecretTraceRegistry } from '@/executor/utils/resolved-secret-trace-registry' diff --git a/apps/sim/lib/knowledge/__integration__/gitlab-live.integration.ts b/apps/sim/lib/knowledge/__integration__/gitlab-live.integration.ts index e515978e548..2aabfd9fdc4 100644 --- a/apps/sim/lib/knowledge/__integration__/gitlab-live.integration.ts +++ b/apps/sim/lib/knowledge/__integration__/gitlab-live.integration.ts @@ -86,9 +86,9 @@ import { listKnowledgeChunks } from '@/lib/knowledge/application/chunks' import { updateKnowledgeConnectorAccess } from '@/lib/knowledge/application/connector-access' import { updateKnowledgeConnector } from '@/lib/knowledge/application/connectors' import { readKnowledgeDocument } from '@/lib/knowledge/application/documents' -import { readSearchDocument } from '@/lib/knowledge/application/read-search-document' import { searchKnowledge } from '@/lib/knowledge/application/search' import { executeSync } from '@/lib/knowledge/connectors/sync-engine' +import { readSearchDocument } from '@/lib/sim-search/indexed/documents/read-search-document' import * as storage from '@/lib/uploads/core/storage-service' import { downloadFileFromUrl } from '@/lib/uploads/utils/file-utils.server' import { PATCH as updateConnectorRoute } from '@/app/api/knowledge/[id]/connectors/[connectorId]/route' diff --git a/apps/sim/lib/knowledge/__integration__/gmail-member.integration.ts b/apps/sim/lib/knowledge/__integration__/gmail-member.integration.ts index 02bebeabca4..1d12d4cdd5b 100644 --- a/apps/sim/lib/knowledge/__integration__/gmail-member.integration.ts +++ b/apps/sim/lib/knowledge/__integration__/gmail-member.integration.ts @@ -24,6 +24,11 @@ import { eq, inArray } from 'drizzle-orm' import { afterAll, afterEach, beforeAll, beforeEach, describe, expect, it, vi } from 'vitest' const counters = vi.hoisted(() => ({ embeddedTexts: 0 })) +/** This suite covers indexed organization search, which is dormant unless Live Search is off. */ +vi.mock('@/lib/core/config/env-flags', async (importOriginal) => ({ + ...(await importOriginal()), + isLiveEnterpriseSearchEnabled: false, +})) vi.mock('@/lib/embeddings', async () => ({ ...(await import('@/lib/embeddings/client')), assertKnowledgeEmbeddingCapacity: async () => {}, diff --git a/apps/sim/lib/knowledge/__integration__/google-calendar-member.integration.ts b/apps/sim/lib/knowledge/__integration__/google-calendar-member.integration.ts index f00ea34e458..1ddb3d3407c 100644 --- a/apps/sim/lib/knowledge/__integration__/google-calendar-member.integration.ts +++ b/apps/sim/lib/knowledge/__integration__/google-calendar-member.integration.ts @@ -24,6 +24,11 @@ import { generateId } from '@sim/utils/id' import { and, eq, inArray, isNull } from 'drizzle-orm' import { afterAll, beforeAll, beforeEach, describe, expect, it, vi } from 'vitest' +/** This suite covers indexed organization search, which is dormant unless Live Search is off. */ +vi.mock('@/lib/core/config/env-flags', async (importOriginal) => ({ + ...(await importOriginal()), + isLiveEnterpriseSearchEnabled: false, +})) vi.mock('@/lib/embeddings', async () => ({ ...(await import('@/lib/embeddings/client')), assertKnowledgeEmbeddingCapacity: async () => {}, diff --git a/apps/sim/lib/knowledge/__integration__/jira-member.integration.ts b/apps/sim/lib/knowledge/__integration__/jira-member.integration.ts index a8f970c6b87..0df85cd89da 100644 --- a/apps/sim/lib/knowledge/__integration__/jira-member.integration.ts +++ b/apps/sim/lib/knowledge/__integration__/jira-member.integration.ts @@ -25,6 +25,11 @@ import { generateId } from '@sim/utils/id' import { and, eq, inArray, isNull } from 'drizzle-orm' import { afterAll, beforeAll, beforeEach, describe, expect, it, vi } from 'vitest' +/** This suite covers indexed organization search, which is dormant unless Live Search is off. */ +vi.mock('@/lib/core/config/env-flags', async (importOriginal) => ({ + ...(await importOriginal()), + isLiveEnterpriseSearchEnabled: false, +})) vi.mock('@/lib/embeddings', async () => ({ ...(await import('@/lib/embeddings/client')), assertKnowledgeEmbeddingCapacity: async () => {}, diff --git a/apps/sim/lib/knowledge/__integration__/kb-block-search.integration.ts b/apps/sim/lib/knowledge/__integration__/kb-block-search.integration.ts index 22490631337..06d16343b9c 100644 --- a/apps/sim/lib/knowledge/__integration__/kb-block-search.integration.ts +++ b/apps/sim/lib/knowledge/__integration__/kb-block-search.integration.ts @@ -11,12 +11,12 @@ import { } from '@/lib/knowledge/__integration__/seed-source-access-fixture' import { createKnowledgeAccessProvider } from '@/lib/knowledge/access/scope' import { - forgetProjectionFilled, resolvePermittedDocuments, retrieveKnowledgeSearch, VECTOR_PROBE_DOCUMENT_LIMIT, } from '@/lib/knowledge/search/queries' import { embeddingVectorValues } from '@/lib/knowledge/vector-columns' +import { forgetProjectionFilled } from '@/lib/sim-search/indexed/retrieval' describe('API-key KB block fan-out', () => { const ids = createKnowledgeAclFixtureIds() diff --git a/apps/sim/lib/knowledge/__integration__/knowledge-projection.integration.ts b/apps/sim/lib/knowledge/__integration__/knowledge-projection.integration.ts index 7f12a726888..a4f6743c434 100644 --- a/apps/sim/lib/knowledge/__integration__/knowledge-projection.integration.ts +++ b/apps/sim/lib/knowledge/__integration__/knowledge-projection.integration.ts @@ -41,8 +41,13 @@ import { and, eq, inArray, isNull, sql } from 'drizzle-orm' import postgres from 'postgres' import { afterAll, beforeAll, beforeEach, describe, expect, it, vi } from 'vitest' +/** This suite covers indexed organization search, which is dormant unless Live Search is off. */ +vi.mock('@/lib/core/config/env-flags', async (importOriginal) => ({ + ...(await importOriginal()), + isLiveEnterpriseSearchEnabled: false, +})) /** The TINQL `resolveTinKeywordQuery` renders for `fixture`: its `english` stem, quoted. */ -vi.mock('@/lib/knowledge/search/tin-keyword', () => ({ +vi.mock('@/lib/sim-search/indexed/retrieval/tin-keyword', () => ({ resolveTinKeywordQuery: async () => '"fixtur"', })) @@ -76,11 +81,11 @@ import type { import { leaseTransaction } from '@/lib/knowledge/connectors/sync-lock' import { executeKeywordSearch, - forgetProjectionFilled, handleVectorOnlySearch, liveSourceAccessFor, } from '@/lib/knowledge/search/queries' import { GITHUB_INSTALLATION_PROVIDER_ID } from '@/lib/oauth/github-installation-types' +import { forgetProjectionFilled } from '@/lib/sim-search/indexed/retrieval' const ids = createKnowledgeAclFixtureIds() const connectorId = generateId() diff --git a/apps/sim/lib/knowledge/__integration__/organization-mcp-search.integration.ts b/apps/sim/lib/knowledge/__integration__/organization-mcp-search.integration.ts index 76aab3ba344..b12e41b5a0b 100644 --- a/apps/sim/lib/knowledge/__integration__/organization-mcp-search.integration.ts +++ b/apps/sim/lib/knowledge/__integration__/organization-mcp-search.integration.ts @@ -85,14 +85,14 @@ import { confluencePageAcl } from '@/lib/knowledge/access/confluence-permissions import { listKnowledgeChunks } from '@/lib/knowledge/application/chunks' import { readKnowledgeDocument } from '@/lib/knowledge/application/documents' import { listKnowledgeBaseCatalog } from '@/lib/knowledge/application/knowledge-bases' -import { readIndexedKnowledgeDocument } from '@/lib/knowledge/application/read-indexed-document' import { prepareSearchSource } from '@/lib/knowledge/application/sim-search' -import { searchScopedKnowledge } from '@/lib/knowledge/application/workspace-search' import { createContentSyncLease } from '@/lib/knowledge/connectors/sync-lock' import { addDocument, persistDocumentAcls } from '@/lib/knowledge/connectors/sync-persistence' import { processDocumentAsync } from '@/lib/knowledge/documents/service' import { getSearchMcpUrl } from '@/lib/knowledge/mcp/urls' import { replaceKnowledgeEmbeddingSecretProvenanceInTx } from '@/lib/knowledge/secret-provenance' +import { readIndexedKnowledgeDocument } from '@/lib/sim-search/indexed/documents/read-indexed-document' +import { searchScopedKnowledge } from '@/lib/sim-search/indexed/search/scoped-search' import { DELETE, GET, POST } from '@/app/api/mcp/search/organizations/[organizationId]/route' import { ResolvedSecretTraceRegistry } from '@/executor/utils/resolved-secret-trace-registry' diff --git a/apps/sim/lib/knowledge/__integration__/provider-processing-recovery.integration.ts b/apps/sim/lib/knowledge/__integration__/provider-processing-recovery.integration.ts index c653e69eac0..7acce090da3 100644 --- a/apps/sim/lib/knowledge/__integration__/provider-processing-recovery.integration.ts +++ b/apps/sim/lib/knowledge/__integration__/provider-processing-recovery.integration.ts @@ -25,6 +25,11 @@ import { and, eq, inArray, sql } from 'drizzle-orm' import { afterAll, afterEach, beforeAll, beforeEach, describe, expect, it, vi } from 'vitest' const fixtureStorage = vi.hoisted(() => ({ root: '' })) +/** This suite covers indexed organization search, which is dormant unless Live Search is off. */ +vi.mock('@/lib/core/config/env-flags', async (importOriginal) => ({ + ...(await importOriginal()), + isLiveEnterpriseSearchEnabled: false, +})) vi.mock('@/lib/uploads/core/setup.server', () => ({ get UPLOAD_DIR_SERVER() { return fixtureStorage.root @@ -49,7 +54,6 @@ import { seedKnowledgeMemberFixture, } from '@/lib/knowledge/__integration__/seed-source-access-fixture' import { searchKnowledge } from '@/lib/knowledge/application/search' -import { searchScopedKnowledge } from '@/lib/knowledge/application/workspace-search' import { materializeDocumentAcls, recordMemberObservations, @@ -61,6 +65,7 @@ import { knowledgeDocumentProcessingOutboxHandlers } from '@/lib/knowledge/docum import { assertDocumentProcessingPayload } from '@/lib/knowledge/documents/processing-payload' import * as providerContinuation from '@/lib/knowledge/documents/processing-provider-continuation' import { processDocumentsWithQueue } from '@/lib/knowledge/documents/service' +import { searchScopedKnowledge } from '@/lib/sim-search/indexed/search/scoped-search' const PNG = Buffer.from( 'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAwMCAO+jRZkAAAAASUVORK5CYII=', diff --git a/apps/sim/lib/knowledge/__integration__/read-indexed-document.integration.ts b/apps/sim/lib/knowledge/__integration__/read-indexed-document.integration.ts index a4cb6ff0afa..300929aa778 100644 --- a/apps/sim/lib/knowledge/__integration__/read-indexed-document.integration.ts +++ b/apps/sim/lib/knowledge/__integration__/read-indexed-document.integration.ts @@ -21,6 +21,11 @@ import { eq, inArray } from 'drizzle-orm' import { afterAll, beforeAll, describe, expect, it, vi } from 'vitest' const fixtures = vi.hoisted(() => ({ storageRoot: '' })) +/** This suite covers indexed organization search, which is dormant unless Live Search is off. */ +vi.mock('@/lib/core/config/env-flags', async (importOriginal) => ({ + ...(await importOriginal()), + isLiveEnterpriseSearchEnabled: false, +})) vi.mock('@/lib/uploads/core/setup.server', () => ({ get UPLOAD_DIR_SERVER() { return fixtures.storageRoot @@ -45,13 +50,13 @@ import { seedKnowledgeAclFixture, } from '@/lib/knowledge/__integration__/seed-source-access-fixture' import { confluencePageAcl } from '@/lib/knowledge/access/confluence-permissions' -import { - type ReadIndexedKnowledgeDocumentInput, - readIndexedKnowledgeDocument, -} from '@/lib/knowledge/application/read-indexed-document' import { createContentSyncLease } from '@/lib/knowledge/connectors/sync-lock' import { addDocument, persistDocumentAcls } from '@/lib/knowledge/connectors/sync-persistence' import { processDocumentAsync } from '@/lib/knowledge/documents/service' +import { + type ReadIndexedKnowledgeDocumentInput, + readIndexedKnowledgeDocument, +} from '@/lib/sim-search/indexed/documents/read-indexed-document' import { ResolvedSecretTraceRegistry } from '@/executor/utils/resolved-secret-trace-registry' describe('indexed document references', () => { diff --git a/apps/sim/lib/knowledge/__integration__/stored-document-recovery.integration.ts b/apps/sim/lib/knowledge/__integration__/stored-document-recovery.integration.ts index dc7b1dcecd0..85d48d53194 100644 --- a/apps/sim/lib/knowledge/__integration__/stored-document-recovery.integration.ts +++ b/apps/sim/lib/knowledge/__integration__/stored-document-recovery.integration.ts @@ -26,6 +26,11 @@ const fixture = vi.hoisted(() => ({ listRuns: vi.fn(), batchTrigger: vi.fn(), })) +/** This suite covers indexed organization search, which is dormant unless Live Search is off. */ +vi.mock('@/lib/core/config/env-flags', async (importOriginal) => ({ + ...(await importOriginal()), + isLiveEnterpriseSearchEnabled: false, +})) vi.mock('@/lib/core/config/trigger-runtime', () => ({ isInsideTriggerRun: () => fixture.useTrigger, })) @@ -78,7 +83,6 @@ import { seedKnowledgeAclFixture, } from '@/lib/knowledge/__integration__/seed-source-access-fixture' import { searchKnowledge } from '@/lib/knowledge/application/search' -import { searchScopedKnowledge } from '@/lib/knowledge/application/workspace-search' import { createContentSyncLease } from '@/lib/knowledge/connectors/sync-lock' import { addDocument } from '@/lib/knowledge/connectors/sync-persistence' import { sweepStuckDocuments } from '@/lib/knowledge/connectors/sync-primitives' @@ -99,6 +103,7 @@ import { retryDocumentProcessing, } from '@/lib/knowledge/documents/service' import { MAX_PROCESSING_ATTEMPTS, QUEUED_DISPATCH_GRACE_MS } from '@/lib/knowledge/documents/types' +import { searchScopedKnowledge } from '@/lib/sim-search/indexed/search/scoped-search' import type { SyncResult } from '@/connectors/types' const fixtures: ReturnType[] = [] diff --git a/apps/sim/lib/knowledge/__integration__/unfilled-projection-source.integration.ts b/apps/sim/lib/knowledge/__integration__/unfilled-projection-source.integration.ts index 163f1fcea0d..f2c660bbf32 100644 --- a/apps/sim/lib/knowledge/__integration__/unfilled-projection-source.integration.ts +++ b/apps/sim/lib/knowledge/__integration__/unfilled-projection-source.integration.ts @@ -27,8 +27,13 @@ import { isRecordLike } from '@sim/utils/object' import { eq, inArray, sql } from 'drizzle-orm' import { afterAll, afterEach, beforeAll, beforeEach, describe, expect, it, vi } from 'vitest' +/** This suite covers indexed organization search, which is dormant unless Live Search is off. */ +vi.mock('@/lib/core/config/env-flags', async (importOriginal) => ({ + ...(await importOriginal()), + isLiveEnterpriseSearchEnabled: false, +})) /** The TINQL `resolveTinKeywordQuery` renders for `fixture`: its `english` stem, quoted. */ -vi.mock('@/lib/knowledge/search/tin-keyword', () => ({ +vi.mock('@/lib/sim-search/indexed/retrieval/tin-keyword', () => ({ resolveTinKeywordQuery: async () => '"fixtur"', })) @@ -44,11 +49,11 @@ import type { } from '@/lib/knowledge/access/types' import { executeKeywordSearch, - forgetProjectionFilled, handleVectorOnlySearch, liveSourceAccessFor, } from '@/lib/knowledge/search/queries' import { GITHUB_INSTALLATION_PROVIDER_ID } from '@/lib/oauth/github-installation-types' +import { forgetProjectionFilled } from '@/lib/sim-search/indexed/retrieval' const ids = createKnowledgeAclFixtureIds() const connectorId = generateId() diff --git a/apps/sim/lib/knowledge/api/route-policies.test.ts b/apps/sim/lib/knowledge/api/route-policies.test.ts index 25a4c789890..2f5cf0cc6aa 100644 --- a/apps/sim/lib/knowledge/api/route-policies.test.ts +++ b/apps/sim/lib/knowledge/api/route-policies.test.ts @@ -12,7 +12,11 @@ import { WorkspaceApiKeyScopeAuthorizationError, } from '@/lib/core/application' import { OrchestrationError } from '@/lib/core/orchestration/types' -import { v2KnowledgeErrorPolicies } from '@/lib/knowledge/api/route-policies' +import { + internalKnowledgeErrorPolicies, + v2KnowledgeErrorPolicies, +} from '@/lib/knowledge/api/route-policies' +import { SearchIndexDormantError } from '@/lib/sim-search/indexed/gate' describe('v2 knowledge error policies', () => { it.each([ @@ -70,3 +74,23 @@ describe('v2 knowledge error policies', () => { }) }) }) + +describe('dormant search index projection', () => { + const message = 'This search index is inactive; use Sim Search.' + + it('refuses the internal search as a conflict with the dormant message', () => { + expect(internalKnowledgeErrorPolicies.search.project(new SearchIndexDormantError())).toEqual({ + status: 409, + body: { error: message }, + headers: undefined, + }) + }) + + it('refuses the public search as a conflict with the dormant message', async () => { + const response = v2KnowledgeErrorPolicies.concealKnowledgeBaseUsageAuthorization.render( + new SearchIndexDormantError() + ) + expect(response?.status).toBe(409) + expect(await response?.json()).toEqual({ error: { code: 'CONFLICT', message } }) + }) +}) diff --git a/apps/sim/lib/knowledge/api/route-policies.ts b/apps/sim/lib/knowledge/api/route-policies.ts index 5f046e65909..e4255f0ceb8 100644 --- a/apps/sim/lib/knowledge/api/route-policies.ts +++ b/apps/sim/lib/knowledge/api/route-policies.ts @@ -16,6 +16,7 @@ import { KnowledgeUsageLimitExceededError } from '@/lib/knowledge/application/bi import { KnowledgeDocumentNotReadyError } from '@/lib/knowledge/application/chunk-errors' import { KnowledgeSearchProvenanceUnavailableError } from '@/lib/knowledge/application/search' import { KnowledgeDocumentUnsupportedMediaTypeError } from '@/lib/knowledge/application/upload-sessions' +import { SearchIndexDormantError } from '@/lib/sim-search/indexed/gate' import { v2Error } from '@/app/api/v2/lib/response' function internalKnowledgeErrorPolicy(unhandledMessage: string): InternalErrorPolicy { @@ -47,6 +48,10 @@ const internalKnowledgeSearchErrorPolicy: InternalErrorPolicy = { if (error instanceof KnowledgeSearchProvenanceUnavailableError) { return internalErrorResponse(422, { error: error.message }) } + /** The knowledge base exists and is readable; its state is what refuses the search. */ + if (error instanceof SearchIndexDormantError) { + return internalErrorResponse(409, { error: error.message }) + } return internalOrchestrationErrorPolicy.project(error) }, unhandled: () => internalErrorResponse(500, { error: 'Failed to perform vector search' }), @@ -124,6 +129,9 @@ const v2KnowledgeUsageErrorPolicy = { if (error instanceof KnowledgeSearchProvenanceUnavailableError) { return v2Error('CONFLICT', error.message) } + if (error instanceof SearchIndexDormantError) { + return v2Error('CONFLICT', error.message) + } return v2OrchestrationErrorPolicy.render(error) }, } satisfies V2ErrorPolicy diff --git a/apps/sim/lib/knowledge/application/search.test.ts b/apps/sim/lib/knowledge/application/search.test.ts index e3e8081b4a5..68d46c8e2dc 100644 --- a/apps/sim/lib/knowledge/application/search.test.ts +++ b/apps/sim/lib/knowledge/application/search.test.ts @@ -3,8 +3,8 @@ */ import { member } from '@sim/db/schema' -import { queueTableRows, resetDbChainMock } from '@sim/testing' -import { beforeEach, describe, expect, it, vi } from 'vitest' +import { queueTableRows, resetDbChainMock, resetEnvFlagsMock, setEnvFlags } from '@sim/testing' +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' import { OrchestrationError } from '@/lib/core/orchestration/types' const mocks = vi.hoisted(() => ({ @@ -115,6 +115,7 @@ vi.mock('@/lib/knowledge/secret-provenance', () => ({ })) import { searchKnowledge } from '@/lib/knowledge/application/search' +import { SearchIndexDormantError } from '@/lib/sim-search/indexed/gate' const workspace = { workspaceId: 'workspace-1', @@ -344,6 +345,62 @@ describe('knowledge search application use case', () => { expect(mocks.executeSearch).not.toHaveBeenCalled() }) + describe('while indexed organization search is dormant', () => { + const principal = { kind: 'session', userId: 'user-1', sessionId: 'session-1' } as const + beforeEach(() => setEnvFlags({ isLiveEnterpriseSearchEnabled: true })) + afterEach(resetEnvFlagsMock) + + it.each([ + ['an organization', { workspaceId: null, organizationId: 'org-canonical' }], + ['a workspace', {}], + ])( + 'refuses %s search index named by id before any spend or retrieval', + async (_owner, owner) => { + mocks.getKnowledgeBase.mockResolvedValue({ + ...knowledgeBase, + ...owner, + isSearchIndex: true, + }) + queueTableRows(member, [{ role: 'member' }]) + const search = searchKnowledge.execute({ + principal, + input: { knowledgeBaseIds: ['knowledge-1'], query: 'answer', topK: 5 }, + }) + await expect(search).rejects.toBeInstanceOf(SearchIndexDormantError) + await expect(search).rejects.toThrow('This search index is inactive; use Sim Search.') + expect(mocks.resolveBilling).not.toHaveBeenCalled() + expect(mocks.generateEmbedding).not.toHaveBeenCalled() + expect(mocks.executeSearch).not.toHaveBeenCalled() + expect(mocks.recordActivity).not.toHaveBeenCalled() + } + ) + + it('refuses a batch that names a search index beside a workspace knowledge base', async () => { + mocks.getKnowledgeBases.mockResolvedValue([ + knowledgeBase, + { ...knowledgeBase, id: 'knowledge-2', isSearchIndex: true }, + ]) + await expect( + searchKnowledge.execute({ + principal, + input: { knowledgeBaseIds: ['knowledge-1', 'knowledge-2'], query: 'answer', topK: 5 }, + }) + ).rejects.toBeInstanceOf(SearchIndexDormantError) + expect(mocks.executeSearch).not.toHaveBeenCalled() + }) + + it('searches a workspace knowledge base exactly as before', async () => { + const result = await searchKnowledge.execute({ + principal, + input: { knowledgeBaseIds: ['knowledge-1'], query: 'answer', topK: 5 }, + }) + expect(result.results).toHaveLength(1) + expect(mocks.executeSearch).toHaveBeenCalledWith( + expect.objectContaining({ knowledgeBaseIds: ['knowledge-1'], searchIndexOnly: false }) + ) + }) + }) + it('authorizes every canonical knowledge base before billing and search', async () => { const result = await searchKnowledge.execute({ principal: { kind: 'session', userId: 'user-1', sessionId: 'session-1' }, diff --git a/apps/sim/lib/knowledge/application/search.ts b/apps/sim/lib/knowledge/application/search.ts index 706071b894b..b01ad8d9eed 100644 --- a/apps/sim/lib/knowledge/application/search.ts +++ b/apps/sim/lib/knowledge/application/search.ts @@ -56,6 +56,7 @@ import { import { getDocumentTagDefinitionsByKnowledgeBaseIds } from '@/lib/knowledge/tags/service' import type { DocumentTagDefinition } from '@/lib/knowledge/tags/types' import type { StructuredFilter } from '@/lib/knowledge/types' +import { assertSearchIndexesActive } from '@/lib/sim-search/indexed/gate' import { estimateTokenCount } from '@/lib/tokenization/estimators' import { ResolvedSecretTraceRegistry } from '@/executor/utils/resolved-secret-trace-registry' import { getRerankModelPricing } from '@/providers/models' @@ -296,6 +297,8 @@ export async function runKnowledgeSearch({ scopeKind: context.organizationId ? 'organization' : 'workspace', knowledgeBaseCount: context.knowledgeBases.length, }) + /** A search index is readable only while indexed organization search is on; nothing is spent first. */ + assertSearchIndexesActive(context.knowledgeBases) input.signal?.throwIfAborted() const requestId = generateRequestId() const hasQuery = Boolean(input.query?.trim()) diff --git a/apps/sim/lib/knowledge/connectors/indexing-policy.ts b/apps/sim/lib/knowledge/connectors/indexing-policy.ts index 37b680847a2..bbb61df298f 100644 --- a/apps/sim/lib/knowledge/connectors/indexing-policy.ts +++ b/apps/sim/lib/knowledge/connectors/indexing-policy.ts @@ -1,13 +1,13 @@ import { knowledgeBase } from '@sim/db/schema' import { eq } from 'drizzle-orm' -import { isLiveEnterpriseSearchEnabled } from '@/lib/core/config/env-flags' +import { isIndexedOrgSearchEnabled } from '@/lib/sim-search/indexed/gate' /** Federated Search keeps source configuration but does not crawl content into a knowledge base. */ export function requiresConnectorIndexing(isSearchIndex?: boolean | null): boolean { - return !isLiveEnterpriseSearchEnabled || isSearchIndex !== true + return isIndexedOrgSearchEnabled() || isSearchIndex !== true } /** Keeps federated sources out of bounded indexing scheduler pages. */ export function connectorIndexingCondition() { - return isLiveEnterpriseSearchEnabled ? eq(knowledgeBase.isSearchIndex, false) : undefined + return isIndexedOrgSearchEnabled() ? undefined : eq(knowledgeBase.isSearchIndex, false) } diff --git a/apps/sim/lib/knowledge/mcp/server.protocol.test.ts b/apps/sim/lib/knowledge/mcp/server.protocol.test.ts index c10b8609894..13c9ae45ba4 100644 --- a/apps/sim/lib/knowledge/mcp/server.protocol.test.ts +++ b/apps/sim/lib/knowledge/mcp/server.protocol.test.ts @@ -19,7 +19,7 @@ vi.mock('@/lib/api/server/routes/v2-json-route', () => ({ vi.mock('@/lib/knowledge/application/search', () => ({ searchKnowledge: { execute: mocks.indexedSearch }, })) -vi.mock('@/lib/knowledge/application/read-indexed-document', () => ({ +vi.mock('@/lib/sim-search/indexed', () => ({ readIndexedKnowledgeDocument: { execute: mocks.indexedRead }, })) vi.mock('@/lib/sim-search/live/application', () => ({ diff --git a/apps/sim/lib/knowledge/mcp/server.test.ts b/apps/sim/lib/knowledge/mcp/server.test.ts index 7d887b3545f..e76d44ab0fa 100644 --- a/apps/sim/lib/knowledge/mcp/server.test.ts +++ b/apps/sim/lib/knowledge/mcp/server.test.ts @@ -51,7 +51,7 @@ vi.mock('@/lib/api/server/routes/v2-json-route', () => ({ vi.mock('@/lib/knowledge/application/search', () => ({ searchKnowledge: { execute: mocks.search }, })) -vi.mock('@/lib/knowledge/application/read-indexed-document', () => ({ +vi.mock('@/lib/sim-search/indexed', () => ({ readIndexedKnowledgeDocument: { execute: mocks.read }, })) vi.mock('@/lib/sim-search/live/application', () => ({ diff --git a/apps/sim/lib/knowledge/mcp/server.ts b/apps/sim/lib/knowledge/mcp/server.ts index 818472e2dca..b8c8b822204 100644 --- a/apps/sim/lib/knowledge/mcp/server.ts +++ b/apps/sim/lib/knowledge/mcp/server.ts @@ -15,13 +15,11 @@ import { import type { V2ApiKeyAuthContext } from '@/lib/api/server/routes/v2-api-key-auth' import { v2RateLimits } from '@/lib/api/server/routes/v2-json-route' import type { ApplicationOperation } from '@/lib/core/application' -import { isLiveEnterpriseSearchEnabled } from '@/lib/core/config/env-flags' import type { ResourceScope } from '@/lib/core/resource-scope' import { afterResponse } from '@/lib/core/utils/after-response' import { getBaseUrl } from '@/lib/core/utils/urls' import { organizationSearchChatOperation } from '@/lib/knowledge/application/chat-operations' import { knowledgeOperations } from '@/lib/knowledge/application/operations' -import { readIndexedKnowledgeDocument } from '@/lib/knowledge/application/read-indexed-document' import { searchKnowledge } from '@/lib/knowledge/application/search' import { recordOrganizationSearchMcpActivity, @@ -29,6 +27,8 @@ import { } from '@/lib/knowledge/mcp/activity' import { createKnowledgeDocumentCitation, liveCitationId } from '@/lib/knowledge/search/citation' import { toolError } from '@/lib/mcp/tool-result' +import { readIndexedKnowledgeDocument } from '@/lib/sim-search/indexed' +import { isIndexedOrgSearchEnabled } from '@/lib/sim-search/indexed/gate' import { readLiveDocument, searchLiveKnowledge } from '@/lib/sim-search/live/application' import { v2CaughtOrchestrationError } from '@/app/api/v2/lib/response' import { projectResolvedSecretModelContent } from '@/executor/utils/resolved-secret-content-projection' @@ -73,6 +73,8 @@ export function createKnowledgeMcpServer(context: KnowledgeMcpContext): McpServe const scope: ResourceScope = { kind: 'organization', organizationId } const principal = auth.principal const server = new McpServer({ name: 'Sim Search', version: '1.0.0' }) + /** Live Search serves both tools unless indexed organization search is on. */ + const liveSearch = !isIndexedOrgSearchEnabled() async function execute( toolName: SearchMcpActivityInput['toolName'], @@ -145,15 +147,15 @@ export function createKnowledgeMcpServer(context: KnowledgeMcpContext): McpServe 'search', { title: 'Search', - description: isLiveEnterpriseSearchEnabled + description: liveSearch ? 'Search this organization’s sources through their live APIs, within your access and the admin’s source settings. Use source and date filters to narrow results, or nativeQueries for provider queries and pagination. Inspect live.accounts for provider status and continuation cursors, and live.guidance for query syntax. Results are candidates, not proof of complete coverage. Use read_document with the exact returned documentId for context and cite citationUrl when available.' : 'Search accessible passages in this organization’s Search index. Use source (for example, jira), modifiedAfter (an ISO timestamp), or documentIds to narrow results. Results are candidates; score is similarity, not answer confidence. Use read_document for context and cite citationUrl.', - inputSchema: isLiveEnterpriseSearchEnabled ? liveSearchMcpSchema : searchMcpSchema, + inputSchema: liveSearch ? liveSearchMcpSchema : searchMcpSchema, annotations: READ_ONLY, }, async (input: unknown, extra: { signal: AbortSignal }) => execute('search', knowledgeOperations.search, extra.signal, async (registry, signal) => { - if (isLiveEnterpriseSearchEnabled) { + if (liveSearch) { const { query, topK, nativeQueries, ...filters } = liveSearchMcpSchema.parse(input) const result = await searchLiveKnowledge.execute({ principal, @@ -242,12 +244,10 @@ export function createKnowledgeMcpServer(context: KnowledgeMcpContext): McpServe 'read_document', { title: 'Read document', - description: isLiveEnterpriseSearchEnabled + description: liveSearch ? 'Read a live document using the exact documentId returned by search. Access and the admin’s source settings are checked again on every read. When hasMore is true, pass next.startChunkIndex and next.startOffset with the same documentId to continue. Cite citationUrl when available.' : 'Read an indexed document by documentId from search or its original URL. URLs must match an accessible indexed source; this tool does not browse the web. Set aroundChunkIndex to a search hit’s chunkIndex for nearby context, or use offset for sequential pages. When pagination.hasMore is true, continue with pagination.offset + pagination.limit. Cite citationUrl. Documents still indexing return metadata only.', - inputSchema: isLiveEnterpriseSearchEnabled - ? readLiveDocumentMcpSchema - : readDocumentMcpSchema, + inputSchema: liveSearch ? readLiveDocumentMcpSchema : readDocumentMcpSchema, annotations: READ_ONLY, }, async (raw: unknown, extra: { signal: AbortSignal }) => @@ -256,7 +256,7 @@ export function createKnowledgeMcpServer(context: KnowledgeMcpContext): McpServe knowledgeOperations.readDocument, extra.signal, async (registry, signal) => { - if (isLiveEnterpriseSearchEnabled) { + if (liveSearch) { const input = readLiveDocumentMcpSchema.parse(raw) const result = await readLiveDocument.execute({ principal, diff --git a/apps/sim/lib/knowledge/projection/run.test.ts b/apps/sim/lib/knowledge/projection/run.test.ts index 876ca3d41fa..546bd4062de 100644 --- a/apps/sim/lib/knowledge/projection/run.test.ts +++ b/apps/sim/lib/knowledge/projection/run.test.ts @@ -1,6 +1,7 @@ /** * @vitest-environment node */ +import { resetEnvFlagsMock, setEnvFlags } from '@sim/testing' import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' const mocks = vi.hoisted(() => ({ @@ -33,6 +34,7 @@ const drained = { settled: 1, deferred: 0, pages: 2, written: 3, remaining: fals describe('runKnowledgeProjectionPass', () => { beforeEach(() => { vi.clearAllMocks() + resetEnvFlagsMock() mocks.runProjection.mockResolvedValue(drained) mocks.isFeatureEnabled.mockResolvedValue(false) mocks.marks.mockReturnValue(2) @@ -46,6 +48,10 @@ describe('runKnowledgeProjectionPass', () => { remaining: false, }) expect(mocks.runProjection).toHaveBeenCalledTimes(2) + expect(mocks.runProjection).toHaveBeenCalledWith( + expect.anything(), + expect.objectContaining({ includeTin: true }) + ) expect(mocks.isFeatureEnabled).toHaveBeenCalledWith('knowledge-projection-fill') expect(mocks.markUnfilled).not.toHaveBeenCalled() expect(mocks.end).toHaveBeenCalledTimes(2) @@ -99,10 +105,26 @@ describe('runKnowledgeProjectionPass', () => { .mockResolvedValueOnce({ marked: 0, cursor: null }) const result = await runKnowledgeProjectionPass({ budgetMs: 60_000 }) expect(result).toMatchObject({ filled: 2, remaining: false }) - expect(mocks.markUnfilled).toHaveBeenNthCalledWith(1, expect.anything(), undefined) - expect(mocks.markUnfilled).toHaveBeenNthCalledWith(2, expect.anything(), { - projection: 0, - afterId: 'row-2', + expect(mocks.markUnfilled).toHaveBeenNthCalledWith(1, expect.anything(), undefined, { + includeSearchIndexes: true, + }) + expect(mocks.markUnfilled).toHaveBeenNthCalledWith( + 2, + expect.anything(), + { projection: 0, afterId: 'row-2' }, + { includeSearchIndexes: true } + ) + }) + + it('writes no Tin and passes over search-index rows while indexed search is dormant', async () => { + setEnvFlags({ isLiveEnterpriseSearchEnabled: true }) + mocks.isFeatureEnabled.mockResolvedValue(true) + mocks.markUnfilled.mockResolvedValueOnce({ marked: 0, cursor: null }) + await runKnowledgeProjectionPass({ budgetMs: 60_000 }) + for (const [, options] of mocks.runProjection.mock.calls) + expect(options).toMatchObject({ includeTin: false }) + expect(mocks.markUnfilled).toHaveBeenCalledWith(expect.anything(), undefined, { + includeSearchIndexes: false, }) }) diff --git a/apps/sim/lib/knowledge/projection/run.ts b/apps/sim/lib/knowledge/projection/run.ts index 853bd8e5ce2..4ab2e8b5972 100644 --- a/apps/sim/lib/knowledge/projection/run.ts +++ b/apps/sim/lib/knowledge/projection/run.ts @@ -9,6 +9,7 @@ import { createLogger } from '@sim/logger' import postgres, { type Sql } from 'postgres' import { env, envNumber } from '@/lib/core/config/env' import { isFeatureEnabled } from '@/lib/core/config/feature-flags' +import { isIndexedOrgSearchEnabled } from '@/lib/sim-search/indexed/gate' const logger = createLogger('KnowledgeProjectionPass') @@ -52,6 +53,10 @@ export interface KnowledgeProjectionPassResult extends KnowledgeProjectionProgre * read the same oldest marks and split them at those locks. A round ends when every worker found * nothing more it could take; the pass goes on while rounds settle documents, and `remaining` * reports marks it left or a fill it did not finish, so the caller can schedule another. + * + * While indexed organization search is dormant the pass writes no Tin keyword rows, which only + * that search reads, and the fill passes over search-index rows: rewriting them would spend index + * writes on content no search serves. Workspace knowledge bases are projected either way. */ export async function runKnowledgeProjectionPass(options: { budgetMs: number @@ -71,6 +76,7 @@ export async function runKnowledgeProjectionPass(options: { ) ) const deadline = Date.now() + options.budgetMs + const indexedOrgSearch = isIndexedOrgSearchEnabled() const result: KnowledgeProjectionPassResult = { settled: 0, deferred: 0, @@ -88,7 +94,10 @@ export async function runKnowledgeProjectionPass(options: { /** Every worker finishes its document before a failure ends the pass, so none is cut off. */ const outcomes = await Promise.allSettled( workers.map((session) => - runKnowledgeProjection(session, { budgetMs: Math.max(0, deadline - Date.now()) }) + runKnowledgeProjection(session, { + budgetMs: Math.max(0, deadline - Date.now()), + includeTin: indexedOrgSearch, + }) ) ) const round: KnowledgeProjectionProgress[] = [] @@ -108,7 +117,9 @@ export async function runKnowledgeProjectionPass(options: { } /** The fill starts only inside the budget, and what it marks is left for a round to settle. */ if (fillCursor === null || Date.now() >= deadline) break - const fill = await markUnfilledProjectionDocuments(sessions[0], fillCursor) + const fill = await markUnfilledProjectionDocuments(sessions[0], fillCursor, { + includeSearchIndexes: indexedOrgSearch, + }) result.filled += fill.marked fillCursor = fill.cursor if (fill.marked > 0) result.remaining = true diff --git a/apps/sim/lib/knowledge/search/queries.test.ts b/apps/sim/lib/knowledge/search/queries.test.ts index 8b81e7b5506..137c39316e9 100644 --- a/apps/sim/lib/knowledge/search/queries.test.ts +++ b/apps/sim/lib/knowledge/search/queries.test.ts @@ -17,7 +17,7 @@ const { mockResolveTinKeywordQuery } = vi.hoisted(() => ({ mockResolveTinKeywordQuery: vi.fn<() => Promise>(async () => null), })) -vi.mock('@/lib/knowledge/search/tin-keyword', () => ({ +vi.mock('@/lib/sim-search/indexed/retrieval/tin-keyword', () => ({ resolveTinKeywordQuery: mockResolveTinKeywordQuery, })) @@ -35,7 +35,6 @@ import { import type { SearchStage } from '@/lib/knowledge/search/diagnostics' import { executeKeywordSearch, - forgetProjectionFilled, forgetSearchReach, fuseByReciprocalRank, getStructuredTagFilters, @@ -56,6 +55,7 @@ import { import { RRF_K } from '@/lib/knowledge/search/recency' import { forgetIndexedVectorSources } from '@/lib/knowledge/search/source-vector-indexes' import type { StructuredFilter } from '@/lib/knowledge/types' +import { forgetProjectionFilled } from '@/lib/sim-search/indexed/retrieval' /** * The builder only reads `embeddingTable[tagSlot]`, so a slot-to-name map stands diff --git a/apps/sim/lib/knowledge/search/queries.ts b/apps/sim/lib/knowledge/search/queries.ts index 5362d51d4cd..49b359adf59 100644 --- a/apps/sim/lib/knowledge/search/queries.ts +++ b/apps/sim/lib/knowledge/search/queries.ts @@ -1,5 +1,4 @@ import { db } from '@sim/db' -import { SOURCE_ACL_PROJECTIONS, type SourceAclProjection } from '@sim/db/knowledge-projection' import { document, embedding, @@ -55,7 +54,6 @@ import { workspaceSearchFilterConditions } from '@/lib/knowledge/search/filter-c import type { WorkspaceSearchFilters } from '@/lib/knowledge/search/filters' import { applyRecencyBoost, RRF_K } from '@/lib/knowledge/search/recency' import { indexedVectorSources } from '@/lib/knowledge/search/source-vector-indexes' -import { resolveTinKeywordQuery } from '@/lib/knowledge/search/tin-keyword' import { coerceTagFilterValue, escapeLikePattern, @@ -67,6 +65,8 @@ import { embeddingCandidateDistance, embeddingDistance, } from '@/lib/knowledge/vector-columns' +import { isIndexedOrgSearchEnabled } from '@/lib/sim-search/indexed/gate' +import { isProjectionFilled, resolveTinKeywordQuery } from '@/lib/sim-search/indexed/retrieval' const logger = createLogger('KnowledgeSearchQueries') @@ -110,60 +110,12 @@ function onRowWalkScanTuples( : Number(CANDIDATE_HNSW_MAX_SCAN_TUPLES) } -/** How long a fully filled projection is taken on trust before its unfilled rows are looked for again. */ -const PROJECTION_FILLED_TTL_MS = 60_000 - /** - * Whether the ranking projection still holds rows the source and ACL fill has not reached. Read off the - * unfilled-rows index in milliseconds and remembered briefly: the answer only ever changes once. - * - * The read asks for the last unfilled row by id, not whether one exists: an `EXISTS` drops its - * order and limit, and while most rows are unfilled the planner expects a sequential scan to - * meet one at once, then walks the whole projection when the unfilled rows sit past the filled - * ones. Ordered by id and capped at one row, the read can only be the partial index, whose - * last entry is the row the fill reaches last. + * Whether the search-index-only retrieval strategies (the projection-fill probe and Tin keyword + * ranking) may run: every base is a search index and indexed organization search is on. */ -const projectionFilled = new LRUCache< - SourceAclProjection, - boolean, - { budget: SearchBudget | undefined; stage: SearchStage } ->({ - max: SOURCE_ACL_PROJECTIONS.length, - ttl: PROJECTION_FILLED_TTL_MS, - /** - * The read that misses the cache is the search's own, under its budget like every other read - * of the leg, and the searches that miss together share it. A read that fails is not - * remembered: it answers unfilled, the slower and safe form, and the next search reads again. - */ - fetchMethod: async (projection, _stale, { context }) => { - const table = projection === 'embedding_search' ? embeddingSearch : embeddingKeywordTin - try { - const [row] = await runSearchQuery(context.budget, context.stage, (executor) => - executor.execute<{ unfilled: boolean }>(sql` - SELECT ( - SELECT ${table.id} FROM ${table} WHERE ${table.acl} IS NULL - ORDER BY ${table.id} DESC LIMIT 1 - ) IS NOT NULL AS unfilled`) - ) - return !row?.unfilled - } catch { - return undefined - } - }, -}) - -/** Whether every row of the projection carries its mirrored source and ACL; unknown counts as not yet. */ -async function isProjectionFilled( - projection: SourceAclProjection, - stage: SearchStage, - budget: SearchBudget | undefined -): Promise { - return (await projectionFilled.fetch(projection, { context: { budget, stage } })) ?? false -} - -/** Forgets whether the projections were filled; the memo is per process and otherwise expires on its own. */ -export function forgetProjectionFilled(): void { - projectionFilled.clear() +function usesIndexedRetrieval(searchIndexOnly: boolean | undefined): boolean { + return searchIndexOnly === true && isIndexedOrgSearchEnabled() } /** @@ -1834,7 +1786,7 @@ async function selectVectorResults(params: SearchParams): Promise ({ vi.mock('@/lib/mothership/chat/organization-chats', () => ({ authorizeOrganizationChatDelegation: { execute: mocks.authorizeChat }, })) -vi.mock('@/lib/knowledge/application/workspace-search', () => ({ +vi.mock('@/lib/sim-search/indexed', () => ({ searchOrganizationKnowledge: { get operation() { return knowledgeOperations.search @@ -26,8 +26,6 @@ vi.mock('@/lib/knowledge/application/workspace-search', () => ({ }, execute: mocks.search, }, -})) -vi.mock('@/lib/knowledge/application/read-search-document', () => ({ readSearchDocument: { get operation() { return knowledgeOperations.readDocument diff --git a/apps/sim/lib/mothership/tools/server/knowledge/workspace-search.ts b/apps/sim/lib/mothership/tools/server/knowledge/workspace-search.ts index 8dff592231e..429644913c6 100644 --- a/apps/sim/lib/mothership/tools/server/knowledge/workspace-search.ts +++ b/apps/sim/lib/mothership/tools/server/knowledge/workspace-search.ts @@ -7,11 +7,6 @@ import { import { isLiveEnterpriseSearchEnabled } from '@/lib/core/config/env-flags' import { getBaseUrl } from '@/lib/core/utils/urls' import { EmbeddingConfigurationError } from '@/lib/embeddings/configuration-error' -import { readSearchDocument } from '@/lib/knowledge/application/read-search-document' -import { - searchOrganizationKnowledge, - searchWorkspaceKnowledge, -} from '@/lib/knowledge/application/workspace-search' import { sourceAuthor } from '@/lib/knowledge/search/author' import { SearchDeadlineError } from '@/lib/knowledge/search/budget' import { createKnowledgeDocumentCitation, liveCitationId } from '@/lib/knowledge/search/citation' @@ -31,6 +26,12 @@ import { } from '@/lib/mothership/application/execute-knowledge-use-case' import type { BaseServerTool, ServerToolContext } from '@/lib/mothership/tools/server/base-tool' import { connectorDisplayName } from '@/lib/sim-search/connectors' +import { + readSearchDocument, + searchOrganizationKnowledge, + searchWorkspaceKnowledge, +} from '@/lib/sim-search/indexed' +import { isIndexedOrgSearchEnabled } from '@/lib/sim-search/indexed/gate' import { readLiveDocument, searchLiveKnowledge } from '@/lib/sim-search/live/application' import { projectResolvedSecretModelContent } from '@/executor/utils/resolved-secret-content-projection' @@ -78,7 +79,7 @@ export const searchWorkspaceServerTool: BaseServerTool = { resultSecretRegistry: registry, signal: context?.abortSignal, } as const - if (isLiveEnterpriseSearchEnabled) { + if (!isIndexedOrgSearchEnabled()) { const nativeProjection = projectResolvedSecretModelContent( nativeQueries ?? [], registry @@ -237,7 +238,7 @@ export const readDocumentServerTool: BaseServerTool = { const input = readDocumentInputSchema.parse(raw) const registry = context?.resolvedSecretTraceRegistry if (!registry) throw new Error('Knowledge result provenance is unavailable') - if (isLiveEnterpriseSearchEnabled) { + if (!isIndexedOrgSearchEnabled()) { const liveInput = { ...input, filters: intersectWorkspaceSearchFilters( diff --git a/apps/sim/lib/sim-search/indexed/README.md b/apps/sim/lib/sim-search/indexed/README.md new file mode 100644 index 00000000000..a4ef6389152 --- /dev/null +++ b/apps/sim/lib/sim-search/indexed/README.md @@ -0,0 +1,51 @@ +# Indexed organization search (dormant) + +This directory holds the indexed backend for Sim Search: retrieval over `is_search_index` knowledge bases that organization and workspace connectors crawl into, ranked from the embedding projections, with Tin keyword ranking where the database provides it. Live Search (`../live/`) replaced it. **This code is dormant.** It stays in the tree so it can be switched back on, but no request reaches it in a default deployment. + +Ordinary workspace knowledge bases, the Knowledge block, the embedding projections, the projector, connector member and ACL machinery, and source vector indexes are shared with other features. They live outside this directory and behave the same whichever backend is selected. + +## The gate + +One switch decides: `isIndexedOrgSearchEnabled()` in `gate.ts`. It is the inverse of the Live Search backend selector, `SIM_SEARCH_LIVE`, which defaults to `true`. Indexed search is on only when a deployment sets `SIM_SEARCH_LIVE=false`. + +While the gate is off: + +- The internal search route, the MCP search and read tools, and Sim's `search_workspace` and `read_document` tools serve Live Search. +- The indexed document page (`/o/[organizationId]/knowledge/[knowledgeBaseId]/[documentId]`) is not found. +- A search that names a search-index knowledge base by id (the Knowledge block, internal, v1 and v2 knowledge search) is refused with `SearchIndexDormantError`, a `409` reading "This search index is inactive; use Sim Search." Workspace knowledge bases are searched as before. +- Connector content syncs and member syncs skip search-index knowledge bases (`lib/knowledge/connectors/indexing-policy.ts`). +- The projector writes no `embedding_keyword_tin` rows, and its source and ACL fill (`knowledge-projection-fill`) passes over search-index rows. Every other projection of a marked document, search index or not, is still written, so workspace knowledge bases project as before. +- Shared retrieval (`lib/knowledge/search/queries.ts`) does not run the search-index-only strategies below. + +## Layout + +- `gate.ts` is the switch, the dormant error, and `assertSearchIndexesActive`. Anything may import it. +- `index.ts` is the use-case barrel: `searchScopedKnowledge`, `searchOrganizationKnowledge`, `searchWorkspaceKnowledge`, `readSearchDocument`, and `readIndexedKnowledgeDocument`. +- `search/scoped-search.ts` resolves an owner's search index and runs the shared knowledge search over it. +- `documents/` reads indexed passages for Sim's `read_document` tool, the document page, and the MCP `read_document` tool. +- `retrieval/` is the barrel for the search-index-only retrieval strategies that shared retrieval calls: Tin keyword ranking (`tin-keyword.ts`, `tin-query.ts`, flag `knowledge-tin-keyword`) and the projection-fill probe (`projection-fill.ts`). It is separate from `index.ts` because the use cases depend on shared retrieval, which depends on these. + +`scripts/check-indexed-org-search-boundary.ts` (`bun run check:indexed-org-search-boundary`, part of `check:audits`) enforces the edge. Outside this directory, only the entry files it allowlists import a barrel, each of them must import the gate, and no file imports past a barrel. Tests are exempt. + +## The projection-fill probe + +Search-index retrieval asks whether a projection still has rows the source and ACL fill has not reached. Filled rows are decided on the row alone and permit a wider vector walk. The probe reads the last unfilled row from the partial index over rows with a null `acl`. It spends at most 250 ms of the leg's budget, and a filled answer is remembered for 60 seconds. An unknown answer, whether from a failure or an exhausted cap, counts as unfilled and is remembered for 5 seconds. A probe cancelled by its own search is not remembered. + +## Re-enabling + +1. Confirm the Tin objects still exist, and restore them where they were removed: the `tin` extension, `knowledge_tin_stream`, `knowledge_tin_base_token`, and `knowledge_tin_membership_key`, the `embedding_keyword_tin_sync` and `knowledge_base_keyword_tin_sync` triggers, and a valid `embedding_keyword_tin_content_idx`. Script migrations `0019_tin_keyword_projection` and `0024_knowledge_projection_async` install them. +2. Backfill `embedding_keyword_tin` for every search-index knowledge base, then build the Tin index. +3. Set `SIM_SEARCH_LIVE=false` and deploy. The container entrypoint (`apps/sim/bootstrap.ts`) mirrors it to `NEXT_PUBLIC_SIM_SEARCH_LIVE` for the client. +4. Re-enable and resync organization connectors on search-index knowledge bases. Content and member syncs resume on their own once the gate is on. Paused sources need to be resumed, and a full resync refreshes content that went stale while the gate was off. +5. Turn on `knowledge-projection-fill` so the fill reaches search-index rows again, and `knowledge-tin-keyword` once the Tin index is valid. +6. Optionally run `apps/sim/scripts/prewarm-search-projection.ts` after the backfill. + +## Database objects it depends on + +- `knowledge_base.is_search_index`, `document.acl`, `document.connector_id`, and `knowledge_connector`. +- `embedding` and the projections: `embedding_search` (vectors plus mirrored `connector_id` and `acl`), `embedding_keyword_search` (GIN keyword ranking), and `embedding_keyword_tin` (Tin keyword ranking, search-index rows only). +- The partial indexes over unfilled projection rows (`acl IS NULL`) that the fill and the probe read. +- `knowledge_projection_dirty` and the projector triggers that mark it. +- The Tin extension objects listed under re-enabling. + +No schema or migration belongs to this directory. Every object above is shared or owned by `packages/db`. diff --git a/apps/sim/lib/knowledge/application/read-indexed-document.ts b/apps/sim/lib/sim-search/indexed/documents/read-indexed-document.ts similarity index 100% rename from apps/sim/lib/knowledge/application/read-indexed-document.ts rename to apps/sim/lib/sim-search/indexed/documents/read-indexed-document.ts diff --git a/apps/sim/lib/knowledge/application/read-search-document.test.ts b/apps/sim/lib/sim-search/indexed/documents/read-search-document.test.ts similarity index 99% rename from apps/sim/lib/knowledge/application/read-search-document.test.ts rename to apps/sim/lib/sim-search/indexed/documents/read-search-document.test.ts index ad3b3831013..9970856f047 100644 --- a/apps/sim/lib/knowledge/application/read-search-document.test.ts +++ b/apps/sim/lib/sim-search/indexed/documents/read-search-document.test.ts @@ -32,7 +32,7 @@ vi.mock('@/lib/execution/durable-secret-provenance', () => ({ importDurableSecretProvenance: mocks.importProvenance, })) -import { readSearchDocument } from '@/lib/knowledge/application/read-search-document' +import { readSearchDocument } from '@/lib/sim-search/indexed/documents/read-search-document' import { ResolvedSecretTraceRegistry } from '@/executor/utils/resolved-secret-trace-registry' const principal = { kind: 'session', userId: 'reader', sessionId: 'session' } as const diff --git a/apps/sim/lib/knowledge/application/read-search-document.ts b/apps/sim/lib/sim-search/indexed/documents/read-search-document.ts similarity index 100% rename from apps/sim/lib/knowledge/application/read-search-document.ts rename to apps/sim/lib/sim-search/indexed/documents/read-search-document.ts diff --git a/apps/sim/lib/sim-search/indexed/gate.test.ts b/apps/sim/lib/sim-search/indexed/gate.test.ts new file mode 100644 index 00000000000..a1fa3ae592a --- /dev/null +++ b/apps/sim/lib/sim-search/indexed/gate.test.ts @@ -0,0 +1,39 @@ +/** + * @vitest-environment node + */ +import { resetEnvFlagsMock, setEnvFlags } from '@sim/testing' +import { beforeEach, describe, expect, it } from 'vitest' +import { + assertSearchIndexesActive, + isIndexedOrgSearchEnabled, + SearchIndexDormantError, +} from '@/lib/sim-search/indexed/gate' + +describe('indexed organization search gate', () => { + beforeEach(resetEnvFlagsMock) + + it('is dormant while Live Search is the backend', () => { + setEnvFlags({ isLiveEnterpriseSearchEnabled: true }) + expect(isIndexedOrgSearchEnabled()).toBe(false) + }) + + it('is on only when Live Search is turned off', () => { + setEnvFlags({ isLiveEnterpriseSearchEnabled: false }) + expect(isIndexedOrgSearchEnabled()).toBe(true) + }) + + it('refuses any search index while dormant and never a workspace knowledge base', () => { + setEnvFlags({ isLiveEnterpriseSearchEnabled: true }) + expect(() => + assertSearchIndexesActive([{ isSearchIndex: false }, { isSearchIndex: true }]) + ).toThrow(SearchIndexDormantError) + expect(() => + assertSearchIndexesActive([{ isSearchIndex: false }, { isSearchIndex: null }, {}]) + ).not.toThrow() + }) + + it('admits search indexes while indexed organization search is on', () => { + setEnvFlags({ isLiveEnterpriseSearchEnabled: false }) + expect(() => assertSearchIndexesActive([{ isSearchIndex: true }])).not.toThrow() + }) +}) diff --git a/apps/sim/lib/sim-search/indexed/gate.ts b/apps/sim/lib/sim-search/indexed/gate.ts new file mode 100644 index 00000000000..e6ad4677a82 --- /dev/null +++ b/apps/sim/lib/sim-search/indexed/gate.ts @@ -0,0 +1,36 @@ +import { isLiveEnterpriseSearchEnabled } from '@/lib/core/config/env-flags' + +/** What a caller is told when it names a search index while indexed organization search is off. */ +export const SEARCH_INDEX_DORMANT_MESSAGE = 'This search index is inactive; use Sim Search.' + +/** + * The single switch for indexed organization search: retrieval over `is_search_index` knowledge + * bases and the crawling that fills them. It is the inverse of the Live Search backend selector + * (`SIM_SEARCH_LIVE`, on by default), so indexed search is dormant unless a deployment sets + * `SIM_SEARCH_LIVE=false`. Read on every call rather than captured, so a caller always observes + * the current selector. + */ +export function isIndexedOrgSearchEnabled(): boolean { + return !isLiveEnterpriseSearchEnabled +} + +/** A search named a search-index knowledge base while indexed organization search is dormant. */ +export class SearchIndexDormantError extends Error { + constructor() { + super(SEARCH_INDEX_DORMANT_MESSAGE) + this.name = 'SearchIndexDormantError' + } +} + +/** + * Refuses a search over any search-index knowledge base while indexed organization search is + * dormant. Its rows are neither crawled nor projected in that state, so answering from them would + * return stale content under current permissions. Workspace knowledge bases are never refused. + */ +export function assertSearchIndexesActive( + knowledgeBases: ReadonlyArray<{ isSearchIndex?: boolean | null }> +): void { + if (isIndexedOrgSearchEnabled()) return + if (knowledgeBases.some((knowledgeBase) => knowledgeBase.isSearchIndex === true)) + throw new SearchIndexDormantError() +} diff --git a/apps/sim/lib/sim-search/indexed/index.ts b/apps/sim/lib/sim-search/indexed/index.ts new file mode 100644 index 00000000000..0cffbe7a9e2 --- /dev/null +++ b/apps/sim/lib/sim-search/indexed/index.ts @@ -0,0 +1,13 @@ +/** + * Dormant indexed organization search: the use cases that search and read `is_search_index` + * knowledge bases. Every caller outside this directory is an allowlisted entry point that checks + * `isIndexedOrgSearchEnabled()` from `@/lib/sim-search/indexed/gate` before it reaches these; see + * `scripts/check-indexed-org-search-boundary.ts` and this directory's README. + */ +export { readIndexedKnowledgeDocument } from '@/lib/sim-search/indexed/documents/read-indexed-document' +export { readSearchDocument } from '@/lib/sim-search/indexed/documents/read-search-document' +export { + searchOrganizationKnowledge, + searchScopedKnowledge, + searchWorkspaceKnowledge, +} from '@/lib/sim-search/indexed/search/scoped-search' diff --git a/apps/sim/lib/sim-search/indexed/retrieval/index.ts b/apps/sim/lib/sim-search/indexed/retrieval/index.ts new file mode 100644 index 00000000000..286c1645fe9 --- /dev/null +++ b/apps/sim/lib/sim-search/indexed/retrieval/index.ts @@ -0,0 +1,11 @@ +/** + * Search-index-only retrieval strategies the shared retrieval layer + * (`lib/knowledge/search/queries.ts`) runs while indexed organization search is on. Kept apart + * from the use-case barrel because the use cases depend on that retrieval layer, which depends on + * these. + */ +export { + forgetProjectionFilled, + isProjectionFilled, +} from '@/lib/sim-search/indexed/retrieval/projection-fill' +export { resolveTinKeywordQuery } from '@/lib/sim-search/indexed/retrieval/tin-keyword' diff --git a/apps/sim/lib/sim-search/indexed/retrieval/projection-fill.test.ts b/apps/sim/lib/sim-search/indexed/retrieval/projection-fill.test.ts new file mode 100644 index 00000000000..decebc71670 --- /dev/null +++ b/apps/sim/lib/sim-search/indexed/retrieval/projection-fill.test.ts @@ -0,0 +1,90 @@ +/** + * @vitest-environment node + */ +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' +import { SearchBudget } from '@/lib/knowledge/search/budget' +import { + forgetProjectionFilled, + isProjectionFilled, + PROJECTION_FILLED_PROBE_BUDGET_MS, +} from '@/lib/sim-search/indexed/retrieval/projection-fill' + +const LEG_BUDGET_MS = 8000 + +describe('projection fill probe', () => { + beforeEach(() => forgetProjectionFilled()) + afterEach(() => vi.restoreAllMocks()) + + it('spends at most its capped share of the leg budget', async () => { + vi.spyOn(performance, 'now').mockReturnValue(1000) + const deadlines: number[] = [] + vi.spyOn(SearchBudget.prototype, 'query').mockImplementation(async function ( + this: SearchBudget + ) { + deadlines.push(this.deadline) + return [{ unfilled: false }] + } as SearchBudget['query']) + const budget = new SearchBudget('vector', 1000 + LEG_BUDGET_MS) + await expect( + isProjectionFilled('embedding_search', 'vector.projection_filled', budget) + ).resolves.toBe(true) + expect(deadlines).toEqual([1000 + PROJECTION_FILLED_PROBE_BUDGET_MS]) + expect(budget.timedOut).toBe(false) + }) + + it('remembers a failed probe as unfilled rather than probing again on every search', async () => { + const query = vi + .spyOn(SearchBudget.prototype, 'query') + .mockRejectedValue(new Error('connection reset')) + const budget = new SearchBudget('keyword', performance.now() + LEG_BUDGET_MS) + await expect( + isProjectionFilled('embedding_keyword_tin', 'keyword.projection_filled', budget) + ).resolves.toBe(false) + await expect( + isProjectionFilled('embedding_keyword_tin', 'keyword.projection_filled', budget) + ).resolves.toBe(false) + expect(query).toHaveBeenCalledOnce() + }) + + it('does not hold a later search past its own share while another search probes', async () => { + let answerFirst: (rows: Array<{ unfilled: boolean }>) => void = () => {} + vi.spyOn(SearchBudget.prototype, 'query').mockImplementation( + () => + new Promise((resolve) => { + answerFirst = resolve as typeof answerFirst + }) as ReturnType + ) + const first = isProjectionFilled( + 'embedding_search', + 'vector.projection_filled', + new SearchBudget('vector', performance.now() + LEG_BUDGET_MS) + ) + const started = performance.now() + await expect( + isProjectionFilled( + 'embedding_search', + 'vector.projection_filled', + new SearchBudget('vector', performance.now() + 20) + ) + ).resolves.toBe(false) + expect(performance.now() - started).toBeLessThan(PROJECTION_FILLED_PROBE_BUDGET_MS) + answerFirst([{ unfilled: false }]) + await expect(first).resolves.toBe(true) + }) + + it('does not remember a probe its own search cancelled', async () => { + const query = vi + .spyOn(SearchBudget.prototype, 'query') + .mockRejectedValue(new DOMException('aborted', 'AbortError')) + const controller = new AbortController() + controller.abort() + const budget = new SearchBudget('vector', performance.now() + LEG_BUDGET_MS, controller.signal) + await expect( + isProjectionFilled('embedding_search', 'vector.projection_filled', budget) + ).resolves.toBe(false) + await expect( + isProjectionFilled('embedding_search', 'vector.projection_filled', budget) + ).resolves.toBe(false) + expect(query).toHaveBeenCalledTimes(2) + }) +}) diff --git a/apps/sim/lib/sim-search/indexed/retrieval/projection-fill.ts b/apps/sim/lib/sim-search/indexed/retrieval/projection-fill.ts new file mode 100644 index 00000000000..c66e803084d --- /dev/null +++ b/apps/sim/lib/sim-search/indexed/retrieval/projection-fill.ts @@ -0,0 +1,99 @@ +import { SOURCE_ACL_PROJECTIONS, type SourceAclProjection } from '@sim/db/knowledge-projection' +import { embeddingKeywordTin, embeddingSearch } from '@sim/db/schema' +import { sleep } from '@sim/utils/helpers' +import { sql } from 'drizzle-orm' +import { LRUCache } from 'lru-cache' +import { runSearchQuery, type SearchBudget } from '@/lib/knowledge/search/budget' +import type { SearchStage } from '@/lib/knowledge/search/diagnostics' + +/** How long a fully filled projection is taken on trust before its unfilled rows are looked for again. */ +const PROJECTION_FILLED_TTL_MS = 60_000 + +/** + * The most of a leg's budget the probe may spend. The probe is one index read that answers in + * milliseconds when the partial index serves it; a read slower than this is fighting a cold cache + * or a busy database, and waiting longer would spend the leg's ranking time on an optimization. + * An unanswered probe costs only the slower plan, so a small cap loses nothing. + */ +export const PROJECTION_FILLED_PROBE_BUDGET_MS = 250 + +/** + * How long an unanswered probe is remembered as unfilled. Short enough that a recovered database + * is asked again within seconds, long enough that searches arriving during an outage do not each + * spend their own probe budget rediscovering it. + */ +export const PROJECTION_FILLED_UNKNOWN_TTL_MS = 5_000 + +/** + * Whether the ranking projection still holds rows the source and ACL fill has not reached. Read off the + * unfilled-rows index in milliseconds and remembered briefly: the answer only ever changes once. + * + * The read asks for the last unfilled row by id, not whether one exists: an `EXISTS` drops its + * order and limit, and while most rows are unfilled the planner expects a sequential scan to + * meet one at once, then walks the whole projection when the unfilled rows sit past the filled + * ones. Ordered by id and capped at one row, the read can only be the partial index, whose + * last entry is the row the fill reaches last. + */ +const projectionFilled = new LRUCache< + SourceAclProjection, + boolean, + { budget: SearchBudget | undefined; stage: SearchStage } +>({ + max: SOURCE_ACL_PROJECTIONS.length, + ttl: PROJECTION_FILLED_TTL_MS, + /** + * The read that misses the cache is the search's own, capped to a small share of its budget, + * and the searches that miss together share it. A read that fails or runs out of that share + * answers unfilled, the slower and safe form, and that answer is remembered briefly so the + * searches behind it do not each pay for the same failure. A read cut short by its own + * search's cancellation learned nothing about the projection and is not remembered. + */ + fetchMethod: async (projection, _stale, { context, options }) => { + const table = projection === 'embedding_search' ? embeddingSearch : embeddingKeywordTin + try { + const [row] = await runSearchQuery( + context.budget?.capped(PROJECTION_FILLED_PROBE_BUDGET_MS), + context.stage, + (executor) => + executor.execute<{ unfilled: boolean }>(sql` + SELECT ( + SELECT ${table.id} FROM ${table} WHERE ${table.acl} IS NULL + ORDER BY ${table.id} DESC LIMIT 1 + ) IS NOT NULL AS unfilled`) + ) + return !row?.unfilled + } catch { + if (context.budget?.signal?.aborted) return undefined + options.ttl = PROJECTION_FILLED_UNKNOWN_TTL_MS + return false + } + }, +}) + +/** + * Whether every row of the projection carries its mirrored source and ACL; unknown counts as not yet. + * + * Searches that miss the cache together share the first one's read, which is capped to that + * search's share. Each caller still waits no longer than its own share, or its own deadline if + * nearer, and reads an unanswered probe as unfilled: a caller that joined late, with less of its + * leg left, never waits on another search's timetable. + */ +export async function isProjectionFilled( + projection: SourceAclProjection, + stage: SearchStage, + budget: SearchBudget | undefined +): Promise { + const answer = projectionFilled.fetch(projection, { context: { budget, stage } }) + if (!budget) return (await answer) ?? false + const waitMs = Math.max( + 0, + Math.min(PROJECTION_FILLED_PROBE_BUDGET_MS, budget.deadline - performance.now()) + ) + const unanswered = sleep(waitMs).then(() => undefined) + return (await Promise.race([answer.catch(() => undefined), unanswered])) ?? false +} + +/** Forgets whether the projections were filled; the memo is per process and otherwise expires on its own. */ +export function forgetProjectionFilled(): void { + projectionFilled.clear() +} diff --git a/apps/sim/lib/knowledge/search/tin-keyword-readiness.test.ts b/apps/sim/lib/sim-search/indexed/retrieval/tin-keyword-readiness.test.ts similarity index 92% rename from apps/sim/lib/knowledge/search/tin-keyword-readiness.test.ts rename to apps/sim/lib/sim-search/indexed/retrieval/tin-keyword-readiness.test.ts index f42b4bc2cfb..911a7304bd6 100644 --- a/apps/sim/lib/knowledge/search/tin-keyword-readiness.test.ts +++ b/apps/sim/lib/sim-search/indexed/retrieval/tin-keyword-readiness.test.ts @@ -8,7 +8,7 @@ vi.mock('@/lib/core/config/feature-flags', () => ({ isFeatureEnabled: vi.fn(async () => true), })) -import { resolveTinKeywordQuery } from '@/lib/knowledge/search/tin-keyword' +import { resolveTinKeywordQuery } from '@/lib/sim-search/indexed/retrieval/tin-keyword' /** Its own file, so the process-wide readiness cache starts empty. */ it('stays on the GIN projection while the Tin index is incomplete, and remembers that', async () => { diff --git a/apps/sim/lib/knowledge/search/tin-keyword.test.ts b/apps/sim/lib/sim-search/indexed/retrieval/tin-keyword.test.ts similarity index 96% rename from apps/sim/lib/knowledge/search/tin-keyword.test.ts rename to apps/sim/lib/sim-search/indexed/retrieval/tin-keyword.test.ts index ac82444c449..9d625d94f9a 100644 --- a/apps/sim/lib/knowledge/search/tin-keyword.test.ts +++ b/apps/sim/lib/sim-search/indexed/retrieval/tin-keyword.test.ts @@ -13,7 +13,7 @@ vi.mock('@/lib/core/config/feature-flags', () => ({ })) import { SearchBudget, SearchDeadlineError } from '@/lib/knowledge/search/budget' -import { resolveTinKeywordQuery } from '@/lib/knowledge/search/tin-keyword' +import { resolveTinKeywordQuery } from '@/lib/sim-search/indexed/retrieval/tin-keyword' /** Readiness is cached per process; the incomplete-index case lives in its own file, where the cache starts empty. */ describe('resolveTinKeywordQuery', () => { diff --git a/apps/sim/lib/knowledge/search/tin-keyword.ts b/apps/sim/lib/sim-search/indexed/retrieval/tin-keyword.ts similarity index 97% rename from apps/sim/lib/knowledge/search/tin-keyword.ts rename to apps/sim/lib/sim-search/indexed/retrieval/tin-keyword.ts index 534cb42f094..91ccef45911 100644 --- a/apps/sim/lib/knowledge/search/tin-keyword.ts +++ b/apps/sim/lib/sim-search/indexed/retrieval/tin-keyword.ts @@ -5,7 +5,7 @@ import { 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' +import { tinQueryFromTsquery } from '@/lib/sim-search/indexed/retrieval/tin-query' const logger = createLogger('TinKeywordSearch') diff --git a/apps/sim/lib/knowledge/search/tin-query.test.ts b/apps/sim/lib/sim-search/indexed/retrieval/tin-query.test.ts similarity index 95% rename from apps/sim/lib/knowledge/search/tin-query.test.ts rename to apps/sim/lib/sim-search/indexed/retrieval/tin-query.test.ts index db9053d3a84..7cbc465c7d5 100644 --- a/apps/sim/lib/knowledge/search/tin-query.test.ts +++ b/apps/sim/lib/sim-search/indexed/retrieval/tin-query.test.ts @@ -2,7 +2,7 @@ * @vitest-environment node */ import { describe, expect, it } from 'vitest' -import { tinQueryFromTsquery } from '@/lib/knowledge/search/tin-query' +import { tinQueryFromTsquery } from '@/lib/sim-search/indexed/retrieval/tin-query' /** Inputs are `websearch_to_tsquery('english', …)::text` exactly as PostgreSQL renders them. */ describe('tinQueryFromTsquery', () => { diff --git a/apps/sim/lib/knowledge/search/tin-query.ts b/apps/sim/lib/sim-search/indexed/retrieval/tin-query.ts similarity index 100% rename from apps/sim/lib/knowledge/search/tin-query.ts rename to apps/sim/lib/sim-search/indexed/retrieval/tin-query.ts diff --git a/apps/sim/lib/knowledge/application/workspace-search.activity.test.ts b/apps/sim/lib/sim-search/indexed/search/scoped-search.activity.test.ts similarity index 99% rename from apps/sim/lib/knowledge/application/workspace-search.activity.test.ts rename to apps/sim/lib/sim-search/indexed/search/scoped-search.activity.test.ts index dc933d694fa..a0186cc224b 100644 --- a/apps/sim/lib/knowledge/application/workspace-search.activity.test.ts +++ b/apps/sim/lib/sim-search/indexed/search/scoped-search.activity.test.ts @@ -51,7 +51,7 @@ vi.mock('@/lib/knowledge/application/search', () => ({ import { searchOrganizationKnowledge, searchScopedKnowledge, -} from '@/lib/knowledge/application/workspace-search' +} from '@/lib/sim-search/indexed/search/scoped-search' const principal = { kind: 'session', userId: 'reader', sessionId: 'session' } as const const input = { organizationId: 'org', query: 'policy', topK: 20, surface: 'slack' } as const diff --git a/apps/sim/lib/knowledge/application/workspace-search.test.ts b/apps/sim/lib/sim-search/indexed/search/scoped-search.test.ts similarity index 97% rename from apps/sim/lib/knowledge/application/workspace-search.test.ts rename to apps/sim/lib/sim-search/indexed/search/scoped-search.test.ts index 1d632f6535c..e846a9a9776 100644 --- a/apps/sim/lib/knowledge/application/workspace-search.test.ts +++ b/apps/sim/lib/sim-search/indexed/search/scoped-search.test.ts @@ -36,7 +36,7 @@ vi.mock('@/lib/knowledge/application/search', () => ({ afterKnowledgeSearch: mocks.afterSearch, })) -import { searchWorkspaceKnowledge } from '@/lib/knowledge/application/workspace-search' +import { searchWorkspaceKnowledge } from '@/lib/sim-search/indexed/search/scoped-search' const principal = { kind: 'session', userId: 'reader', sessionId: 'session' } as const const input = { workspaceId: 'workspace', query: 'orion', topK: 20, filters: { source: 'slack' } } diff --git a/apps/sim/lib/knowledge/application/workspace-search.ts b/apps/sim/lib/sim-search/indexed/search/scoped-search.ts similarity index 100% rename from apps/sim/lib/knowledge/application/workspace-search.ts rename to apps/sim/lib/sim-search/indexed/search/scoped-search.ts diff --git a/apps/sim/lib/sim-search/live/README.md b/apps/sim/lib/sim-search/live/README.md index 6a9ccff8660..19822255b4b 100644 --- a/apps/sim/lib/sim-search/live/README.md +++ b/apps/sim/lib/sim-search/live/README.md @@ -1,6 +1,6 @@ # Federated Search access and connector behavior -This describes the live enterprise-search path. Credential Groups and ordinary knowledge-base indexing retain their existing behavior. Live Search is enabled by default; only an explicit `SIM_SEARCH_LIVE=false` selects the legacy indexed backend. Live Search sources do not create content-indexing jobs, and queued content or persisted-directory jobs stop before crawling, embedding, or building ACL snapshots. Ordinary KB jobs remain enabled. Administrators can still maintain GitLab CSV grants, and request-time source permission checks remain required. +This describes the live enterprise-search path. Credential Groups and ordinary knowledge-base indexing retain their existing behavior. Live Search is enabled by default; only an explicit `SIM_SEARCH_LIVE=false` selects the legacy indexed backend, which is kept dormant in `../indexed/` (see its README). Live Search sources do not create content-indexing jobs, and queued content or persisted-directory jobs stop before crawling, embedding, or building ACL snapshots. Ordinary KB jobs remain enabled. Administrators can still maintain GitLab CSV grants, and request-time source permission checks remain required. ## Admin and member surfaces diff --git a/apps/sim/scripts/dormant-org-search/README.md b/apps/sim/scripts/dormant-org-search/README.md new file mode 100644 index 00000000000..5ff222549ff --- /dev/null +++ b/apps/sim/scripts/dormant-org-search/README.md @@ -0,0 +1,271 @@ +# Dormant organization search: database recovery runbook + +Organization-level indexed search ("Sim Search" over an organization search index, `knowledge_base.is_search_index = true`) is dormant: live search answers those queries instead. An organization search index can hold most of the rows in `embedding`, `embedding_search` and `embedding_keyword_search`, and all of `embedding_keyword_tin`. That inflates the shared vector index (`embedding_search_512_cosine_hnsw_idx`) past memory and slows every workspace knowledge base search. + +This runbook removes that data safely and in a way you can repeat and resume, then gives the space back. Run the steps in order. Every script **only reads unless `--execute` is passed**, logs what it would do, and is safe to re-run. + +Placeholders: `` is the search index's `knowledge_base.id`; `` is the connection string of the role that owns the knowledge tables (the migrations role). The scripts read `MIGRATION_DATABASE_URL`, then `DATABASE_URL`. The deletion script uses the app's database client, which reads only `DATABASE_URL`. So set `DATABASE_URL` to the migrations role for every step. + +```sh +export DATABASE_URL='' +cd +``` + +Run long steps in `tmux`/`screen` from a host close to the database, off-peak. + +| Script | Step | Writes with `--execute` | +|---|---|---| +| `disable-tin-projection.ts` | 1 | drops the three Tin sync triggers, truncates `embedding_keyword_tin` | +| `restore-tin-projection.ts` | rollback | reinstalls those triggers from migrations `0019` + `0024`, optional backfill | +| `delete-search-index-documents.ts` | 2 | deletes the search index's connector documents and chunks, queues storage cleanup, resets connector cursors | +| `maintenance.ts` | 3 | `REINDEX INDEX CONCURRENTLY`, `VACUUM (VERBOSE, ANALYZE)` | + +## 0. Preconditions + +Do not start until all of these hold: + +1. **The release with live search as the default is deployed** everywhere, and old app instances are drained. No deployed code path should read the organization index for a user query. +2. **The organization's connectors are paused.** Their processing queue is already cancelled. Step 2 refuses to run otherwise. Check: + ```sql + SELECT id, connector_type, status, member_sync_status, + sync_lock_token IS NOT NULL AS content_sync_running, + member_sync_lock_token IS NOT NULL AS member_sync_running, + deleted_at, detached_at + FROM knowledge_connector WHERE knowledge_base_id = ''; + ``` + Every row that is not deleted should be `paused` (or `disabled`), with neither sync running and `detached_at` NULL. +3. **The fill does not reach the index.** A release that carries the indexed search gate skips search-index rows in the `knowledge-projection-fill` pass while indexed search is dormant. On an older release, turn the flag off; otherwise the fill keeps marking documents of the index for the projector while you delete them. +4. The knowledge base is the one you mean: + ```sql + SELECT id, name, organization_id, is_search_index, deleted_at + FROM knowledge_base WHERE id = ''; + ``` +5. Record a baseline to compare against later: + ```sql + SELECT relname, n_live_tup, n_dead_tup, pg_size_pretty(pg_total_relation_size(relid)) AS size, + last_autovacuum, last_vacuum + FROM pg_stat_user_tables + WHERE relname IN ('embedding', 'embedding_search', 'embedding_keyword_search', + 'embedding_keyword_tin', 'document', 'knowledge_projection_dirty'); + SELECT count(*) FROM knowledge_projection_dirty; + SELECT count(*) FROM outbox_event + WHERE status = 'pending' AND event_type = 'knowledge.document.storage.cleanup'; + ``` + `bun apps/sim/scripts/dormant-org-search/maintenance.ts` (no flags) prints all of this plus index sizes, running vacuums and index builds. +6. A recent point-in-time-recovery window exists. Deleted documents come back only by resyncing the connectors or restoring from backup. + +## 1. Stop maintaining the Tin keyword projection + +`embedding_keyword_tin` is the BM25 (`tin` access method) projection. Only organization search indexes are projected into it. Three triggers keep it current inside writers' transactions: + +| Trigger | Table | Installed by | +|---|---|---| +| `embedding_keyword_tin_sync` | `embedding` | `0019_tin_keyword_projection`, re-guarded by `0024_knowledge_projection_async` | +| `knowledge_base_keyword_tin_sync` | `knowledge_base` | `0019` | +| `embedding_keyword_tin_source_acl_set` | `embedding_keyword_tin` | `0021`, re-guarded by `0024` | + +The script drops exactly these three and then runs `TRUNCATE embedding_keyword_tin`, all in one transaction. The table, its indexes, and the Tin SQL functions stay, so the schema and any code that references them remain valid. + +```sh +bun apps/sim/scripts/dormant-org-search/disable-tin-projection.ts # dry run: lists triggers, size +bun apps/sim/scripts/dormant-org-search/disable-tin-projection.ts --execute +``` + +- **Locks.** `DROP TRIGGER` and `TRUNCATE` take ACCESS EXCLUSIVE locks on `embedding`, `knowledge_base` and `embedding_keyword_tin`. A queued ACCESS EXCLUSIVE request blocks every later reader of that table while it waits. Each attempt therefore waits at most `3s` (`lock_timeout`) and then retries with backoff, for up to `--retry-budget-minutes` (default 20). Once it has the locks, it finishes in well under a second: truncate unlinks files and does not delete rows one by one. +- **The document ACL fan-out is left as it is.** `sync_projection_source_acl()` (the `projection_source_acl_sync` trigger on `document`, from `0024`) updates `embedding_search` and `embedding_keyword_tin` inline, in one function body. Rewriting it would fork `0024`'s SQL. Against an empty table, its Tin `UPDATE` is one probe of `embedding_keyword_tin_document_idx` that finds nothing. +- **The knowledge projector writes no Tin rows while indexed search is dormant.** An older release, or one with `SIM_SEARCH_LIVE=false`, still writes Tin for marked documents of any search index whenever the Tin functions exist. If rows reappear (`SELECT EXISTS (SELECT 1 FROM embedding_keyword_tin)`), another search index is being written; find it with `SELECT knowledge_base_id, count(*) FROM embedding_keyword_tin GROUP BY 1`. +- **Verify:** + ```sql + SELECT tgname, tgrelid::regclass FROM pg_trigger + WHERE tgname IN ('embedding_keyword_tin_sync', 'knowledge_base_keyword_tin_sync', + 'embedding_keyword_tin_source_acl_set'); -- expect no rows + SELECT pg_size_pretty(pg_total_relation_size('embedding_keyword_tin')); -- a few pages + ``` +- **Do not** run `0019_tin_keyword_projection.ts` directly or `db:push` against this database afterwards. Both re-install the Tin triggers and backfill the whole projection. Recorded script migrations do not run again on deploy, so normal releases leave this state alone. A future script migration that re-runs `installKnowledgeProjectionAsync` would bring back only `embedding_keyword_tin_source_acl_set`, which fires only on Tin writes and is harmless. + +### Re-enable Tin + +```sh +bun apps/sim/scripts/dormant-org-search/restore-tin-projection.ts # dry run +bun apps/sim/scripts/dormant-org-search/restore-tin-projection.ts --execute # triggers only +bun apps/sim/scripts/dormant-org-search/restore-tin-projection.ts --execute --backfill # + refill +``` + +The restore calls `installProjection` from `0019` and then `installKnowledgeProjectionAsync` from `0024`, the same functions the migrations use. The Tin triggers come back exactly as the migrations define them, including the deferred-projection `WHEN` guard. The script then checks that all three exist. Triggers only project chunks written after they are installed, so the projection needs a **backfill** (`--backfill` runs `0019`'s paced keyset `backfillProjection`). With a large search index, backfilling into the live Tin index is slow. For a large refill, drop `embedding_keyword_tin_content_idx`, backfill, then run `bun packages/db/script-migrations/0019_tin_keyword_projection.ts` to rebuild the index concurrently. + +## 2. Delete the search index's documents and chunks + +```sh +# Dry run: guard checks, then the first 3 pages with their chunk counts. Nothing is written. +bun apps/sim/scripts/dormant-org-search/delete-search-index-documents.ts --knowledge-base-id= + +# A short first execution, to measure throughput and watch the database: +bun apps/sim/scripts/dormant-org-search/delete-search-index-documents.ts --knowledge-base-id= --execute --max-pages=5 + +# The full run (stop any time with Ctrl-C; each page is its own transaction): +DB_APP_NAME=sim-dormant-org-search-delete bun apps/sim/scripts/dormant-org-search/delete-search-index-documents.ts --knowledge-base-id= --execute +``` + +Exit codes: `0` means it reached the end, or stopped at `--max-pages`. The last log line then gives the `afterId` to pass as `--after-id`. `1` means it failed (it logs the cursor it reached). `2` means a precondition refused the run; nothing after the check was written. + +### What each page does + +1. **Guard, re-read before every page.** The knowledge base must exist and have `is_search_index = true`. Every connector that is not deleted must be `paused` or `disabled`, and neither its content sync nor its member sync may hold a lease. A detached connector is refused outright, because its worker is turning those same documents into standalone uploads. If a connector resumes mid-run, the run stops before its next page. +2. **Storage backpressure.** If at least `--storage-cleanup-ceiling` (default 2000) `knowledge.document.storage.cleanup` outbox events are pending, the run waits 30 s and checks again until the outbox worker catches up. +3. **Next page.** The run reads the next `--page-size` (default 200) documents with `connector_id IS NOT NULL`, in id order after the cursor. +4. **Chunks.** It deletes the page's chunks in transactions of at most `--chunk-batch-size` (default 1000), each with `lock_timeout` 5 s and `statement_timeout` 60 s (`--lock-timeout-ms`, `--statement-timeout-ms`), pausing `--pause-ms` (default 250) after each. +5. **Documents, in one transaction.** This mirrors the app's connector cleanup worker (`lib/knowledge/connectors/deletion.ts`). It share-locks the knowledge base and re-checks `is_search_index`, then locks the page's documents `FOR UPDATE` so a late indexing commit can't slip in. It checks that no chunk remains; if one does, it deletes chunks again, up to 3 passes, then fails. Next it queues the storage cleanup intents with the app's own `enqueueKnowledgeStorageCleanup`, and deletes the documents. +6. **Retries.** A transient failure (lock or statement timeout, deadlock or serialization conflict, lost connection) retries that step in place with backoff, up to 8 times. Each step commits entirely or not at all, so a retry repeats nothing that committed. + +### Why this is cheap for triggers, and why it does not flood the projector + +Deleting fires **no row trigger** on these tables. The trigger inventory on `document`, `embedding`, `knowledge_base` and the projections has only `INSERT`/`UPDATE` triggers: + +- the synchronous projection triggers `embedding_search_sync`, `embedding_keyword_search_sync`, and the dropped `embedding_keyword_tin_sync` +- the projection marks `embedding_projection_mark_insert` and `embedding_projection_mark_update` +- the source/ACL triggers `projection_source_acl_sync` and `*_source_acl_set` +- the secret-provenance demotions + +As `0024` documents, "a delete marks nothing": the projections' rows go through `ON DELETE CASCADE` foreign keys in the deleting statement. For each chunk, the cascade is a primary-key delete in `embedding_search`, `embedding_keyword_search`, `embedding_keyword_tin` (empty after step 1) and `embedding_secret_provenance`. For each document, it is a primary-key delete in `knowledge_projection_dirty`, `knowledge_document_observation` and `document_secret_provenance`. So **no `knowledge_projection_dirty` marks are written**, and setting `sim.projection_mode` would change nothing, since it only guards insert and update triggers. The run logs the mark count before and after as a check. + +Row deletes do not touch the HNSW or GIN indexes; their entries become dead until vacuum. That is why step 3 exists. + +### Storage objects + +Stored objects are cleaned up the way the app does it, not with raw bucket deletes. Each deleted document whose `file_url` names a `kb/` or `knowledge-base/` key with an active `workspace_files` ownership binding gets a `knowledge.document.storage.cleanup` outbox event, committed in the same transaction as the delete. The app's outbox worker then deletes the object. It re-checks ownership, content version and that no other document references the key, and soft-deletes the binding. Watch progress: + +```sql +SELECT status, count(*) FROM outbox_event +WHERE event_type = 'knowledge.document.storage.cleanup' AND created_at > '' +GROUP BY status; +``` + +**Left behind, by design** (the same as any app deletion): + +- Objects of documents whose key has no active ownership binding. The run logs `Cannot queue knowledge storage cleanup without an ownership binding` with the document id. Clean these up only through the repair process in `lib/knowledge/documents/storage-cleanup.md` ("Previously orphaned files"). +- Documents without a stored object (remote `file_url`) have nothing to clean. +- Events that exhaust their 48 attempts stay as `status = 'dead_letter'` with their last error, for recovery by an operator. + +Source-connected documents are not metered as uploaded storage, so there is no storage-usage ledger to adjust. **Standalone uploads** (`connector_id IS NULL`) are billed storage. They are never selected here, and the final log line reports whether any remain; delete those through the app. + +### What is kept, and why the connector cursors are reset + +The **knowledge base row and its paused connectors are kept**: their configuration, credentials, members, permission snapshots and sync history. Re-enabling is then a matter of resuming the connectors, with nothing to set up again. + +Every page that deletes documents also clears each stopped connector's listing state, in the same transaction, so a run stopped partway never leaves a connector that would skip what was already deleted. These are the columns the app clears when a connector must list everything again (an access-mode switch or a source change), plus the directory checkpoint and each member's retry time: + +- `lastSyncAt`, `lastSyncDocCount`, `listingCheckpoint`, `directoryCheckpoint` and `memberTombstoneCursor` on `knowledge_connector` +- `listingCheckpoint`, `changeCursor`, `memberSyncedThrough`, `lastCompleteListingAt`, `lastListedCount` and `nextAttemptAt` (so every member is due) on its `knowledge_connector_member` rows + +Without this, a resumed connector would sync incrementally from its old cursor ("changed since last sync"), skip directory reconciliation behind a `complete` checkpoint, or leave members waiting on a future retry. It would never re-list the documents deleted here, and the index would stay silently incomplete. Partition work rows (`knowledge_connector_partition`) belong to the old listing generation, and the next full listing replaces them. `--no-connector-reset` skips the reset. When a run finishes with no connector documents left, it resets once more to catch rows a page could not (for example a member added mid-run). + +Each deleting transaction also re-decides the guard while holding the knowledge base and its connectors `FOR SHARE`. Resuming a connector or claiming a sync updates the connector row, so it waits for the page in flight to commit, and the next page refuses. + +### Duration and monitoring + +Throughput depends on chunk count, WAL volume and autovacuum. Measure it: the `--max-pages=5` run logs `elapsedMs` per page. Remaining time is about `(documents left / 200) × seconds per page`. Get the document count with `SELECT count(*) FROM document WHERE knowledge_base_id = ''`. + +While it runs, watch: + +- Replication lag and WAL generation. If lag grows, raise `--pause-ms` or lower `--chunk-batch-size`. +- `pg_stat_activity` for lock waits caused by the scripts. The postgres.js scripts connect as `application_name = 'sim-dormant-org-search-ops'`; the deletion runs on the app's database client, so start it with `DB_APP_NAME=sim-dormant-org-search-delete` to tell its sessions apart. +- Search latency: it should not get worse. Deletes do not touch the indexes. +- The storage cleanup backlog, which the ceiling bounds. + +Autovacuum will start on the heavily deleted tables during the run. That is expected; step 3 handles what it cannot finish. + +### Re-running and resuming + +Re-running with the same arguments is always safe. Deleted documents are gone, so a fresh run from the start sees only what is left. `--after-id` only skips re-reading the already emptied range. A second run over an empty index reads one empty page, clears the already cleared cursors again (idempotent), and exits `0`. + +## 3. Give the space back: reindex, then vacuum + +After step 2, most entries in the listed indexes point at dead tuples. **Reindex first, then vacuum.** Vacuuming an HNSW index repairs the graph around every deleted element, so autovacuum on a heavily deleted `embedding_search` can run for a very long time (pgvector's docs: "Vacuuming can take a while for HNSW indexes. Speed it up by reindexing first"). `REINDEX INDEX CONCURRENTLY` builds a fresh index from live rows only, without blocking reads or writes. The vacuum afterwards then has almost nothing to repair. + +```sh +bun apps/sim/scripts/dormant-org-search/maintenance.ts # report +bun apps/sim/scripts/dormant-org-search/maintenance.ts --reindex=all # dry run: the plan +bun apps/sim/scripts/dormant-org-search/maintenance.ts --reindex=embedding_search_512_cosine_hnsw_idx --execute --maintenance-work-mem=4GB +bun apps/sim/scripts/dormant-org-search/maintenance.ts --reindex=embedding_search_acl_unfilled_idx --execute +bun apps/sim/scripts/dormant-org-search/maintenance.ts --reindex=embedding_keyword_search_content_idx --execute +bun apps/sim/scripts/dormant-org-search/maintenance.ts --reindex=emb_content_fts_idx --execute +bun apps/sim/scripts/dormant-org-search/maintenance.ts --reindex=embedding_search_document_lookup_idx --execute +``` + +`--reindex=all` runs those five in that order. Any other index of a maintenance table can be named (for example another HNSW width, or `emb_doc_id_idx`); the report lists sizes so you can decide. Run one index at a time and check the report between them. + +- The session sets `lock_timeout = 0` and `statement_timeout = 0`, because a cancelled concurrent reindex leaves an invalid `_ccnew` copy behind. The script drops such a leftover before rebuilding that index. +- `--maintenance-work-mem`: HNSW builds are much faster when the graph fits (pgvector reports `hnsw graph no longer fits into maintenance_work_mem` otherwise). Size it against the *new*, live-only index and the server's free memory. Don't exhaust the server's memory. `--parallel-workers=N` sets `max_parallel_maintenance_workers`. +- Disk: the new index is built next to the old one until the swap. +- Progress: `pg_stat_progress_create_index` (also in the report). +- A running non-wraparound autovacuum on the same table is cancelled by the reindex's lock request, as intended. A **wraparound** autovacuum (`pg_stat_activity.query LIKE '%to prevent wraparound%'`) does not yield. Let it finish, or cancel it and vacuum manually right after. + +Then vacuum, **one table at a time**, largest first, off-peak: + +```sh +bun apps/sim/scripts/dormant-org-search/maintenance.ts --vacuum=embedding_search --execute --maintenance-work-mem=1GB +bun apps/sim/scripts/dormant-org-search/maintenance.ts --vacuum=embedding --execute --maintenance-work-mem=1GB +bun apps/sim/scripts/dormant-org-search/maintenance.ts --vacuum=embedding_keyword_search --execute +bun apps/sim/scripts/dormant-org-search/maintenance.ts --vacuum=embedding_secret_provenance --execute +bun apps/sim/scripts/dormant-org-search/maintenance.ts --vacuum=document --execute +bun apps/sim/scripts/dormant-org-search/maintenance.ts --vacuum=knowledge_document_observation --execute +bun apps/sim/scripts/dormant-org-search/maintenance.ts --vacuum=document_secret_provenance --execute +``` + +A manual `VACUUM (VERBOSE, ANALYZE)` runs without autovacuum's cost-delay throttling and with the session's `maintenance_work_mem`. A larger value holds more dead tuple ids per index pass, so it needs fewer passes over every index. This is why a manual run finishes when autovacuum did not. It also refreshes planner statistics, which matter because the tables just shrank. Watch `pg_stat_progress_vacuum`: `index_vacuum_count` above 1 means `maintenance_work_mem` was too small for one pass. `VACUUM` does not shrink files; freed pages are reused by later writes. Do **not** use `VACUUM FULL`: it takes an ACCESS EXCLUSIVE lock and rewrites the table. + +Finally, warm the rebuilt projection into cache: `bun apps/sim/scripts/prewarm-search-projection.ts`. + +### Verify + +```sql +-- Index sizes: the listed indexes should have shrunk roughly in proportion to the rows removed. +SELECT c.relname, pg_size_pretty(pg_relation_size(c.oid)), i.indisvalid +FROM pg_index i JOIN pg_class c ON c.oid = i.indexrelid +WHERE c.relname IN ('embedding_search_512_cosine_hnsw_idx', 'embedding_search_acl_unfilled_idx', + 'embedding_keyword_search_content_idx', 'emb_content_fts_idx', + 'embedding_search_document_lookup_idx'); +-- No invalid leftovers: +SELECT c.relname FROM pg_index i JOIN pg_class c ON c.oid = i.indexrelid WHERE NOT i.indisvalid; +-- Dead tuples back near zero, last_vacuum set: +SELECT relname, n_live_tup, n_dead_tup, last_vacuum, last_autovacuum +FROM pg_stat_user_tables +WHERE relname IN ('embedding', 'embedding_search', 'embedding_keyword_search', 'document'); +-- Nothing left of the index: +SELECT count(*) FROM embedding WHERE knowledge_base_id = ''; +SELECT count(*) FROM embedding_search WHERE knowledge_base_id = ''; +``` + +**Search latency.** Pick an ordinary workspace knowledge base, ``, and compare against the same query from before step 2: + +```sql +EXPLAIN (ANALYZE, BUFFERS) +SELECT id FROM embedding_search +WHERE knowledge_base_id = '' AND enabled +ORDER BY vector_512 <=> (SELECT vector_512 FROM embedding_search + WHERE knowledge_base_id = '' + AND vector_512 IS NOT NULL LIMIT 1) +LIMIT 20; +``` + +Expect far fewer shared buffers read and a shorter execution time. The application's knowledge search latency metrics are the final check. + +## Re-enabling organization indexed search later + +1. Restore Tin (`restore-tin-projection.ts --execute`) if keyword ranking over the index should use it. +2. Resume the connectors from the product. Their cursors were reset, so they list and index the whole source again. This costs about as much as the original indexing (provider API quota, embedding spend, and the index growth this runbook reversed). +3. Once the resync has written the chunks, backfill Tin (`--backfill`, or rebuild the index as described above). + +## Local testing + +The unit tests (`*.test.ts` here) mock the database and run with the normal suite. `dormant-org-search.integration.ts` runs the real scripts against a disposable local PostgreSQL database. It is not part of CI; run it by hand: + +```sh +createdb sim_acl_test_dormant_ops +psql -d sim_acl_test_dormant_ops -c 'CREATE EXTENSION vector; CREATE EXTENSION btree_gin; CREATE EXTENSION pg_trgm;' +(cd packages/db && DATABASE_URL=postgresql://localhost:5432/sim_acl_test_dormant_ops bun run ./scripts/migrate.ts) +cd apps/sim && KNOWLEDGE_ACL_TEST_DATABASE_URL=postgresql://localhost:5432/sim_acl_test_dormant_ops \ + bunx vitest run --mode integration scripts/dormant-org-search/dormant-org-search.integration.ts +``` + +The test installs `0019`'s Tin functions and triggers where they are missing (the `tin` extension is not needed for them), and removes them again afterwards. diff --git a/apps/sim/scripts/dormant-org-search/cli.ts b/apps/sim/scripts/dormant-org-search/cli.ts new file mode 100644 index 00000000000..081ce7ee012 --- /dev/null +++ b/apps/sim/scripts/dormant-org-search/cli.ts @@ -0,0 +1,50 @@ +import { resolveMigrationDatabaseUrl } from '@sim/db/script-migrations/database-url' +import postgres, { type Sql } from 'postgres' + +/** + * Whether a run may write. Every script in this directory reads and reports unless the operator + * passes `--execute`; `--dry-run` is accepted for explicitness, and passing both is refused rather + * than resolved in either direction. + */ +export function resolveExecuteFlag(flags: { execute?: boolean; 'dry-run'?: boolean }): boolean { + if (flags.execute && flags['dry-run']) { + throw new Error('Pass either --execute or --dry-run, not both') + } + return flags.execute === true +} + +/** Parses a positive integer flag, refusing anything a typo could turn into an unbounded run. */ +export function parsePositiveInteger(name: string, value: string | undefined, fallback: number) { + if (value === undefined) return fallback + const parsed = Number(value) + if (!Number.isSafeInteger(parsed) || parsed < 1) { + throw new Error(`--${name} must be a positive integer, got ${value}`) + } + return parsed +} + +/** Parses a non-negative integer flag, such as a pause in milliseconds. */ +export function parseNonNegativeInteger(name: string, value: string | undefined, fallback: number) { + if (value === undefined) return fallback + const parsed = Number(value) + if (!Number.isSafeInteger(parsed) || parsed < 0) { + throw new Error(`--${name} must be a non-negative integer, got ${value}`) + } + return parsed +} + +/** + * One dedicated connection on the migrations role (`MIGRATION_DATABASE_URL`, else `DATABASE_URL`), + * the role that owns the knowledge tables and can therefore drop their triggers, truncate, reindex + * and vacuum them. One connection keeps session settings on the connection that uses them. + */ +export function connectMigrationRole(): Sql { + const url = resolveMigrationDatabaseUrl() + if (!url) throw new Error('MIGRATION_DATABASE_URL or DATABASE_URL is required') + return postgres(url, { + max: 1, + max_lifetime: null, + onnotice: () => undefined, + connection: { application_name: 'sim-dormant-org-search-ops' }, + }) +} diff --git a/apps/sim/scripts/dormant-org-search/delete-search-index-documents.test.ts b/apps/sim/scripts/dormant-org-search/delete-search-index-documents.test.ts new file mode 100644 index 00000000000..d707e705891 --- /dev/null +++ b/apps/sim/scripts/dormant-org-search/delete-search-index-documents.test.ts @@ -0,0 +1,60 @@ +/** + * @vitest-environment node + */ +import { describe, expect, it } from 'vitest' +import { parseDeleteSearchIndexDocumentsArgs } from '@/scripts/dormant-org-search/delete-search-index-documents' + +describe('parseDeleteSearchIndexDocumentsArgs', () => { + it('defaults to a dry run with the documented bounds', () => { + expect(parseDeleteSearchIndexDocumentsArgs(['--knowledge-base-id=kb-1'])).toEqual({ + knowledgeBaseId: 'kb-1', + execute: false, + pageSize: 200, + chunkBatchSize: 1_000, + pauseMs: 250, + maxPages: undefined, + afterId: '', + storageCleanupCeiling: 2_000, + resetConnectors: true, + lockTimeoutMs: 5_000, + statementTimeoutMs: 60_000, + }) + }) + + it('writes only with --execute', () => { + expect( + parseDeleteSearchIndexDocumentsArgs(['--knowledge-base-id=kb-1', '--execute']).execute + ).toBe(true) + expect( + parseDeleteSearchIndexDocumentsArgs(['--knowledge-base-id=kb-1', '--dry-run']).execute + ).toBe(false) + expect(() => + parseDeleteSearchIndexDocumentsArgs(['--knowledge-base-id=kb-1', '--execute', '--dry-run']) + ).toThrow('not both') + }) + + it('parses resume and bound flags', () => { + expect( + parseDeleteSearchIndexDocumentsArgs([ + '--knowledge-base-id=kb-1', + '--max-pages=5', + '--after-id=doc-9', + '--pause-ms=0', + '--no-connector-reset', + ]) + ).toMatchObject({ maxPages: 5, afterId: 'doc-9', pauseMs: 0, resetConnectors: false }) + }) + + it('refuses a missing base, unknown flags and non-positive bounds', () => { + expect(() => parseDeleteSearchIndexDocumentsArgs([])).toThrow('--knowledge-base-id') + expect(() => + parseDeleteSearchIndexDocumentsArgs(['--knowledge-base-id=kb-1', '--force']) + ).toThrow() + expect(() => + parseDeleteSearchIndexDocumentsArgs(['--knowledge-base-id=kb-1', '--page-size=0']) + ).toThrow('positive integer') + expect(() => + parseDeleteSearchIndexDocumentsArgs(['--knowledge-base-id=kb-1', '--max-pages=1.5']) + ).toThrow('positive integer') + }) +}) diff --git a/apps/sim/scripts/dormant-org-search/delete-search-index-documents.ts b/apps/sim/scripts/dormant-org-search/delete-search-index-documents.ts new file mode 100644 index 00000000000..af549142121 --- /dev/null +++ b/apps/sim/scripts/dormant-org-search/delete-search-index-documents.ts @@ -0,0 +1,145 @@ +#!/usr/bin/env bun + +/** + * Deletes every connector-owned document and chunk of one organization search index, keeping the + * knowledge base row and its (paused) connectors, then clears the connectors' listing cursors so a + * resumed connector lists its source from scratch. + * + * Usage: + * DATABASE_URL= bun apps/sim/scripts/dormant-org-search/delete-search-index-documents.ts \ + * --knowledge-base-id= # dry run: guard + first pages + * DATABASE_URL= bun apps/sim/scripts/dormant-org-search/delete-search-index-documents.ts \ + * --knowledge-base-id= --execute [--max-pages=50] [--after-id=] + * + * Options: + * --page-size=200 documents per page + * --chunk-batch-size=1000 chunks per delete transaction + * --pause-ms=250 pause after every committed transaction + * --max-pages=N stop after N pages (dry run default 3; unbounded when executing) + * --after-id= resume after this document id + * --storage-cleanup-ceiling=2000 pending storage cleanup events before the run waits + * --lock-timeout-ms=5000 per transaction + * --statement-timeout-ms=60000 per statement + * --no-connector-reset keep the connectors' listing cursors + * + * Exit codes: 0 finished or stopped at --max-pages (resume with the logged --after-id), 1 failed, + * 2 refused by a precondition (checked before every page; nothing after the refusal is written). + */ + +import { parseArgs } from 'node:util' +import { createLogger } from '@sim/logger' +import { toError } from '@sim/utils/errors' +import { generateShortId } from '@sim/utils/id' +import { + parseNonNegativeInteger, + parsePositiveInteger, + resolveExecuteFlag, +} from '@/scripts/dormant-org-search/cli' +import { + DEFAULT_CHUNK_BATCH_SIZE, + DEFAULT_DOCUMENT_PAGE_SIZE, + DEFAULT_PAUSE_MS, + DEFAULT_STORAGE_CLEANUP_CEILING, + deleteSearchIndexDocuments, + type SearchIndexDeletionOptions, + SearchIndexDeletionRefused, +} from '@/scripts/dormant-org-search/search-index-deletion' + +const logger = createLogger('DeleteSearchIndexDocuments') + +/** Refused by a precondition, before the run or before one of its pages. */ +export const EXIT_REFUSED = 2 + +export interface DeleteSearchIndexDocumentsArgs + extends Omit { + lockTimeoutMs: number + statementTimeoutMs: number +} + +/** Parses the command line; every mutation requires `--execute`. */ +export function parseDeleteSearchIndexDocumentsArgs( + argv: readonly string[] +): DeleteSearchIndexDocumentsArgs { + const { values } = parseArgs({ + args: [...argv], + options: { + 'knowledge-base-id': { type: 'string' }, + execute: { type: 'boolean' }, + 'dry-run': { type: 'boolean' }, + 'page-size': { type: 'string' }, + 'chunk-batch-size': { type: 'string' }, + 'pause-ms': { type: 'string' }, + 'max-pages': { type: 'string' }, + 'after-id': { type: 'string' }, + 'storage-cleanup-ceiling': { type: 'string' }, + 'lock-timeout-ms': { type: 'string' }, + 'statement-timeout-ms': { type: 'string' }, + 'no-connector-reset': { type: 'boolean' }, + }, + strict: true, + }) + const knowledgeBaseId = values['knowledge-base-id']?.trim() + if (!knowledgeBaseId) throw new Error('--knowledge-base-id is required') + return { + knowledgeBaseId, + execute: resolveExecuteFlag(values), + pageSize: parsePositiveInteger('page-size', values['page-size'], DEFAULT_DOCUMENT_PAGE_SIZE), + chunkBatchSize: parsePositiveInteger( + 'chunk-batch-size', + values['chunk-batch-size'], + DEFAULT_CHUNK_BATCH_SIZE + ), + pauseMs: parseNonNegativeInteger('pause-ms', values['pause-ms'], DEFAULT_PAUSE_MS), + maxPages: + values['max-pages'] === undefined + ? undefined + : parsePositiveInteger('max-pages', values['max-pages'], 1), + afterId: values['after-id'] ?? '', + storageCleanupCeiling: parsePositiveInteger( + 'storage-cleanup-ceiling', + values['storage-cleanup-ceiling'], + DEFAULT_STORAGE_CLEANUP_CEILING + ), + resetConnectors: values['no-connector-reset'] !== true, + lockTimeoutMs: parsePositiveInteger('lock-timeout-ms', values['lock-timeout-ms'], 5_000), + statementTimeoutMs: parsePositiveInteger( + 'statement-timeout-ms', + values['statement-timeout-ms'], + 60_000 + ), + } +} + +async function main(): Promise { + const args = parseDeleteSearchIndexDocumentsArgs(process.argv.slice(2)) + /** Imported after parsing, so an argument error is reported before the app's database client requires DATABASE_URL. */ + const { drizzleSearchIndexDeletionStore } = await import( + '@/scripts/dormant-org-search/search-index-deletion-store' + ) + const { lockTimeoutMs, statementTimeoutMs, ...options } = args + try { + await deleteSearchIndexDocuments( + drizzleSearchIndexDeletionStore({ lockTimeoutMs, statementTimeoutMs }), + { ...options, requestId: `dormant-org-search:${generateShortId()}` } + ) + return 0 + } catch (error) { + if (error instanceof SearchIndexDeletionRefused) { + logger.error('Refused by a precondition; no further page was written', { + reasons: error.reasons, + }) + return EXIT_REFUSED + } + throw error + } +} + +if (import.meta.main) { + main().then( + (code) => process.exit(code), + (error) => { + logger.error('Search index deletion failed', toError(error)) + process.exit(1) + } + ) +} diff --git a/apps/sim/scripts/dormant-org-search/disable-tin-projection.test.ts b/apps/sim/scripts/dormant-org-search/disable-tin-projection.test.ts new file mode 100644 index 00000000000..41993587f11 --- /dev/null +++ b/apps/sim/scripts/dormant-org-search/disable-tin-projection.test.ts @@ -0,0 +1,77 @@ +/** + * @vitest-environment node + */ +import { describe, expect, it } from 'vitest' +import { + disableTinProjection, + TIN_TRIGGERS, + type TinProjectionDatabase, + type TinProjectionState, +} from '@/scripts/dormant-org-search/disable-tin-projection' + +/** A database whose Tin triggers and rows change only through the calls under test. */ +function fakeDatabase(initial: Partial = {}) { + const state: TinProjectionState = { + triggers: [...TIN_TRIGGERS], + hasRows: true, + estimatedRows: 1_000, + totalBytes: 8_192, + functionsInstalled: true, + ...initial, + } + const database: TinProjectionDatabase & { drops: number } = { + drops: 0, + readState: async () => ({ ...state, triggers: [...state.triggers] }), + dropTriggersAndTruncate: async () => { + database.drops += 1 + state.triggers = [] + state.hasRows = false + state.totalBytes = 0 + }, + } + return { database, state } +} + +describe('disableTinProjection', () => { + it('drops nothing in a dry run', async () => { + const { database } = fakeDatabase() + const result = await disableTinProjection(database, { execute: false }) + expect(database.drops).toBe(0) + expect(result).toMatchObject({ executed: false, after: null }) + expect(result.before.triggers.map((trigger) => trigger.name)).toEqual([ + 'embedding_keyword_tin_sync', + 'knowledge_base_keyword_tin_sync', + 'embedding_keyword_tin_source_acl_set', + ]) + }) + + it('drops the triggers and truncates when executing', async () => { + const { database } = fakeDatabase() + const result = await disableTinProjection(database, { execute: true }) + expect(database.drops).toBe(1) + expect(result.after).toMatchObject({ triggers: [], hasRows: false }) + }) + + it('does nothing when the triggers are gone and the table is empty', async () => { + const { database } = fakeDatabase({ triggers: [], hasRows: false, estimatedRows: -1 }) + const result = await disableTinProjection(database, { execute: true }) + expect(database.drops).toBe(0) + expect(result.executed).toBe(false) + }) + + it('truncates again when only rows remain', async () => { + const { database } = fakeDatabase({ triggers: [] }) + await disableTinProjection(database, { execute: true }) + expect(database.drops).toBe(1) + }) + + it('fails when a trigger survives the drop', async () => { + const { database, state } = fakeDatabase() + database.dropTriggersAndTruncate = async () => { + state.triggers = [TIN_TRIGGERS[0]] + } + await expect(disableTinProjection(database, { execute: true })).rejects.toThrow( + 'still installed' + ) + }) +}) diff --git a/apps/sim/scripts/dormant-org-search/disable-tin-projection.ts b/apps/sim/scripts/dormant-org-search/disable-tin-projection.ts new file mode 100644 index 00000000000..fd5bcb41a9b --- /dev/null +++ b/apps/sim/scripts/dormant-org-search/disable-tin-projection.ts @@ -0,0 +1,219 @@ +#!/usr/bin/env bun + +/** + * Stops maintaining the Tin keyword projection and empties it, keeping the table, its indexes and + * the Tin functions so the schema and the code that reads it stay valid. + * + * Drops exactly the triggers that write `embedding_keyword_tin` in a writer's transaction: + * - `embedding_keyword_tin_sync` on `embedding` (installed by `0019_tin_keyword_projection`, + * re-guarded by `0024_knowledge_projection_async`) + * - `knowledge_base_keyword_tin_sync` on `knowledge_base` (`0019`) + * - `embedding_keyword_tin_source_acl_set` on `embedding_keyword_tin` (`0021`, re-guarded by `0024`) + * + * then truncates the table, all in one transaction. The document trigger's source and ACL fan-out + * (`sync_projection_source_acl`, `0024`) also updates `embedding_keyword_tin` inline; it is shared + * with `embedding_search`, so it is left as it is and becomes an index probe of an empty table. + * `restore-tin-projection.ts` reinstalls the dropped triggers from the migrations that own them. + * + * Usage: + * DATABASE_URL= bun apps/sim/scripts/dormant-org-search/disable-tin-projection.ts # dry run + * DATABASE_URL= bun apps/sim/scripts/dormant-org-search/disable-tin-projection.ts --execute + * + * Exit codes: 0 done (or nothing to do), 1 failed. + */ + +import { parseArgs } from 'node:util' +import { retryOnLockTimeout } from '@sim/db/scripts/lock-timeout-retry' +import { createLogger } from '@sim/logger' +import { toError } from '@sim/utils/errors' +import type { Sql } from 'postgres' +import { + connectMigrationRole, + parsePositiveInteger, + resolveExecuteFlag, +} from '@/scripts/dormant-org-search/cli' + +const logger = createLogger('DisableTinProjection') + +/** The triggers that write the Tin projection in the writer's transaction, in lock order. */ +export const TIN_TRIGGERS = [ + { name: 'embedding_keyword_tin_sync', table: 'embedding' }, + { name: 'knowledge_base_keyword_tin_sync', table: 'knowledge_base' }, + { name: 'embedding_keyword_tin_source_acl_set', table: 'embedding_keyword_tin' }, +] as const + +export type TinTrigger = (typeof TIN_TRIGGERS)[number] + +/** + * How long each attempt may queue for the table locks. `DROP TRIGGER` and `TRUNCATE` take ACCESS + * EXCLUSIVE locks, and a queued ACCESS EXCLUSIVE request blocks every later reader of the table, + * including search's reads of `knowledge_base`, for as long as it waits; the retries below find a + * quiet moment instead of waiting long once. + */ +const LOCK_TIMEOUT = '3s' +const DEFAULT_RETRY_BUDGET_MINUTES = 20 +const RETRY_BACKOFF = { baseMs: 2_000, maxMs: 30_000 } as const + +export interface TinProjectionState { + /** Tin triggers present now, of {@link TIN_TRIGGERS}. */ + triggers: TinTrigger[] + /** Whether the projection holds any row, read directly. */ + hasRows: boolean + /** + * Planner estimate of the projection's rows, since exact counting would read the whole table. + * It is -1 until the table is first analyzed and lags a truncate until the next analyze. + */ + estimatedRows: number + totalBytes: number + /** Whether `0019` installed the Tin functions here; absent where the database has no `tin`. */ + functionsInstalled: boolean +} + +export interface TinProjectionDatabase { + readState(): Promise + /** Drops {@link TIN_TRIGGERS} and truncates the projection in one transaction. */ + dropTriggersAndTruncate(): Promise +} + +export interface DisableTinProjectionResult { + executed: boolean + before: TinProjectionState + after: TinProjectionState | null +} + +/** + * Reports the projection's state and, when `execute` is set and there is anything to remove, + * drops the triggers and truncates the table. Idempotent: a second run finds nothing to drop and + * an empty table, and truncates again only if something wrote rows meanwhile. + */ +export async function disableTinProjection( + database: TinProjectionDatabase, + options: { execute: boolean } +): Promise { + const before = await database.readState() + logger.info('Tin projection state', { + triggers: before.triggers.map((trigger) => `${trigger.name} ON ${trigger.table}`), + hasRows: before.hasRows, + estimatedRows: before.estimatedRows, + totalBytes: before.totalBytes, + functionsInstalled: before.functionsInstalled, + }) + if (before.triggers.length === 0 && !before.hasRows) { + logger.info('Nothing to do: no Tin triggers and an empty projection') + return { executed: false, before, after: null } + } + if (!options.execute) { + logger.info('Dry run: would drop the triggers above and TRUNCATE embedding_keyword_tin', { + drop: before.triggers.map((trigger) => trigger.name), + }) + return { executed: false, before, after: null } + } + await database.dropTriggersAndTruncate() + const after = await database.readState() + logger.info('Tin projection disabled', { + remainingTriggers: after.triggers.map((trigger) => trigger.name), + hasRows: after.hasRows, + totalBytes: after.totalBytes, + }) + if (after.triggers.length > 0) { + throw new Error('Tin triggers are still installed after the drop') + } + if (after.hasRows) { + logger.warn( + 'The projection gained rows after the truncate: a release that predates the indexed search gate, or runs with SIM_SEARCH_LIVE=false, still projects Tin rows for marked search-index documents' + ) + } + return { executed: true, before, after } +} + +/** The postgres.js implementation, on the migrations role's single connection. */ +export function postgresTinProjectionDatabase( + sql: Sql, + options: { retryBudgetMs: number } +): TinProjectionDatabase { + return { + async readState() { + const present = await sql>` + SELECT t.tgname AS name, c.relname AS table + FROM pg_trigger t JOIN pg_class c ON c.oid = t.tgrelid + WHERE NOT t.tgisinternal + AND t.tgname IN ${sql(TIN_TRIGGERS.map((trigger) => trigger.name))}` + const triggers = TIN_TRIGGERS.filter((trigger) => + present.some((row) => row.name === trigger.name && row.table === trigger.table) + ) + const [size] = await sql< + Array<{ estimated_rows: string | null; total_bytes: string | null; functions: boolean }> + >` + SELECT c.reltuples::bigint AS estimated_rows, + pg_total_relation_size(c.oid)::bigint AS total_bytes, + to_regprocedure('knowledge_tin_stream(tsvector)') IS NOT NULL AS functions + FROM pg_class c WHERE c.oid = to_regclass('embedding_keyword_tin')` + const [heap] = await sql>` + SELECT EXISTS (SELECT 1 FROM embedding_keyword_tin) AS has_rows` + return { + triggers, + hasRows: Boolean(heap?.has_rows), + estimatedRows: Number(size?.estimated_rows ?? 0), + totalBytes: Number(size?.total_bytes ?? 0), + functionsInstalled: Boolean(size?.functions), + } + }, + async dropTriggersAndTruncate() { + await retryOnLockTimeout( + () => + sql.begin(async (tx) => { + await tx.unsafe(`SET LOCAL lock_timeout = '${LOCK_TIMEOUT}'`) + for (const trigger of TIN_TRIGGERS) { + await tx.unsafe(`DROP TRIGGER IF EXISTS ${trigger.name} ON ${trigger.table}`) + } + await tx.unsafe('TRUNCATE embedding_keyword_tin') + }), + { + budgetMs: options.retryBudgetMs, + backoff: RETRY_BACKOFF, + onRetry: ({ attempt, delayMs }) => + logger.warn('Tin drop waited out its lock timeout; retrying', { + attempt, + retryInMs: Math.round(delayMs), + }), + } + ) + }, + } +} + +async function main(): Promise { + const { values } = parseArgs({ + options: { + execute: { type: 'boolean' }, + 'dry-run': { type: 'boolean' }, + 'retry-budget-minutes': { type: 'string' }, + }, + strict: true, + }) + const execute = resolveExecuteFlag(values) + const retryBudgetMinutes = parsePositiveInteger( + 'retry-budget-minutes', + values['retry-budget-minutes'], + DEFAULT_RETRY_BUDGET_MINUTES + ) + const sql = connectMigrationRole() + try { + await disableTinProjection( + postgresTinProjectionDatabase(sql, { retryBudgetMs: retryBudgetMinutes * 60_000 }), + { execute } + ) + } finally { + await sql.end() + } +} + +if (import.meta.main) { + main().then( + () => process.exit(0), + (error) => { + logger.error('Disabling the Tin projection failed', toError(error)) + process.exit(1) + } + ) +} diff --git a/apps/sim/scripts/dormant-org-search/dormant-org-search.integration.ts b/apps/sim/scripts/dormant-org-search/dormant-org-search.integration.ts new file mode 100644 index 00000000000..333a333d749 --- /dev/null +++ b/apps/sim/scripts/dormant-org-search/dormant-org-search.integration.ts @@ -0,0 +1,411 @@ +/** + * Real PostgreSQL coverage for the dormant organization search runbook: the Tin trigger drop and + * restore, and the search index deletion with its storage intents, connector reset, and the + * guarantee that deleting marks nothing for the projector. Runs only by hand against a disposable + * local database (see README.md); it is deliberately not listed in CI. + */ +import { db } from '@sim/db' +import { + document, + embedding, + embeddingKeywordSearch, + embeddingKeywordTin, + embeddingSearch, + knowledgeBase, + knowledgeConnector, + knowledgeProjectionDirty, + organization, + outboxEvent, + user, + workspace, + workspaceFiles, +} from '@sim/db/schema' +import { installProjection } from '@sim/db/script-migrations/0019_tin_keyword_projection' +import { generateId } from '@sim/utils/id' +import { and, count, eq, inArray, sql } from 'drizzle-orm' +import postgres from 'postgres' +import { afterAll, beforeAll, describe, expect, it } from 'vitest' +import { + createKnowledgeAclFixtureIds, + seedKnowledgeAclFixture, +} from '@/lib/knowledge/__integration__/seed-source-access-fixture' +import { KNOWLEDGE_STORAGE_CLEANUP_EVENT } from '@/lib/knowledge/documents/storage-cleanup' +import { + disableTinProjection, + postgresTinProjectionDatabase, +} from '@/scripts/dormant-org-search/disable-tin-projection' +import { + migrationTinRestoreSteps, + restoreTinProjection, +} from '@/scripts/dormant-org-search/restore-tin-projection' +import { + deleteSearchIndexDocuments, + SearchIndexDeletionRefused, +} from '@/scripts/dormant-org-search/search-index-deletion' +import { drizzleSearchIndexDeletionStore } from '@/scripts/dormant-org-search/search-index-deletion-store' + +const ids = createKnowledgeAclFixtureIds() +const indexId = generateId() +const indexConnectorId = generateId() +const vector = Array.from({ length: 1536 }, (_, index) => (index % 7) / 10) +const DOCUMENTS = 7 +const CHUNKS = 3 +const timeouts = { lockTimeoutMs: 5_000, statementTimeoutMs: 30_000 } +const sleep = async () => undefined + +let pg: postgres.Sql +let tinFunctionsExisted = false +const indexDocumentIds: string[] = [] +const workspaceDocumentId = generateId() +const bindingIds: string[] = [] + +function chunk(documentId: string, knowledgeBaseId: string, chunkIndex: number) { + return { + id: generateId(), + documentId, + knowledgeBaseId, + chunkIndex, + chunkHash: `fixture-hash-${chunkIndex}`, + content: `dormant search fixture chunk ${chunkIndex}`, + contentLength: 30, + tokenCount: 5, + startOffset: 0, + endOffset: 30, + embeddingModel: 'text-embedding-3-small', + embedding: vector, + } +} + +async function rowsOf(knowledgeBaseId: string) { + const counts = await Promise.all( + [embedding, embeddingSearch, embeddingKeywordSearch, embeddingKeywordTin].map(async (table) => { + const [row] = await db + .select({ rows: count() }) + .from(table) + .where(eq(table.knowledgeBaseId, knowledgeBaseId)) + return row.rows + }) + ) + const [docs] = await db + .select({ rows: count() }) + .from(document) + .where(eq(document.knowledgeBaseId, knowledgeBaseId)) + return { + documents: docs.rows, + chunks: counts[0], + vector: counts[1], + keyword: counts[2], + tin: counts[3], + } +} + +async function triggerNames() { + const rows = await pg>` + SELECT tgname AS name, pg_get_triggerdef(oid) AS definition FROM pg_trigger + WHERE NOT tgisinternal AND tgname IN + ('embedding_keyword_tin_sync', 'knowledge_base_keyword_tin_sync', 'embedding_keyword_tin_source_acl_set') + ORDER BY tgname` + return rows +} + +beforeAll(async () => { + pg = postgres(process.env.DATABASE_URL as string, { max: 1, onnotice: () => undefined }) + const [existing] = await pg>` + SELECT to_regprocedure('knowledge_tin_stream(tsvector)') IS NOT NULL AS present` + tinFunctionsExisted = Boolean(existing?.present) + /** + * `0019` installs the Tin triggers only where the `tin` extension exists; their functions are + * plain SQL and PL/pgSQL, so the triggers are installed here directly to exercise the drop. + */ + await installProjection(pg) + + await seedKnowledgeAclFixture(ids) + await db.insert(knowledgeBase).values({ + id: indexId, + userId: ids.aliceId, + organizationId: ids.organizationId, + name: 'Dormant search fixture', + isSearchIndex: true, + }) + await db.insert(knowledgeConnector).values({ + id: indexConnectorId, + knowledgeBaseId: indexId, + connectorType: 'google_drive', + sourceConfig: {}, + accessMode: 'admin', + status: 'active', + lastSyncAt: new Date(), + listingCheckpoint: { cursor: 'fixture' }, + }) + for (let index = 0; index < DOCUMENTS; index++) { + const documentId = `dormant-${String(index).padStart(2, '0')}-${generateId()}` + indexDocumentIds.push(documentId) + const key = `kb/${generateId()}.txt` + /** Every other document has a stored object with an organization-owned binding. */ + const stored = index % 2 === 0 + if (stored) { + const bindingId = generateId() + bindingIds.push(bindingId) + await db.insert(workspaceFiles).values({ + id: bindingId, + key, + userId: ids.aliceId, + organizationId: ids.organizationId, + context: 'knowledge-base', + originalName: `${documentId}.txt`, + contentType: 'text/plain', + sizeBytes: 12, + }) + } + await db.insert(document).values({ + id: documentId, + knowledgeBaseId: indexId, + connectorId: indexConnectorId, + filename: `${documentId}.txt`, + fileUrl: stored + ? `http://localhost:3000/api/files/serve/${key}?context=knowledge-base` + : 'https://fixture.test/remote', + storageKey: stored ? key : null, + fileSize: 12, + mimeType: 'text/plain', + processingStatus: 'completed', + }) + await db + .insert(embedding) + .values( + Array.from({ length: CHUNKS }, (_, chunkIndex) => chunk(documentId, indexId, chunkIndex)) + ) + } + await db.insert(document).values({ + id: workspaceDocumentId, + knowledgeBaseId: ids.knowledgeBaseId, + connectorId: ids.connectorId, + filename: 'workspace.txt', + fileUrl: 'https://fixture.test/workspace', + fileSize: 12, + mimeType: 'text/plain', + processingStatus: 'completed', + }) + await db.insert(embedding).values(chunk(workspaceDocumentId, ids.knowledgeBaseId, 0)) + /** Seeding marks documents; the deletion under test must add none. */ + await db.delete(knowledgeProjectionDirty) +}, 60_000) + +afterAll(async () => { + await db + .delete(outboxEvent) + .where( + and( + eq(outboxEvent.eventType, KNOWLEDGE_STORAGE_CLEANUP_EVENT), + inArray(sql`${outboxEvent.payload}::jsonb ->> 'fileId'`, bindingIds) + ) + ) + await db.delete(knowledgeBase).where(eq(knowledgeBase.id, indexId)) + await db.delete(workspace).where(eq(workspace.id, ids.workspaceId)) + await db.delete(organization).where(eq(organization.id, ids.organizationId)) + await db.delete(user).where(inArray(user.id, [ids.aliceId, ids.bobId])) + if (!tinFunctionsExisted) { + await pg.unsafe(`DROP TRIGGER IF EXISTS embedding_keyword_tin_sync ON embedding; + DROP TRIGGER IF EXISTS knowledge_base_keyword_tin_sync ON knowledge_base; + DROP FUNCTION IF EXISTS sync_embedding_keyword_tin(); + DROP FUNCTION IF EXISTS sync_knowledge_base_keyword_tin(); + DROP FUNCTION IF EXISTS knowledge_tin_stream(tsvector); + DROP FUNCTION IF EXISTS knowledge_tin_base_token(text); + DROP FUNCTION IF EXISTS knowledge_tin_membership_key(text);`) + } + await pg.end() +}) + +describe('dormant organization search runbook in PostgreSQL', () => { + it('drops the Tin triggers, truncates the projection, and leaves the other projections alone', async () => { + expect(await rowsOf(indexId)).toEqual({ + documents: DOCUMENTS, + chunks: DOCUMENTS * CHUNKS, + vector: DOCUMENTS * CHUNKS, + keyword: DOCUMENTS * CHUNKS, + tin: DOCUMENTS * CHUNKS, + }) + const database = postgresTinProjectionDatabase(pg, { retryBudgetMs: 30_000 }) + + const dryRun = await disableTinProjection(database, { execute: false }) + expect(dryRun.before.triggers).toHaveLength(3) + expect((await rowsOf(indexId)).tin).toBe(DOCUMENTS * CHUNKS) + + const result = await disableTinProjection(database, { execute: true }) + expect(result.after).toMatchObject({ triggers: [], hasRows: false }) + expect(await triggerNames()).toEqual([]) + expect(await rowsOf(indexId)).toMatchObject({ tin: 0, vector: DOCUMENTS * CHUNKS }) + + /** A new chunk in the search index no longer reaches Tin. */ + const late = chunk(indexDocumentIds[0], indexId, 99) + await db.insert(embedding).values(late) + expect((await rowsOf(indexId)).tin).toBe(0) + await db.delete(embedding).where(eq(embedding.id, late.id)) + await db.delete(knowledgeProjectionDirty) + }) + + it('refuses while a connector is active, and writes nothing', async () => { + const store = drizzleSearchIndexDeletionStore(timeouts) + await expect( + deleteSearchIndexDocuments(store, { + knowledgeBaseId: indexId, + execute: true, + requestId: 'fixture', + sleep, + }) + ).rejects.toBeInstanceOf(SearchIndexDeletionRefused) + await expect( + deleteSearchIndexDocuments(store, { + knowledgeBaseId: ids.knowledgeBaseId, + execute: true, + requestId: 'fixture', + sleep, + }) + ).rejects.toThrow('not an organization search index') + /** A deleting transaction re-decides the guard itself, so a resume between pages cannot race it. */ + await expect( + store.deleteChunkBatch(indexId, indexDocumentIds.slice(0, 1), 10) + ).rejects.toBeInstanceOf(SearchIndexDeletionRefused) + await expect( + store.deleteDocuments(indexId, indexDocumentIds.slice(0, 1), 'fixture', true) + ).rejects.toBeInstanceOf(SearchIndexDeletionRefused) + expect((await rowsOf(indexId)).documents).toBe(DOCUMENTS) + }) + + it('deletes the search index in resumable pages, queues storage cleanup, and marks nothing', async () => { + await db + .update(knowledgeConnector) + .set({ + status: 'paused', + lastSyncAt: new Date(), + listingCheckpoint: { cursor: 'fixture' }, + directoryCheckpoint: { phase: 'complete' }, + }) + .where(eq(knowledgeConnector.id, indexConnectorId)) + const store = drizzleSearchIndexDeletionStore(timeouts) + const options = { + knowledgeBaseId: indexId, + requestId: 'fixture', + pageSize: 3, + chunkBatchSize: 4, + pauseMs: 0, + sleep, + } + + const dryRun = await deleteSearchIndexDocuments(store, { + ...options, + execute: false, + maxPages: 10, + }) + expect(dryRun).toMatchObject({ pages: 3, chunksCounted: DOCUMENTS * CHUNKS, done: true }) + expect((await rowsOf(indexId)).documents).toBe(DOCUMENTS) + + const first = await deleteSearchIndexDocuments(store, { + ...options, + execute: true, + maxPages: 1, + }) + expect(first).toMatchObject({ pages: 1, documentsDeleted: 3, done: false }) + expect((await rowsOf(indexId)).documents).toBe(DOCUMENTS - 3) + /** A run stopped after one page already leaves the connector listing from scratch on resume. */ + const [partial] = await db + .select() + .from(knowledgeConnector) + .where(eq(knowledgeConnector.id, indexConnectorId)) + expect(partial).toMatchObject({ + lastSyncAt: null, + listingCheckpoint: null, + directoryCheckpoint: null, + }) + + const rest = await deleteSearchIndexDocuments(store, { + ...options, + execute: true, + afterId: first.afterId, + }) + expect(rest).toMatchObject({ + documentsDeleted: DOCUMENTS - 3, + done: true, + connectorsReset: { connectors: 0, members: 0 }, + standaloneDocumentsRemain: false, + projectionMarks: { before: 0, after: 0 }, + }) + expect(first.chunksDeleted + rest.chunksDeleted).toBe(DOCUMENTS * CHUNKS) + expect(first.storageCleanupQueued + rest.storageCleanupQueued).toBe(bindingIds.length) + expect(await rowsOf(indexId)).toEqual({ + documents: 0, + chunks: 0, + vector: 0, + keyword: 0, + tin: 0, + }) + expect(await rowsOf(ids.knowledgeBaseId)).toMatchObject({ documents: 1, chunks: 1 }) + const [marks] = await db.select({ rows: count() }).from(knowledgeProjectionDirty) + expect(marks.rows).toBe(0) + + const events = await db + .select({ payload: outboxEvent.payload }) + .from(outboxEvent) + .where( + and( + eq(outboxEvent.eventType, KNOWLEDGE_STORAGE_CLEANUP_EVENT), + eq(outboxEvent.status, 'pending'), + inArray(sql`${outboxEvent.payload}::jsonb ->> 'fileId'`, bindingIds) + ) + ) + expect( + events.map((event) => (event.payload as { organizationId: string }).organizationId) + ).toEqual(bindingIds.map(() => ids.organizationId)) + + const [connectorRow] = await db + .select() + .from(knowledgeConnector) + .where(eq(knowledgeConnector.id, indexConnectorId)) + expect(connectorRow).toMatchObject({ + status: 'paused', + lastSyncAt: null, + listingCheckpoint: null, + directoryCheckpoint: null, + }) + const [kbRow] = await db.select().from(knowledgeBase).where(eq(knowledgeBase.id, indexId)) + expect(kbRow).toMatchObject({ isSearchIndex: true, deletedAt: null }) + + const again = await deleteSearchIndexDocuments(store, { ...options, execute: true }) + expect(again).toMatchObject({ pages: 0, documentsDeleted: 0, done: true }) + }) + + it('restores exactly the dropped Tin triggers, guarded, and backfills', async () => { + const database = postgresTinProjectionDatabase(pg, { retryBudgetMs: 30_000 }) + const result = await restoreTinProjection(database, migrationTinRestoreSteps(pg), { + execute: true, + backfill: true, + }) + expect(result).toMatchObject({ executed: true, backfilled: 0 }) + const triggers = await triggerNames() + expect(triggers.map((trigger) => trigger.name)).toEqual([ + 'embedding_keyword_tin_source_acl_set', + 'embedding_keyword_tin_sync', + 'knowledge_base_keyword_tin_sync', + ]) + for (const trigger of triggers.filter( + (row) => row.name !== 'knowledge_base_keyword_tin_sync' + )) { + expect(trigger.definition).toContain('sim.projection_mode') + } + + /** A synchronous chunk write in a search index reaches Tin again. */ + const documentId = generateId() + await db.insert(document).values({ + id: documentId, + knowledgeBaseId: indexId, + connectorId: indexConnectorId, + filename: 'resumed.txt', + fileUrl: 'https://fixture.test/resumed', + fileSize: 12, + mimeType: 'text/plain', + processingStatus: 'completed', + }) + await db.insert(embedding).values(chunk(documentId, indexId, 0)) + expect((await rowsOf(indexId)).tin).toBe(1) + }) +}) diff --git a/apps/sim/scripts/dormant-org-search/maintenance.test.ts b/apps/sim/scripts/dormant-org-search/maintenance.test.ts new file mode 100644 index 00000000000..6c8225efd22 --- /dev/null +++ b/apps/sim/scripts/dormant-org-search/maintenance.test.ts @@ -0,0 +1,120 @@ +/** + * @vitest-environment node + */ +import { describe, expect, it } from 'vitest' +import { + type MaintenanceDatabase, + planMaintenance, + REINDEX_TARGETS, + runMaintenance, +} from '@/scripts/dormant-org-search/maintenance' + +const INDEX_TABLES: Record = { + embedding_search_512_cosine_hnsw_idx: 'embedding_search', + embedding_search_acl_unfilled_idx: 'embedding_search', + embedding_keyword_search_content_idx: 'embedding_keyword_search', + emb_content_fts_idx: 'embedding', + embedding_search_document_lookup_idx: 'embedding_search', + embedding_search_cosine_hnsw_idx: 'embedding_search', + user_pkey: 'user', +} + +function fakeDatabase(leftovers: Record = {}) { + const ran: string[] = [] + const database: MaintenanceDatabase = { + report: async () => ({}), + indexTable: async (index) => INDEX_TABLES[index] ?? null, + invalidReindexLeftovers: async (index) => leftovers[index] ?? [], + run: async (statement) => { + ran.push(statement) + }, + } + return { database, ran } +} + +describe('planMaintenance', () => { + it('reports only when no action is requested', () => { + expect(planMaintenance({})).toEqual({ + action: { kind: 'report' }, + settings: [], + statements: [], + }) + }) + + it('reindexes every listed index concurrently, without lock or statement timeouts', () => { + const plan = planMaintenance({ reindex: 'all', maintenanceWorkMem: '2GB', parallelWorkers: 4 }) + expect(plan.statements).toEqual( + REINDEX_TARGETS.map((name) => `REINDEX INDEX CONCURRENTLY ${name}`) + ) + expect(plan.settings).toEqual([ + 'SET lock_timeout = 0', + 'SET statement_timeout = 0', + "SET maintenance_work_mem = '2GB'", + 'SET max_parallel_maintenance_workers = 4', + ]) + }) + + it('refuses malformed names, tables outside the allowlist and malformed memory settings', () => { + expect(() => planMaintenance({ reindex: 'idx; DROP TABLE x' })).toThrow('Not an index name') + expect(() => planMaintenance({ vacuum: 'user; DROP TABLE x' })).toThrow( + 'Not a maintenance table' + ) + expect(() => planMaintenance({ vacuum: 'embedding', maintenanceWorkMem: "1GB'; --" })).toThrow( + 'maintenance-work-mem' + ) + expect(() => planMaintenance({ reindex: 'all', vacuum: 'embedding' })).toThrow('not both') + }) + + it('vacuums exactly one table', () => { + expect(planMaintenance({ vacuum: 'embedding_search' }).statements).toEqual([ + 'VACUUM (VERBOSE, ANALYZE) embedding_search', + ]) + }) +}) + +describe('runMaintenance', () => { + it('refuses an index outside the maintenance tables before running anything', async () => { + const { database, ran } = fakeDatabase() + for (const index of ['user_pkey', 'missing_idx']) { + await expect( + runMaintenance(database, planMaintenance({ reindex: index }), { execute: true }) + ).rejects.toThrow('Not a maintenance index') + } + expect(ran).toEqual([]) + }) + + it('accepts another index of a maintenance table by name', async () => { + const { database, ran } = fakeDatabase() + await runMaintenance( + database, + planMaintenance({ reindex: 'embedding_search_cosine_hnsw_idx' }), + { + execute: true, + } + ) + expect(ran.at(-1)).toBe('REINDEX INDEX CONCURRENTLY embedding_search_cosine_hnsw_idx') + }) + + it('runs nothing in a dry run', async () => { + const { database, ran } = fakeDatabase() + await runMaintenance(database, planMaintenance({ reindex: 'all' }), { execute: false }) + expect(ran).toEqual([]) + }) + + it('drops an interrupted reindex leftover before rebuilding that index', async () => { + const { database, ran } = fakeDatabase({ + embedding_search_512_cosine_hnsw_idx: ['embedding_search_512_cosine_hnsw_idx_ccnew'], + }) + await runMaintenance( + database, + planMaintenance({ reindex: 'embedding_search_512_cosine_hnsw_idx' }), + { execute: true } + ) + expect(ran).toEqual([ + 'SET lock_timeout = 0', + 'SET statement_timeout = 0', + 'DROP INDEX CONCURRENTLY IF EXISTS embedding_search_512_cosine_hnsw_idx_ccnew', + 'REINDEX INDEX CONCURRENTLY embedding_search_512_cosine_hnsw_idx', + ]) + }) +}) diff --git a/apps/sim/scripts/dormant-org-search/maintenance.ts b/apps/sim/scripts/dormant-org-search/maintenance.ts new file mode 100644 index 00000000000..8fc329781d4 --- /dev/null +++ b/apps/sim/scripts/dormant-org-search/maintenance.ts @@ -0,0 +1,293 @@ +#!/usr/bin/env bun + +/** + * Post-deletion maintenance for the knowledge search tables: a health report, `REINDEX INDEX + * CONCURRENTLY` of the indexes the deletion left mostly dead, and a manual `VACUUM (VERBOSE, + * ANALYZE)` one table at a time. + * + * Reindex before vacuuming: vacuuming an HNSW index repairs the graph around every deleted element, + * which can keep autovacuum on a heavily deleted `embedding_search` running for a very long time, + * while a concurrent reindex builds a fresh graph from the live rows only and leaves the vacuum + * almost nothing to repair. + * + * Usage: + * DATABASE_URL= bun apps/sim/scripts/dormant-org-search/maintenance.ts # report only + * ... maintenance.ts --reindex=embedding_search_512_cosine_hnsw_idx # dry run: prints the plan + * ... maintenance.ts --reindex=embedding_search_512_cosine_hnsw_idx --execute [--maintenance-work-mem=4GB] + * ... maintenance.ts --reindex=all --execute + * ... maintenance.ts --vacuum=embedding_search --execute [--maintenance-work-mem=1GB] + * + * Exit codes: 0 done, 1 failed or refused. + */ + +import { parseArgs } from 'node:util' +import { createLogger } from '@sim/logger' +import { toError } from '@sim/utils/errors' +import type { Sql } from 'postgres' +import { + connectMigrationRole, + parsePositiveInteger, + resolveExecuteFlag, +} from '@/scripts/dormant-org-search/cli' + +const logger = createLogger('DormantOrgSearchMaintenance') + +/** + * The indexes a search index deletion leaves mostly dead, in the order to rebuild them: the vector + * index search walks first, then the keyword and lookup indexes. `--reindex=all` rebuilds these; + * any other index of a {@link VACUUM_TARGETS} table can be named explicitly. + */ +export const REINDEX_TARGETS = [ + 'embedding_search_512_cosine_hnsw_idx', + 'embedding_search_acl_unfilled_idx', + 'embedding_keyword_search_content_idx', + 'emb_content_fts_idx', + 'embedding_search_document_lookup_idx', +] as const + +/** The tables the deletion wrote dead tuples into, largest first. */ +export const VACUUM_TARGETS = [ + 'embedding_search', + 'embedding', + 'embedding_keyword_search', + 'embedding_secret_provenance', + 'document', + 'knowledge_document_observation', + 'document_secret_provenance', + 'embedding_keyword_tin', +] as const + +/** Tables the report covers. */ +const REPORT_TABLES = [...VACUUM_TARGETS, 'knowledge_projection_dirty', 'outbox_event'] as const + +/** Accepts `1GB`, `512MB` and the like, the only shapes interpolated into `SET`. */ +const MEMORY_SETTING = /^[1-9][0-9]{0,5}(kB|MB|GB)$/ + +/** An unquoted lower-case identifier, the only index name shape interpolated into `REINDEX`. */ +const INDEX_NAME = /^[a-z_][a-z0-9_]{0,62}$/ + +export type MaintenanceAction = + | { kind: 'report' } + | { kind: 'reindex'; indexes: string[] } + | { kind: 'vacuum'; table: string } + +export interface MaintenancePlan { + action: MaintenanceAction + /** Session settings applied before the statements. */ + settings: string[] + statements: string[] +} + +/** + * Validates targets and renders the statements: a table must be listed, and an index name must be + * a plain identifier, which {@link runMaintenance} then checks belongs to a listed table before + * anything runs. A vacuum takes exactly one table: running two at once doubles the I/O and memory + * the database has to absorb. + */ +export function planMaintenance(input: { + reindex?: string + vacuum?: string + maintenanceWorkMem?: string + parallelWorkers?: number +}): MaintenancePlan { + if (input.reindex && input.vacuum) throw new Error('Pass --reindex or --vacuum, not both') + if (input.maintenanceWorkMem && !MEMORY_SETTING.test(input.maintenanceWorkMem)) { + throw new Error( + `--maintenance-work-mem must look like 1GB or 512MB, got ${input.maintenanceWorkMem}` + ) + } + /** + * `REINDEX CONCURRENTLY` and `VACUUM` wait for older transactions; a lock or statement timeout + * would cancel them partway, and a cancelled concurrent reindex leaves an invalid copy behind. + */ + const settings = [ + 'SET lock_timeout = 0', + 'SET statement_timeout = 0', + ...(input.maintenanceWorkMem + ? [`SET maintenance_work_mem = '${input.maintenanceWorkMem}'`] + : []), + ...(input.parallelWorkers !== undefined + ? [`SET max_parallel_maintenance_workers = ${input.parallelWorkers}`] + : []), + ] + if (input.reindex) { + const requested = + input.reindex === 'all' + ? [...REINDEX_TARGETS] + : input.reindex.split(',').map((name) => name.trim()) + const malformed = requested.filter((name) => !INDEX_NAME.test(name)) + if (malformed.length > 0) throw new Error(`Not an index name: ${malformed.join(', ')}`) + return { + action: { kind: 'reindex', indexes: requested }, + settings, + statements: requested.map((name) => `REINDEX INDEX CONCURRENTLY ${name}`), + } + } + if (input.vacuum) { + if (!(VACUUM_TARGETS as readonly string[]).includes(input.vacuum)) { + throw new Error( + `Not a maintenance table: ${input.vacuum}; expected one of ${VACUUM_TARGETS.join(', ')}` + ) + } + return { + action: { kind: 'vacuum', table: input.vacuum }, + settings, + statements: [`VACUUM (VERBOSE, ANALYZE) ${input.vacuum}`], + } + } + return { action: { kind: 'report' }, settings: [], statements: [] } +} + +export interface MaintenanceDatabase { + report(): Promise> + /** The table an index belongs to, or `null` when no such index exists. */ + indexTable(index: string): Promise + /** Invalid leftovers of an interrupted concurrent reindex of `index` (`_ccnew`, `_ccnew1`, ...). */ + invalidReindexLeftovers(index: string): Promise + run(statement: string): Promise +} + +/** + * Prints the report, and when `execute` is set runs the plan: each reindex first drops any invalid + * copy an interrupted run left, which would otherwise be maintained on every write while serving + * no query. The report is printed again at the end. + */ +export async function runMaintenance( + database: MaintenanceDatabase, + plan: MaintenancePlan, + options: { execute: boolean } +): Promise<{ executed: string[] }> { + logger.info('Knowledge search tables before maintenance', await database.report()) + if (plan.action.kind === 'report') return { executed: [] } + if (plan.action.kind === 'reindex') { + for (const index of plan.action.indexes) { + const table = await database.indexTable(index) + if (!table || !(VACUUM_TARGETS as readonly string[]).includes(table)) { + throw new Error( + `Not a maintenance index: ${index}; expected an index of ${VACUUM_TARGETS.join(', ')}` + ) + } + } + } + if (!options.execute) { + logger.info('Dry run: would run', { settings: plan.settings, statements: plan.statements }) + return { executed: [] } + } + const executed: string[] = [] + for (const setting of plan.settings) await database.run(setting) + for (const [position, statement] of plan.statements.entries()) { + if (plan.action.kind === 'reindex') { + const index = plan.action.indexes[position] + for (const leftover of await database.invalidReindexLeftovers(index)) { + const drop = `DROP INDEX CONCURRENTLY IF EXISTS ${leftover}` + logger.warn('Dropping an invalid index left by an interrupted reindex', { leftover }) + await database.run(drop) + executed.push(drop) + } + } + const startedAt = Date.now() + logger.info('Running', { statement }) + await database.run(statement) + executed.push(statement) + logger.info('Finished', { statement, elapsedMs: Date.now() - startedAt }) + } + logger.info('Knowledge search tables after maintenance', await database.report()) + return { executed } +} + +/** The postgres.js implementation on one reserved session, so `SET` applies to what follows it. */ +export function postgresMaintenanceDatabase(sql: Sql): MaintenanceDatabase { + return { + async report() { + const tables = await sql` + SELECT s.relname AS table, s.n_live_tup AS live_rows, s.n_dead_tup AS dead_rows, + pg_size_pretty(pg_total_relation_size(s.relid)) AS total_size, + s.last_vacuum, s.last_autovacuum, s.last_analyze, s.last_autoanalyze + FROM pg_stat_user_tables s + WHERE s.relname IN ${sql([...REPORT_TABLES])} + ORDER BY pg_total_relation_size(s.relid) DESC` + const indexes = await sql` + SELECT c.relname AS index, t.relname AS table, pg_size_pretty(pg_relation_size(c.oid)) AS size, + i.indisvalid AS valid + FROM pg_index i + JOIN pg_class c ON c.oid = i.indexrelid + JOIN pg_class t ON t.oid = i.indrelid + WHERE t.relname IN ${sql([...REPORT_TABLES])} + AND (c.relname IN ${sql([...REINDEX_TARGETS])} OR c.relname LIKE '%hnsw%' + OR c.relname LIKE '%ccnew%' OR NOT i.indisvalid) + ORDER BY pg_relation_size(c.oid) DESC` + const vacuums = await sql` + SELECT p.pid, c.relname AS table, p.phase, p.heap_blks_total, p.heap_blks_scanned, + p.index_vacuum_count, now() - a.xact_start AS running_for + FROM pg_stat_progress_vacuum p + JOIN pg_class c ON c.oid = p.relid + LEFT JOIN pg_stat_activity a ON a.pid = p.pid` + const builds = await sql` + SELECT p.pid, c.relname AS index, p.phase, p.blocks_total, p.blocks_done, + p.tuples_total, p.tuples_done + FROM pg_stat_progress_create_index p + LEFT JOIN pg_class c ON c.oid = p.index_relid` + return { tables, indexes, vacuums, builds } + }, + async indexTable(index) { + const [row] = await sql>` + SELECT t.relname AS table + FROM pg_index i + JOIN pg_class c ON c.oid = i.indexrelid + JOIN pg_class t ON t.oid = i.indrelid + JOIN pg_namespace n ON n.oid = c.relnamespace + WHERE c.relname = ${index} AND n.nspname = current_schema()` + return row?.table ?? null + }, + async invalidReindexLeftovers(index) { + const rows = await sql>` + SELECT c.relname AS name + FROM pg_index i JOIN pg_class c ON c.oid = i.indexrelid + WHERE NOT i.indisvalid AND c.relname ~ ${`^${index}_ccnew[0-9]*$`}` + return rows.map((row) => row.name) + }, + async run(statement) { + await sql.unsafe(statement) + }, + } +} + +async function main(): Promise { + const { values } = parseArgs({ + options: { + execute: { type: 'boolean' }, + 'dry-run': { type: 'boolean' }, + reindex: { type: 'string' }, + vacuum: { type: 'string' }, + 'maintenance-work-mem': { type: 'string' }, + 'parallel-workers': { type: 'string' }, + }, + strict: true, + }) + const execute = resolveExecuteFlag(values) + const plan = planMaintenance({ + reindex: values.reindex, + vacuum: values.vacuum, + maintenanceWorkMem: values['maintenance-work-mem'], + parallelWorkers: + values['parallel-workers'] === undefined + ? undefined + : parsePositiveInteger('parallel-workers', values['parallel-workers'], 1), + }) + const sql = connectMigrationRole() + try { + await runMaintenance(postgresMaintenanceDatabase(sql), plan, { execute }) + } finally { + await sql.end() + } +} + +if (import.meta.main) { + main().then( + () => process.exit(0), + (error) => { + logger.error('Maintenance failed', toError(error)) + process.exit(1) + } + ) +} diff --git a/apps/sim/scripts/dormant-org-search/restore-tin-projection.test.ts b/apps/sim/scripts/dormant-org-search/restore-tin-projection.test.ts new file mode 100644 index 00000000000..05d7262a27a --- /dev/null +++ b/apps/sim/scripts/dormant-org-search/restore-tin-projection.test.ts @@ -0,0 +1,109 @@ +/** + * @vitest-environment node + */ +import { describe, expect, it, vi } from 'vitest' + +vi.mock('@sim/db/script-migrations/0019_tin_keyword_projection', () => ({ + installProjection: vi.fn(), + backfillProjection: vi.fn(), +})) +vi.mock('@sim/db/script-migrations/0024_knowledge_projection_async', () => ({ + installKnowledgeProjectionAsync: vi.fn(), +})) + +import { + TIN_TRIGGERS, + type TinProjectionDatabase, + type TinProjectionState, +} from '@/scripts/dormant-org-search/disable-tin-projection' +import { + restoreTinProjection, + type TinRestoreSteps, +} from '@/scripts/dormant-org-search/restore-tin-projection' + +/** A database whose Tin triggers and rows change only through the calls under test. */ +function fakeDatabase(initial: Partial = {}) { + const state: TinProjectionState = { + triggers: [...TIN_TRIGGERS], + hasRows: true, + estimatedRows: 1_000, + totalBytes: 8_192, + functionsInstalled: true, + ...initial, + } + const database: TinProjectionDatabase & { drops: number } = { + drops: 0, + readState: async () => ({ ...state, triggers: [...state.triggers] }), + dropTriggersAndTruncate: async () => { + database.drops += 1 + state.triggers = [] + state.hasRows = false + state.totalBytes = 0 + }, + } + return { database, state } +} + +function fakeSteps(onInstall: () => void) { + const calls: string[] = [] + const steps: TinRestoreSteps = { + installProjection: async () => { + calls.push('0019 installProjection') + }, + installKnowledgeProjectionAsync: async () => { + calls.push('0024 installKnowledgeProjectionAsync') + onInstall() + }, + backfillProjection: async () => { + calls.push('0019 backfillProjection') + return 42 + }, + } + return { steps, calls } +} + +describe('restoreTinProjection', () => { + it('runs nothing in a dry run', async () => { + const { database } = fakeDatabase({ triggers: [], hasRows: false }) + const { steps, calls } = fakeSteps(() => undefined) + const result = await restoreTinProjection(database, steps, { execute: false, backfill: true }) + expect(calls).toEqual([]) + expect(result).toEqual({ executed: false, backfilled: null }) + }) + + it('reinstalls from 0019 then re-guards with 0024, and backfills only when asked', async () => { + const { database, state } = fakeDatabase({ triggers: [], hasRows: false }) + const { steps, calls } = fakeSteps(() => { + state.triggers = [...TIN_TRIGGERS] + }) + expect(await restoreTinProjection(database, steps, { execute: true, backfill: false })).toEqual( + { executed: true, backfilled: null } + ) + expect(calls).toEqual(['0019 installProjection', '0024 installKnowledgeProjectionAsync']) + + calls.length = 0 + expect(await restoreTinProjection(database, steps, { execute: true, backfill: true })).toEqual({ + executed: true, + backfilled: 42, + }) + expect(calls.at(-1)).toBe('0019 backfillProjection') + }) + + it('fails when a trigger is still missing after the install', async () => { + const { database } = fakeDatabase({ triggers: [] }) + const { steps, calls } = fakeSteps(() => undefined) + await expect( + restoreTinProjection(database, steps, { execute: true, backfill: true }) + ).rejects.toThrow('missing after restore') + expect(calls).not.toContain('0019 backfillProjection') + }) + + it('refuses where the Tin functions were never installed', async () => { + const { database } = fakeDatabase({ triggers: [], functionsInstalled: false }) + const { steps, calls } = fakeSteps(() => undefined) + await expect( + restoreTinProjection(database, steps, { execute: true, backfill: false }) + ).rejects.toThrow('not installed') + expect(calls).toEqual([]) + }) +}) diff --git a/apps/sim/scripts/dormant-org-search/restore-tin-projection.ts b/apps/sim/scripts/dormant-org-search/restore-tin-projection.ts new file mode 100644 index 00000000000..3a5eb5fa7a7 --- /dev/null +++ b/apps/sim/scripts/dormant-org-search/restore-tin-projection.ts @@ -0,0 +1,163 @@ +#!/usr/bin/env bun + +/** + * Reinstalls exactly the Tin triggers `disable-tin-projection.ts` dropped, from the migrations that + * own them rather than a copy of their SQL: + * + * 1. `installProjection` from `0019_tin_keyword_projection` re-creates the Tin functions and the + * `embedding_keyword_tin_sync` and `knowledge_base_keyword_tin_sync` triggers. + * 2. `installKnowledgeProjectionAsync` from `0024_knowledge_projection_async` re-guards + * `embedding_keyword_tin_sync` with the deferred-projection `WHEN` and re-creates + * `embedding_keyword_tin_source_acl_set`, and re-installs the (unchanged, idempotent) marking. + * + * Triggers write rows only for chunks written after they exist, so the projection is refilled with + * `--backfill`, which runs the same keyset backfill `0019` runs (`backfillProjection`). It inserts + * row by row into the live Tin index; for a large refill it is faster to drop the index first and + * let `bun packages/db/script-migrations/0019_tin_keyword_projection.ts` rebuild it afterwards. + * + * Usage: + * DATABASE_URL= bun apps/sim/scripts/dormant-org-search/restore-tin-projection.ts # dry run + * DATABASE_URL= bun apps/sim/scripts/dormant-org-search/restore-tin-projection.ts --execute [--backfill] + * + * Exit codes: 0 done, 1 failed or refused. + */ + +import { parseArgs } from 'node:util' +import { + backfillProjection, + installProjection, +} from '@sim/db/script-migrations/0019_tin_keyword_projection' +import { installKnowledgeProjectionAsync } from '@sim/db/script-migrations/0024_knowledge_projection_async' +import { retryOnLockTimeout } from '@sim/db/scripts/lock-timeout-retry' +import { createLogger } from '@sim/logger' +import { toError } from '@sim/utils/errors' +import type { Sql } from 'postgres' +import { connectMigrationRole, resolveExecuteFlag } from '@/scripts/dormant-org-search/cli' +import { + postgresTinProjectionDatabase, + TIN_TRIGGERS, + type TinProjectionDatabase, +} from '@/scripts/dormant-org-search/disable-tin-projection' + +const logger = createLogger('RestoreTinProjection') + +const INSTALL_RETRY = { + budgetMs: 20 * 60_000, + backoff: { baseMs: 2_000, maxMs: 30_000 }, +} as const + +/** The migration steps the restore runs, injectable so the order is testable without a database. */ +export interface TinRestoreSteps { + installProjection(): Promise + installKnowledgeProjectionAsync(): Promise + backfillProjection(): Promise +} + +export interface RestoreTinProjectionResult { + executed: boolean + backfilled: number | null +} + +/** + * Refuses where `tin` is not installed, since the Tin triggers would then write rows no index + * serves; otherwise reinstalls the triggers in migration order and optionally backfills. + */ +export async function restoreTinProjection( + database: TinProjectionDatabase, + steps: TinRestoreSteps, + options: { execute: boolean; backfill: boolean } +): Promise { + const before = await database.readState() + const missing = TIN_TRIGGERS.filter( + (trigger) => !before.triggers.some((present) => present.name === trigger.name) + ) + logger.info('Tin projection state', { + installed: before.triggers.map((trigger) => trigger.name), + missing: missing.map((trigger) => trigger.name), + hasRows: before.hasRows, + functionsInstalled: before.functionsInstalled, + }) + if (!before.functionsInstalled) { + throw new Error( + 'The Tin functions are not installed here; run 0019_tin_keyword_projection directly instead' + ) + } + if (!options.execute) { + logger.info('Dry run: would reinstall the Tin triggers from 0019 and 0024', { + backfill: options.backfill, + }) + return { executed: false, backfilled: null } + } + await steps.installProjection() + await steps.installKnowledgeProjectionAsync() + const after = await database.readState() + const stillMissing = TIN_TRIGGERS.filter( + (trigger) => !after.triggers.some((present) => present.name === trigger.name) + ) + if (stillMissing.length > 0) { + throw new Error( + `Tin triggers missing after restore: ${stillMissing.map((trigger) => trigger.name).join(', ')}` + ) + } + logger.info('Tin triggers reinstalled', { + triggers: after.triggers.map((trigger) => trigger.name), + }) + if (!options.backfill) { + logger.warn( + 'Existing chunks are not projected until a backfill runs; rerun with --execute --backfill' + ) + return { executed: true, backfilled: null } + } + const backfilled = await steps.backfillProjection() + logger.info('Tin projection backfilled', { rows: backfilled }) + return { executed: true, backfilled } +} + +/** The migrations' own functions, on the migrations role's connection. */ +export function migrationTinRestoreSteps(sql: Sql): TinRestoreSteps { + return { + installProjection: () => + retryOnLockTimeout(() => installProjection(sql), { + ...INSTALL_RETRY, + onRetry: ({ attempt, delayMs }) => + logger.warn('Tin trigger install waited out its lock timeout; retrying', { + attempt, + retryInMs: Math.round(delayMs), + }), + }), + installKnowledgeProjectionAsync: () => installKnowledgeProjectionAsync(sql), + backfillProjection: () => backfillProjection(sql), + } +} + +async function main(): Promise { + const { values } = parseArgs({ + options: { + execute: { type: 'boolean' }, + 'dry-run': { type: 'boolean' }, + backfill: { type: 'boolean' }, + }, + strict: true, + }) + const execute = resolveExecuteFlag(values) + const sql = connectMigrationRole() + try { + await restoreTinProjection( + postgresTinProjectionDatabase(sql, { retryBudgetMs: INSTALL_RETRY.budgetMs }), + migrationTinRestoreSteps(sql), + { execute, backfill: values.backfill === true } + ) + } finally { + await sql.end() + } +} + +if (import.meta.main) { + main().then( + () => process.exit(0), + (error) => { + logger.error('Restoring the Tin projection failed', toError(error)) + process.exit(1) + } + ) +} diff --git a/apps/sim/scripts/dormant-org-search/search-index-deletion-store.ts b/apps/sim/scripts/dormant-org-search/search-index-deletion-store.ts new file mode 100644 index 00000000000..7f2e9882af1 --- /dev/null +++ b/apps/sim/scripts/dormant-org-search/search-index-deletion-store.ts @@ -0,0 +1,357 @@ +import { db } from '@sim/db' +import { + document, + embedding, + knowledgeBase, + knowledgeConnector, + knowledgeConnectorMember, + knowledgeProjectionDirty, + outboxEvent, +} from '@sim/db/schema' +import { and, asc, count, eq, gt, inArray, isNotNull, isNull, or, sql } from 'drizzle-orm' +import type { DbTransaction } from '@/lib/db/types' +import { + enqueueKnowledgeStorageCleanup, + KNOWLEDGE_STORAGE_CLEANUP_EVENT, +} from '@/lib/knowledge/documents/storage-cleanup' +import { + evaluateDeletionGuard, + SearchIndexDeletionRefused, + type SearchIndexDeletionStore, + STOPPED_CONNECTOR_STATUSES, +} from '@/scripts/dormant-org-search/search-index-deletion' + +/** Bounds for one deletion transaction: short lock waits, and a cap on any one statement. */ +export interface DeletionTimeouts { + lockTimeoutMs: number + statementTimeoutMs: number +} + +async function enterBoundedTransaction(tx: DbTransaction, timeouts: DeletionTimeouts) { + await tx.execute( + sql`SELECT set_config('lock_timeout', ${`${timeouts.lockTimeoutMs}ms`}, true), + set_config('statement_timeout', ${`${timeouts.statementTimeoutMs}ms`}, true)` + ) +} + +/** + * Re-decides the deletion guard inside a deleting transaction, holding the base and its + * connectors `FOR SHARE` until commit. A resume or a sync claim updates the connector row, so it + * waits for this page to commit and the next page's guard refuses it: no page can delete + * documents under a sync that started after the page-level guard ran. + */ +async function lockGuard(tx: DbTransaction, knowledgeBaseId: string) { + const [base] = await tx + .select({ + id: knowledgeBase.id, + isSearchIndex: knowledgeBase.isSearchIndex, + deletedAt: knowledgeBase.deletedAt, + workspaceId: knowledgeBase.workspaceId, + organizationId: knowledgeBase.organizationId, + userId: knowledgeBase.userId, + }) + .from(knowledgeBase) + .where(eq(knowledgeBase.id, knowledgeBaseId)) + .for('share') + .limit(1) + const connectors = await tx + .select({ + id: knowledgeConnector.id, + status: knowledgeConnector.status, + syncLockToken: knowledgeConnector.syncLockToken, + memberSyncLockToken: knowledgeConnector.memberSyncLockToken, + deletedAt: knowledgeConnector.deletedAt, + detachedAt: knowledgeConnector.detachedAt, + }) + .from(knowledgeConnector) + .where(eq(knowledgeConnector.knowledgeBaseId, knowledgeBaseId)) + .orderBy(asc(knowledgeConnector.id)) + .for('share') + const reasons = evaluateDeletionGuard( + base ?? null, + connectors.map((row) => ({ + id: row.id, + status: row.status, + syncLockHeld: row.syncLockToken !== null, + memberSyncLockHeld: row.memberSyncLockToken !== null, + deletedAt: row.deletedAt, + detachedAt: row.detachedAt, + })) + ) + if (!base || reasons.length > 0) throw new SearchIndexDeletionRefused(reasons) + return base +} + +/** + * Makes every stopped connector of the base list its sources from scratch when it resumes: the + * same columns the app clears when a connector must list everything again (an access mode + * switch, a source change), plus the directory checkpoint, and every member made due. Runs in + * each deleting transaction, so a run stopped partway never leaves a connector whose cursors + * would skip the documents already deleted; rows already reset are left alone. + */ +async function resetCursors(tx: DbTransaction, knowledgeBaseId: string, now: Date) { + const connectors = await tx + .update(knowledgeConnector) + .set({ + lastSyncAt: null, + lastSyncDocCount: null, + listingCheckpoint: null, + directoryCheckpoint: null, + memberTombstoneCursor: null, + updatedAt: now, + }) + .where( + and( + eq(knowledgeConnector.knowledgeBaseId, knowledgeBaseId), + isNull(knowledgeConnector.deletedAt), + isNull(knowledgeConnector.detachedAt), + inArray(knowledgeConnector.status, [...STOPPED_CONNECTOR_STATUSES]), + isNull(knowledgeConnector.syncLockToken), + isNull(knowledgeConnector.memberSyncLockToken), + or( + isNotNull(knowledgeConnector.lastSyncAt), + isNotNull(knowledgeConnector.lastSyncDocCount), + isNotNull(knowledgeConnector.listingCheckpoint), + isNotNull(knowledgeConnector.directoryCheckpoint), + isNotNull(knowledgeConnector.memberTombstoneCursor) + ) + ) + ) + .returning({ id: knowledgeConnector.id }) + const stopped = tx + .select({ id: knowledgeConnector.id }) + .from(knowledgeConnector) + .where( + and( + eq(knowledgeConnector.knowledgeBaseId, knowledgeBaseId), + isNull(knowledgeConnector.deletedAt), + isNull(knowledgeConnector.detachedAt), + inArray(knowledgeConnector.status, [...STOPPED_CONNECTOR_STATUSES]) + ) + ) + const members = await tx + .update(knowledgeConnectorMember) + .set({ + listingCheckpoint: null, + changeCursor: null, + memberSyncedThrough: null, + lastCompleteListingAt: null, + lastListedCount: null, + nextAttemptAt: null, + updatedAt: now, + }) + .where( + and( + inArray(knowledgeConnectorMember.connectorId, stopped), + or( + isNotNull(knowledgeConnectorMember.listingCheckpoint), + isNotNull(knowledgeConnectorMember.changeCursor), + isNotNull(knowledgeConnectorMember.memberSyncedThrough), + isNotNull(knowledgeConnectorMember.lastCompleteListingAt), + isNotNull(knowledgeConnectorMember.lastListedCount), + isNotNull(knowledgeConnectorMember.nextAttemptAt) + ) + ) + ) + .returning({ id: knowledgeConnectorMember.id }) + return { connectors: connectors.length, members: members.length } +} + +/** + * The deletion store on the app's database client, so storage cleanup intents are queued by the + * app's own `enqueueKnowledgeStorageCleanup` in the deleting transaction, and deleted by the app's + * outbox worker with its ownership, content-version and reference checks. + */ +export function drizzleSearchIndexDeletionStore( + timeouts: DeletionTimeouts +): SearchIndexDeletionStore { + return { + async loadKnowledgeBase(knowledgeBaseId) { + const [row] = await db + .select({ + id: knowledgeBase.id, + isSearchIndex: knowledgeBase.isSearchIndex, + organizationId: knowledgeBase.organizationId, + deletedAt: knowledgeBase.deletedAt, + }) + .from(knowledgeBase) + .where(eq(knowledgeBase.id, knowledgeBaseId)) + .limit(1) + return row ?? null + }, + + async listConnectors(knowledgeBaseId) { + const rows = await db + .select({ + id: knowledgeConnector.id, + status: knowledgeConnector.status, + syncLockToken: knowledgeConnector.syncLockToken, + memberSyncLockToken: knowledgeConnector.memberSyncLockToken, + deletedAt: knowledgeConnector.deletedAt, + detachedAt: knowledgeConnector.detachedAt, + }) + .from(knowledgeConnector) + .where(eq(knowledgeConnector.knowledgeBaseId, knowledgeBaseId)) + .orderBy(asc(knowledgeConnector.id)) + return rows.map((row) => ({ + id: row.id, + status: row.status, + syncLockHeld: row.syncLockToken !== null, + memberSyncLockHeld: row.memberSyncLockToken !== null, + deletedAt: row.deletedAt, + detachedAt: row.detachedAt, + })) + }, + + async nextDocumentPage(knowledgeBaseId, afterId, limit) { + const rows = await db + .select({ id: document.id }) + .from(document) + .where( + and( + eq(document.knowledgeBaseId, knowledgeBaseId), + isNotNull(document.connectorId), + gt(document.id, afterId) + ) + ) + .orderBy(asc(document.id)) + .limit(limit) + return rows.map((row) => row.id) + }, + + async countChunks(documentIds) { + const [row] = await db + .select({ chunks: count() }) + .from(embedding) + .where(inArray(embedding.documentId, [...documentIds])) + return Number(row?.chunks ?? 0) + }, + + async deleteChunkBatch(knowledgeBaseId, documentIds, limit) { + return db.transaction(async (tx) => { + await enterBoundedTransaction(tx, timeouts) + await lockGuard(tx, knowledgeBaseId) + /** + * Only chunks of documents the connectors still own: a document detachment converted to a + * standalone upload after the page was read keeps its chunks, and the page skips it. + */ + const owned = tx + .select({ id: document.id }) + .from(document) + .where( + and( + inArray(document.id, [...documentIds]), + eq(document.knowledgeBaseId, knowledgeBaseId), + isNotNull(document.connectorId) + ) + ) + const batch = tx + .select({ id: embedding.id }) + .from(embedding) + .where(inArray(embedding.documentId, owned)) + .limit(limit) + const deleted = await tx + .delete(embedding) + .where(inArray(embedding.id, batch)) + .returning({ id: embedding.id }) + return deleted.length + }) + }, + + async deleteDocuments(knowledgeBaseId, documentIds, requestId, resetConnectors) { + return db.transaction(async (tx) => { + await enterBoundedTransaction(tx, timeouts) + const owner = await lockGuard(tx, knowledgeBaseId) + /** Locks the documents against a late indexing commit between the chunk check and the delete. */ + const docs = await tx + .select({ id: document.id, fileUrl: document.fileUrl }) + .from(document) + .where( + and( + inArray(document.id, [...documentIds]), + eq(document.knowledgeBaseId, knowledgeBaseId), + isNotNull(document.connectorId) + ) + ) + .orderBy(asc(document.id)) + .for('update') + if (docs.length === 0) { + return { kind: 'deleted', deleted: 0, storageCleanupQueued: 0 } as const + } + const ids = docs.map((doc) => doc.id) + const [remaining] = await tx + .select({ id: embedding.id }) + .from(embedding) + .where(inArray(embedding.documentId, ids)) + .limit(1) + if (remaining) return { kind: 'chunks-remain' } as const + const queued = await enqueueKnowledgeStorageCleanup( + tx, + docs.map((doc) => ({ + ...doc, + workspaceId: owner.workspaceId, + organizationId: owner.organizationId, + userId: owner.userId, + })), + requestId + ) + const deleted = await tx + .delete(document) + .where(inArray(document.id, ids)) + .returning({ id: document.id }) + if (resetConnectors) await resetCursors(tx, knowledgeBaseId, new Date()) + return { + kind: 'deleted', + deleted: deleted.length, + storageCleanupQueued: queued.length, + } as const + }) + }, + + async pendingStorageCleanup(cap) { + const pending = db + .select({ one: sql`1` }) + .from(outboxEvent) + .where( + and( + eq(outboxEvent.status, 'pending'), + eq(outboxEvent.eventType, KNOWLEDGE_STORAGE_CLEANUP_EVENT) + ) + ) + .limit(cap) + .as('pending') + const [row] = await db.select({ pending: count() }).from(pending) + return Number(row?.pending ?? 0) + }, + + async projectionMarkCount() { + const [row] = await db.select({ marks: count() }).from(knowledgeProjectionDirty) + return Number(row?.marks ?? 0) + }, + + async hasConnectorDocuments(knowledgeBaseId) { + const [row] = await db + .select({ id: document.id }) + .from(document) + .where(and(eq(document.knowledgeBaseId, knowledgeBaseId), isNotNull(document.connectorId))) + .limit(1) + return Boolean(row) + }, + + async hasStandaloneDocuments(knowledgeBaseId) { + const [row] = await db + .select({ id: document.id }) + .from(document) + .where(and(eq(document.knowledgeBaseId, knowledgeBaseId), isNull(document.connectorId))) + .limit(1) + return Boolean(row) + }, + + async resetConnectorCursors(knowledgeBaseId) { + return db.transaction(async (tx) => { + await enterBoundedTransaction(tx, timeouts) + return resetCursors(tx, knowledgeBaseId, new Date()) + }) + }, + } +} diff --git a/apps/sim/scripts/dormant-org-search/search-index-deletion.test.ts b/apps/sim/scripts/dormant-org-search/search-index-deletion.test.ts new file mode 100644 index 00000000000..6db1a12adea --- /dev/null +++ b/apps/sim/scripts/dormant-org-search/search-index-deletion.test.ts @@ -0,0 +1,321 @@ +/** + * @vitest-environment node + */ +import { beforeEach, describe, expect, it, vi } from 'vitest' +import { + deleteSearchIndexDocuments, + evaluateDeletionGuard, + type SearchIndexConnector, + SearchIndexDeletionRefused, + type SearchIndexDeletionStore, + type SearchIndexKnowledgeBase, +} from '@/scripts/dormant-org-search/search-index-deletion' + +const KB_ID = 'kb-search-index' + +function connector(overrides: Partial = {}): SearchIndexConnector { + return { + id: 'connector-1', + status: 'paused', + syncLockHeld: false, + memberSyncLockHeld: false, + deletedAt: null, + detachedAt: null, + ...overrides, + } +} + +const searchIndex: SearchIndexKnowledgeBase = { + id: KB_ID, + isSearchIndex: true, + organizationId: 'org-1', + deletedAt: null, +} + +/** An in-memory knowledge base: documents with chunk counts, plus the calls the run made. */ +class FakeStore implements SearchIndexDeletionStore { + knowledgeBase: SearchIndexKnowledgeBase | null = searchIndex + connectors: SearchIndexConnector[] = [connector()] + documents = new Map() + pending = 0 + marks = 3 + mutations: string[] = [] + /** Chunks a late writer adds to a document the first time its page is deleted. */ + lateChunks = new Map() + guardReads = 0 + + constructor(documents: number, chunksPerDocument = 3) { + for (let index = 0; index < documents; index++) { + this.documents.set(`doc-${String(index).padStart(4, '0')}`, { + chunks: chunksPerDocument, + connector: true, + }) + } + } + + async loadKnowledgeBase() { + this.guardReads += 1 + return this.knowledgeBase + } + async listConnectors() { + return this.connectors + } + async nextDocumentPage(_kb: string, afterId: string, limit: number) { + return [...this.documents.entries()] + .filter(([id, doc]) => doc.connector && id > afterId) + .map(([id]) => id) + .sort() + .slice(0, limit) + } + async countChunks(ids: readonly string[]) { + return ids.reduce((sum, id) => sum + (this.documents.get(id)?.chunks ?? 0), 0) + } + async deleteChunkBatch(_knowledgeBaseId: string, ids: readonly string[], limit: number) { + this.mutations.push('deleteChunkBatch') + let deleted = 0 + for (const id of ids) { + const doc = this.documents.get(id) + if (!doc) continue + const take = Math.min(doc.chunks, limit - deleted) + doc.chunks -= take + deleted += take + if (deleted === limit) break + } + return deleted + } + async deleteDocuments(_kb: string, ids: readonly string[], _requestId: string, _reset: boolean) { + this.mutations.push('deleteDocuments') + for (const id of ids) { + const late = this.lateChunks.get(id) + if (late) { + this.lateChunks.delete(id) + this.documents.get(id)!.chunks += late + } + } + if (ids.some((id) => (this.documents.get(id)?.chunks ?? 0) > 0)) { + return { kind: 'chunks-remain' } as const + } + let deleted = 0 + for (const id of ids) if (this.documents.delete(id)) deleted += 1 + return { kind: 'deleted', deleted, storageCleanupQueued: deleted } as const + } + async pendingStorageCleanup(cap: number) { + return Math.min(this.pending, cap) + } + async projectionMarkCount() { + return this.marks + } + async hasConnectorDocuments() { + return [...this.documents.values()].some((doc) => doc.connector) + } + async hasStandaloneDocuments() { + return [...this.documents.values()].some((doc) => !doc.connector) + } + async resetConnectorCursors() { + this.mutations.push('resetConnectorCursors') + return { connectors: this.connectors.length, members: 0 } + } +} + +const sleep = vi.fn(async () => undefined) + +function run(store: FakeStore, options: Partial[1]>) { + return deleteSearchIndexDocuments(store, { + knowledgeBaseId: KB_ID, + execute: false, + requestId: 'test-request', + pauseMs: 0, + sleep, + ...options, + }) +} + +describe('evaluateDeletionGuard', () => { + it('accepts a search index whose connectors are all stopped', () => { + expect( + evaluateDeletionGuard(searchIndex, [ + connector(), + connector({ id: 'connector-2', status: 'disabled' }), + ]) + ).toEqual([]) + }) + + it('refuses a missing base and a base that is not a search index', () => { + expect(evaluateDeletionGuard(null, [])).toEqual(['knowledge base not found']) + expect(evaluateDeletionGuard({ ...searchIndex, isSearchIndex: false }, [])).toEqual([ + expect.stringContaining('not an organization search index'), + ]) + expect(evaluateDeletionGuard({ ...searchIndex, organizationId: null }, [])).toEqual([ + expect.stringContaining('workspace search index'), + ]) + }) + + it.each(['active', 'pending', 'syncing', 'error'])('refuses a %s connector', (status) => { + expect(evaluateDeletionGuard(searchIndex, [connector({ status })])).toEqual([ + expect.stringContaining(`is ${status}`), + ]) + }) + + it('refuses a paused connector that still holds either sync lease', () => { + expect( + evaluateDeletionGuard(searchIndex, [ + connector({ syncLockHeld: true, memberSyncLockHeld: true }), + ]) + ).toHaveLength(2) + }) + + it('refuses a detached connector and ignores a deleted one', () => { + expect( + evaluateDeletionGuard(searchIndex, [ + connector({ id: 'detached', status: 'active', detachedAt: new Date() }), + connector({ id: 'deleted', status: 'active', deletedAt: new Date() }), + ]) + ).toEqual([expect.stringContaining('connector detached is detached')]) + }) +}) + +describe('deleteSearchIndexDocuments', () => { + beforeEach(() => { + sleep.mockClear() + }) + + it('refuses before reading any page when the guard fails', async () => { + const store = new FakeStore(5) + store.connectors = [connector({ status: 'active' })] + const nextPage = vi.spyOn(store, 'nextDocumentPage') + await expect(run(store, { execute: true })).rejects.toBeInstanceOf(SearchIndexDeletionRefused) + expect(nextPage).not.toHaveBeenCalled() + expect(store.mutations).toEqual([]) + }) + + it('only reads in a dry run, walking at most three pages by default', async () => { + const store = new FakeStore(10) + const summary = await run(store, { pageSize: 2 }) + expect(store.mutations).toEqual([]) + expect(store.documents.size).toBe(10) + expect(summary).toMatchObject({ + executed: false, + pages: 3, + documentsSeen: 6, + chunksCounted: 18, + afterId: 'doc-0005', + done: false, + connectorsReset: null, + }) + }) + + it('reports the end of a dry run without resetting connectors', async () => { + const store = new FakeStore(3) + const summary = await run(store, { pageSize: 2, maxPages: 10 }) + expect(summary.done).toBe(true) + expect(summary.connectorsReset).toBeNull() + expect(store.mutations).toEqual([]) + }) + + it('deletes chunks in bounded batches before each page of documents, then resets connectors', async () => { + const store = new FakeStore(5, 4) + const summary = await run(store, { execute: true, pageSize: 2, chunkBatchSize: 3 }) + expect(store.documents.size).toBe(0) + expect(summary).toMatchObject({ + executed: true, + pages: 3, + documentsDeleted: 5, + chunksDeleted: 20, + storageCleanupQueued: 5, + afterId: 'doc-0004', + done: true, + connectorsReset: { connectors: 1, members: 0 }, + standaloneDocumentsRemain: false, + }) + const pageMutations = store.mutations.slice(0, store.mutations.indexOf('deleteDocuments') + 1) + expect(pageMutations).toEqual([ + 'deleteChunkBatch', + 'deleteChunkBatch', + 'deleteChunkBatch', + 'deleteDocuments', + ]) + expect(store.mutations.at(-1)).toBe('resetConnectorCursors') + }) + + it('stops at --max-pages and resumes from the reported cursor', async () => { + const store = new FakeStore(5) + const first = await run(store, { execute: true, pageSize: 2, maxPages: 1 }) + expect(first).toMatchObject({ pages: 1, documentsDeleted: 2, done: false, afterId: 'doc-0001' }) + expect(store.mutations).not.toContain('resetConnectorCursors') + + const second = await run(store, { execute: true, pageSize: 2, afterId: first.afterId }) + expect(second).toMatchObject({ documentsDeleted: 3, done: true }) + expect(store.documents.size).toBe(0) + }) + + it('does not reset connectors when documents before the resume cursor remain', async () => { + const store = new FakeStore(4) + const summary = await run(store, { execute: true, pageSize: 2, afterId: 'doc-0001' }) + expect(summary.done).toBe(true) + expect(summary.connectorsReset).toBeNull() + expect(store.documents.size).toBe(2) + }) + + it('never selects standalone uploads and reports that they remain', async () => { + const store = new FakeStore(2) + store.documents.set('doc-upload', { chunks: 2, connector: false }) + const summary = await run(store, { execute: true }) + expect(store.documents.has('doc-upload')).toBe(true) + expect(summary.standaloneDocumentsRemain).toBe(true) + expect(summary.connectorsReset).toEqual({ connectors: 1, members: 0 }) + }) + + it('deletes chunks a late writer committed after the chunk pass', async () => { + const store = new FakeStore(2) + store.lateChunks.set('doc-0000', 2) + const summary = await run(store, { execute: true }) + expect(summary.documentsDeleted).toBe(2) + expect(summary.chunksDeleted).toBe(8) + }) + + it('re-reads the guard before every page and stops when a connector resumes', async () => { + const store = new FakeStore(6) + const deleteDocuments = store.deleteDocuments.bind(store) + vi.spyOn(store, 'deleteDocuments').mockImplementation(async (...args) => { + const outcome = await deleteDocuments(...args) + store.connectors = [connector({ status: 'active' })] + return outcome + }) + await expect(run(store, { execute: true, pageSize: 2 })).rejects.toBeInstanceOf( + SearchIndexDeletionRefused + ) + expect(store.documents.size).toBe(4) + }) + + it('waits while the storage cleanup backlog is at its ceiling', async () => { + const store = new FakeStore(1) + store.pending = 50 + const pending = vi.spyOn(store, 'pendingStorageCleanup') + let checks = 0 + pending.mockImplementation(async () => { + checks += 1 + return checks <= 3 ? 50 : 0 + }) + await run(store, { execute: true, storageCleanupCeiling: 50, backpressureWaitMs: 1234 }) + expect(sleep).toHaveBeenCalledWith(1234) + expect(store.documents.size).toBe(0) + }) + + it('retries a transient failure in place and fails on a permanent one', async () => { + const store = new FakeStore(1) + const deleteChunkBatch = store.deleteChunkBatch.bind(store) + const lockTimeout = Object.assign(new Error('canceling statement due to lock timeout'), { + code: '55P03', + }) + vi.spyOn(store, 'deleteChunkBatch') + .mockRejectedValueOnce(lockTimeout) + .mockImplementation(deleteChunkBatch) + const summary = await run(store, { execute: true }) + expect(summary.documentsDeleted).toBe(1) + + const failing = new FakeStore(1) + vi.spyOn(failing, 'deleteChunkBatch').mockRejectedValue(new Error('permission denied')) + await expect(run(failing, { execute: true })).rejects.toThrow('permission denied') + expect(failing.documents.size).toBe(1) + }) +}) diff --git a/apps/sim/scripts/dormant-org-search/search-index-deletion.ts b/apps/sim/scripts/dormant-org-search/search-index-deletion.ts new file mode 100644 index 00000000000..a3cb0b6b68b --- /dev/null +++ b/apps/sim/scripts/dormant-org-search/search-index-deletion.ts @@ -0,0 +1,399 @@ +import { createLogger } from '@sim/logger' +import { getTransientDatabaseFailure } from '@sim/utils/errors' +import { sleep as defaultSleep } from '@sim/utils/helpers' +import { backoffWithJitter } from '@sim/utils/retry' + +const logger = createLogger('SearchIndexDeletion') + +/** Documents per page: each page's documents are deleted, with their storage intents, in one transaction. */ +export const DEFAULT_DOCUMENT_PAGE_SIZE = 200 + +/** + * Chunks deleted per transaction. Each chunk's delete cascades, by foreign key, to its rows in + * `embedding_search`, `embedding_keyword_search`, `embedding_keyword_tin` and + * `embedding_secret_provenance`, so a batch writes about five times this many row deletions. + * Matches the app's connector cleanup worker. + */ +export const DEFAULT_CHUNK_BATCH_SIZE = 1_000 + +/** Pause after every committed transaction, so the run shares the database with search. */ +export const DEFAULT_PAUSE_MS = 250 + +/** Pages a dry run walks unless `--max-pages` says otherwise; a dry run only reads. */ +export const DEFAULT_DRY_RUN_MAX_PAGES = 3 + +/** + * Pending storage cleanup events above which the run waits for the outbox worker before queueing + * more: every deleted document with a stored object queues one event, and the worker deletes one + * object per event. + */ +export const DEFAULT_STORAGE_CLEANUP_CEILING = 2_000 + +/** Wait between checks while the storage cleanup backlog is above its ceiling. */ +export const DEFAULT_BACKPRESSURE_WAIT_MS = 30_000 + +/** Consecutive transient failures (lock or statement timeout, conflict, lost connection) one step may take. */ +export const DEFAULT_TRANSIENT_RETRIES = 8 + +/** How many times a page re-deletes chunks that a late writer committed after its chunk pass. */ +const MAX_CHUNK_PASSES_PER_PAGE = 3 + +const RETRY_BACKOFF = { baseMs: 2_000, maxMs: 60_000 } as const + +/** Connector statuses under which neither the content nor the member sync engine may run. */ +export const STOPPED_CONNECTOR_STATUSES = ['paused', 'disabled'] as const + +export interface SearchIndexKnowledgeBase { + id: string + isSearchIndex: boolean + organizationId: string | null + deletedAt: Date | null +} + +export interface SearchIndexConnector { + id: string + status: string + syncLockHeld: boolean + memberSyncLockHeld: boolean + deletedAt: Date | null + detachedAt: Date | null +} + +export type DeleteDocumentsOutcome = + | { kind: 'deleted'; deleted: number; storageCleanupQueued: number } + | { kind: 'chunks-remain' } + +export interface SearchIndexDeletionStore { + loadKnowledgeBase(knowledgeBaseId: string): Promise + listConnectors(knowledgeBaseId: string): Promise + /** Connector-owned document ids after `afterId`, in id order. */ + nextDocumentPage(knowledgeBaseId: string, afterId: string, limit: number): Promise + /** Chunks of these documents; read only by a dry run. */ + countChunks(documentIds: readonly string[]): Promise + /** Deletes at most `limit` chunks of these documents in one transaction; returns how many. */ + deleteChunkBatch( + knowledgeBaseId: string, + documentIds: readonly string[], + limit: number + ): Promise + /** + * In one transaction: share-locks the knowledge base and its connectors and re-decides the + * guard, locks the documents, and, when none has a chunk left, queues their storage cleanup and + * deletes them. With `resetConnectors`, the same transaction resets the connectors' listing + * cursors, so a run stopped partway never leaves a connector that would skip deleted documents. + */ + deleteDocuments( + knowledgeBaseId: string, + documentIds: readonly string[], + requestId: string, + resetConnectors: boolean + ): Promise + /** Pending storage cleanup events, counted up to `cap`. */ + pendingStorageCleanup(cap: number): Promise + projectionMarkCount(): Promise + hasConnectorDocuments(knowledgeBaseId: string): Promise + hasStandaloneDocuments(knowledgeBaseId: string): Promise + /** + * Clears the listing cursors of the base's stopped connectors and of their members, so a resumed + * connector lists its source from scratch. Returns the connectors and members reset. + */ + resetConnectorCursors(knowledgeBaseId: string): Promise<{ connectors: number; members: number }> +} + +export interface SearchIndexDeletionOptions { + knowledgeBaseId: string + execute: boolean + requestId: string + pageSize?: number + chunkBatchSize?: number + pauseMs?: number + /** Stop after this many pages; unbounded for an executing run, {@link DEFAULT_DRY_RUN_MAX_PAGES} for a dry run. */ + maxPages?: number + /** Resume after this document id. */ + afterId?: string + storageCleanupCeiling?: number + backpressureWaitMs?: number + transientRetries?: number + /** Reset the connectors' listing cursors once no connector-owned document remains. */ + resetConnectors?: boolean + sleep?: (ms: number) => Promise +} + +export interface SearchIndexDeletionSummary { + executed: boolean + pages: number + documentsSeen: number + documentsDeleted: number + chunksDeleted: number + /** Chunks counted by a dry run on the pages it walked. */ + chunksCounted: number + storageCleanupQueued: number + /** The last document id reached; pass it as `--after-id` to resume. */ + afterId: string + /** Whether every connector-owned document after the starting cursor was reached. */ + done: boolean + connectorsReset: { connectors: number; members: number } | null + standaloneDocumentsRemain: boolean | null + projectionMarks: { before: number; after: number } +} + +/** A precondition failed, before the run or before one of its pages; nothing after the check is written. */ +export class SearchIndexDeletionRefused extends Error { + constructor(readonly reasons: string[]) { + super(`Refusing to delete search index documents: ${reasons.join('; ')}`) + this.name = 'SearchIndexDeletionRefused' + } +} + +/** + * Why deleting this base's documents is unsafe now, or an empty list. The base must be an + * organization-owned search index; a workspace search index is out of this tool's scope. Every + * connector that can still write documents — not deleted and not detached — must be stopped, and no sync of either engine may hold its lease: a running sync + * would write documents behind the cursor. A detached connector is refused outright, because its + * worker is converting the same documents into standalone uploads. + */ +export function evaluateDeletionGuard( + knowledgeBase: SearchIndexKnowledgeBase | null, + connectors: readonly SearchIndexConnector[] +): string[] { + if (!knowledgeBase) return ['knowledge base not found'] + const reasons: string[] = [] + if (!knowledgeBase.isSearchIndex) { + reasons.push('knowledge base is not an organization search index (is_search_index = false)') + } else if (!knowledgeBase.organizationId) { + reasons.push('knowledge base is a workspace search index, not an organization one') + } + for (const connector of connectors) { + if (connector.deletedAt) continue + if (connector.detachedAt) { + reasons.push(`connector ${connector.id} is detached and still releasing documents`) + continue + } + if (!(STOPPED_CONNECTOR_STATUSES as readonly string[]).includes(connector.status)) { + reasons.push(`connector ${connector.id} is ${connector.status}, not paused or disabled`) + } + if (connector.syncLockHeld) reasons.push(`connector ${connector.id} holds a content sync lease`) + if (connector.memberSyncLockHeld) { + reasons.push(`connector ${connector.id} holds a member sync lease`) + } + } + return reasons +} + +/** + * Runs `step`, retrying transient database failures in place with backoff. A step either commits + * whole or not at all, so a retry repeats nothing that committed. + */ +async function withTransientRetry( + label: string, + step: () => Promise, + retries: number, + sleep: (ms: number) => Promise +): Promise { + for (let attempt = 1; ; attempt++) { + try { + return await step() + } catch (error) { + const failure = getTransientDatabaseFailure(error) + if (!failure || attempt > retries) throw error + const delayMs = backoffWithJitter(attempt, null, RETRY_BACKOFF) + logger.warn('Transient database failure; retrying', { + step: label, + failure, + attempt, + retryInMs: Math.round(delayMs), + }) + await sleep(delayMs) + } + } +} + +async function assertGuard(store: SearchIndexDeletionStore, knowledgeBaseId: string) { + const [knowledgeBase, connectors] = await Promise.all([ + store.loadKnowledgeBase(knowledgeBaseId), + store.listConnectors(knowledgeBaseId), + ]) + const reasons = evaluateDeletionGuard(knowledgeBase, connectors) + if (reasons.length > 0) throw new SearchIndexDeletionRefused(reasons) + return { knowledgeBase: knowledgeBase as SearchIndexKnowledgeBase, connectors } +} + +/** + * Deletes a search index's connector-owned documents and their chunks in resumable pages. + * + * Each page takes the next documents after the cursor in id order, deletes their chunks in + * bounded transactions, then deletes the documents in one transaction that also commits their + * storage cleanup intents, as the app's connector cleanup worker does. A delete marks nothing in + * `knowledge_projection_dirty` and fires no projection trigger: the projections' rows go by + * foreign-key cascade in the deleting statement. The guard is re-read before every page, so + * resuming a connector mid-run stops the run before its next page. + * + * Standalone uploads (`connector_id IS NULL`) are never selected; they carry billed storage that + * only the app's document deletion settles, and the summary reports whether any remain. + */ +export async function deleteSearchIndexDocuments( + store: SearchIndexDeletionStore, + options: SearchIndexDeletionOptions +): Promise { + const pageSize = options.pageSize ?? DEFAULT_DOCUMENT_PAGE_SIZE + const chunkBatchSize = options.chunkBatchSize ?? DEFAULT_CHUNK_BATCH_SIZE + const pauseMs = options.pauseMs ?? DEFAULT_PAUSE_MS + const maxPages = options.maxPages ?? (options.execute ? undefined : DEFAULT_DRY_RUN_MAX_PAGES) + const ceiling = options.storageCleanupCeiling ?? DEFAULT_STORAGE_CLEANUP_CEILING + const backpressureWaitMs = options.backpressureWaitMs ?? DEFAULT_BACKPRESSURE_WAIT_MS + const retries = options.transientRetries ?? DEFAULT_TRANSIENT_RETRIES + const sleep = options.sleep ?? defaultSleep + const { knowledgeBaseId, execute } = options + const pause = async () => { + if (pauseMs > 0) await sleep(pauseMs) + } + const retry = (label: string, step: () => Promise) => + withTransientRetry(label, step, retries, sleep) + + const { knowledgeBase, connectors } = await assertGuard(store, knowledgeBaseId) + const marksBefore = await store.projectionMarkCount() + logger.info(execute ? 'Deleting search index documents' : 'Dry run: nothing will be written', { + knowledgeBaseId, + knowledgeBaseDeleted: Boolean(knowledgeBase.deletedAt), + connectors: connectors.map((connector) => ({ + id: connector.id, + status: connector.status, + deleted: Boolean(connector.deletedAt), + })), + afterId: options.afterId ?? '', + pageSize, + chunkBatchSize, + maxPages: maxPages ?? 'unbounded', + projectionMarks: marksBefore, + pendingStorageCleanup: await store.pendingStorageCleanup(ceiling), + }) + + const summary: SearchIndexDeletionSummary = { + executed: execute, + pages: 0, + documentsSeen: 0, + documentsDeleted: 0, + chunksDeleted: 0, + chunksCounted: 0, + storageCleanupQueued: 0, + afterId: options.afterId ?? '', + done: false, + connectorsReset: null, + standaloneDocumentsRemain: null, + projectionMarks: { before: marksBefore, after: marksBefore }, + } + const startedAt = Date.now() + + while (maxPages === undefined || summary.pages < maxPages) { + if (summary.pages > 0) await assertGuard(store, knowledgeBaseId) + if (execute) { + for (;;) { + const pending = await store.pendingStorageCleanup(ceiling) + if (pending < ceiling) break + logger.info('Storage cleanup backlog at its ceiling; waiting for the outbox worker', { + pending, + ceiling, + waitMs: backpressureWaitMs, + }) + await sleep(backpressureWaitMs) + } + } + const documentIds = await retry('read document page', () => + store.nextDocumentPage(knowledgeBaseId, summary.afterId, pageSize) + ) + if (documentIds.length === 0) { + summary.done = true + break + } + summary.documentsSeen += documentIds.length + const lastId = documentIds[documentIds.length - 1] + + if (!execute) { + const chunks = await store.countChunks(documentIds) + summary.chunksCounted += chunks + summary.pages += 1 + summary.afterId = lastId + logger.info('Dry run page', { + page: summary.pages, + documents: documentIds.length, + chunks, + afterId: lastId, + }) + continue + } + + let outcome: DeleteDocumentsOutcome = { kind: 'chunks-remain' } + for ( + let pass = 0; + pass < MAX_CHUNK_PASSES_PER_PAGE && outcome.kind === 'chunks-remain'; + pass++ + ) { + for (;;) { + const deleted = await retry('delete chunk batch', () => + store.deleteChunkBatch(knowledgeBaseId, documentIds, chunkBatchSize) + ) + summary.chunksDeleted += deleted + if (deleted === 0) break + await pause() + if (deleted < chunkBatchSize) break + } + outcome = await retry('delete documents', () => + store.deleteDocuments( + knowledgeBaseId, + documentIds, + options.requestId, + options.resetConnectors !== false + ) + ) + } + if (outcome.kind === 'chunks-remain') { + throw new Error( + `Chunks kept appearing for documents after ${summary.afterId || '(start)'}; is a writer still indexing this knowledge base?` + ) + } + summary.documentsDeleted += outcome.deleted + summary.storageCleanupQueued += outcome.storageCleanupQueued + summary.pages += 1 + summary.afterId = lastId + logger.info('Page deleted', { + page: summary.pages, + documents: outcome.deleted, + storageCleanupQueued: outcome.storageCleanupQueued, + totalDocuments: summary.documentsDeleted, + totalChunks: summary.chunksDeleted, + afterId: lastId, + elapsedMs: Date.now() - startedAt, + }) + await pause() + } + + if (summary.done) { + summary.standaloneDocumentsRemain = await store.hasStandaloneDocuments(knowledgeBaseId) + if (summary.standaloneDocumentsRemain) { + logger.warn( + 'Standalone uploads remain in this knowledge base; delete them through the app so their storage is settled' + ) + } + if (execute && options.resetConnectors !== false) { + if (await store.hasConnectorDocuments(knowledgeBaseId)) { + logger.warn( + 'Connector documents remain before the starting cursor; rerun without --after-id before connectors are reset' + ) + } else { + summary.connectorsReset = await store.resetConnectorCursors(knowledgeBaseId) + logger.info('Connector listing cursors reset', summary.connectorsReset) + } + } + } + summary.projectionMarks.after = await store.projectionMarkCount() + logger.info( + summary.done + ? 'Search index deletion reached the end' + : 'Search index deletion stopped; resume with --after-id', + { + ...summary, + elapsedMs: Date.now() - startedAt, + } + ) + return summary +} diff --git a/package.json b/package.json index e279e523822..62c9cd396dc 100644 --- a/package.json +++ b/package.json @@ -44,6 +44,7 @@ "check:api-validation:strict": "bun run scripts/check-api-validation-contracts.ts --check --enforce-boundary-baseline", "check:realtime-prune": "bun run scripts/check-realtime-prune-graph.ts", "check:tool-request-boundary": "bun run scripts/check-tool-request-boundary.ts", + "check:indexed-org-search-boundary": "bun run scripts/check-indexed-org-search-boundary.ts", "check:actorless-executor-operations": "bun run scripts/check-actorless-executor-operations.ts", "check:permission-group-enforcement": "bun run scripts/check-permission-group-enforcement.ts", "check:capability-subject": "bun run scripts/check-capability-subject.ts", diff --git a/packages/db/knowledge-projection.test.ts b/packages/db/knowledge-projection.test.ts index 1989277b42f..5db231776fa 100644 --- a/packages/db/knowledge-projection.test.ts +++ b/packages/db/knowledge-projection.test.ts @@ -161,6 +161,25 @@ describe('runKnowledgeProjection', () => { ]) }) + it('writes no Tin rows when the caller leaves Tin out, and projects everything else', async () => { + const state = database({ + marks: new Map([ + ['content', { generation: 1, content: true }], + ['acl', { generation: 1, content: false }], + ]), + }) + const { sql, trace } = fakeSql(state) + await expect(runKnowledgeProjection(sql, { includeTin: false })).resolves.toMatchObject({ + settled: 2, + remaining: false, + }) + expect(trace.pages.map((page) => `${page.documentId}:${page.projection}`)).toEqual([ + 'content:embedding_search', + 'content:embedding_keyword_search', + 'acl:embedding_search', + ]) + }) + it('pages a document by chunk rows until a page comes back short', async () => { const state = database({ marks: new Map([['doc', { generation: 1, content: false }]]), @@ -334,4 +353,35 @@ describe('markUnfilledProjectionDocuments', () => { markUnfilledProjectionDocuments(done.sql, { projection: 1, afterId: 'tin-1' }) ).resolves.toEqual({ marked: 0, cursor: null }) }) + + it('reads search-index rows only when they are included', async () => { + const { sql, statements } = fillSql(0, [{ marked: 1, last_id: 'row-3' }]) + await markUnfilledProjectionDocuments(sql) + expect(statements[0]?.text).not.toContain('is_search_index') + }) + + it('passes over search-index rows and never reads Tin when they are excluded', async () => { + const { sql, statements } = fillSql(0, [ + { marked: 0, last_id: 'row-9' }, + { marked: 0, last_id: null }, + ]) + await expect( + markUnfilledProjectionDocuments(sql, undefined, { includeSearchIndexes: false }) + ).resolves.toEqual({ marked: 0, cursor: { projection: 0, afterId: 'row-9' } }) + expect(statements[0]?.text).toContain('FROM embedding_search WHERE acl IS NULL') + expect(statements[0]?.text).toContain( + 'SELECT 1 FROM knowledge_base k WHERE k.id = u.knowledge_base_id AND k.is_search_index' + ) + await expect( + markUnfilledProjectionDocuments( + sql, + { projection: 0, afterId: 'row-9' }, + { includeSearchIndexes: false } + ) + ).resolves.toEqual({ marked: 0, cursor: null }) + expect(statements).toHaveLength(2) + expect(statements.some((statement) => statement.text.includes('embedding_keyword_tin'))).toBe( + false + ) + }) }) diff --git a/packages/db/knowledge-projection.ts b/packages/db/knowledge-projection.ts index 802874bd3ee..6357b8f8b0b 100644 --- a/packages/db/knowledge-projection.ts +++ b/packages/db/knowledge-projection.ts @@ -244,6 +244,12 @@ export interface KnowledgeProjectionOptions { */ budgetMs?: number pageSize?: number + /** + * Whether the Tin keyword projection is written. It holds only search-index rows, so the + * application turns it off while indexed organization search is dormant; the other projections + * are written either way. Defaults to true, and Tin is still skipped where it is not installed. + */ + includeTin?: boolean /** Called after each page commits, for tests that interleave writes with a run. */ onPage?: (page: { documentId: string @@ -439,9 +445,10 @@ export async function runKnowledgeProjection( ): Promise { const deadline = options.budgetMs === undefined ? Number.POSITIVE_INFINITY : Date.now() + options.budgetMs - const projections = (await tinInstalled(sql)) - ? KNOWLEDGE_PROJECTIONS - : KNOWLEDGE_PROJECTIONS.filter((projection) => projection !== 'embedding_keyword_tin') + const projections = + options.includeTin !== false && (await tinInstalled(sql)) + ? KNOWLEDGE_PROJECTIONS + : KNOWLEDGE_PROJECTIONS.filter((projection) => projection !== 'embedding_keyword_tin') const totals = { pages: 0, written: 0 } const skipped: string[] = [] let settled = 0 @@ -467,6 +474,14 @@ export async function runKnowledgeProjection( return { settled, deferred: skipped.length, ...totals, remaining: true } } +export interface MarkUnfilledProjectionOptions { + /** + * Whether rows of search-index knowledge bases are filled. Defaults to true; the application + * turns it off while indexed organization search is dormant. + */ + includeSearchIndexes?: boolean +} + /** Unfilled rows the fill reads per round, from each projection's unfilled-rows index. */ const FILL_SCAN_ROWS = 2_000 @@ -490,24 +505,39 @@ export const FILL_MARK_CEILING = 100 * removes or rewrites its rows itself, and the next pass reads whatever is still unfilled. Returns * how many documents were marked and the id to continue after, or `null` once every projection's * unfilled rows have been read. + * + * With `includeSearchIndexes` off, rows of search-index knowledge bases are passed over like rows + * whose document is gone, so the fill never rewrites them, and the Tin projection, which holds + * only search-index rows, is not read at all. The package cannot see the application's switch, so + * the caller decides. */ export async function markUnfilledProjectionDocuments( sql: Sql, - cursor: { projection: number; afterId: string } = { projection: 0, afterId: '' } + cursor: { projection: number; afterId: string } = { projection: 0, afterId: '' }, + options: MarkUnfilledProjectionOptions = {} ): Promise<{ marked: number; cursor: { projection: number; afterId: string } | null }> { + const includeSearchIndexes = options.includeSearchIndexes ?? true const [{ outstanding }] = await sql>` SELECT count(*)::int AS outstanding FROM knowledge_projection_dirty` if (outstanding >= FILL_MARK_CEILING) return { marked: 0, cursor } for (let index = cursor.projection; index < SOURCE_ACL_PROJECTIONS.length; index++) { const projection = SOURCE_ACL_PROJECTIONS[index] + if (!includeSearchIndexes && projection === 'embedding_keyword_tin') continue const afterId = index === cursor.projection ? cursor.afterId : '' + /** Search-index rows are read and passed over, so the cursor still moves past them. */ + const outsideSearchIndexes = includeSearchIndexes + ? '' + : ` + AND NOT EXISTS ( + SELECT 1 FROM knowledge_base k WHERE k.id = u.knowledge_base_id AND k.is_search_index + )` const [row] = await sql.unsafe>( `WITH unfilled AS MATERIALIZED ( - SELECT id, document_id FROM ${projection} WHERE acl IS NULL AND id > $1 + SELECT id, document_id${includeSearchIndexes ? '' : ', knowledge_base_id'} FROM ${projection} WHERE acl IS NULL AND id > $1 ORDER BY id LIMIT ${FILL_SCAN_ROWS} ), documents AS MATERIALIZED ( SELECT u.document_id, min(u.id) AS first_id FROM unfilled u - WHERE EXISTS (SELECT 1 FROM document d WHERE d.id = u.document_id) + WHERE EXISTS (SELECT 1 FROM document d WHERE d.id = u.document_id)${outsideSearchIndexes} GROUP BY u.document_id ), chosen AS MATERIALIZED ( SELECT document_id, first_id FROM documents ORDER BY first_id LIMIT $2 diff --git a/scripts/check-indexed-org-search-boundary.test.ts b/scripts/check-indexed-org-search-boundary.test.ts new file mode 100644 index 00000000000..bceab765348 --- /dev/null +++ b/scripts/check-indexed-org-search-boundary.test.ts @@ -0,0 +1,166 @@ +import { describe, expect, it } from 'vitest' +import { + BARREL_ENTRY_FILES, + findBoundaryViolations, + findStaleEntries, + GATE_SPECIFIER, + RETRIEVAL_BARREL, + USE_CASE_BARREL, +} from './check-indexed-org-search-boundary' + +const ENTRY = 'apps/sim/lib/knowledge/mcp/server.ts' +const RETRIEVAL_ENTRY = 'apps/sim/lib/knowledge/search/queries.ts' +const GATE_IMPORT = `import { isIndexedOrgSearchEnabled } from '${GATE_SPECIFIER}'` +const GATED = `${GATE_IMPORT}\nif (!isIndexedOrgSearchEnabled()) throw new Error('dormant')` + +describe('indexed organization search boundary audit', () => { + it('lets any file read the gate', () => { + expect( + findBoundaryViolations([ + { file: 'apps/sim/lib/knowledge/connectors/indexing-policy.ts', source: GATE_IMPORT }, + ]) + ).toEqual([]) + }) + + it('lets an allowlisted entry import its barrel beside the gate', () => { + expect( + findBoundaryViolations([ + { + file: ENTRY, + source: `import { readIndexedKnowledgeDocument } from '${USE_CASE_BARREL}'\n${GATED}`, + }, + { + file: RETRIEVAL_ENTRY, + source: `import { isProjectionFilled } from '${RETRIEVAL_BARREL}'\n${GATED}`, + }, + ]) + ).toEqual([]) + }) + + it('rejects a barrel import from a file that is not allowlisted', () => { + expect( + findBoundaryViolations([ + { + file: 'apps/sim/lib/knowledge/application/chat.ts', + source: `import { searchScopedKnowledge } from '${USE_CASE_BARREL}'\n${GATE_IMPORT}`, + }, + ]) + ).toEqual([ + expect.objectContaining({ + file: 'apps/sim/lib/knowledge/application/chat.ts', + line: 1, + reason: `is not an allowlisted entry point for ${USE_CASE_BARREL}`, + }), + ]) + }) + + it('rejects an entry importing the barrel it is not allowlisted for', () => { + expect( + findBoundaryViolations([ + { + file: ENTRY, + source: `import { isProjectionFilled } from '${RETRIEVAL_BARREL}'\n${GATED}`, + }, + ]) + ).toEqual([expect.objectContaining({ file: ENTRY, specifier: RETRIEVAL_BARREL })]) + }) + + it('rejects deep imports, re-exports, dynamic imports, and relative paths into the directory', () => { + const violations = findBoundaryViolations([ + { + file: ENTRY, + source: [ + GATE_IMPORT, + "import { readSearchDocument } from '@/lib/sim-search/indexed/documents/read-search-document'", + "export { resolveTinKeywordQuery } from '@/lib/sim-search/indexed/retrieval/tin-keyword'", + "const scoped = await import('@/lib/sim-search/indexed/search/scoped-search')", + ].join('\n'), + }, + { + file: 'apps/sim/lib/sim-search/connectors.ts', + source: "import { searchScopedKnowledge } from './indexed/search/scoped-search'", + }, + ]) + expect(violations.map(({ file, line }) => `${file}:${line}`)).toEqual([ + `${ENTRY}:2`, + `${ENTRY}:3`, + `${ENTRY}:4`, + 'apps/sim/lib/sim-search/connectors.ts:1', + ]) + }) + + it.each([ + ['without the gate', `import { readIndexedKnowledgeDocument } from '${USE_CASE_BARREL}'`], + [ + 'importing the gate but never calling it', + `import { readIndexedKnowledgeDocument } from '${USE_CASE_BARREL}'\n${GATE_IMPORT}`, + ], + ])('rejects an entry reaching in %s', (_case, source) => { + expect(findBoundaryViolations([{ file: ENTRY, source }])).toEqual([ + expect.objectContaining({ + file: ENTRY, + reason: 'reaches dormant indexed search without calling isIndexedOrgSearchEnabled()', + }), + ]) + }) + + it('catches a dynamic import whose target is held in a variable', () => { + expect( + findBoundaryViolations([ + { + file: 'apps/sim/lib/other/loader.ts', + source: `const target = '${USE_CASE_BARREL}'\nexport const load = () => import(target)`, + }, + ]) + ).toEqual([ + expect.objectContaining({ + file: 'apps/sim/lib/other/loader.ts', + specifier: USE_CASE_BARREL, + reason: `is not an allowlisted entry point for ${USE_CASE_BARREL}`, + }), + ]) + }) + + it('accepts a gate called under an aliased import', () => { + expect( + findBoundaryViolations([ + { + file: ENTRY, + source: `import { readIndexedKnowledgeDocument } from '${USE_CASE_BARREL}'\nimport { isIndexedOrgSearchEnabled as indexedOn } from '${GATE_SPECIFIER}'\nexport const on = indexedOn()`, + }, + ]) + ).toEqual([]) + }) + + it('ignores tests, the dormant directory itself, and text that only names the module', () => { + expect( + findBoundaryViolations([ + { + file: 'apps/sim/lib/knowledge/search/queries.test.ts', + source: "import { x } from '@/lib/sim-search/indexed/retrieval/tin-keyword'", + }, + { + file: 'apps/sim/lib/sim-search/indexed/index.ts', + source: "export { x } from '@/lib/sim-search/indexed/search/scoped-search'", + }, + { + file: 'apps/sim/lib/sim-search/live/application.ts', + source: + "/** See '@/lib/sim-search/indexed/search/scoped-search'. */\nconst a = 'indexed'", + }, + ]) + ).toEqual([]) + }) + + it('reports allowlisted entries that no longer import their barrel', () => { + const importing = Object.entries(BARREL_ENTRY_FILES).flatMap(([barrel, entries]) => + Object.keys(entries).map((file) => ({ file, source: `import { x } from '${barrel}'` })) + ) + expect(findStaleEntries(importing)).toEqual([]) + expect( + findStaleEntries( + importing.map((entry) => (entry.file === ENTRY ? { file: ENTRY, source: '' } : entry)) + ) + ).toEqual([`${ENTRY} (${USE_CASE_BARREL})`]) + }) +}) diff --git a/scripts/check-indexed-org-search-boundary.ts b/scripts/check-indexed-org-search-boundary.ts new file mode 100644 index 00000000000..7a016426a61 --- /dev/null +++ b/scripts/check-indexed-org-search-boundary.ts @@ -0,0 +1,295 @@ +#!/usr/bin/env bun +/** + * Keeps the dormant indexed organization search behind its gate. + * + * `apps/sim/lib/sim-search/indexed/` holds the retrieval over `is_search_index` knowledge bases + * that Live Search replaced. It stays in the tree for a possible re-enable, switched by + * `isIndexedOrgSearchEnabled()` in `@/lib/sim-search/indexed/gate`. Dormant code is only safe + * while nothing reaches it except through that switch, so this audit pins the module edge: + * + * - `@/lib/sim-search/indexed/gate` is the switch itself and may be imported from anywhere. + * - `@/lib/sim-search/indexed` (the use cases) and `@/lib/sim-search/indexed/retrieval` (the + * search-index-only retrieval strategies) are imported only by the entry files allowlisted + * below, each of which must also import the gate so it can decide before calling in. + * - Nothing outside the directory reaches past those barrels into a file. + * + * Tests are exempt: they exercise the dormant code directly, whatever the switch says. + * + * Parsed with the TypeScript AST, so comments and strings naming the module are not imports. + * + * Usage: bun run scripts/check-indexed-org-search-boundary.ts + */ +import type { Dirent } from 'node:fs' +import { readdirSync, readFileSync } from 'node:fs' +import path from 'node:path' +import { fileURLToPath } from 'node:url' +import ts from '@typescript/typescript6' + +const ROOT = path.resolve(path.dirname(fileURLToPath(import.meta.url)), '..') +const APP_ROOT = 'apps/sim' +const INDEXED_DIR = 'apps/sim/lib/sim-search/indexed' + +export const GATE_SPECIFIER = '@/lib/sim-search/indexed/gate' +export const USE_CASE_BARREL = '@/lib/sim-search/indexed' +export const RETRIEVAL_BARREL = '@/lib/sim-search/indexed/retrieval' + +/** + * The only files allowed to import each barrel, and why. Keep these small: every entry is a + * surface that can run dormant code, and each must check the gate first. + */ +export const BARREL_ENTRY_FILES: Record> = { + [USE_CASE_BARREL]: { + 'apps/sim/app/api/knowledge/search/route.ts': + 'selects the indexed route over the live route only while the gate is on', + 'apps/sim/app/o/[organizationId]/knowledge/[knowledgeBaseId]/[documentId]/page.tsx': + 'indexed citation page; not found while the gate is off', + 'apps/sim/lib/knowledge/mcp/server.ts': + 'registers indexed search and read tools only while the gate is on', + 'apps/sim/lib/mothership/tools/server/knowledge/workspace-search.ts': + 'Sim search and read tools take the indexed branch only while the gate is on', + }, + [RETRIEVAL_BARREL]: { + 'apps/sim/lib/knowledge/search/queries.ts': + 'shared retrieval; runs the search-index-only strategies only while the gate is on', + }, +} + +const SKIP_DIRS = new Set([ + 'node_modules', + '.next', + '.turbo', + 'dist', + 'coverage', + '__integration__', +]) + +export interface BoundaryViolation { + file: string + line: number + specifier: string + reason: string +} + +export interface SourceFile { + /** Repository-relative path with forward slashes. */ + file: string + source: string +} + +/** Whether a repository-relative path is a test the audit leaves alone. */ +export function isTestFile(file: string): boolean { + return /\.(test|integration)\.tsx?$/.test(file) || file.split('/').includes('__integration__') +} + +/** The module specifiers a file loads at runtime or for types, with their lines. */ +export function moduleSpecifiers( + file: string, + source: string +): Array<{ specifier: string; line: number }> { + const sourceFile = ts.createSourceFile(file, source, ts.ScriptTarget.Latest, false) + const found: Array<{ specifier: string; line: number }> = [] + const record = (node: ts.Node, specifier: string) => { + const { line } = sourceFile.getLineAndCharacterOfPosition(node.getStart(sourceFile)) + found.push({ specifier, line: line + 1 }) + } + const consumed = new Set() + const visit = (node: ts.Node): void => { + if ( + (ts.isImportDeclaration(node) || ts.isExportDeclaration(node)) && + node.moduleSpecifier && + ts.isStringLiteralLike(node.moduleSpecifier) + ) { + record(node, node.moduleSpecifier.text) + consumed.add(node.moduleSpecifier) + } else if (ts.isCallExpression(node)) { + const isDynamicImport = node.expression.kind === ts.SyntaxKind.ImportKeyword + const isRequire = ts.isIdentifier(node.expression) && node.expression.text === 'require' + const argument = node.arguments[0] + if ((isDynamicImport || isRequire) && argument && ts.isStringLiteralLike(argument)) { + record(node, argument.text) + consumed.add(argument) + } + } else if (ts.isStringLiteralLike(node) && !consumed.has(node)) { + /** + * Any other string naming a module is a specifier too: a dynamic import or `require` of a + * variable is only as safe as the strings that variable can hold, so the string is where + * the reach into the dormant directory is caught. + */ + record(node, node.text) + } + ts.forEachChild(node, visit) + } + visit(sourceFile) + return found +} + +/** The gate function whose call an entry point must make before reaching indexed search. */ +const GATE_FUNCTION = 'isIndexedOrgSearchEnabled' + +/** + * Whether the file calls the gate, under any local name its import gave it. An entry only + * importing the gate proves nothing: it must ask it. That the call guards each use is left to + * review of the allowlist, which names every entry explicitly. + */ +export function callsGate(file: string, source: string): boolean { + const sourceFile = ts.createSourceFile(file, source, ts.ScriptTarget.Latest, false) + const localNames = new Set() + for (const statement of sourceFile.statements) { + if ( + !ts.isImportDeclaration(statement) || + !ts.isStringLiteralLike(statement.moduleSpecifier) || + normalizeSpecifier(file, statement.moduleSpecifier.text) !== GATE_SPECIFIER + ) + continue + const bindings = statement.importClause?.namedBindings + if (!bindings || !ts.isNamedImports(bindings)) continue + for (const element of bindings.elements) { + if ((element.propertyName ?? element.name).text === GATE_FUNCTION) { + localNames.add(element.name.text) + } + } + } + if (localNames.size === 0) return false + let called = false + const visit = (node: ts.Node): void => { + if (called) return + if ( + ts.isCallExpression(node) && + ts.isIdentifier(node.expression) && + localNames.has(node.expression.text) + ) { + called = true + return + } + ts.forEachChild(node, visit) + } + visit(sourceFile) + return called +} + +/** The `@/`-rooted form of a specifier, resolving a relative one against the importing file. */ +function normalizeSpecifier(file: string, specifier: string): string | null { + if (specifier.startsWith('@/')) return specifier + if (!specifier.startsWith('.')) return null + const resolved = path.posix.normalize(path.posix.join(path.posix.dirname(file), specifier)) + return resolved.startsWith(`${APP_ROOT}/`) ? `@/${resolved.slice(APP_ROOT.length + 1)}` : null +} + +/** Every violation across `files`, which are the non-test sources outside the indexed directory. */ +export function findBoundaryViolations(files: readonly SourceFile[]): BoundaryViolation[] { + const violations: BoundaryViolation[] = [] + const importedBarrels = new Map>() + const sources = new Map() + for (const { file, source } of files) { + if (file.startsWith(`${INDEXED_DIR}/`) || isTestFile(file)) continue + /** Any specifier reaching the directory, absolute or relative, names its `indexed` segment. */ + if (!source.includes('indexed')) continue + for (const { specifier, line } of moduleSpecifiers(file, source)) { + const normalized = normalizeSpecifier(file, specifier) + if ( + !normalized || + (normalized !== USE_CASE_BARREL && !normalized.startsWith(`${USE_CASE_BARREL}/`)) + ) + continue + if (normalized === GATE_SPECIFIER) continue + const entries = BARREL_ENTRY_FILES[normalized] + if (!entries) { + violations.push({ + file, + line, + specifier, + reason: `reaches into dormant indexed search; import ${USE_CASE_BARREL} or ${RETRIEVAL_BARREL}`, + }) + continue + } + if (!(file in entries)) { + violations.push({ + file, + line, + specifier, + reason: `is not an allowlisted entry point for ${normalized}`, + }) + continue + } + sources.set(file, source) + const barrels = importedBarrels.get(file) ?? new Set() + barrels.add(normalized) + importedBarrels.set(file, barrels) + } + } + for (const file of importedBarrels.keys()) { + if (!callsGate(file, sources.get(file) ?? '')) { + violations.push({ + file, + line: 1, + specifier: GATE_SPECIFIER, + reason: `reaches dormant indexed search without calling ${GATE_FUNCTION}()`, + }) + } + } + return violations +} + +/** Allowlist entries that no longer import their barrel, so the list cannot outlive its callers. */ +export function findStaleEntries(files: readonly SourceFile[]): string[] { + const byFile = new Map(files.map((entry) => [entry.file, entry])) + const stale: string[] = [] + for (const [barrel, entries] of Object.entries(BARREL_ENTRY_FILES)) { + for (const file of Object.keys(entries)) { + const entry = byFile.get(file) + const imports = entry + ? moduleSpecifiers(entry.file, entry.source).some( + ({ specifier }) => normalizeSpecifier(entry.file, specifier) === barrel + ) + : false + if (!imports) stale.push(`${file} (${barrel})`) + } + } + return stale +} + +function walk(dir: string, out: string[] = []): string[] { + let entries: Dirent[] + try { + entries = readdirSync(dir, { withFileTypes: true }) + } catch { + throw new Error(`check-indexed-org-search-boundary: "${dir}" does not exist`) + } + for (const entry of entries) { + if (SKIP_DIRS.has(entry.name)) continue + const full = path.join(dir, entry.name) + if (entry.isDirectory()) walk(full, out) + else if (/\.(ts|tsx)$/.test(entry.name)) out.push(full) + } + return out +} + +function main() { + const files = walk(path.join(ROOT, APP_ROOT)) + .map((full) => path.relative(ROOT, full).split(path.sep).join('/')) + .filter((file) => !isTestFile(file)) + .map((file) => ({ file, source: readFileSync(path.join(ROOT, file), 'utf8') })) + const violations = findBoundaryViolations(files) + const stale = findStaleEntries(files) + if (violations.length === 0 && stale.length === 0) { + console.log( + `✓ check-indexed-org-search-boundary: ${files.length} files, dormant indexed search reached only through its gate` + ) + process.exit(0) + } + console.error('✗ check-indexed-org-search-boundary: dormant indexed search escaped its gate\n') + for (const violation of violations) { + console.error(` ${violation.file}:${violation.line} ${violation.specifier}`) + console.error(` ${violation.reason}`) + } + for (const entry of stale) { + console.error(` ${entry}\n allowlisted but no longer imports the barrel; remove it`) + } + console.error( + '\n Indexed organization search is dormant. Reach it only from an allowlisted entry point\n' + + ' that checks isIndexedOrgSearchEnabled() first; see apps/sim/lib/sim-search/indexed/README.md.\n' + ) + process.exit(1) +} + +if (import.meta.main) main()