improvement(knowledge): instrument the search source overview and bound its probes (#7957)

* improvement(knowledge): instrument the search source overview and bound its probes

The source overview backs the Sim Search indexing banner and source filter, but
it emitted no telemetry, so a slow request could only be observed as a duration
on the route. It now shares the search diagnostics trace, timing availability
resolution, the provider read, and each per-batch indexing and searchable probe,
and records how many access batches and live source proofs a request consumed.

The access batch generator is instrumented in the same trace, separating the
connector discovery query from the live source proof that reaches the provider
over the network, so a slow batch says which half was slow for every caller.

The searchable probe also stops repeating: one searchable document is the whole
answer, yet the bounded existence query ran again for every remaining batch.

* fix(knowledge): report the free first access predicate separately

The batch counter started at one to include the ordinary predicate, which is
yielded without a connector query or a live source proof, while the field
documented every batch as costing both. The count now says what it measures,
and it is recorded before the first yield so a caller that stops consuming
there still reports the predicate it used.

Author: waleedlatif1

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

@@ -1,4 +1,5 @@
 import type { AuthorizingUseCase } from '@/lib/core/application'
+import type { ResourceOwner } from '@/lib/core/resource-scope'
 import type { knowledgeOperations } from '@/lib/knowledge/application/operations'
 import type { SearchKnowledgeInput } from '@/lib/knowledge/application/search'
 import {
@@ -43,3 +44,29 @@ export function instrumentSearchUseCase<
       ),
   }
 }
+
+/**
+ * Times the source overview, whose cost is access batching and live source proof rather than
+ * retrieval. It shares the search trace so one log line explains a slow Sim Search surface.
+ */
+export function instrumentSourceOverviewUseCase<I extends ResourceOwner, R>(
+  useCase: AuthorizingUseCase<typeof knowledgeOperations.readSearchSourceOverview, I, R>
+): AuthorizingUseCase<typeof knowledgeOperations.readSearchSourceOverview, I, R> {
+  return {
+    ...useCase,
+    execute: (args) =>
+      withSearchDiagnostics(
+        {
+          operation: 'read_search_source_overview',
+          principalKind: args.principal.kind,
+          /** Derived without `resourceScopeFromOwner`, which throws before authorization runs. */
+          scopeKind: args.input.organizationId
+            ? 'organization'
+            : args.input.workspaceId
+              ? 'workspace'
+              : undefined,
+        },
+        () => measureSearchStage('source_overview', () => useCase.execute(args))
+      ),
+  }
+}

apps/sim/lib/knowledge/application/search-source-overview.test.ts

