improvement(knowledge): embed only after every prerequisite holds, refuse a contradicting owner, count the shared fill read

Author: waleedlatif1

apps/sim/lib/knowledge/__integration__/kb-block-search.integration.ts

@@ -132,18 +132,20 @@ describe('API-key KB block fan-out', () => {
           statements.filter((query) => query.includes(fragment))
         /**
          * Every statement runs under the leg's deadline: the candidate search applies it with the
-         * scan settings in one statement, and the probe, the exact ranking, the rerank and
-         * hydration each open with one of their own.
+         * scan settings in one statement, and the probe, the exact ranking and hydration each
+         * open with one of their own. The projection-fill read is shared by the searches that
+         * miss its memo together, so it appears once.
          */
-        expect(matching('statement_timeout')).toHaveLength(bases.length * 5)
+        expect(matching('statement_timeout')).toHaveLength(bases.length * 4 + 1)
         /**
          * A scope this small leaves the bounded traversal short of its candidate limit, so every
          * search probes once and rescues once — never a widening retry loop.
          */
         expect(matching('hnsw.iterative_scan')).toHaveLength(bases.length)
         expect(matching('AS visible')).toHaveLength(bases.length)
         expect(matching(') + 0 LIMIT')).toHaveLength(bases.length)
-        expect(matching('"embedding_search"."id" = ANY(')).toHaveLength(bases.length)
+        /** The walk carries each candidate's identities, so a filled projection reads no page. */
+        expect(matching('"embedding_search"."id" = ANY(')).toHaveLength(0)
         /** The probe enumerates visible documents and reports saturation; it never ranks them. */
         expect(
           statements.filter(

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

@@ -390,51 +390,42 @@ export async function runKnowledgeSearch({
     : undefined
   const resultSecretRegistry = preparedRegistry ?? input.resultSecretRegistry
   input.signal?.throwIfAborted()
-  /** The embedding is requested the moment admission passes, beside the scope and defaults reads. */
-  const admittedEmbedding = async () => {
-    const billingAttribution = await admit()
-    input.signal?.throwIfAborted()
-    const queryEmbedding = hasQuery
-      ? await measureSearchStage('embedding', () =>
-          runWithKnowledgeModelInputProvenance(resultSecretRegistry, () =>
-            generateSearchEmbedding(
-              input.query!,
-              embeddingTarget!,
-              context.workspaceId,
-              input.signal
-            )
-          )
-        )
-      : null
-    return { billingAttribution, queryEmbedding }
-  }
-  const [access, searchDefaults, admitted, tagDefinitions, rerankerCredential] = await Promise.all([
-    measureSearchStage('access_scope', () => context.access.get()),
-    measureSearchStage('defaults', () =>
-      resolveKnowledgeSearchDefaults({
-        workspaceId: context.workspaceId,
-        organizationId: context.organizationId,
+  const [access, searchDefaults, billingAttribution, tagDefinitions, rerankerCredential] =
+    await Promise.all([
+      measureSearchStage('access_scope', () => context.access.get()),
+      measureSearchStage('defaults', () =>
+        resolveKnowledgeSearchDefaults({
+          workspaceId: context.workspaceId,
+          organizationId: context.organizationId,
 
-        /** The signed-in person, if any; never the billing owner or a key's creator. */
-        userId: resolvePrincipalSubjectUserId(principal) ?? undefined,
-        requestedMode: input.searchMode,
-      })
-    ),
-    admittedEmbedding(),
-    /** The tag names the results are labelled with depend on the bases alone. */
-    filters.length === 0
-      ? measureSearchStage('tag_definitions', () =>
-          getDocumentTagDefinitionsByKnowledgeBaseIds(knowledgeBaseIds)
-        )
-      : Promise.resolve(definitionsByKnowledgeBase),
-    /** A surface may ask to rerank; without a key for the workspace or the platform there is nothing to ask. */
-    input.rerankerEnabled && hasQuery
-      ? hasRerankerCredential(context.workspaceId, input.rerankerApiKey)
-      : false,
-  ])
-  const { billingAttribution, queryEmbedding } = admitted
+          /** The signed-in person, if any; never the billing owner or a key's creator. */
+          userId: resolvePrincipalSubjectUserId(principal) ?? undefined,
+          requestedMode: input.searchMode,
+        })
+      ),
+      admit(),
+      /** The tag names the results are labelled with depend on the bases alone. */
+      filters.length === 0
+        ? measureSearchStage('tag_definitions', () =>
+            getDocumentTagDefinitionsByKnowledgeBaseIds(knowledgeBaseIds)
+          )
+        : Promise.resolve(definitionsByKnowledgeBase),
+      /** A surface may ask to rerank; without a key for the workspace or the platform there is nothing to ask. */
+      input.rerankerEnabled && hasQuery
+        ? hasRerankerCredential(context.workspaceId, input.rerankerApiKey)
+        : false,
+    ])
   definitionsByKnowledgeBase = tagDefinitions
   input.signal?.throwIfAborted()
+  /** Requested only once every prerequisite held: a search refused for any reason spends no model call. */
+  const queryEmbedding = hasQuery
+    ? await measureSearchStage('embedding', () =>
+        runWithKnowledgeModelInputProvenance(resultSecretRegistry, () =>
+          generateSearchEmbedding(input.query!, embeddingTarget!, context.workspaceId, input.signal)
+        )
+      )
+    : null
+  input.signal?.throwIfAborted()
   annotateSearchDiagnostics({
     accessScopeKind: access.kind,
     searchMode: searchDefaults.searchMode,

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

@@ -1,5 +1,6 @@
 import type { Principal } from '@sim/auth/principal'
 import { resolvePrincipalSubjectUserId } from '@sim/auth/principal'
+import { OrchestrationError } from '@/lib/core/orchestration/types'
 import { type ResourceOwner, resourceScopeFromOwner } from '@/lib/core/resource-scope'
 import { requireOrganizationSearchAvailable } from '@/lib/knowledge/access/availability'
 import { defineAuthorizedKnowledgeUseCase } from '@/lib/knowledge/application/authorized-knowledge-use-case'
@@ -112,6 +113,13 @@ function defineScopedSearchUseCase<
         ...surface.searchInput(input, context),
         knowledgeBaseIds: [index.id],
       }
+      /** An owner the request asserts is the one that was resolved, or the request names none. */
+      if (
+        (searchInput.organizationId && searchInput.organizationId !== context.organizationId) ||
+        (searchInput.workspaceId && searchInput.workspaceId !== context.workspaceId)
+      ) {
+        throw new OrchestrationError('not_found', 'Knowledge base not found')
+      }
       validateKnowledgeSearchInput(searchInput)
       return runKnowledgeSearch({
         principal,