fix(knowledge): retain model input provenance (#6540)

Author: TheodoreSpeaks

apps/sim/lib/api/contracts/knowledge/search.test.ts

@@ -0,0 +1,22 @@
+import { describe, expect, it } from 'vitest'
+import { internalKnowledgeSearchBodySchema } from '@/lib/api/contracts/knowledge/search'
+import { RESOLVED_SECRET_PROVENANCE_FIELD } from '@/lib/execution/private-tool-metadata'
+
+describe('internal Knowledge search contract', () => {
+  it('retains the private model-input provenance envelope for boundary validation', () => {
+    const provenance = {
+      version: 1 as const,
+      complete: true,
+      entries: [{ name: 'QUERY_SECRET', encryptedValue: 'encrypted-query-secret' }],
+      scope: { userId: 'workflow-owner', workspaceId: 'workspace-1' },
+    }
+
+    expect(
+      internalKnowledgeSearchBodySchema.parse({
+        knowledgeBaseIds: ['knowledge-base-1'],
+        query: 'search query',
+        [RESOLVED_SECRET_PROVENANCE_FIELD]: provenance,
+      })
+    ).toMatchObject({ [RESOLVED_SECRET_PROVENANCE_FIELD]: provenance })
+  })
+})

apps/sim/lib/api/contracts/knowledge/search.ts

@@ -1,5 +1,7 @@
 import { z } from 'zod'
+import { resolvedSecretTraceProvenanceSchema } from '@/lib/api/contracts/primitives'
 import { defineRouteContract } from '@/lib/api/contracts/types'
+import { RESOLVED_SECRET_PROVENANCE_FIELD } from '@/lib/execution/private-tool-metadata'
 import { DEFAULT_RERANKER_MODEL, rerankerModelSchema } from '@/lib/knowledge/reranker-models'
 
 export const knowledgeSearchTagFilterSchema = z.object({
@@ -92,6 +94,7 @@ export const internalKnowledgeSearchBodySchema = z.intersection(
   z.object({
     workflowId: z.string().optional(),
     skipUsageBilling: z.boolean().optional(),
+    [RESOLVED_SECRET_PROVENANCE_FIELD]: resolvedSecretTraceProvenanceSchema.optional(),
   })
 )
 

apps/sim/lib/knowledge/model-input-provenance.test.ts

@@ -3,6 +3,7 @@
  */
 import { encryptionMockFns, environmentUtilsMockFns, resetEnvironmentUtilsMock } from '@sim/testing'
 import { afterEach, describe, expect, it, vi } from 'vitest'
+import { internalKnowledgeSearchBodySchema } from '@/lib/api/contracts/knowledge/search'
 import { PRIVATE_MODEL_INPUT_PROVENANCE_HEADER } from '@/lib/execution/model-input-provenance'
 import {
   RESOLVED_SECRET_PROVENANCE_FIELD,
@@ -86,6 +87,25 @@ describe('Knowledge model input provenance', () => {
     expect(environmentUtilsMockFns.mockGetEffectiveEnvironmentSnapshot).not.toHaveBeenCalled()
   })
 
+  it('accepts a verified envelope after the internal route contract parses it', async () => {
+    const body = internalKnowledgeSearchBodySchema.parse({
+      knowledgeBaseIds: ['knowledge-base-1'],
+      ...verifiedPayload(),
+    })
+
+    const result = await prepareKnowledgeModelInputProvenance({
+      headers: verifiedHeaders(),
+      payload: body,
+      isInternalRequest: true,
+      userId: 'user-1',
+      workspaceId: 'workspace-1',
+      modelInput: body.query,
+    })
+
+    expect(result.success).toBe(true)
+    expect(result.success && result.registry?.isComplete()).toBe(true)
+  })
+
   it('does not activate an authenticated entry absent from the exact model input', async () => {
     environmentUtilsMockFns.mockGetEffectiveEnvironmentSnapshot.mockResolvedValue({
       personalEncrypted: { TOKEN: 'encrypted-token' },