@@ -0,0 +1,76 @@
+/** @vitest-environment node */
+import { document, knowledgeConnector, member } from '@sim/db/schema'
+import { dbChainMockFns, queueTableRows, resetDbChainMock } from '@sim/testing'
+import { sql } from 'drizzle-orm'
+import { beforeEach, describe, expect, it, vi } from 'vitest'
+
+const mocks = vi.hoisted(() => ({
+  context: vi.fn(),
+  availability: vi.fn(),
+  provider: vi.fn(),
+  batches: vi.fn(),
+}))
+vi.mock('@/lib/knowledge/application/contexts', () => ({
+  resolveKnowledgeOwnerContext: mocks.context,
+}))
+vi.mock('@/lib/knowledge/access/availability', () => ({
+  resolveKnowledgeAccessAvailability: mocks.availability,
+}))
+vi.mock('@/lib/knowledge/access/scope', () => ({
+  createKnowledgeAccessProvider: mocks.provider,
+}))
+vi.mock('@/lib/knowledge/read-access', () => ({
+  knowledgeReadAccessBatches: mocks.batches,
+}))
+
+import { readSearchSourceOverview } from '@/lib/knowledge/application/search-source-overview'
+
+const principal = { kind: 'session', userId: 'reader', sessionId: 'session' } as const
+const input = { organizationId: 'org-1', workspaceId: null }
+
+/** Authorization reads the caller's membership with its own single-row bound, before any probe. */
+const AUTHORIZATION_SINGLE_ROW_READS = 1
+
+/** Every other single-row read in this use case is the searchable existence probe. */
+const searchableProbeCount = () =>
+  dbChainMockFns.limit.mock.calls.filter(([rows]) => rows === 1).length -
+  AUTHORIZATION_SINGLE_ROW_READS
+
+function yieldBatches(count: number) {
+  mocks.batches.mockImplementation(async function* () {
+    for (let index = 0; index < count; index += 1) yield sql`batch-${sql.raw(String(index))}`
+  })
+}
+
+beforeEach(() => {
+  vi.clearAllMocks()
+  resetDbChainMock()
+  mocks.context.mockResolvedValue(input)
+  mocks.availability.mockResolvedValue({ memberScoped: true, sourceMirrored: false })
+  mocks.provider.mockReturnValue({ get: async () => ({ kind: 'user' }) })
+})
+
+describe('readSearchSourceOverview', () => {
+  it('stops probing for a searchable document once one batch has found one', async () => {
+    yieldBatches(3)
+    queueTableRows(member, [{ role: 'owner' }])
+    queueTableRows(knowledgeConnector, [{ connectorType: 'gmail' }])
+    queueTableRows(document, [{ id: 'doc-1' }])
+
+    const result = await readSearchSourceOverview.execute({ principal, input })
+
+    expect(result.hasSearchableDocuments).toBe(true)
+    expect(searchableProbeCount()).toBe(1)
+  })
+
+  it('keeps probing every batch while no searchable document has been found', async () => {
+    yieldBatches(3)
+    queueTableRows(member, [{ role: 'owner' }])
+    queueTableRows(knowledgeConnector, [{ connectorType: 'gmail' }])
+
+    const result = await readSearchSourceOverview.execute({ principal, input })
+
+    expect(result.hasSearchableDocuments).toBe(false)
+    expect(searchableProbeCount()).toBe(3)
+  })
+})

apps/sim/lib/knowledge/application/search-source-overview.ts

@@ -9,8 +9,10 @@ import { createKnowledgeAccessProvider } from '@/lib/knowledge/access/scope'
 import { defineAuthorizedKnowledgeUseCase } from '@/lib/knowledge/application/authorized-knowledge-use-case'
 import { resolveKnowledgeOwnerContext } from '@/lib/knowledge/application/contexts'
 import { knowledgeOperations } from '@/lib/knowledge/application/operations'
+import { instrumentSourceOverviewUseCase } from '@/lib/knowledge/application/search-diagnostics'
 import { MAX_SEARCH_SOURCE_PROVIDER_TYPES } from '@/lib/knowledge/constants'
 import { knowledgeReadAccessBatches } from '@/lib/knowledge/read-access'
+import { annotateSearchDiagnostics, measureSearchStage } from '@/lib/knowledge/search/diagnostics'
 import { searchIntegrationAccessCondition } from '@/lib/knowledge/search/integration-policy'
 import { CONNECTOR_META_REGISTRY } from '@/connectors/registry'
 
@@ -54,120 +56,144 @@ export async function listConfiguredSearchProviderTypes(owner: ResourceOwner): P
 }
 
 /** Provider-level existence probes keep setup cards independent of the loaded source pages. */
-export const readSearchSourceOverview = defineAuthorizedKnowledgeUseCase({
-  operation: knowledgeOperations.readSearchSourceOverview,
-  resolveContext: ({ input }: { input: ResourceOwner }) => resolveKnowledgeOwnerContext(input),
-  async execute({ principal, context }): Promise<SearchSourceOverview> {
-    const availability = await resolveKnowledgeAccessAvailability(context)
-    const access = createKnowledgeAccessProvider(principal, context)
-    const providerTypes = searchProviderTypes()
-    const configured = configuredSearchSourceCondition(context, providerTypes)
-    const available = or(
-      availability.memberScoped ? eq(knowledgeConnector.accessMode, 'members') : undefined,
-      availability.sourceMirrored
-        ? and(
-            eq(knowledgeConnector.accessMode, 'admin'),
-            inArray(
-              knowledgeConnector.connectorType,
-              providerTypes.filter(
-                (type) =>
-                  availability.memberScoped || !CONNECTOR_META_REGISTRY[type].requiresMemberIdentity
+export const readSearchSourceOverview = instrumentSourceOverviewUseCase(
+  defineAuthorizedKnowledgeUseCase({
+    operation: knowledgeOperations.readSearchSourceOverview,
+    resolveContext: ({ input }: { input: ResourceOwner }) => resolveKnowledgeOwnerContext(input),
+    async execute({ principal, context }): Promise<SearchSourceOverview> {
+      const availability = await measureSearchStage('source_overview.availability', () =>
+        resolveKnowledgeAccessAvailability(context)
+      )
+      const access = createKnowledgeAccessProvider(principal, context)
+      const providerTypes = searchProviderTypes()
+      const configured = configuredSearchSourceCondition(context, providerTypes)
+      const available = or(
+        availability.memberScoped ? eq(knowledgeConnector.accessMode, 'members') : undefined,
+        availability.sourceMirrored
+          ? and(
+              eq(knowledgeConnector.accessMode, 'admin'),
+              inArray(
+                knowledgeConnector.connectorType,
+                providerTypes.filter(
+                  (type) =>
+                    availability.memberScoped ||
+                    !CONNECTOR_META_REGISTRY[type].requiresMemberIdentity
+                )
               )
             )
-          )
-        : undefined
-    )
-    const syncingEnabled = and(
-      available,
-      searchIntegrationAccessCondition(),
-      notInArray(knowledgeConnector.status, ['paused', 'disabled']),
-      or(
-        eq(knowledgeConnector.accessMode, 'admin'),
-        notInArray(knowledgeConnector.memberSyncStatus, ['disabled'])
+          : undefined
+      )
+      const syncingEnabled = and(
+        available,
+        searchIntegrationAccessCondition(),
+        notInArray(knowledgeConnector.status, ['paused', 'disabled']),
+        or(
+          eq(knowledgeConnector.accessMode, 'admin'),
+          notInArray(knowledgeConnector.memberSyncStatus, ['disabled'])
+        )
       )
-    )
-    const documentConditions = and(
-      eq(document.connectorId, knowledgeConnector.id),
-      eq(document.knowledgeBaseId, knowledgeConnector.knowledgeBaseId),
-      eq(document.enabled, true),
-      eq(document.userExcluded, false),
-      isNull(document.archivedAt),
-      isNull(document.deletedAt)
-    )
-    const providers = await configuredProvidersQuery()
-      .where(configured)
-      .limit(MAX_SEARCH_SOURCE_PROVIDER_TYPES)
-    const indexingTypes = new Set<string>()
-    let hasSearchableDocuments = false
-    for await (const accessCondition of knowledgeReadAccessBatches(access, [
-      configured,
-      available,
-      documentConditions,
-    ])) {
-      const readableDocument = and(documentConditions, accessCondition)
-      const [indexing, searchable] = await Promise.all([
-        availability.memberScoped || availability.sourceMirrored
-          ? configuredProvidersQuery()
-              .where(
-                and(
-                  configured,
-                  syncingEnabled,
-                  or(
-                    inArray(knowledgeConnector.status, ['pending', 'syncing']),
-                    and(
-                      eq(knowledgeConnector.accessMode, 'members'),
-                      inArray(knowledgeConnector.memberSyncStatus, ['pending', 'running'])
-                    ),
-                    exists(
-                      db
-                        .select({ id: document.id })
-                        .from(document)
-                        .where(
+      const documentConditions = and(
+        eq(document.connectorId, knowledgeConnector.id),
+        eq(document.knowledgeBaseId, knowledgeConnector.knowledgeBaseId),
+        eq(document.enabled, true),
+        eq(document.userExcluded, false),
+        isNull(document.archivedAt),
+        isNull(document.deletedAt)
+      )
+      const providers = await measureSearchStage('source_overview.providers', () =>
+        configuredProvidersQuery().where(configured).limit(MAX_SEARCH_SOURCE_PROVIDER_TYPES)
+      )
+      annotateSearchDiagnostics({ configuredProviderCount: providers.length })
+      const indexingTypes = new Set<string>()
+      let searchableProbes = 0
+      let hasSearchableDocuments = false
+      for await (const accessCondition of knowledgeReadAccessBatches(access, [
+        configured,
+        available,
+        documentConditions,
+      ])) {
+        const readableDocument = and(documentConditions, accessCondition)
+        const probesSources: boolean = availability.memberScoped || availability.sourceMirrored
+        /** One searchable document is the whole answer, so later batches skip the probe entirely. */
+        const probesSearchable: boolean = probesSources && !hasSearchableDocuments
+        if (probesSearchable) searchableProbes += 1
+        /** Annotated so the searchable probe's guard does not infer through its own result. */
+        const [indexing, searchable]: [{ connectorType: string }[], { id: string }[]] =
+          await Promise.all([
+            probesSources
+              ? measureSearchStage('source_overview.indexing', () =>
+                  configuredProvidersQuery()
+                    .where(
+                      and(
+                        configured,
+                        syncingEnabled,
+                        or(
+                          inArray(knowledgeConnector.status, ['pending', 'syncing']),
                           and(
-                            readableDocument,
-                            inArray(document.processingStatus, ['pending', 'processing'])
+                            eq(knowledgeConnector.accessMode, 'members'),
+                            inArray(knowledgeConnector.memberSyncStatus, ['pending', 'running'])
+                          ),
+                          exists(
+                            db
+                              .select({ id: document.id })
+                              .from(document)
+                              .where(
+                                and(
+                                  readableDocument,
+                                  inArray(document.processingStatus, ['pending', 'processing'])
+                                )
+                              )
                           )
                         )
+                      )
                     )
-                  )
+                    .limit(MAX_SEARCH_SOURCE_PROVIDER_TYPES)
                 )
-              )
-              .limit(MAX_SEARCH_SOURCE_PROVIDER_TYPES)
-          : [],
-        availability.memberScoped || availability.sourceMirrored
-          ? db
-              .select({ id: document.id })
-              .from(document)
-              .innerJoin(knowledgeConnector, eq(knowledgeConnector.id, document.connectorId))
-              .innerJoin(knowledgeBase, eq(knowledgeBase.id, knowledgeConnector.knowledgeBaseId))
-              .where(
-                and(
-                  configured,
-                  available,
-                  readableDocument,
-                  eq(document.processingStatus, 'completed'),
-                  exists(
-                    db
-                      .select({ id: embedding.id })
-                      .from(embedding)
-                      .where(
-                        and(eq(embedding.documentId, document.id), eq(embedding.enabled, true))
+              : [],
+            probesSearchable
+              ? measureSearchStage('source_overview.searchable', () =>
+                  db
+                    .select({ id: document.id })
+                    .from(document)
+                    .innerJoin(knowledgeConnector, eq(knowledgeConnector.id, document.connectorId))
+                    .innerJoin(
+                      knowledgeBase,
+                      eq(knowledgeBase.id, knowledgeConnector.knowledgeBaseId)
+                    )
+                    .where(
+                      and(
+                        configured,
+                        available,
+                        readableDocument,
+                        eq(document.processingStatus, 'completed'),
+                        exists(
+                          db
+                            .select({ id: embedding.id })
+                            .from(embedding)
+                            .where(
+                              and(
+                                eq(embedding.documentId, document.id),
+                                eq(embedding.enabled, true)
+                              )
+                            )
+                        )
                       )
-                  )
+                    )
+                    .limit(1)
                 )
-              )
-              .limit(1)
-          : [],
-      ])
-      for (const provider of indexing) indexingTypes.add(provider.connectorType)
-      hasSearchableDocuments ||= searchable.length > 0
-    }
-    return {
-      providers: providers.map(({ connectorType }) => ({
-        connectorType,
-        isSyncing: indexingTypes.has(connectorType),
-      })),
-      hasSearchableDocuments,
-    }
-  },
-})
+              : [],
+          ])
+        for (const provider of indexing) indexingTypes.add(provider.connectorType)
+        hasSearchableDocuments ||= searchable.length > 0
+      }
+      annotateSearchDiagnostics({ searchableProbeCount: searchableProbes })
+      return {
+        providers: providers.map(({ connectorType }) => ({
+          connectorType,
+          isSyncing: indexingTypes.has(connectorType),
+        })),
+        hasSearchableDocuments,
+      }
+    },
+  })
+)